Multi-Image Generation with In-Context LoRA

This repo aims to generate and evaluate coherent 2D multi-view scenes (multiple images with intrinsic relationships, such as different viewpoints of the same object or scene) using minimal training data and limited compute. The approach builds on In-Context LoRA, a method to adapt diffusion transformer (DiT) models to multi-image outputs without changing the model architecture. The key idea is to concatenate multiple images into one larger image during training, use a joint caption describing all sub-images, and fine-tune on a small dataset . The limitation of the paper is that only qualitative visual evaluation has been performed. This repo aims to implement a quantitative approach for evaluating In-Context LoRA vs Base Flux.1 model. Essentialy, we are trying to answer: Does LoRA preserves object identity without duplicating views?

Project status

The 4x4 grid composite images have been automatically captioned with Gemini 3.5 Flash. This run over every eligible four-view instance under data/ produced 423 accepted training pairs and retained 106 ambiguous composites for abstention analysis. Another model can be used instead of Gemini 3.5 Flash for captioning. The minimal Study 1 pilot split, training config, paired generation harness, and blinded scorecard are implemented in configs/study1_pilot.yaml The gold benchmark, model comparison, confidence calibration, and human review described in PLANS.md detail future work.

Latest dataset build

Item	Result
Input root	data/
Model	gemini-3.5-flash
Eligible four-view instances	529
Selected source images	2,116
Accepted image/caption pairs	423
Abstention composites	106
Invalid cached responses after repair	0
Unexpected processing failures	0

Generated data is written to training/composites_4view_grid_all, with ambiguous examples in training/composites_4view_grid_all_abstention.

Current structured-data examples

The following contact sheets sample accepted 2×2 composites together with their structured captions. The first shows two examples:

The second shows four examples:

If you want to display another deterministic contact sheet with the captioned composite images RUN:

.venv/bin/python -m src.sanity_check \
  --input-dir training/composites_4view_grid_all \
  --output sanity_check.png \
  --count 4 \
  --columns 4 \
  --seed 17

Research questions

• Does LoRA preserves object identity without duplicating views?
• Which Gemini model annotates viewpoints best?

1. Dataset Preparation (Multi-View Image Sets)

Collect or curate groups containing multiple views of the same object or scene. MVImgNet is one suitable source. The current dataset build found 529 eligible instances and deterministically selected four spaced views from each; structured annotation accepted 423 of the resulting composites. The original small-scale experiments used 126 selected images, while the current Study 1 pilot uses the larger accepted dataset described above.

2. Automatic Caption Generation for Multi-Image Scenes

python -m src.dataset_builder generates captioned 2x2 composites from an MVImgNet-style directory. Gemini credentials are loaded only when captioning starts. Set GOOGLE_API_KEY in your environment or a local .env file:

export GOOGLE_API_KEY=your-key
python -m src.dataset_builder \
  --objects-dir data/ \
  --category-file mvimgnet_category.txt \
  --output-dir training/composites_4view_grid \
  --abstention-dir training/composites_4view_grid_abstention \
  --cache-dir .gemini_cache \
  --model gemini-3.5-flash \
  --limit 63 \
  --tile-width 512 \
  --tile-height 512

Gemini 3.5 Flash (gemini-3.5-flash) returns a JSON annotation constrained by the MultiviewAnnotation Pydantic output schema rather than a free-form caption. Each tile receives an absolute horizontal viewpoint, side, vertical angle, framing, visible features, and confidence. Python validates the annotation and renders the final caption deterministically.

Annotations containing an indeterminate viewpoint field are not written into the LoRA training directory. Their composites and annotations are retained in the abstention directory for evaluating whether the vision-language model declines ambiguous orientation judgments. If --abstention-dir is omitted, it defaults to a sibling named <output-dir>_abstention.

The cache consists of versioned JSON envelopes. Each entry preserves the exact raw model text, parsed annotation, validation result, model and prompt versions, latency, composite-image hash, and hashes of the four source images. Old free-form .txt cache entries are left untouched and are not reused by the structured pipeline.

For hundreds of composites, use the resumable Gemini Batch API workflow instead of the synchronous builder. Preparation creates keyed JSONL chunks, submission records each job before continuing, and collection downloads, validates, caches, and renders the completed dataset:

python -m src.batch_dataset_builder prepare \
  --objects-dir data \
  --category-file data/mvimgnet_category.txt \
  --output-dir training/composites_4view_grid_all \
  --abstention-dir training/composites_4view_grid_all_abstention \
  --cache-dir .gemini_cache \
  --work-dir .gemini_batch/all_data \
  --model gemini-3.5-flash
python -m src.batch_dataset_builder submit --work-dir .gemini_batch/all_data
python -m src.batch_dataset_builder collect --work-dir .gemini_batch/all_data --wait

The commands are resumable. state.json records every submitted job ID, successful responses are stored in the versioned cache, and rerunning prepare excludes valid cached annotations while retrying invalid ones. The July 2026 run used the paths in the example above and produced the counts reported in Latest dataset build.

For each image set, we need a single descriptive caption that encompasses all views/images. Writing these by hand is possible but to ensure scalability and consistency, we can automate caption generation using multimodal models: The captioning pipeline defaults to Gemini 3.5 Flash (gemini-3.5-flash) for each image set. Override it explicitly with --model when running another model. The prompt works best when the images have already been concatenated into a single composite image (also cheaper).

2. Fine-tuning In-Context LoRA on the FLUX model

To train the model, we will use the In-Context LoRA approach on a diffusion transformer model like FLUX.

LoRA (Low-Rank Adaptation) inserts trainable low-rank weight matrices into the model (typically into attention layers) and freezes the original model weights. This drastically reduces the number of parameters that need updating (and thus memory usage), making it feasible to train on a single high-end RTX 4090 NVIDIA GPU. I trained a LoRA adapter on the FLUX model using the prepared multi-image dataset using the repository AI-toolkit and used as inspiration the In-Context LoRA config file.

The repository retains the original config_4views.yaml and now includes the reproducible pilot config at configs/study1_pilot.yaml.

RUN Fine-tuning with LoRA

You can remove the Wandb login (I use it for monitoring the loss). You need to login into HuggingFace to accept the FLUX model license usage.

../ai-toolkit/.venv/bin/wandb login
../ai-toolkit/.venv/bin/huggingface-cli login

WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 \
  ../ai-toolkit/.venv/bin/python \
  ../ai-toolkit/run.py \
  configs/study1_pilot.yaml

Development and reproducibility

For CPU-only preprocessing development and tests, install the lightweight project dependencies rather than the CUDA training stack:

python3.11 -m venv .venv
.venv/bin/python -m pip install -e ".[test]"
.venv/bin/python -m ruff check .
.venv/bin/python -m pytest

requirements.txt remains the ** CUDA 12.1** environment for model training. Install it only on a compatible GPU system. To run the local Gradio demo, supply the checkpoint path explicitly:

python app.py --lora-model models/4views.safetensors

To renumber an existing directory of image/caption pairs:

python -m src.rename_files training/composites_4view_grid --start 1

Study 1 Pilot Fine-Tuning LoRA

The minimal pilot uses a deterministic, source-instance-level 90/10 split of the 423 accepted pairs. Seed 17 produces 381 training pairs and 42 holdout pairs. The split command verifies matching PNG/TXT stems and rejects source-instance, composite image SHA-256, or selected source-image SHA-256 overlap across partitions:

.venv/bin/python -m src.study1_split

This writes training/study1_pilot/train, training/study1_pilot/holdout, and training/study1_pilot/split_manifest.jsonl. The manifest records every source instance, destination path, split parameter, composite hash, and selected source image hashes. Re-running the command is idempotent when the existing files match; it refuses to overwrite a different file or accept stale extra pairs.

Training uses the neighboring AI Toolkit checkout already expected by this project. From this repository root, launch the 500-step pilot explicitly with:

python3.11 ../ai-toolkit/run.py configs/study1_pilot.yaml

The config trains one rank-16/alpha-16 LoRA and samples two fixed holdout captions every 100 steps. This repository never launches that GPU command automatically.

After selecting the desired AI Toolkit checkpoint, generate the controlled base/LoRA pairs with the eight synthetic prompts and two fixed seeds:

.venv/bin/python -m evaluation.generate_pairs \
  --lora output/study1_pilot/study1_pilot.safetensors \
  --cpu-offload

The generator produces 32 grids: eight prompts × seeds 1001 and 1002 × base and LoRA. Both conditions use guidance 3.5, 20 steps, and 1024×1024 resolution. It recreates the CPU-backed torch generator immediately before every condition, leaves the LoRA unfused, and writes all settings, hashes, and output paths to evaluation/outputs/study1_pilot/generation_manifest.jsonl.

Create randomized A/B image copies, a blank scoring CSV, and a separate condition key with:

.venv/bin/python -m evaluation.create_blinded_scorecard \
  --generation-manifest evaluation/outputs/study1_pilot/generation_manifest.jsonl \
  --output-dir evaluation/outputs/study1_pilot_blinded \
  --blind-seed 17

Give raters scorecard.csv and the adjacent images/ directory, but withhold blind_key.jsonl until scoring is complete. Run the focused CPU checks with:

.venv/bin/python -m pytest tests/test_study1_split.py tests/test_evaluation_pairing.py

Implemented

• Deterministic four-view selection and 2×2 composite construction.
• Duplicate source-image detection and hashing.
• Gemini model selection through --model.
• Pydantic structured annotation schema.
• Controlled viewpoint labels and explicit indeterminate abstention.
• Raw-response caching with model, prompt, latency, and image metadata.
• Deterministic caption rendering.
• Separate accepted and abstention manifests.
• Resumable Batch API preparation, submission, polling, and collection.
• Keyed JSONL chunks with persisted job IDs and per-request result matching.
• Gemini-compatible flattened structured-output schema.
• Full Gemini 3.5 Flash batch run over all 529 eligible instances in data/.
• Existing FLUX LoRA checkpoint
• Unit tests for annotation, cache, rendering, preprocessing, and batch payloads.
• Seed-17 instance-level Study 1 pilot split with leakage checks.
• Minimal 500-step Study 1 LoRA config with fixed holdout monitor prompts.
• Paired base-FLUX/LoRA generation manifest and blinded scoring CSV tooling.

Next steps

• Review the 106 abstentions and spot-check the 423 accepted annotations before treating the generated output as training data.
• Build and adjudicate the 100-composite gold benchmark.
• Benchmark the four selected Gemini models.
• Calibrate the low-confidence threshold.
• Add benchmark, strict-pose, and appearance-only dataset exports.
• Verify that identity captions never contain pose labels or indeterminate.
• Train pose-conditioned and appearance-only LoRAs with three seeds each.
• Run the Study 1 pilot training, paired generation, and blinded human scoring.
• Add DINOv2, DreamSim, LPIPS, and perceptual-hash supporting metrics.

Try the Demo on Hugging Face :hugging_face:

Here is the Hugging Face Model Card and the Hugging Face Demo Space.

Historical example outputs :image:

These outputs and free-form captions predate the current structured annotation pipeline. They are retained as historical examples and are not the final research caption format.

 [FOUR-VIEWS] a red desk lamp from multiple views;[TOP-LEFT] This photo shows a 45-degree angle of desk lamp;[TOP-RIGHT] This photo shows a high-angle shot of the lamp; [BOTTOM-LEFT] Here is a side view shot of lamp; [BOTTOM-RIGHT] The back view of the desk lamp.

[FOUR-VIEWS] a bedroom from multiple views;[TOP-LEFT] This photo shows a 45-degree angle of the bedroom;[TOP-RIGHT] This photo shows a high-angle shot of the bedroom; [BOTTOM-LEFT] Here is a side view shot of bedroom; [BOTTOM-RIGHT] A low angle view of the bedroom.