Adjust third tab

app.py

@@ -1,46 +1,21 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionPipeline
 from PIL import Image
 import qrcode
 from qrcode.constants import ERROR_CORRECT_H
 
-# ========= Stable Diffusion (prompt-only) =========
+# --- shared device/dtype ---
 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
 
-sd_pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    torch_dtype=dtype
-).to(device)
+# --- 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)
 
-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:
+# --- lazy ControlNet loader (canny) ---
+_cn = {"pipe": None, "canny": None}
+def _load_cn():
+    if _cn["pipe"] is None:
         from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
         from controlnet_aux import CannyDetector
         controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=dtype)
@@ -52,61 +27,60 @@ def _load_controlnet():
         ).to(device)
         pipe.enable_attention_slicing()
         pipe.enable_vae_slicing()
-        _cn_loaded["pipe"] = pipe
-        _cn_loaded["canny"] = CannyDetector()
-    return _cn_loaded["pipe"], _cn_loaded["canny"]
+        _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)
 
-def stylize_qr(prompt: str, qr_image: Image.Image, steps: int, guidance: float, seed: int, canny_low: int, canny_high: int):
-    if qr_image is None:
-        raise gr.Error("Please provide a QR image (use the QR Maker tab first or upload one).")
-    pipe, canny = _load_controlnet()
-    # ensure 512x512 for speed/quality
-    qr_img = qr_image.convert("RGB").resize((512, 512), Image.NEAREST)
-    edges = canny(qr_img, low_threshold=int(canny_low), high_threshold=int(canny_high))
+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)
 
+    # 2) Canny edges
+    pipe, canny = _load_cn()
+    edges = canny(qr_512, 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):
-            return run()
-    return run()
-
-# ========= UI =========
-with gr.Blocks() as demo:
-    gr.Markdown("## Stable Diffusion + QR Code + ControlNet (step by step)")
+            out = run()
+    else:
+        out = run()
+    return out, qr_512, edges
 
-    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)
+# ====== 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")
 
-    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)
+    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("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")
-        qr_in    = gr.Image(label="QR Input (use output from QR Maker or upload)", type="pil")
-        out_styl = gr.Image(label="Stylized QR")
-        gr.Button("Stylize").click(stylize_qr, [s_prompt, qr_in, s_steps, s_cfg, s_seed, canny_l, canny_h], out_styl)
-
-if __name__ == "__main__":
-    demo.launch()
+    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]
+    )