Testing img2img

app.py

@@ -8,6 +8,7 @@ from qrcode.constants import ERROR_CORRECT_H
 from diffusers import (
     StableDiffusionPipeline,
     StableDiffusionControlNetPipeline,
+    StableDiffusionControlNetImg2ImgPipeline,   # NEW: img2img pipeline
     ControlNetModel,
     DPMSolverMultistepScheduler,
 )
@@ -16,7 +17,7 @@ from diffusers import (
 os.environ.setdefault("MPLCONFIGDIR", "/tmp/mpl")
 
 MODEL_ID = "runwayml/stable-diffusion-v1-5"
-CN_QRMON = "monster-labs/control_v1p_sd15_qrcode_monster"  # v2 on the repo
+CN_QRMON = "monster-labs/control_v1p_sd15_qrcode_monster"
 DTYPE = torch.float16
 
 # ---------- helpers ----------
@@ -41,7 +42,7 @@ def normalize_color(c):
     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 v2. :contentReference[oaicite:1]{index=1}
+    # Mid-gray background improves blending & scan rate with QR-Monster.
     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,7 +52,7 @@ def make_qr(url="http://www.mybirdfire.com", size=768, border=12, back_color="#8
     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 (simple post 'repair')."""
+    """Gently push ControlNet-required blacks/whites for scannability."""
     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)))
@@ -65,7 +66,15 @@ def enforce_qr_contrast(stylized: Image.Image, qr_img: Image.Image, strength: fl
 
 # ---------- lazy pipelines (CPU-offloaded for ZeroGPU) ----------
 _SD = None
-_CN = None
+_CN_TXT2IMG = None
+_CN_IMG2IMG = None
+
+def _base_scheduler_for(pipe):
+    pipe.scheduler = DPMSolverMultistepScheduler.from_config(
+        pipe.scheduler.config, use_karras_sigmas=True, algorithm_type="dpmsolver++"
+    )
+    pipe.enable_attention_slicing(); pipe.enable_vae_slicing(); pipe.enable_model_cpu_offload()
+    return pipe
 
 def get_sd_pipe():
     global _SD
@@ -77,16 +86,13 @@ def get_sd_pipe():
             use_safetensors=True,
             low_cpu_mem_usage=True,
         )
-        pipe.scheduler = DPMSolverMultistepScheduler.from_config(
-            pipe.scheduler.config, use_karras_sigmas=True, algorithm_type="dpmsolver++"
-        )
-        pipe.enable_attention_slicing(); pipe.enable_vae_slicing(); pipe.enable_model_cpu_offload()
-        _SD = pipe
+        _SD = _base_scheduler_for(pipe)
     return _SD
 
-def get_qrmon_pipe():
-    global _CN
-    if _CN is None:
+def get_qrmon_txt2img_pipe():
+    """(kept for completeness; not used in the two-stage flow)"""
+    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,
@@ -96,12 +102,24 @@ def get_qrmon_pipe():
             use_safetensors=True,
             low_cpu_mem_usage=True,
         )
-        pipe.scheduler = DPMSolverMultistepScheduler.from_config(
-            pipe.scheduler.config, use_karras_sigmas=True, algorithm_type="dpmsolver++"
+        _CN_TXT2IMG = _base_scheduler_for(pipe)
+    return _CN_TXT2IMG
+
+def get_qrmon_img2img_pipe():
+    """This is the pipeline we want for stage B."""
+    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,
         )
-        pipe.enable_attention_slicing(); pipe.enable_vae_slicing(); pipe.enable_model_cpu_offload()
-        _CN = pipe
-    return _CN
+        _CN_IMG2IMG = _base_scheduler_for(pipe)
+    return _CN_IMG2IMG
 
 # ---------- ZeroGPU tasks ----------
 @spaces.GPU(duration=120)
@@ -127,48 +145,55 @@ def txt2img(prompt: str, negative: str, steps: int, cfg: float, width: int, heig
 @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):
-    pipe = get_qrmon_pipe()
+               qr_weight: float, repair_strength: float, feather: float, seed: int,
+               denoise: float = 0.45):
     s = snap8(size)
-    qr_img = make_qr(url=url, size=s, border=int(border), back_color=back_color, blur_radius=float(blur))
 
+    # --- 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))
 
-    # Tip from the article: don't stuff "QR code" into the prompt; let ControlNet shape it. :contentReference[oaicite:2]{index=2}
     if torch.cuda.is_available(): torch.cuda.empty_cache()
     gc.collect()
+    with torch.autocast(device_type="cuda", dtype=DTYPE):
+        base = sd(
+            prompt=str(style_prompt),                      # don't include "QR code" here
+            negative_prompt=str(negative or ""),
+            num_inference_steps=max(int(steps)//2, 12),
+            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: ControlNet img2img (QR Monster) ---
+    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:
-            # diffusers ≥ 0.30.x uses `image=` for ControlNet’s conditioning input
-            out = pipe(
-                prompt=str(style_prompt),
-                negative_prompt=str(negative or ""),
-                image=qr_img,  # <-- use `image`
-                controlnet_conditioning_scale=float(qr_weight),
-                num_inference_steps=int(steps),
-                guidance_scale=float(cfg),
-                width=s, height=s,
-                generator=gen,
-            )
-        except TypeError:
-            # fallback for older versions that still expect `control_image=`
-            out = pipe(
-                prompt=str(style_prompt),
-                negative_prompt=str(negative or ""),
-                control_image=qr_img,
-                controlnet_conditioning_scale=float(qr_weight),
-                num_inference_steps=int(steps),
-                guidance_scale=float(cfg),
-                width=s, height=s,
-                generator=gen,
-            )
+        out = pipe(
+            prompt=str(style_prompt),
+            negative_prompt=str(negative or ""),
+            image=base,                                   # init image (img2img)
+            control_image=qr_img,                         # control image (QR)
+            strength=float(denoise),                      # 0.3–0.6 keeps composition
+            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,
+        )
 
     img = out.images[0]
     img = enforce_qr_contrast(img, qr_img, strength=float(repair_strength), feather=float(feather))
-    return img, qr_img
+    return img, qr_img, base
 
 # ---------- UI ----------
 with gr.Blocks() as demo:
@@ -185,26 +210,28 @@ with gr.Blocks() as demo:
         out_img = gr.Image(label="Image", interactive=False)
         gr.Button("Generate").click(txt2img, [prompt, negative, steps, cfg, width, height, seed], out_img)
 
-    with gr.Tab("QR Code Stylizer (ControlNet Monster)"):
+    with gr.Tab("QR Code Stylizer (ControlNet Monster — two-stage)"):
         url       = gr.Textbox(label="URL/Text", value="http://www.mybirdfire.com")
         s_prompt  = gr.Textbox(label="Style prompt (no 'QR code' needed)", value="baroque palace interior, intricate roots, dramatic lighting, ultra detailed")
         s_negative= gr.Textbox(label="Negative prompt", value="lowres, low contrast, blurry, jpeg artifacts, worst quality, watermark, text")
         size      = gr.Slider(384, 1024, value=768, step=64, label="Canvas (px)")
-        steps2    = gr.Slider(10, 50, value=28, step=1, label="Steps")
+        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")
         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")
         ctrl_img  = gr.Image(label="Control QR used")
+        base_img  = gr.Image(label="Base art (Stage A)")
         gr.Button("Stylize QR").click(
             qr_stylize,
-            [url, s_prompt, s_negative, steps2, cfg2, size, border, back_col, blur, qr_w, repair, feather, seed2],
-            [final_img, ctrl_img]
+            [url, s_prompt, s_negative, steps2, cfg2, size, border, back_col, blur, qr_w, repair, feather, seed2, denoise],
+            [final_img, ctrl_img, base_img]
         )
 
 if __name__ == "__main__":