T-DIG policies

LoRA adapters for the prompt-expansion policies of T-DIG: Text-conditioned Diversity for Image Generation (under review).

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 builds it. Generation, with vLLM:

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.

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

Citation information will be added after review.