Hugging Face's logo

Spaces:
JohanDL
/
Text_to_SVG_GRPO
Running

App Files Community
Text_to_SVG_GRPO
File size: 4,643 Bytes 
fe5070c
 
 
be0cea3
fe5070c
 
 
 
 
 
be0cea3
fe5070c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9b179e0
 
 
00ff9a0
9b179e0
 
 
 
 
 
 
 
 
 
 
00ff9a0
9b179e0
fe5070c
9b179e0
fe5070c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
00ff9a0
 
fe5070c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
 
#!/usr/bin/env python3
# svg_compare_gradio.py
# ------------------------------------------------------------
import spaces
import re, os, torch, cairosvg, lpips, clip, gradio as gr
from io import BytesIO
from pathlib import Path
from PIL import Image
from transformers import BitsAndBytesConfig, AutoTokenizer
import gradio as gr


# ---------- paths YOU may want to edit ----------------------
ADAPTER_DIR = "unsloth_trained_weights/checkpoint-1700"  # LoRA ckpt
BASE_MODEL  = "Qwen/Qwen2.5-Coder-7B-Instruct"
MAX_NEW     = 512
DEVICE      = "cuda" # if torch.cuda.is_available() else "cpu"

# ---------- utils -------------------------------------------
SVG_PAT = re.compile(r"<svg[^>]*>.*?</svg>", re.S | re.I)
def extract_svg(txt:str):
    m = list(SVG_PAT.finditer(txt))
    return m[-1].group(0) if m else None                       # last match βœ”

def svg2pil(svg:str):
    try:
        png = cairosvg.svg2png(bytestring=svg.encode())
        return Image.open(BytesIO(png)).convert("RGB")
    except Exception:
        return None

# ---------- backbone loaders (CLIP + LPIPS) -----------------
_CLIP,_PREP,_LP=None,None,None
@spaces.GPU
def _load_backbones():
    global _CLIP,_PREP,_LP
    if _CLIP is None:
        _CLIP,_PREP = clip.load("ViT-L/14", device=DEVICE); _CLIP.eval()
    if _LP is None:
        _LP = lpips.LPIPS(net="vgg").to(DEVICE).eval()

@spaces.GPU
@torch.no_grad()
def fused_sim(a:Image.Image,b:Image.Image,Ξ±=.5):
    _load_backbones()
    ta,tb = _PREP(a).unsqueeze(0).to(DEVICE), _PREP(b).unsqueeze(0).to(DEVICE)
    fa = _CLIP.encode_image(ta); fa/=fa.norm(dim=-1,keepdim=True)
    fb = _CLIP.encode_image(tb); fb/=fb.norm(dim=-1,keepdim=True)
    clip_sim=(([email protected]).item()+1)/2
    lp_sim = 1 - _LP(ta,tb,normalize=True).item()
    return Ξ±*clip_sim + (1-Ξ±)*lp_sim

# ---------- load models once at startup ---------------------
@spaces.GPU
def load_models():
    from unsloth import FastLanguageModel
    global base, tok, lora
    bnb_cfg = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_use_double_quant=True)
    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, tok, lora = load_models()

def build_prompt(desc:str):
    msgs=[{"role":"system","content":"You are an SVG illustrator."},
          {"role":"user",
           "content":f"ONLY reply with a valid, complete <svg>…</svg> file that depicts: {desc}"}]
    return tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)

@spaces.GPU
@torch.no_grad()
def draw(model, desc:str):
    prompt = build_prompt(desc)
    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)
    img = svg2pil(svg) if svg else None
    return img, svg or "(no SVG found)"

# ---------- gradio interface --------------------------------
def compare(desc):
    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")
    gr.Markdown(
        "Type an image **description** (e.g. *a purple forest at dusk*). "
        "Click **Generate** to see what the base model and your fine-tuned LoRA produce."
    )
    inp = gr.Textbox(label="Description", placeholder="a purple forest at dusk")
    btn = gr.Button("Generate")
    with gr.Row():
        out_base = gr.Image(label="Base model", type="pil")
        out_lora = gr.Image(label="LoRA-tuned model", type="pil")
    sim_lbl = gr.Markdown()
    with gr.Accordion("βš™οΈ  Raw SVG code", open=False):
        svg_base_box = gr.Textbox(label="Base SVG", lines=6)
        svg_lora_box = gr.Textbox(label="LoRA SVG", lines=6)
    btn.click(compare, inp, [out_base, out_lora, sim_lbl, svg_base_box, svg_lora_box])

demo.launch()