Testing img2img

app.py

@@ -7,17 +7,18 @@ import qrcode
 from qrcode.constants import ERROR_CORRECT_H
 from diffusers import (
     StableDiffusionPipeline,
-    StableDiffusionControlNetPipeline,         # TXT2IMG (Method 1)
-    StableDiffusionControlNetImg2ImgPipeline,  # two-stage img2img
+    StableDiffusionControlNetPipeline,
+    StableDiffusionControlNetImg2ImgPipeline,   # for Hi-Res Fix
     ControlNetModel,
     DPMSolverMultistepScheduler,
 )
 
-# Optional: silence matplotlib cache warning in Spaces
+# Quiet matplotlib cache warning on Spaces
 os.environ.setdefault("MPLCONFIGDIR", "/tmp/mpl")
 
 MODEL_ID = "runwayml/stable-diffusion-v1-5"
-CN_QRMON  = "monster-labs/control_v1p_sd15_qrcode_monster"
+# You can swap to a QR-Pattern-v2 repo if you know one on HF.
+CN_QRMON = "monster-labs/control_v1p_sd15_qrcode_monster"
 DTYPE = torch.float16
 
 # ---------- helpers ----------
@@ -41,8 +42,11 @@ def normalize_color(c):
         return s
     return "white"
 
-def make_qr(url="http://www.mybirdfire.com", size=768, border=12, back_color="#808080", blur_radius=1.2):
-    # Mid-gray background improves blending & scan rate with QR-Monster.
+def make_qr(url="http://www.mybirdfire.com", size=768, border=12, back_color="#FFFFFF", blur_radius=0.0):
+    """
+    IMPORTANT for Method 1: give ControlNet a sharp, black-on-WHITE QR.
+    (No blur. Pixel-perfect.)
+    """
     qr = qrcode.QRCode(version=None, error_correction=ERROR_CORRECT_H, box_size=10, border=int(border))
     qr.add_data(url.strip()); qr.make(fit=True)
     img = qr.make_image(fill_color="black", back_color=normalize_color(back_color)).convert("RGB")
@@ -51,16 +55,16 @@ def make_qr(url="http://www.mybirdfire.com", size=768, border=12, back_color="#8
         img = img.filter(ImageFilter.GaussianBlur(radius=float(blur_radius)))
     return img
 
-def enforce_qr_contrast(stylized: Image.Image, qr_img: Image.Image, strength: float = 0.6, feather: float = 1.0) -> Image.Image:
-    """Gently push ControlNet-required blacks/whites for scannability."""
+def enforce_qr_contrast(stylized: Image.Image, qr_img: Image.Image, strength: float = 0.0, feather: float = 1.0) -> Image.Image:
+    """Optional gentle repair. Default OFF for Method 1."""
     if strength <= 0: return stylized
     q = qr_img.convert("L")
     black_mask = q.point(lambda p: 255 if p < 128 else 0).filter(ImageFilter.GaussianBlur(radius=float(feather)))
     black = np.asarray(black_mask, dtype=np.float32) / 255.0
     white = 1.0 - black
     s = np.asarray(stylized.convert("RGB"), dtype=np.float32) / 255.0
-    s = s * (1.0 - float(strength) * black[..., None])                    # deepen blacks
-    s = s + (1.0 - s) * (float(strength) * 0.85 * white[..., None])       # lift whites
+    s = s * (1.0 - float(strength) * black[..., None])
+    s = s + (1.0 - s) * (float(strength) * 0.85 * white[..., None])
     s = np.clip(s, 0.0, 1.0)
     return Image.fromarray((s * 255.0).astype(np.uint8), mode="RGB")
 
@@ -80,43 +84,29 @@ def get_sd_pipe():
     global _SD
     if _SD is None:
         pipe = StableDiffusionPipeline.from_pretrained(
-            MODEL_ID,
-            torch_dtype=DTYPE,
-            safety_checker=None,
-            use_safetensors=True,
-            low_cpu_mem_usage=True,
+            MODEL_ID, torch_dtype=DTYPE, safety_checker=None, use_safetensors=True, low_cpu_mem_usage=True
         )
         _SD = _base_scheduler_for(pipe)
     return _SD
 
 def get_qrmon_txt2img_pipe():
-    """Method 1 (TXT2IMG): SD + ControlNet QR-Monster, no init image, only conditioning image."""
     global _CN_TXT2IMG
     if _CN_TXT2IMG is None:
         cn = ControlNetModel.from_pretrained(CN_QRMON, torch_dtype=DTYPE, use_safetensors=True)
         pipe = StableDiffusionControlNetPipeline.from_pretrained(
-            MODEL_ID,
-            controlnet=cn,
-            torch_dtype=DTYPE,
-            safety_checker=None,
-            use_safetensors=True,
-            low_cpu_mem_usage=True,
+            MODEL_ID, controlnet=cn, torch_dtype=DTYPE, safety_checker=None,
+            use_safetensors=True, low_cpu_mem_usage=True
         )
         _CN_TXT2IMG = _base_scheduler_for(pipe)
     return _CN_TXT2IMG
 
 def get_qrmon_img2img_pipe():
-    """Two-stage B: SD img2img with ControlNet QR-Monster (kept so you can compare)."""
     global _CN_IMG2IMG
     if _CN_IMG2IMG is None:
         cn = ControlNetModel.from_pretrained(CN_QRMON, torch_dtype=DTYPE, use_safetensors=True)
         pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-            MODEL_ID,
-            controlnet=cn,
-            torch_dtype=DTYPE,
-            safety_checker=None,
-            use_safetensors=True,
-            low_cpu_mem_usage=True,
+            MODEL_ID, controlnet=cn, torch_dtype=DTYPE, safety_checker=None,
+            use_safetensors=True, low_cpu_mem_usage=True
         )
         _CN_IMG2IMG = _base_scheduler_for(pipe)
     return _CN_IMG2IMG
@@ -142,196 +132,114 @@ def txt2img(prompt: str, negative: str, steps: int, cfg: float, width: int, heig
         )
     return out.images[0]
 
-# ---- Method 1: TXT2IMG ControlNet (no init image; QR as conditioning only) ----
+# -------- Method 1: QR control model in text-to-image (+ optional Hi-Res Fix) --------
 @spaces.GPU(duration=120)
-def qr_txt2img(url: str, style_prompt: str, negative: str, steps: int, cfg: float,
-               size: int, border: int, back_color: str, blur: float,
-               qr_weight: float, start: float, end: float, seed: int,
+def qr_txt2img(url: str, style_prompt: str, negative: str,
+               steps: int, cfg: float, size: int, border: int,
+               qr_weight: float, seed: int,
+               use_hires: bool, hires_upscale: float, hires_strength: float,
                repair_strength: float, feather: float):
+
     s = snap8(size)
-    qr_img = make_qr(url=url, size=s, border=int(border), back_color=back_color, blur_radius=float(blur))
 
+    # Control image: crisp black-on-white QR
+    qr_img = make_qr(url=url, size=s, border=int(border), back_color="#FFFFFF", blur_radius=0.0)
+
+    # Seed / generator
     if int(seed) < 0:
         seed = random.randint(0, 2**31 - 1)
     gen = torch.Generator(device="cuda").manual_seed(int(seed))
 
+    # --- Stage A: txt2img with ControlNet (the actual "Method 1")
     pipe = get_qrmon_txt2img_pipe()
     if torch.cuda.is_available(): torch.cuda.empty_cache()
     gc.collect()
-    with torch.autocast(device_type="cuda", dtype=DTYPE):
-        try:
-            out = pipe(
-                prompt=str(style_prompt),
-                negative_prompt=str(negative or ""),
-                image=qr_img,  # ControlNet conditioning
-                controlnet_conditioning_scale=float(qr_weight),
-                control_guidance_start=float(start),
-                control_guidance_end=float(end),
-                num_inference_steps=int(steps),
-                guidance_scale=float(cfg),
-                width=s, height=s,
-                generator=gen,
-            )
-        except TypeError:
-            # Fallback for older diffusers param names
-            out = pipe(
-                prompt=str(style_prompt),
-                negative_prompt=str(negative or ""),
-                control_image=qr_img,
-                controlnet_conditioning_scale=float(qr_weight),
-                controlnet_start=float(start),
-                controlnet_end=float(end),
-                num_inference_steps=int(steps),
-                guidance_scale=float(cfg),
-                width=s, height=s,
-                generator=gen,
-            )
-
-    img = out.images[0]
-    img = enforce_qr_contrast(img, qr_img, strength=float(repair_strength), feather=float(feather))
-    return img, qr_img
-
-# ---- Two-stage (your previous Method-1 variant using IMG2IMG) ----
-@spaces.GPU(duration=120)
-def qr_stylize(url: str, style_prompt: str, negative: str, steps: int, cfg: float,
-               size: int, border: int, back_color: str, blur: float,
-               qr_weight: float, repair_strength: float, feather: float, seed: int,
-               denoise: float = 0.45):
-    s = snap8(size)
-
-    # Stage A: base art (txt2img)
-    sd = get_sd_pipe()
-    if int(seed) < 0:
-        seed = random.randint(0, 2**31 - 1)
-    gen = torch.Generator(device="cuda").manual_seed(int(seed))
 
-    if torch.cuda.is_available(): torch.cuda.empty_cache()
-    gc.collect()
     with torch.autocast(device_type="cuda", dtype=DTYPE):
-        base = sd(
+        # diffusers ≥ 0.30.x uses `image=` for control image
+        out = pipe(
             prompt=str(style_prompt),
             negative_prompt=str(negative or ""),
-            num_inference_steps=max(int(steps)//2, 12),
+            image=qr_img,
+            controlnet_conditioning_scale=float(qr_weight),   # ~1.0–1.2 works well
+            control_guidance_start=0.0,                       # "Balanced" feel
+            control_guidance_end=1.0,
+            num_inference_steps=int(steps),
             guidance_scale=float(cfg),
             width=s, height=s,
             generator=gen,
-        ).images[0]
-
-    # control image (QR)
-    qr_img = make_qr(url=url, size=s, border=int(border),
-                     back_color=back_color, blur_radius=float(blur))
-
-    # Stage B: img2img with ControlNet QR
-    pipe = get_qrmon_img2img_pipe()
-    if torch.cuda.is_available(): torch.cuda.empty_cache()
-    gc.collect()
-    with torch.autocast(device_type="cuda", dtype=DTYPE):
-        try:
-            out = pipe(
-                prompt=str(style_prompt),
-                negative_prompt=str(negative or ""),
-                image=base,            # init image
-                image_guidance_scale=None,
-                control_image=qr_img,  # QR conditioning
-                strength=float(denoise),
-                controlnet_conditioning_scale=float(qr_weight),
-                control_guidance_start=0.05,
-                control_guidance_end=0.95,
-                num_inference_steps=int(steps),
-                guidance_scale=float(cfg),
-                width=s, height=s,
-                generator=gen,
-            )
-        except TypeError:
-            out = pipe(
+        )
+    lowres = out.images[0]
+
+    # --- Optional Stage B: Hi-Res Fix (img2img with same QR)
+    final = lowres
+    if use_hires:
+        up = max(1.0, min(2.0, float(hires_upscale)))
+        W = snap8(int(s * up)); H = W
+        pipe2 = get_qrmon_img2img_pipe()
+        if torch.cuda.is_available(): torch.cuda.empty_cache()
+        gc.collect()
+        with torch.autocast(device_type="cuda", dtype=DTYPE):
+            out2 = pipe2(
                 prompt=str(style_prompt),
                 negative_prompt=str(negative or ""),
-                image=base,
-                control_image=qr_img,
-                strength=float(denoise),
+                image=lowres,                      # init image
+                control_image=qr_img,              # same QR
+                strength=float(hires_strength),    # ~0.7 like "Hires Fix"
                 controlnet_conditioning_scale=float(qr_weight),
-                controlnet_start=0.05,
-                controlnet_end=0.95,
+                control_guidance_start=0.0,
+                control_guidance_end=1.0,
                 num_inference_steps=int(steps),
                 guidance_scale=float(cfg),
-                width=s, height=s,
+                width=W, height=H,
                 generator=gen,
             )
+        final = out2.images[0]
 
-    img = out.images[0]
-    img = enforce_qr_contrast(img, qr_img, strength=float(repair_strength), feather=float(feather))
-    return img, qr_img, base
+    final = enforce_qr_contrast(final, qr_img, strength=float(repair_strength), feather=float(feather))
+    return final, lowres, qr_img
 
 # ---------- UI ----------
 with gr.Blocks() as demo:
-    gr.Markdown("# ZeroGPU Stable Diffusion + AI QR Codes")
+    gr.Markdown("# ZeroGPU • SD1.5 + AI QR (Method 1)")
 
-    with gr.Tab("Text → Image"):
-        prompt  = gr.Textbox(label="Prompt", value="a cozy reading nook, warm sunlight, cinematic lighting, highly detailed")
-        negative = gr.Textbox(label="Negative (optional)", value="lowres, blurry, watermark, text")
-        steps   = gr.Slider(8, 40, value=28, step=1, label="Steps")
+    with gr.Tab("Plain Text → Image"):
+        prompt  = gr.Textbox(label="Prompt", value="Japanese painting, mountains")
+        negative = gr.Textbox(label="Negative (optional)", value="ugly, disfigured, low quality, blurry, nsfw")
+        steps   = gr.Slider(8, 40, value=20, step=1, label="Steps")
         cfg     = gr.Slider(1.0, 12.0, value=7.0, step=0.5, label="CFG")
-        width   = gr.Slider(256, 1024, value=640, step=16, label="Width")
-        height  = gr.Slider(256, 1024, value=640, step=16, label="Height")
+        width   = gr.Slider(256, 1024, value=512, step=16, label="Width")
+        height  = gr.Slider(256, 1024, value=512, step=16, label="Height")
         seed    = gr.Number(value=-1, precision=0, label="Seed (-1 random)")
         out_img = gr.Image(label="Image", interactive=False)
         gr.Button("Generate").click(txt2img, [prompt, negative, steps, cfg, width, height, seed], out_img)
 
-    # ---- Method 1: TXT2IMG ControlNet ----
-    with gr.Tab("QR (Method 1 — TXT2IMG)"):
-        url_m1      = gr.Textbox(label="URL/Text", value="http://www.mybirdfire.com")
-        prompt_m1   = gr.Textbox(label="Style prompt (no 'QR code' needed)",
-                                 value="epic phoenix in flames, dramatic lighting, detailed, 8k")
-        neg_m1      = gr.Textbox(label="Negative prompt",
-                                 value="lowres, low contrast, blurry, jpeg artifacts, worst quality, bad anatomy, extra digits")
-        size_m1     = gr.Slider(384, 1024, value=768, step=64, label="Canvas (px)")
-        steps_m1    = gr.Slider(10, 60, value=28, step=1, label="Steps")
-        cfg_m1      = gr.Slider(1.0, 12.0, value=6.5, step=0.1, label="CFG")
-        border_m1   = gr.Slider(4, 20, value=12, step=1, label="QR border (quiet zone)")
-        back_m1     = gr.ColorPicker(value="#808080", label="QR background")
-        blur_m1     = gr.Slider(0.0, 3.0, value=1.2, step=0.1, label="Soften control (blur)")
-        weight_m1   = gr.Slider(0.6, 1.6, value=1.2, step=0.05, label="QR control weight")
-        start_m1    = gr.Slider(0.0, 1.0, value=0.05, step=0.01, label="Control start")
-        end_m1      = gr.Slider(0.0, 1.0, value=0.95, step=0.01, label="Control end")
-        seed_m1     = gr.Number(value=-1, precision=0, label="Seed (-1 random)")
-        repair_m1   = gr.Slider(0.0, 1.0, value=0.6, step=0.05, label="Post repair strength")
-        feather_m1  = gr.Slider(0.0, 3.0, value=1.0, step=0.1, label="Repair feather (px)")
-
-        final_m1 = gr.Image(label="Final QR (TXT2IMG)")
-        ctrl_m1  = gr.Image(label="Control QR used")
+    with gr.Tab("Method 1: QR control (txt2img)"):
+        url       = gr.Textbox(label="URL/Text", value="https://example.com")
+        s_prompt  = gr.Textbox(label="Style prompt", value="Japanese painting, mountains, 1girl")
+        s_negative= gr.Textbox(label="Negative prompt", value="ugly, disfigured, low quality, blurry, nsfw")
+        size      = gr.Slider(384, 1024, value=512, step=64, label="Canvas (px)")
+        steps2    = gr.Slider(10, 50, value=20, step=1, label="Steps")
+        cfg2      = gr.Slider(1.0, 12.0, value=7.0, step=0.1, label="CFG")
+        border    = gr.Slider(2, 16, value=4, step=1, label="QR border (quiet zone)")
+        qr_w      = gr.Slider(0.6, 1.6, value=1.1, step=0.05, label="QR control weight")
+        seed2     = gr.Number(value=-1, precision=0, label="Seed (-1 random)")
 
-        gr.Button("Generate (Method 1)").click(
-            qr_txt2img,
-            [url_m1, prompt_m1, neg_m1, steps_m1, cfg_m1, size_m1, border_m1, back_m1, blur_m1,
-             weight_m1, start_m1, end_m1, seed_m1, repair_m1, feather_m1],
-            [final_m1, ctrl_m1]
-        )
+        use_hires = gr.Checkbox(value=True, label="Hi-Res Fix (img2img upscale)")
+        hires_up  = gr.Slider(1.0, 2.0, value=2.0, step=0.25, label="Hi-Res upscale (×)")
+        hires_str = gr.Slider(0.3, 0.9, value=0.7, step=0.05, label="Hi-Res denoise strength")
 
-    # ---- Two-stage (Method-1 variant, IMG2IMG) ----
-    with gr.Tab("QR (Two-stage IMG2IMG)"):
-        url       = gr.Textbox(label="URL/Text", value="http://www.mybirdfire.com")
-        s_prompt  = gr.Textbox(label="Style prompt (no 'QR code' needed)",
-                               value="epic phoenix in flames, dramatic lighting, detailed, 8k")
-        s_negative= gr.Textbox(label="Negative prompt",
-                               value="lowres, low contrast, blurry, jpeg artifacts, worst quality, bad anatomy, extra digits")
-        size      = gr.Slider(384, 1024, value=768, step=64, label="Canvas (px)")
-        steps2    = gr.Slider(10, 60, value=28, step=1, label="Total steps")
-        cfg2      = gr.Slider(1.0, 12.0, value=6.5, step=0.1, label="CFG")
-        border    = gr.Slider(4, 20, value=12, step=1, label="QR border (quiet zone)")
-        back_col  = gr.ColorPicker(value="#808080", label="QR background")
-        blur      = gr.Slider(0.0, 3.0, value=1.2, step=0.1, label="Soften control (blur)")
-        qr_w      = gr.Slider(0.6, 1.6, value=1.2, step=0.05, label="QR control weight")
-        denoise   = gr.Slider(0.2, 0.8, value=0.45, step=0.01, label="Denoising strength (Stage B)")
-        repair    = gr.Slider(0.0, 1.0, value=0.6, step=0.05, label="Post repair strength")
+        repair    = gr.Slider(0.0, 1.0, value=0.0, step=0.05, label="Post repair strength (optional)")
         feather   = gr.Slider(0.0, 3.0, value=1.0, step=0.1, label="Repair feather (px)")
-        seed2     = gr.Number(value=-1, precision=0, label="Seed (-1 random)")
-        final_img = gr.Image(label="Final stylized QR")
+
+        final_img = gr.Image(label="Final (or Hi-Res) image")
+        low_img   = gr.Image(label="Low-res (Stage A) preview")
         ctrl_img  = gr.Image(label="Control QR used")
-        base_img  = gr.Image(label="Base art (Stage A)")
-        gr.Button("Stylize QR (Two-stage)").click(
-            qr_stylize,
-            [url, s_prompt, s_negative, steps2, cfg2, size, border, back_col, blur, qr_w, repair, feather, seed2, denoise],
-            [final_img, ctrl_img, base_img]
+
+        gr.Button("Generate QR Art").click(
+            qr_txt2img,
+            [url, s_prompt, s_negative, steps2, cfg2, size, border, qr_w, seed2, use_hires, hires_up, hires_str, repair, feather],
+            [final_img, low_img, ctrl_img]
         )
 
 if __name__ == "__main__":