Spaces:
Running on Zero
Running on Zero
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import random | |
| import re | |
| import spaces # must precede torch / CUDA-touching imports | |
| import torch | |
| import gradio as gr | |
| from transformers import AutoTokenizer, Qwen3_5ForConditionalGeneration | |
| from peft import PeftModel | |
| from diffusers import Flux2KleinPipeline | |
| BASE_POLICY = "Qwen/Qwen3.5-4B" | |
| POLICY_ADAPTER_REPO = "neulab/t-dig-checkpoints" | |
| POLICY_ADAPTER_SUBFOLDER = "flux2-klein-9b/4b-generate" | |
| RENDERER = "black-forest-labs/FLUX.2-klein-9B" | |
| K = 4 | |
| # T-DIG generation instruction, verbatim (instruction_version tdig-v4-2 from | |
| # neulab/t-dig-data, renderer flux2-klein-9b). word_cap = len(prompt.split()) + 20. | |
| INSTRUCTION_TEMPLATE = """You are given an image-generation prompt and a diversity axis. | |
| Base prompt: {prompt} | |
| Axis: {axis} | |
| Write exactly 4 image-generation prompts which vary the main prompt along the given axis. | |
| RULES, in order of importance: | |
| 1. THE BASE PROMPT'S CHOICES ARE CONSTRAINTS, NOT VARIABLES. Anything it states | |
| explicitly must survive in every prompt: every subject, its count, its kind, and every | |
| attribute it names. If the base prompt says "two men", every prompt says two men -- not | |
| a woman, not three men, not a crowd. | |
| 2. WHERE THE BASE PROMPT ALREADY FIXES SOMETHING, VARY WITHIN IT, NEVER AWAY FROM IT. | |
| The axis can only open up what the prompt left open. Contradicting a stated detail is | |
| the single worst thing you can do here. | |
| 3. Vary ONLY the given axis, and together cover clearly distinct values of it. | |
| 4. KEEP EACH PROMPT UNDER {cap} WORDS. A short prompt keeps the renderer on the base | |
| prompt; every extra description competes with it and pulls the image away from what was | |
| asked. Add the axis change and stop. | |
| 5. Hold everything else as close to the base prompt as you can -- ideally reuse its exact | |
| wording and only substitute the part the axis changes. | |
| Output exactly 4 prompts, each wrapped in <prompt></prompt> tags, and nothing else.""" | |
| def build_instruction(prompt: str, axis: str) -> str: | |
| cap = len(prompt.split()) + 20 | |
| return INSTRUCTION_TEMPLATE.format(prompt=prompt.strip(), axis=axis.strip(), cap=cap) | |
| def parse_prompts(text: str) -> list: | |
| found = re.findall(r"<prompt>(.*?)</prompt>", text, flags=re.S) | |
| prompts = [p.strip().strip('"') for p in found] | |
| prompts = [p for p in prompts if p] | |
| if not prompts: # tag format missed -> fall back to non-empty lines | |
| prompts = [ln.strip() for ln in text.strip().splitlines() if ln.strip()] | |
| return prompts[:K] | |
| print("Loading T-DIG policy: Qwen3.5-4B + LoRA adapter...") | |
| policy_tok = AutoTokenizer.from_pretrained(BASE_POLICY) | |
| policy = Qwen3_5ForConditionalGeneration.from_pretrained( | |
| BASE_POLICY, dtype=torch.bfloat16, attn_implementation="sdpa" | |
| ) | |
| policy = PeftModel.from_pretrained( | |
| policy, POLICY_ADAPTER_REPO, subfolder=POLICY_ADAPTER_SUBFOLDER, torch_device="cpu" | |
| ) | |
| policy = policy.to("cuda").eval() | |
| # Diversity strength alpha in [0, 1]: scale every LoRA B matrix by alpha. | |
| # peft applies (W + scaling * B @ A) with scaling = lora_alpha / r; rescaling the | |
| # module's `scaling` dict interpolates between the untrained base (0) and the | |
| # trained policy (1) without touching weights. | |
| _LORA_SCALING_1 = {} | |
| for _name, _module in policy.named_modules(): | |
| _s = getattr(_module, "scaling", None) | |
| if isinstance(_s, dict) and "default" in _s: | |
| _LORA_SCALING_1[_name] = _s["default"] | |
| def set_diversity_strength(alpha: float): | |
| alpha = float(min(max(alpha, 0.0), 1.0)) | |
| for name, module in policy.named_modules(): | |
| s = getattr(module, "scaling", None) | |
| if isinstance(s, dict) and "default" in s: | |
| module.scaling["default"] = _LORA_SCALING_1[name] * alpha | |
| print("T-DIG policy ready. Loading FLUX.2 [klein] 9B renderer...") | |
| renderer = Flux2KleinPipeline.from_pretrained(RENDERER, torch_dtype=torch.bfloat16) | |
| renderer.to("cuda") | |
| print("FLUX.2 [klein] 9B ready.") | |
| def generate(prompt: str, axis: str, diversity_strength: float = 1.0, seed: int = 0, | |
| randomize_seed: bool = True, progress=gr.Progress(track_tqdm=True)): | |
| """Generate a diverse set of 4 images that vary one prompt along one axis. | |
| The T-DIG policy (LoRA-tuned Qwen3.5-4B) rewrites the base prompt into 4 | |
| axis-conditioned variants, and FLUX.2 [klein] 9B renders one image per variant. | |
| Args: | |
| prompt: the base image-generation prompt. | |
| axis: what should differ between the images, e.g. "the time of day". | |
| diversity_strength: alpha in [0,1]; 0 = untrained Qwen3.5, 1 = full T-DIG policy. | |
| seed: RNG seed for reproducibility. | |
| randomize_seed: draw a fresh seed on every run. | |
| """ | |
| prompt = (prompt or "").strip() | |
| axis = (axis or "").strip() | |
| if not prompt: | |
| raise gr.Error("Please provide a base prompt.") | |
| if not axis: | |
| raise gr.Error("Please provide a diversity axis — what should vary between the images.") | |
| if randomize_seed: | |
| seed = random.randint(0, 2**31 - 1) | |
| seed = int(seed) | |
| # 1) Policy: base prompt + axis -> 4 axis-conditioned prompts | |
| chat = policy_tok.apply_chat_template( | |
| [{"role": "user", "content": build_instruction(prompt, axis)}], | |
| tokenize=False, | |
| add_generation_prompt=True, | |
| enable_thinking=False, | |
| ) | |
| inputs = policy_tok([chat], return_tensors="pt").to("cuda") | |
| set_diversity_strength(diversity_strength) | |
| with torch.no_grad(): | |
| out = policy.generate(**inputs, max_new_tokens=1024, do_sample=True, temperature=1.0) | |
| raw = policy_tok.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True) | |
| variants = parse_prompts(raw) | |
| while len(variants) < K: # degenerate output: fill with what we have | |
| variants.append(variants[-1] if variants else prompt) | |
| # 2) Renderer: one image per variant (sequential to keep VRAM flat) | |
| generator = torch.Generator("cuda").manual_seed(seed) | |
| images = [] | |
| for variant in variants: | |
| images.extend( | |
| renderer( | |
| prompt=variant, | |
| height=1024, | |
| width=1024, | |
| guidance_scale=1.0, | |
| num_inference_steps=4, | |
| generator=generator, | |
| ).images | |
| ) | |
| shown = "\n\n".join(f"**Image {i + 1}:** {v}" for i, v in enumerate(variants)) | |
| return images, shown, seed | |
| CSS = """ | |
| #col-container { max-width: 1200px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo: | |
| gr.Markdown( | |
| """ | |
| # T-DIG: Text-Conditioned Diversity for Image Generation | |
| Give **one prompt** and **one thing that should vary**. A LoRA-tuned Qwen3.5-4B | |
| policy rewrites the prompt into 4 axis-conditioned variants, and FLUX.2 [klein] 9B | |
| renders one image per variant — a set that varies what you asked for and nothing else. | |
| [Paper](https://arxiv.org/abs/2610.29003) · [T-DIG checkpoints](https://huggingface.co/neulab/t-dig-checkpoints) | |
| """ | |
| ) | |
| with gr.Column(elem_id="col-container"): | |
| prompt = gr.Textbox( | |
| label="Base prompt", | |
| placeholder="a baker decorating a square cake", | |
| lines=2, | |
| ) | |
| axis = gr.Textbox( | |
| label="Diversity axis — what should vary between the images?", | |
| placeholder="the theme of the cake", | |
| lines=1, | |
| ) | |
| run = gr.Button("Generate a diverse set", variant="primary") | |
| gallery = gr.Gallery(label="Diverse image set", columns=4, height=340) | |
| variants_md = gr.Markdown() | |
| with gr.Accordion("Advanced settings", open=False): | |
| with gr.Row(): | |
| diversity_strength = gr.Slider( | |
| 0.0, 1.0, value=1.0, step=0.05, | |
| label="Diversity strength (α)", | |
| info="0 = untrained Qwen3.5, 1 = full T-DIG policy", | |
| ) | |
| seed = gr.Number(label="Seed", value=0, precision=0) | |
| randomize_seed = gr.Checkbox(label="Randomize seed", value=True) | |
| gr.Examples( | |
| examples=[ | |
| ["a baker decorating a square cake", "the theme of the cake"], | |
| ["a family of tigers playing in their natural habitat", "an unrealistic way the tigers could be having fun"], | |
| ["A Labrador playing on a riverbank", "colors of the Labrador"], | |
| ["a cozy reading nook with an armchair and floor lamp", "the reading nook decor style"], | |
| ["a handcrafted wooden sailboat at a quiet harbor", "the weather and sky conditions"], | |
| ["a steaming bowl of ramen on a small cafe counter", "who is around the cafe counter"], | |
| ], | |
| inputs=[prompt, axis], | |
| outputs=[gallery, variants_md, seed], | |
| fn=generate, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| run.click( | |
| generate, | |
| inputs=[prompt, axis, diversity_strength, seed, randomize_seed], | |
| outputs=[gallery, variants_md, seed], | |
| api_name="generate", | |
| ) | |
| demo.launch(mcp_server=True) | |