Fix bug

app.py

@@ -1,21 +1,46 @@
 import gradio as gr
 import torch
+from diffusers import StableDiffusionPipeline
 from PIL import Image
 import qrcode
 from qrcode.constants import ERROR_CORRECT_H
 
-# --- shared device/dtype ---
+# ========= Stable Diffusion (prompt-only) =========
 device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
 dtype = torch.float16 if device != "cpu" else torch.float32
 
-# --- prompt-only SD pipe you already loaded as sd_pipe ---
-# from diffusers import StableDiffusionPipeline
-# sd_pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=dtype).to(device)
+sd_pipe = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=dtype
+).to(device)
 
-# --- lazy ControlNet loader (canny) ---
-_cn = {"pipe": None, "canny": None}
-def _load_cn():
-    if _cn["pipe"] is None:
+def sd_generate(prompt, steps, guidance, seed):
+    gen = torch.Generator(device=device).manual_seed(int(seed)) if int(seed) != 0 else None
+    def run():
+        return sd_pipe(prompt, num_inference_steps=int(steps), guidance_scale=float(guidance), generator=gen).images[0]
+    if device in ("cuda", "mps"):
+        with torch.autocast(device):
+            return run()
+    return run()
+
+# ========= QR Maker =========
+def make_qr(url: str = "http://www.mybirdfire.com", size: int = 512, border: int = 2) -> Image.Image:
+    qr = qrcode.QRCode(
+        version=None,
+        error_correction=ERROR_CORRECT_H,
+        box_size=10,
+        border=border
+    )
+    qr.add_data(url.strip())
+    qr.make(fit=True)
+    img = qr.make_image(fill_color="black", back_color="white").convert("RGB")
+    return img.resize((size, size), resample=Image.NEAREST)
+
+# ========= ControlNet Stylizer (prompt + QR) =========
+# lazy-load to speed initial startup
+_cn_loaded = {"pipe": None}
+def _load_controlnet():
+    if _cn_loaded["pipe"] is None:
         from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
         from controlnet_aux import CannyDetector
         controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=dtype)
@@ -27,60 +52,59 @@ def _load_cn():
         ).to(device)
         pipe.enable_attention_slicing()
         pipe.enable_vae_slicing()
-        _cn["pipe"], _cn["canny"] = pipe, CannyDetector()
-    return _cn["pipe"], _cn["canny"]
-
-def _make_qr(url: str, size: int, border: int) -> Image.Image:
-    qr = qrcode.QRCode(version=None, error_correction=ERROR_CORRECT_H, box_size=10, border=border)
-    qr.add_data(url.strip()); qr.make(fit=True)
-    img = qr.make_image(fill_color="black", back_color="white").convert("RGB")
-    return img.resize((size, size), resample=Image.NEAREST)
+        _cn_loaded["pipe"] = pipe
+        _cn_loaded["canny"] = CannyDetector()
+    return _cn_loaded["pipe"], _cn_loaded["canny"]
 
-def stylize_from_url(url: str, size: int, border: int,
-                     prompt: str, steps: int, guidance: float, seed: int,
-                     canny_low: int, canny_high: int):
-    # 1) Make QR
-    qr_img = _make_qr(url, size, border).convert("RGB")
-    qr_512 = qr_img.resize((512, 512), Image.NEAREST)
+def stylize_qr(prompt: str, steps: int, guidance: float, seed: int, canny_low: int, canny_high: int):
+    # Always make QR from URL
+    qr_image = make_qr("http://www.mybirdfire.com", size=512, border=2)
 
-    # 2) Canny edges
-    pipe, canny = _load_cn()
-    edges = canny(qr_512, low_threshold=int(canny_low), high_threshold=int(canny_high))
+    pipe, canny = _load_controlnet()
+    edges = canny(qr_image, low_threshold=int(canny_low), high_threshold=int(canny_high))
 
-    # 3) Generate with ControlNet
     gen = torch.Generator(device=device).manual_seed(int(seed)) if int(seed) != 0 else None
     def run():
-        return pipe(prompt=str(prompt),
-                    image=edges,
-                    num_inference_steps=int(steps),
-                    guidance_scale=float(guidance),
-                    generator=gen).images[0]
+        return pipe(
+            prompt=str(prompt),
+            image=edges,
+            num_inference_steps=int(steps),
+            guidance_scale=float(guidance),
+            generator=gen
+        ).images[0]
     if device in ("cuda", "mps"):
         with torch.autocast(device):
-            out = run()
-    else:
-        out = run()
-    return out, qr_512, edges
+            return run()
+    return run()
 
-# ====== UI: new tab ======
-with gr.Tab("QR Stylizer (ControlNet, auto‑QR)"):
-    url      = gr.Textbox(label="URL/Text", value="http://www.mybirdfire.com")
-    size     = gr.Slider(256, 1024, value=512, step=64, label="QR size (px)")
-    border   = gr.Slider(0, 8, value=2, step=1, label="QR border (quiet zone)")
-    prompt   = gr.Textbox(label="Style Prompt",
-                          value="floral papercut style, high contrast, preserve sharp squares")
-    steps    = gr.Slider(10, 50, value=30, step=1, label="Steps")
-    guidance = gr.Slider(1, 12, value=7.5, step=0.1, label="Guidance Scale")
-    seed     = gr.Number(value=0, label="Seed (0=random)", precision=0)
-    canny_l  = gr.Slider(0, 255, value=100, step=1, label="Canny low")
-    canny_h  = gr.Slider(0, 255, value=200, step=1, label="Canny high")
+# ========= UI =========
+with gr.Blocks() as demo:
+    gr.Markdown("## Stable Diffusion + QR Code + ControlNet (step by step)")
 
-    out_img  = gr.Image(label="Stylized QR")
-    out_qr   = gr.Image(label="Generated QR (input)")
-    out_edge = gr.Image(label="Canny edges (debug)")
+    with gr.Tab("Stable Diffusion (prompt → image)"):
+        prompt = gr.Textbox(label="Prompt", value="A fantasy castle at sunset")
+        steps  = gr.Slider(10, 50, value=30, label="Steps", step=1)
+        cfg    = gr.Slider(1, 12, value=7.5, label="Guidance Scale", step=0.1)
+        seed   = gr.Number(value=0, label="Seed (0 = random)", precision=0)
+        out_sd = gr.Image(label="Generated Image")
+        gr.Button("Generate").click(sd_generate, [prompt, steps, cfg, seed], out_sd)
 
-    gr.Button("Stylize").click(
-        stylize_from_url,
-        inputs=[url, size, border, prompt, steps, guidance, seed, canny_l, canny_h],
-        outputs=[out_img, out_qr, out_edge]
-    )
+    with gr.Tab("QR Maker (mybirdfire)"):
+        url   = gr.Textbox(label="URL/Text", value="http://www.mybirdfire.com")
+        size  = gr.Slider(256, 1024, value=512, step=64, label="Size (px)")
+        quiet = gr.Slider(0, 8, value=2, step=1, label="Border (quiet zone)")
+        out_qr = gr.Image(label="QR Code", type="pil")
+        gr.Button("Generate QR").click(make_qr, [url, size, quiet], out_qr)
+
+    with gr.Tab("QR Stylizer (ControlNet)"):
+        s_prompt = gr.Textbox(label="Style Prompt", value="floral papercut style, high contrast, preserve sharp squares")
+        s_steps  = gr.Slider(10, 50, value=30, label="Steps", step=1)
+        s_cfg    = gr.Slider(1, 12, value=7.5, label="Guidance Scale", step=0.1)
+        s_seed   = gr.Number(value=0, label="Seed (0 = random)", precision=0)
+        canny_l  = gr.Slider(0, 255, value=100, step=1, label="Canny low")
+        canny_h  = gr.Slider(0, 255, value=200, step=1, label="Canny high")
+        out_styl = gr.Image(label="Stylized QR")
+        gr.Button("Stylize").click(stylize_qr, [s_prompt, s_steps, s_cfg, s_seed, canny_l, canny_h], out_styl)
+
+if __name__ == "__main__":
+    demo.launch()