Testing 2 Stable Diffusion

app.py

@@ -6,7 +6,6 @@ import numpy as np
 import qrcode
 from qrcode.constants import ERROR_CORRECT_H
 from diffusers import (
-    StableDiffusionPipeline,
     StableDiffusionControlNetPipeline,
     StableDiffusionControlNetImg2ImgPipeline,   # for Hi-Res Fix
     ControlNetModel,
@@ -16,8 +15,13 @@ from diffusers import (
 # Quiet matplotlib cache warning on Spaces
 os.environ.setdefault("MPLCONFIGDIR", "/tmp/mpl")
 
-MODEL_ID = "runwayml/stable-diffusion-v1-5"
-# You can swap to a QR-Pattern-v2 repo if you know one on HF.
+# ---- base models for the two tabs ----
+BASE_MODELS = {
+    "anything": "andite/anything-v4.5",
+    "dream":    "Lykon/dreamshaper-8",
+}
+
+# ControlNet (QR Monster v2 for SD15)
 CN_QRMON = "monster-labs/control_v1p_sd15_qrcode_monster"
 DTYPE = torch.float16
 
@@ -42,10 +46,9 @@ def normalize_color(c):
         return s
     return "white"
 
-def make_qr(url="http://www.mybirdfire.com", size=768, border=12, back_color="#FFFFFF", blur_radius=0.0):
+def make_qr(url="https://example.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.)
+    IMPORTANT for Method 1: give ControlNet a sharp, black-on-WHITE QR (no blur).
     """
     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)
@@ -69,76 +72,58 @@ def enforce_qr_contrast(stylized: Image.Image, qr_img: Image.Image, strength: fl
     return Image.fromarray((s * 255.0).astype(np.uint8), mode="RGB")
 
 # ---------- lazy pipelines (CPU-offloaded for ZeroGPU) ----------
-_SD = None
-_CN_TXT2IMG = None
-_CN_IMG2IMG = None
+_CN = None                 # shared ControlNet QR Monster
+_CN_TXT2IMG = {}           # per-base-model txt2img pipes
+_CN_IMG2IMG = {}           # per-base-model img2img pipes
 
 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()
+    pipe.enable_attention_slicing()
+    pipe.enable_vae_slicing()
+    pipe.enable_model_cpu_offload()
     return pipe
 
-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
-        )
-        _SD = _base_scheduler_for(pipe)
-    return _SD
+def get_cn():
+    global _CN
+    if _CN is None:
+        _CN = ControlNetModel.from_pretrained(CN_QRMON, torch_dtype=DTYPE, use_safetensors=True)
+    return _CN
 
-def get_qrmon_txt2img_pipe():
-    global _CN_TXT2IMG
-    if _CN_TXT2IMG is None:
-        cn = ControlNetModel.from_pretrained(CN_QRMON, torch_dtype=DTYPE, use_safetensors=True)
+def get_qrmon_txt2img_pipe(model_id: str):
+    if model_id not in _CN_TXT2IMG:
         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=get_cn(),
+            torch_dtype=DTYPE,
+            safety_checker=None,
+            use_safetensors=True,
+            low_cpu_mem_usage=True,
         )
-        _CN_TXT2IMG = _base_scheduler_for(pipe)
-    return _CN_TXT2IMG
+        _CN_TXT2IMG[model_id] = _base_scheduler_for(pipe)
+    return _CN_TXT2IMG[model_id]
 
-def get_qrmon_img2img_pipe():
-    global _CN_IMG2IMG
-    if _CN_IMG2IMG is None:
-        cn = ControlNetModel.from_pretrained(CN_QRMON, torch_dtype=DTYPE, use_safetensors=True)
+def get_qrmon_img2img_pipe(model_id: str):
+    if model_id not in _CN_IMG2IMG:
         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=get_cn(),
+            torch_dtype=DTYPE,
+            safety_checker=None,
+            use_safetensors=True,
+            low_cpu_mem_usage=True,
         )
-        _CN_IMG2IMG = _base_scheduler_for(pipe)
-    return _CN_IMG2IMG
-
-# ---------- ZeroGPU tasks ----------
-@spaces.GPU(duration=120)
-def txt2img(prompt: str, negative: str, steps: int, cfg: float, width: int, height: int, seed: int):
-    pipe = get_sd_pipe()
-    w, h = snap8(width), snap8(height)
-    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):
-        out = pipe(
-            prompt=str(prompt),
-            negative_prompt=str(negative or ""),
-            num_inference_steps=int(steps),
-            guidance_scale=float(cfg),
-            width=w, height=h,
-            generator=gen,
-        )
-    return out.images[0]
+        _CN_IMG2IMG[model_id] = _base_scheduler_for(pipe)
+    return _CN_IMG2IMG[model_id]
 
 # -------- 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,
-               qr_weight: float, seed: int,
-               use_hires: bool, hires_upscale: float, hires_strength: float,
-               repair_strength: float, feather: float):
+def _qr_txt2img_core(model_id: str,
+                     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)
 
@@ -150,19 +135,17 @@ def qr_txt2img(url: str, style_prompt: str, negative: str,
         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()
+    # --- Stage A: txt2img with ControlNet
+    pipe = get_qrmon_txt2img_pipe(model_id)
     if torch.cuda.is_available(): torch.cuda.empty_cache()
     gc.collect()
-
     with torch.autocast(device_type="cuda", dtype=DTYPE):
-        # diffusers ≥ 0.30.x uses `image=` for control image
         out = pipe(
             prompt=str(style_prompt),
             negative_prompt=str(negative or ""),
-            image=qr_img,
-            controlnet_conditioning_scale=float(qr_weight),   # ~1.0–1.2 works well
-            control_guidance_start=0.0,                       # "Balanced" feel
+            image=qr_img,                                   # control image for txt2img
+            controlnet_conditioning_scale=float(qr_weight), # ~1.0–1.2 works well
+            control_guidance_start=0.0,
             control_guidance_end=1.0,
             num_inference_steps=int(steps),
             guidance_scale=float(cfg),
@@ -170,13 +153,14 @@ def qr_txt2img(url: str, style_prompt: str, negative: str,
             generator=gen,
         )
     lowres = out.images[0]
+    lowres = enforce_qr_contrast(lowres, qr_img, strength=float(repair_strength), feather=float(feather))
 
     # --- 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()
+        pipe2 = get_qrmon_img2img_pipe(model_id)
         if torch.cuda.is_available(): torch.cuda.empty_cache()
         gc.collect()
         with torch.autocast(device_type="cuda", dtype=DTYPE):
@@ -199,47 +183,77 @@ def qr_txt2img(url: str, style_prompt: str, negative: str,
     final = enforce_qr_contrast(final, qr_img, strength=float(repair_strength), feather=float(feather))
     return final, lowres, qr_img
 
+# Wrappers for each tab (so Gradio can bind without passing the model id)
+@spaces.GPU(duration=120)
+def qr_txt2img_anything(*args):
+    return _qr_txt2img_core(BASE_MODELS["anything"], *args)
+
+@spaces.GPU(duration=120)
+def qr_txt2img_dream(*args):
+    return _qr_txt2img_core(BASE_MODELS["dream"], *args)
+
 # ---------- UI ----------
 with gr.Blocks() as demo:
-    gr.Markdown("# ZeroGPU • SD1.5 + AI QR (Method 1)")
-
-    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=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)
-
-    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)")
-
-        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")
-
-        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)")
-
-        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")
-
-        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]
+    gr.Markdown("# ZeroGPU • Method 1: QR Control (two base models)")
+
+    # ---- Tab 1: Anything v4.5 (anime/illustration) ----
+    with gr.Tab("Method 1 • Anything v4.5"):
+        url1       = gr.Textbox(label="URL/Text", value="http://www.mybirdfire.com")
+        s_prompt1  = gr.Textbox(label="Style prompt", value="japanese painting, elegant shrine and torii, distant mount fuji, autumn maple trees, warm sunlight, 1girl in kimono, highly detailed, intricate patterns, anime key visual, dramatic composition")
+        s_negative1= gr.Textbox(label="Negative prompt", value="ugly, low quality, blurry, nsfw, watermark, text, low contrast, deformed, extra digits")
+        size1      = gr.Slider(384, 1024, value=512, step=64, label="Canvas (px)")
+        steps1     = gr.Slider(10, 50, value=20, step=1, label="Steps")
+        cfg1       = gr.Slider(1.0, 12.0, value=7.0, step=0.1, label="CFG")
+        border1    = gr.Slider(2, 16, value=4, step=1, label="QR border (quiet zone)")
+        qr_w1      = gr.Slider(0.6, 1.6, value=1.1, step=0.05, label="QR control weight")
+        seed1      = gr.Number(value=-1, precision=0, label="Seed (-1 random)")
+
+        use_hires1 = gr.Checkbox(value=True, label="Hi-Res Fix (img2img upscale)")
+        hires_up1  = gr.Slider(1.0, 2.0, value=2.0, step=0.25, label="Hi-Res upscale (×)")
+        hires_str1 = gr.Slider(0.3, 0.9, value=0.7, step=0.05, label="Hi-Res denoise strength")
+
+        repair1    = gr.Slider(0.0, 1.0, value=0.0, step=0.05, label="Post repair strength (optional)")
+        feather1   = gr.Slider(0.0, 3.0, value=1.0, step=0.1, label="Repair feather (px)")
+
+        final_img1 = gr.Image(label="Final (or Hi-Res) image")
+        low_img1   = gr.Image(label="Low-res (Stage A) preview")
+        ctrl_img1  = gr.Image(label="Control QR used")
+
+        gr.Button("Generate with Anything v4.5").click(
+            qr_txt2img_anything,
+            [url1, s_prompt1, s_negative1, steps1, cfg1, size1, border1, qr_w1, seed1,
+             use_hires1, hires_up1, hires_str1, repair1, feather1],
+            [final_img1, low_img1, ctrl_img1]
+        )
+
+    # ---- Tab 2: DreamShaper (general art/painterly) ----
+    with gr.Tab("Method 1 • DreamShaper 8"):
+        url2       = gr.Textbox(label="URL/Text", value="http://www.mybirdfire.com")
+        s_prompt2  = gr.Textbox(label="Style prompt", value="ornate baroque palace interior, gilded details, chandeliers, volumetric light, ultra detailed, cinematic")
+        s_negative2= gr.Textbox(label="Negative prompt", value="lowres, low contrast, blurry, jpeg artifacts, watermark, text, bad anatomy")
+        size2      = gr.Slider(384, 1024, value=512, step=64, label="Canvas (px)")
+        steps2     = gr.Slider(10, 50, value=24, step=1, label="Steps")
+        cfg2       = gr.Slider(1.0, 12.0, value=6.8, step=0.1, label="CFG")
+        border2    = gr.Slider(2, 16, value=8, step=1, label="QR border (quiet zone)")
+        qr_w2      = gr.Slider(0.6, 1.6, value=1.2, step=0.05, label="QR control weight")
+        seed2      = gr.Number(value=-1, precision=0, label="Seed (-1 random)")
+
+        use_hires2 = gr.Checkbox(value=True, label="Hi-Res Fix (img2img upscale)")
+        hires_up2  = gr.Slider(1.0, 2.0, value=2.0, step=0.25, label="Hi-Res upscale (×)")
+        hires_str2 = gr.Slider(0.3, 0.9, value=0.7, step=0.05, label="Hi-Res denoise strength")
+
+        repair2    = gr.Slider(0.0, 1.0, value=0.0, step=0.05, label="Post repair strength (optional)")
+        feather2   = gr.Slider(0.0, 3.0, value=1.0, step=0.1, label="Repair feather (px)")
+
+        final_img2 = gr.Image(label="Final (or Hi-Res) image")
+        low_img2   = gr.Image(label="Low-res (Stage A) preview")
+        ctrl_img2  = gr.Image(label="Control QR used")
+
+        gr.Button("Generate with DreamShaper 8").click(
+            qr_txt2img_dream,
+            [url2, s_prompt2, s_negative2, steps2, cfg2, size2, border2, qr_w2, seed2,
+             use_hires2, hires_up2, hires_str2, repair2, feather2],
+            [final_img2, low_img2, ctrl_img2]
         )
 
 if __name__ == "__main__":