import spaces # must be imported before torch on ZeroGPU import json import os import random import gradio as gr import torch from diffusers import QwenImage21Pipeline from rewriters import STUDENT_IDS, TEACHER_ID, load_student, load_teacher, size_for, student_rewrite, teacher_rewrite IMAGE_ID = "Qwen/Qwen-Image-2.1" DATASET_ID = "ML-Intern-lab/Qwen-Image-2.1-rewriter-distill" MAX_SEED = 2**31 - 1 ARMS = ["Raw request (no rewriter)", "Official 9B rewriter", "0.8B pocket rewriter", "2B pocket rewriter"] DEFAULT_ARMS = ARMS[1:] HERE = os.path.dirname(os.path.abspath(__file__)) with open(os.path.join(HERE, "sheets.json"), encoding="utf-8") as f: SHEETS = json.load(f) # [{"id","request","lang","category","sheet_url"}] SHEET_CHOICES = [(f"[{s['lang']}] {s['request'][:90]}", s["id"]) for s in SHEETS] OVERVIEW_URL = f"https://huggingface.co/datasets/{DATASET_ID}/resolve/main/renders/sheets/sheet_overview.png" # ----------------------------------------------------------------- models (module scope, eager) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") pipe = QwenImage21Pipeline.from_pretrained(IMAGE_ID, dtype=torch.bfloat16).to(device) teacher_tok, teacher, teacher_system = load_teacher(device) students = {size: load_student(size, device) for size in STUDENT_IDS} def est_seconds(arms, steps, megapixels, teacher_tokens) -> int: t = 10 if ARMS[1] in arms: t += 20 + int(teacher_tokens) * 0.035 # 9B with thinking t += 8 * sum(a in arms for a in ARMS[2:]) # students t += len(arms) * (10 + int(steps) * 0.6 * float(megapixels)) return int(t) def compare_duration(request, arms=DEFAULT_ARMS, steps=28, megapixels=1.0, seed=0, randomize_seed=True, teacher_tokens=3072, *a, **k): return est_seconds(arms, steps, megapixels, teacher_tokens) def _rewrite(arm, request, seed, teacher_tokens): if arm == ARMS[0]: return {"prompt": request, "ratio": "3:2", "parse_ok": True, "tokens": 0, "seconds": 0.0, "thinking": ""} if arm == ARMS[1]: return teacher_rewrite(teacher_tok, teacher, teacher_system, request, seed=seed, max_new_tokens=int(teacher_tokens)) tok, model = students["0.8B" if arm == ARMS[2] else "2B"] return student_rewrite(tok, model, request, seed=seed) def _stats(arm, r, w=None, h=None): if arm == ARMS[0]: s = "No rewrite." else: s = f"{r['tokens']} tokens in {r['seconds']:.1f} s" if r.get("thinking"): s += f" (thinking included)" s += f". Ratio {r['ratio']}." if not r["parse_ok"]: s += " Not valid JSON; raw text used." if w: s += f" Rendered {w}x{h}." return s @spaces.GPU(size="xlarge", duration=compare_duration) def compare(request: str, arms: list = DEFAULT_ARMS, steps: int = 28, megapixels: float = 1.0, seed: int = 0, randomize_seed: bool = True, teacher_tokens: int = 3072): """Rewrite one request with each selected rewriter, then render every rewrite with Qwen-Image 2.1 at the same seed. Yields progressively: prompts first, then one image per arm.""" request = (request or "").strip() if not request: raise gr.Error("Type an image request first.") if not arms: raise gr.Error("Select at least one arm.") if randomize_seed: seed = random.randint(0, MAX_SEED) seed = int(seed) cols = {a: {"image": None, "prompt": "", "stats": "waiting" if a in arms else "not selected"} for a in ARMS} def emit(msg): return tuple([cols[a]["image"] for a in ARMS] + [cols[a]["prompt"] for a in ARMS] + [cols[a]["stats"] for a in ARMS] + [seed, msg]) rewrites = {} for a in ARMS: if a not in arms: continue cols[a]["stats"] = "rewriting..." yield emit(f"Rewriting with {a}...") r = _rewrite(a, request, seed, teacher_tokens) rewrites[a] = r cols[a]["prompt"], cols[a]["stats"] = r["prompt"], _stats(a, r) yield emit(f"Rewrote with {a}.") for a in ARMS: if a not in arms: continue cols[a]["stats"] += " Rendering..." yield emit(f"Rendering {a}...") r = rewrites[a] w, h = size_for(r["ratio"], megapixels) image = pipe(prompt=r["prompt"], width=w, height=h, num_inference_steps=int(steps), generator=torch.Generator(device="cuda").manual_seed(seed)).images[0] cols[a]["image"], cols[a]["stats"] = image, _stats(a, r, w, h) yield emit(f"Rendered {a}.") yield emit(f"Done. Seed {seed}, {int(steps)} steps, about {megapixels} megapixels per image.") def show_sheet(sheet_id): s = next((x for x in SHEETS if x["id"] == sheet_id), None) if not s: return None, "" return s["sheet_url"], f"**Request ({s['lang']}, {s['category']}):** {s['request']}" # ----------------------------------------------------------------- UI CSS = """ #col-container { max-width: 1500px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ INTRO = f""" # ⚖️ Qwen-Image 2.1 Rewriter Arena One request, several prompt rewriters, same seed, side by side. Compare the **official 9B rewriter** ([{TEACHER_ID}](https://huggingface.co/{TEACHER_ID}): 1,700-word system prompt, thinking on, about 1,600 tokens per rewrite) against its two pocket distillations, [0.8B](https://huggingface.co/{STUDENT_IDS['0.8B']}) and [2B](https://huggingface.co/{STUDENT_IDS['2B']}) (no system prompt, no thinking, about 450 tokens), and the raw request with no rewriting. Each rewrite is rendered by [Qwen-Image 2.1](https://huggingface.co/{IMAGE_ID}) at the ratio that rewriter chose. A live run with the 9B arm takes 2 to 3 minutes of GPU time. The **Pre-rendered comparisons** tab shows the 40-request evaluation grid at no cost. Research demo under the Qwen Research License (non-commercial). Built with Qwen. """ ABOUT = """ **What the 40-request evaluation found.** Rendered-text OCR word accuracy: raw request 0.46, official 9B rewriter 0.80, 0.8B student 0.57, 2B student 0.55. The students keep roughly a quarter to a third of the teacher's gain, close to the teacher on Latin text and further behind on Chinese and Japanese. A pairwise vision-model judge could not separate the arms. On text-level checks over 300 held-out requests the students match the teacher on valid JSON and allowed ratios (99 to 100%) and on keeping the user's quoted text (53% and 60% versus 53%), at about 28% of the tokens. **How to read a live run.** Look for the quoted text in each image, whether the ratio suits the subject, and whether the rewrite invented things the request did not ask for. The seed is shared, but each arm renders at its own ratio, so compositions differ. **Fairness notes.** The teacher runs with its official protocol: system prompt, thinking, presence penalty 1.5, sampling at temperature 1.0. The token cap for the teacher is adjustable; its mean is about 1,600 and a few requests need more. The students use the same sampling with no system prompt. All arms use 28 steps here for speed; the evaluation used 40. Data, predictions and all 41 contact sheets: [ML-Intern-lab/Qwen-Image-2.1-rewriter-distill](https://huggingface.co/datasets/ML-Intern-lab/Qwen-Image-2.1-rewriter-distill). The students were trained end to end by ML Intern in HuggingChat for about 16 USD. """ EXAMPLES = [[s["request"]] for s in SHEETS[:12]] with gr.Blocks(title="Qwen-Image 2.1 Rewriter Arena") as demo: with gr.Column(elem_id="col-container"): gr.Markdown(INTRO) with gr.Tabs(): with gr.Tab("Live comparison"): with gr.Row(): request = gr.Textbox(label="Your request (any language)", lines=2, scale=4, placeholder='a man at a bus stop "next bus 15 min"') run = gr.Button("Compare", variant="primary", scale=1) arms = gr.CheckboxGroup(ARMS, value=DEFAULT_ARMS, label="Arms to run") with gr.Accordion("Settings", open=False): steps = gr.Slider(8, 50, value=28, step=1, label="Steps") megapixels = gr.Slider(0.5, 1.5, value=1.0, step=0.25, label="Megapixels per image") teacher_tokens = gr.Slider(1024, 6144, value=3072, step=256, label="Token cap for the 9B rewriter", info="Its rewrite plus thinking averages about 1,600 tokens") seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed (shared by all arms)") randomize_seed = gr.Checkbox(True, label="Randomize seed") status = gr.Markdown("") images, prompts, stats = [], [], [] with gr.Row(): for a in ARMS: with gr.Column(): images.append(gr.Image(label=a, format="png", image_mode="RGBA", height=360)) stats.append(gr.Markdown("")) prompts.append(gr.Textbox(label="Prompt sent to the image model", lines=6, interactive=False)) gr.Examples(examples=EXAMPLES, inputs=[request], label="Held-out evaluation requests (fills the box; press Compare)", cache_examples=False) gr.on([run.click, request.submit], compare, [request, arms, steps, megapixels, seed, randomize_seed, teacher_tokens], images + prompts + stats + [seed, status], api_name="compare") with gr.Tab("Pre-rendered comparisons (no GPU)"): gr.Markdown("The 40 evaluation requests rendered at 40 steps, seed 0, with the raw request, the 9B teacher and both " "students. Columns follow the sheet's own header.") sheet_pick = gr.Dropdown(SHEET_CHOICES, value=SHEET_CHOICES[0][1], label="Request") sheet_caption = gr.Markdown("") sheet_img = gr.Image(label="Contact sheet", height=760) sheet_pick.change(show_sheet, sheet_pick, [sheet_img, sheet_caption], api_visibility="private") demo.load(show_sheet, sheet_pick, [sheet_img, sheet_caption], api_visibility="private") with gr.Accordion("Overview sheet (all 40 requests)", open=False): gr.Image(value=OVERVIEW_URL, label="Overview", height=900) with gr.Accordion("About the numbers and fairness", open=False): gr.Markdown(ABOUT) if __name__ == "__main__": demo.queue(max_size=20).launch(theme=gr.themes.Citrus(), css=CSS)