| --- |
| license: apache-2.0 |
| base_model: |
| - Qwen/Qwen3.5-4B |
| - Qwen/Qwen3.5-9B |
| - Qwen/Qwen3.5-35B-A3B |
| library_name: peft |
| tags: |
| - text-to-image |
| - prompt-expansion |
| - diversity |
| - lora |
| --- |
| |
| # T-DIG policies |
|
|
| LoRA adapters for the prompt-expansion policies of [T-DIG: Text-Conditioned Diversity for Image Generation](https://arxiv.org/abs/2610.29003). |
|
|
| Given a prompt and a request for what should differ, a *generation* policy writes K prompts, one per image, that vary the requested axis and keep everything else in the prompt. An *editing* policy instead writes one edit instruction per image of an existing set. The text-to-image (T2I) model is not changed. Each policy is trained with GRPO against one T2I model; it works with others, but best with its own. |
|
|
| ## Contents |
|
|
| One folder per T2I model and policy size, `<T2I model>/<size>-<generate|edit>/`. Each holds `adapter_model.safetensors`, `adapter_config.json` and `selection.json`. `selection.json` gives the training update and the held-out metrics of that checkpoint. `checkpoints.json` lists all of them. |
|
|
| | T2I model | 4B | 9B | 35B-A3B | |
| |---|---|---|---| |
| | Qwen-Image-2512 | generate, edit | generate, edit, generate-paper | generate | |
| | Stable Diffusion 3.5 Large | generate, edit | generate, edit, generate-paper | — | |
| | Stable Diffusion 3.5 Medium | generate, edit | generate, edit, generate-paper | generate, edit | |
| | FLUX.2 [klein] 9B | generate, edit | generate, edit, generate-paper | generate, edit | |
|
|
| - **`9b-generate-paper`**: the checkpoints behind the paper's results (training update 800 for Qwen-Image and SD3.5 Large, 700 for SD3.5 Medium and FLUX.2 [klein]). |
| - **All other folders**: checkpoints chosen on the 758 held-out (prompt, axis) pairs by a fixed rule: |
| 1. Use checkpoints from update 400 on. |
| 2. Drop any checkpoint that often asks the T2I model to write text into the image, or often adds people the prompt and axis do not mention. |
| 3. Among the rest, take the best VQAScore within 0.04 D_a of the most diverse. |
| |
| `selection_note` marks the runs where no checkpoint met the second step's limits, so the least affected one was taken. |
|
|
| These are the recommended checkpoints for use, and differ in places from the paper's. Training past these points tends to raise automatic scores through such shortcuts rather than better images. |
|
|
| ## Use |
|
|
| The policy expects the instruction the paper trained it on; `tdig.instruction` in the [code repository](https://github.com/simran-khanuja/t-dig) builds it. Generation, with vLLM: |
|
|
| ```python |
| from vllm import LLM, SamplingParams |
| from vllm.lora.request import LoRARequest |
| from transformers import AutoTokenizer |
| from tdig.instruction import instruction # from the t-dig repository |
| |
| base, adapter = "Qwen/Qwen3.5-9B", "flux2-klein-9b/9b-generate" |
| prompt, axis = "a dog playing in a park", "the time of day" |
| text = instruction(prompt, axis, k=4, t2i_model="flux2-klein-9b", word_cap=len(prompt.split()) + 20) |
| tok = AutoTokenizer.from_pretrained(base) |
| chat = tok.apply_chat_template([{"role": "user", "content": text}], tokenize=False, |
| add_generation_prompt=True, enable_thinking=False) |
| llm = LLM(model=base, enable_lora=True, max_lora_rank=32) |
| out = llm.generate([chat], SamplingParams(temperature=1.0, max_tokens=1024), |
| lora_request=LoRARequest("t-dig", 1, adapter))[0].outputs[0].text |
| # out holds 4 <prompt>...</prompt> blocks; generate one image from each with the T2I model. |
| ``` |
|
|
| - **Diversity strength:** scaling the adapter's `lora_B` weights by α between 0 and 1 interpolates between the untrained model (α=0) and the trained policy (α=1). The interactive framework in the code repository serves this as a slider. |
| - **Editing policies:** they also see the K current images. Use `tdig.instruction.edit_instruction` and the framework's editing path. |
|
|
| ## Training |
|
|
| - **Method:** GRPO with set size K=4 and 8 completions per (prompt, axis) pair. |
| - **Adapter:** LoRA rank 32. |
| - **Optimizer:** learning rate 2e-5, no KL term. |
| - **Reward:** |
| - TPIPS-Vendi diversity along the axis, |
| - minus a hinge on VQAScore alignment below the T2I model's own i.i.d. level, |
| - a word-budget penalty, |
| - a small hinge on HPSv3 below the i.i.d. level. |
| - **Data:** about 76k (prompt, axis) pairs from 19.8k prompts, released as [neulab/t-dig-train](https://huggingface.co/datasets/neulab/t-dig-train). |
|
|
| ## Limitations |
|
|
| - **Unrequested additions:** trained policies sometimes add content the user did not ask for (most often extra people) or, for some T2I models, ask for text in the image. `selection.json` reports how often, per checkpoint. |
| - **SD3.5 Medium editing policies** show this the most. |
|
|
| ## Citation |
|
|
| ```bibtex |
| @misc{khanuja2026tdig, |
| title={T-DIG: Text-Conditioned Diversity for Image Generation}, |
| author={Simran Khanuja and Rohit Gandikota and Gaurav Parmar and Graham Neubig}, |
| year={2026}, |
| eprint={2610.29003}, |
| archivePrefix={arXiv}, |
| primaryClass={cs.CV}, |
| url={https://arxiv.org/abs/2610.29003} |
| } |
| ``` |
|
|
