Upload app.py with huggingface_hub

app.py

@@ -136,20 +136,42 @@ print("[dgemma] adapters loaded:", list(_NAME.values()), flush=True)
 model.to("cuda")
 print("[dgemma] model registered on cuda (ZeroGPU will stream on first call)", flush=True)
 
+# The discrete-diffusion backbone denoises a fixed-length canvas whose size is set by the
+# base config (canvas_length=256), independent of the `max_new_tokens` request. Infill masks
+# MUST be sized to this or the canvas-clamp hook mismatches the model's internal tensor.
+def _canvas_length():
+    for obj in (model, getattr(model, "base_model", None)):
+        cfg = getattr(obj, "config", None)
+        v = getattr(cfg, "canvas_length", None)
+        if v:
+            return int(v)
+    return 256
+
+
+CANVAS_LEN = _canvas_length()
+print(f"[dgemma] canvas length = {CANVAS_LEN}", flush=True)
+
 _THINK = re.compile(r"<think>(.*?)</think>", re.S)
 _ANSWER = re.compile(r"<answer>(.*?)</answer>", re.S)
+# Some finetunes emit an OpenAI-Harmony-style channel prefix, e.g. "<|channel|>thought\n..."
+# or a bare "thought\n..." / "analysis\n..." lead-in before the actual answer. Strip it.
+_CHANNEL = re.compile(r"^\s*(?:<\|[^>]*\|>)*\s*(?:thought|analysis|final|channel)\s*\n+", re.I)
 
 
 def _strip_cot(text: str) -> str:
     """Pull the <answer> (and <think> findings) out of a CoT response."""
     t = _THINK.search(text)
     a = _ANSWER.search(text)
-    if not t and not a:
-        return (text.replace("<think>", "").replace("</think>", "")
-                .replace("<answer>", "").replace("</answer>", "").strip())
-    findings = t.group(1).strip() if t else ""
-    impression = a.group(1).strip() if a else ""
-    return (findings + "\n" + impression).strip()
+    if t or a:
+        findings = t.group(1).strip() if t else ""
+        impression = a.group(1).strip() if a else ""
+        return (findings + "\n" + impression).strip()
+    # No explicit CoT tags — clean up channel prefixes and any stray tag tokens.
+    t = re.sub(r"<\|[^>]*\|>", " ", text)
+    t = _CHANNEL.sub("", t)
+    t = (t.replace("<think>", "").replace("</think>", "")
+         .replace("<answer>", "").replace("</answer>", ""))
+    return t.strip()
 
 
 def _prompt_inputs(image, instruction):
@@ -284,15 +306,19 @@ def infill_report(image: Image.Image, template: str, adapter_label: str,
     image = image.convert("RGB")
 
     tok = processor.tokenizer
-    L = int(max_new_tokens)
+    # The model's denoising canvas is always CANVAS_LEN long, regardless of the
+    # max_new_tokens request; the clamp mask must match that exact size.
+    L = CANVAS_LEN
     pad = tok.pad_token_id or 0
 
     # Split around the blank and locate its token span on the canvas.
     before, after = template.split(BLANK, 1)
     before_ids = tok(before, add_special_tokens=False)["input_ids"]
     after_ids = tok(after, add_special_tokens=False)["input_ids"]
-    # Reserve room in the middle for the fill; scale by remaining canvas space.
-    hole_len = max(1, min(48, L - len(before_ids) - len(after_ids) - 2))
+    # Reserve room in the middle for the fill; cap by the user's requested span length
+    # and by the remaining canvas space.
+    max_fill = max(1, min(int(max_new_tokens), 96))
+    hole_len = max(1, min(max_fill, L - len(before_ids) - len(after_ids) - 2))
     h0 = min(len(before_ids), L - 1)
     h1 = min(h0 + hole_len, L - 1)
 
@@ -403,8 +429,8 @@ with gr.Blocks(title="DiffusionGemma Radiology VQA") as demo:
                     with gr.Column():
                         inf_out = gr.Textbox(label="Filled span", lines=8)
                 with gr.Accordion("Advanced", open=False):
-                    inf_tokens = gr.Slider(32, 256, value=MAX_NEW_TOKENS, step=16,
-                                           label="Canvas length (max new tokens)")
+                    inf_tokens = gr.Slider(4, 96, value=48, step=4,
+                                           label="Max fill span (tokens)")
                 gr.Examples(
                     examples=[
                         ["cxr_aorta.png", "The lungs are clear. [BLANK] No pleural effusion is seen.", ADAPTER_LABELS[0], True],