StructGen

This repository contains the dataset utilities, training code, and inference code for StructGen: Disambiguating Multi-Reference Image Generation via Structured Context Modeling.

Overview

StructGen represents each referenced entity or scene as an ordered structured-context entry. The repository provides code to parse the Structured Context Schema, BAGEL-compatible model definitions, and training and inference entrypoints.

Installation

The documented runtime is Python 3.10 with PyTorch 2.5.1 on CUDA/Linux.

conda create -n structgen python=3.10 -y
conda activate structgen
pip install -r requirements.txt

FlashAttention is required by the BAGEL attention kernels and is installed separately after PyTorch and the CUDA build toolchain are available:

pip install flash_attn==2.5.8 --no-build-isolation

Data and Benchmark

The StructGen data package and final 400-case benchmark are available from the StructGen Hugging Face dataset repository.

For MICo-De&Re training, pair the input and target images from the official MICo-150K dataset with StructGen's structured-prompt annotations. Dataset paths and layouts are documented in the reproduction notes.

Training

Configure a BAGEL-compatible base model and the dataset paths described in the reproduction notes. Training uses train/train_structgen.py with the dataset mixture in data/configs/structgen.yaml; the complete distributed launch command is provided in the reproduction notes.

Training writes step-numbered checkpoint directories. Each saved step directory under --checkpoint_dir contains model.safetensors.

Inference

After training, choose a saved model.safetensors checkpoint and prepare a structured-context input file. examples/structured_context_example.json provides a minimal input example.

bash scripts/run_inference.sh \
  --gpu 0 \
  --placeholder_img_prob 1.0 \
  --test_jsonl /path/to/input.jsonl \
  --data_dir /path/to/images \
  --model_path /path/to/BAGEL-7B-MoT \
  --ckpt_path outputs/train/checkpoints/0002000/model.safetensors \
  --output_dir outputs/structured_context_images \
  --seed 42

See the reproduction notes for the Structured Context Schema, inference controls, and validation commands.

License

StructGen's original code is released under the Apache License 2.0. See THIRD_PARTY_NOTICES.md for third-party components; datasets and images remain subject to their source licenses.