Adding initial eval code

app.py

@@ -53,28 +53,46 @@ def fused_sim(a:Image.Image,b:Image.Image,α=.5):
 bnb_cfg = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_use_double_quant=True)
 
 # ---------- load models once at startup ---------------------
-base = None
+_base = None
+# @spaces.GPU
+# def load_models():
+#     from unsloth import FastLanguageModel
+#     global base, tok, lora
+#     if base is None:
+#         print("Loading BASE …")
+#         base, tok = FastLanguageModel.from_pretrained(
+#             BASE_MODEL, max_seq_length=2048,
+#             load_in_4bit=True, quantization_config=bnb_cfg, device_map="auto")
+#         tok.pad_token = tok.eos_token
+
+#         print("Loading LoRA …")
+#         lora, _ = FastLanguageModel.from_pretrained(
+#             ADAPTER_DIR, max_seq_length=2048,
+#             load_in_4bit=True, quantization_config=bnb_cfg, device_map="auto")
+#         print("✔ models loaded")
+
+_base = _lora = _tok = None
+_CLIP = _PREP = _LP = None
+
 @spaces.GPU
-def load_models():
+def ensure_models():
+    """Create base, lora, tok **once per worker**."""
     from unsloth import FastLanguageModel
-    global base, tok, lora
-    if base is None:
-        print("Loading BASE …")
-        base, tok = FastLanguageModel.from_pretrained(
+    global _base, _lora, _tok
+    if _base is None:
+        _base, _tok = FastLanguageModel.from_pretrained(
             BASE_MODEL, max_seq_length=2048,
-            load_in_4bit=True, quantization_config=bnb_cfg, device_map="auto")
-        tok.pad_token = tok.eos_token
-
-        print("Loading LoRA …")
-        lora, _ = FastLanguageModel.from_pretrained(
+            quantization_config=bnb_cfg, device_map="auto")
+        _tok.pad_token = _tok.eos_token
+        _lora, _ = FastLanguageModel.from_pretrained(
             ADAPTER_DIR, max_seq_length=2048,
-            load_in_4bit=True, quantization_config=bnb_cfg, device_map="auto")
-        print("✔ models loaded")
+            quantization_config=bnb_cfg, device_map="auto")
+    return True
 
-@spaces.GPU
-def ensure_models():
-    load_models()
-    return True           # small, pickle-able sentinel
+# @spaces.GPU
+# def ensure_models():
+#     load_models()
+#     return True           # small, pickle-able sentinel
 
 
 def build_prompt(desc:str):
@@ -85,26 +103,37 @@ def build_prompt(desc:str):
 
 @spaces.GPU
 @torch.no_grad()
-def draw(model, desc:str):
+def draw(model_flag, desc):
     ensure_models()
-    prompt = build_prompt(desc)
-    ids = tok(prompt, return_tensors="pt").to(DEVICE)
+    model = _base if model_flag == "base" else _lora
+    prompt = _tok.apply_chat_template(
+        [{"role":"system","content":"You are an SVG illustrator."},
+         {"role":"user",
+          "content":f"ONLY reply with a valid, complete <svg>…</svg> file that depicts: {desc}"}],
+        tokenize=False, add_generation_prompt=True)
+    ids = _tok(prompt, return_tensors="pt").to(DEVICE)
     out = model.generate(**ids, max_new_tokens=MAX_NEW,
                          do_sample=True, temperature=.7, top_p=.8)
-    txt = tok.decode(out[0], skip_special_tokens=True)
-    svg = extract_svg(txt)
+    svg = extract_svg(_tok.decode(out[0], skip_special_tokens=True))
     img = svg2pil(svg) if svg else None
     return img, svg or "(no SVG found)"
 
 # ---------- gradio interface --------------------------------
+#
 def compare(desc):
-    ensure_models()
-    img_base, svg_base = draw(base, desc)
-    img_lora, svg_lora = draw(lora, desc)
-    # sim = (fused_sim(img_lora, img_base) if img_base and img_lora else float("nan"))
-
+    img_b, svg_b = draw("base", desc)
+    img_l, svg_l = draw("lora", desc)
     caption = "Thanks for trying our model 😊\nIf you don't see an image for the base or GRPO model that means it didn't generate a valid SVG!"
-    return img_base, img_lora, caption, svg_base, svg_lora
+    return img_b, img_l, caption, svg_b, svg_l
+
+# def compare(desc):
+#     ensure_models()
+#     img_base, svg_base = draw(base, desc)
+#     img_lora, svg_lora = draw(lora, desc)
+#     # sim = (fused_sim(img_lora, img_base) if img_base and img_lora else float("nan"))
+
+#     caption = "Thanks for trying our model 😊\nIf you don't see an image for the base or GRPO model that means it didn't generate a valid SVG!"
+#     return img_base, img_lora, caption, svg_base, svg_lora
 
 with gr.Blocks(css="body{background:#111;color:#eee}") as demo:
     gr.Markdown("## 🖌️ Qwen-2.5 SVG Generator — base vs GRPO-LoRA")