"""Prompt rewriters for Qwen-Image 2.1: the official 9B teacher and the two pocket students. Teacher protocol follows QwenLM/Qwen-Image-2.1 prompt_rewrite/ (system prompt, thinking on, presence penalty 1.5, JSON answer parsed last-span-first). Student protocol follows the model cards: raw request as the only user turn, no system prompt, thinking disabled. """ import json import time import torch from huggingface_hub import hf_hub_download from transformers import AutoModelForCausalLM, AutoTokenizer, LogitsProcessor, LogitsProcessorList try: import json_repair except ImportError: # pragma: no cover json_repair = None TEACHER_ID = "Qwen/Qwen-Image-2.1-PE-T2I" STUDENT_IDS = { "0.8B": "ML-Intern-lab/Qwen-Image-2.1-PE-T2I-Pocket-0.8B", "2B": "ML-Intern-lab/Qwen-Image-2.1-PE-T2I-Pocket-2B", } # ~1 megapixel sizes per ratio, multiples of 32 (same table as the distillation image eval). RATIO_SIZES = { "1:1": (1024, 1024), "3:2": (1248, 832), "2:3": (832, 1248), "16:9": (1376, 768), "9:16": (768, 1376), "4:3": (1184, 896), "3:4": (896, 1184), "2:1": (1472, 736), "1:2": (736, 1472), "21:9": (1568, 672), "9:21": (672, 1568), "4:5": (928, 1152), "5:4": (1152, 928), "3:1": (1728, 576), "1:3": (576, 1760), } DEFAULT_RATIO = "3:2" SAMPLING = dict(do_sample=True, temperature=1.0, top_p=0.95, top_k=20) def size_for(ratio: str, megapixels: float = 1.0) -> tuple[int, int]: """(width, height) for a ratio, scaled to about `megapixels`, multiples of 32.""" w, h = RATIO_SIZES.get(ratio, RATIO_SIZES[DEFAULT_RATIO]) s = (megapixels * 1_000_000 / (w * h)) ** 0.5 return max(256, int(w * s) // 32 * 32), max(256, int(h * s) // 32 * 32) # ----------------------------------------------------------------------------- parsing (from pe_core.py) def split_thinking(text: str) -> tuple[str, str]: if "" in text: think, _, answer = text.partition("") if "" in think: think = think.partition("")[2] return think.strip(), answer.strip() if "" in text: return text.partition("")[2].strip(), "" return "", text.strip() def _balanced_spans(answer: str) -> list[str]: spans, depth, start = [], 0, -1 for i, ch in enumerate(answer): if ch == "{": if depth == 0: start = i depth += 1 elif ch == "}" and depth: depth -= 1 if depth == 0 and start >= 0: spans.append(answer[start:i + 1]) return spans def _as_obj(candidate: str): try: obj = json.loads(candidate) except json.JSONDecodeError: if json_repair is None: return None obj = json_repair.repair_json(candidate, return_objects=True) if isinstance(obj, list): obj = obj[0] if obj else None return obj if isinstance(obj, dict) else None def parse_answer(answer: str) -> dict: answer = (answer or "").strip() for candidate in reversed(_balanced_spans(answer)): obj = _as_obj(candidate) if obj is None: continue rewritten = obj.get("rewritten_prompt") or obj.get("rewrited_prompt") if not isinstance(rewritten, str) or not rewritten.strip(): continue ratio = str(obj.get("wh_ratio") or "").strip() return {"prompt": rewritten.strip(), "ratio": ratio if ratio in RATIO_SIZES else DEFAULT_RATIO, "ratio_raw": ratio, "parse_ok": True} return {"prompt": answer, "ratio": DEFAULT_RATIO, "ratio_raw": "", "parse_ok": False} class PresencePenalty(LogitsProcessor): """vLLM-style presence penalty: subtract a constant from every already-generated token.""" def __init__(self, penalty: float, prompt_len: int): self.penalty, self.prompt_len = penalty, prompt_len def __call__(self, input_ids, scores): for b in range(input_ids.shape[0]): generated = input_ids[b, self.prompt_len:] if generated.numel(): scores[b, generated.unique()] -= self.penalty return scores # ----------------------------------------------------------------------------- loading def load_student(size: str, device): rid = STUDENT_IDS[size] tok = AutoTokenizer.from_pretrained(rid) model = AutoModelForCausalLM.from_pretrained(rid, dtype=torch.bfloat16).to(device).eval() return tok, model def load_teacher(device): tok = AutoTokenizer.from_pretrained(TEACHER_ID) model = AutoModelForCausalLM.from_pretrained(TEACHER_ID, dtype=torch.bfloat16).to(device).eval() system_prompt = open(hf_hub_download(TEACHER_ID, "system_prompt.txt"), encoding="utf-8").read().strip() return tok, model, system_prompt # ----------------------------------------------------------------------------- rewriting (call inside @spaces.GPU) @torch.inference_mode() def student_rewrite(tok, model, request: str, seed: int = 0, max_new_tokens: int = 1024) -> dict: prompt = tok.apply_chat_template([{"role": "user", "content": request}], add_generation_prompt=True, tokenize=False) ids = tok(prompt, return_tensors="pt").to(model.device) torch.manual_seed(seed) t = time.time() out = model.generate(**ids, max_new_tokens=max_new_tokens, pad_token_id=tok.pad_token_id or tok.eos_token_id, **SAMPLING) gen = out[0, ids["input_ids"].shape[1]:] text = tok.decode(gen, skip_special_tokens=True) _, answer = split_thinking(text) # the template pre-fills an empty think block; strip it if echoed res = parse_answer(answer or text) res.update(seconds=time.time() - t, tokens=int(gen.numel()), thinking="", raw=text) return res @torch.inference_mode() def teacher_rewrite(tok, model, system_prompt: str, request: str, seed: int = 0, max_new_tokens: int = 3072) -> dict: prompt = tok.apply_chat_template( [{"role": "system", "content": system_prompt}, {"role": "user", "content": request}], add_generation_prompt=True, tokenize=False, enable_thinking=True) ids = tok(prompt, return_tensors="pt").to(model.device) prompt_len = ids["input_ids"].shape[1] torch.manual_seed(seed) t = time.time() out = model.generate(**ids, max_new_tokens=max_new_tokens, pad_token_id=tok.eos_token_id, logits_processor=LogitsProcessorList([PresencePenalty(1.5, prompt_len)]), **SAMPLING) gen = out[0, prompt_len:] text = tok.decode(gen, skip_special_tokens=True) thinking, answer = split_thinking(text) res = parse_answer(answer) if not answer: # ran out of tokens inside the think block res.update(prompt=request, parse_ok=False) res.update(seconds=time.time() - t, tokens=int(gen.numel()), thinking=thinking, raw=text) return res