Testing 2 Stable Diffusion

app.py

@@ -7,22 +7,21 @@ import qrcode
 from qrcode.constants import ERROR_CORRECT_H
 from diffusers import (
     StableDiffusionControlNetPipeline,
-    StableDiffusionControlNetImg2ImgPipeline,
-    StableDiffusionPipeline,
+    StableDiffusionControlNetImg2ImgPipeline,   # for Hi-Res Fix
     ControlNetModel,
     DPMSolverMultistepScheduler,
 )
-from huggingface_hub import hf_hub_download
-from huggingface_hub.errors import EntryNotFoundError
 
 # Quiet matplotlib cache warning on Spaces
 os.environ.setdefault("MPLCONFIGDIR", "/tmp/mpl")
 
+# ---- base models for the two tabs ----
 BASE_MODELS = {
-    "counterfeit": "gsdf/Counterfeit-V3.0",   # single-file repo (no model_index.json)
-    "dream":       "Lykon/dreamshaper-8",     # diffusers-style (works with from_pretrained)
+    "stable-diffusion-v1-5": "runwayml/stable-diffusion-v1-5",
+    "dream":    "Lykon/dreamshaper-8",
 }
 
+# ControlNet (QR Monster v2 for SD15)
 CN_QRMON = "monster-labs/control_v1p_sd15_qrcode_monster"
 DTYPE = torch.float16
 
@@ -48,6 +47,9 @@ def normalize_color(c):
     return "white"
 
 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).
+    """
     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")
@@ -57,6 +59,7 @@ def make_qr(url="https://example.com", size=768, border=12, back_color="#FFFFFF"
     return img
 
 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)))
@@ -69,10 +72,9 @@ 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) ----------
-_CN = None
-_CN_TXT2IMG = {}
-_CN_IMG2IMG = {}
-ACTIVE_MODEL_ID = 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(
@@ -83,142 +85,35 @@ def _base_scheduler_for(pipe):
     pipe.enable_model_cpu_offload()
     return pipe
 
-def unload_all_pipes(keep_id=None):
-    for mid in list(_CN_TXT2IMG.keys()):
-        if mid != keep_id:
-            try: _CN_TXT2IMG[mid].to("cpu")
-            except: pass
-            del _CN_TXT2IMG[mid]
-    for mid in list(_CN_IMG2IMG.keys()):
-        if mid != keep_id:
-            try: _CN_IMG2IMG[mid].to("cpu")
-            except: pass
-            del _CN_IMG2IMG[mid]
-    if torch.cuda.is_available(): torch.cuda.empty_cache()
-    gc.collect()
-
 def get_cn():
     global _CN
     if _CN is None:
         _CN = ControlNetModel.from_pretrained(CN_QRMON, torch_dtype=DTYPE, use_safetensors=True)
     return _CN
 
-# ----- single-file fallback for Counterfeit (assemble ControlNet pipelines from components) -----
-CF_CANDIDATES = [
-    "Counterfeit-V3.0_fp16.safetensors",
-    "Counterfeit-V3.0.safetensors",
-    "Counterfeit-V3.0-pruned.safetensors",
-    "Counterfeit-V3.0.ckpt",
-]
-
-def _download_counterfeit_single_file():
-    last_err = None
-    for fname in CF_CANDIDATES:
-        try:
-            path = hf_hub_download(repo_id=BASE_MODELS["counterfeit"], filename=fname)
-            return path
-        except Exception as e:
-            last_err = e
-    raise last_err or RuntimeError("No suitable Counterfeit file found.")
-
-def _build_controlnet_pipes_from_single_file():
-    """Load Counterfeit via single .safetensors/.ckpt and assemble ControlNet txt2img/img2img pipelines."""
-    # 1) base SD pipeline from single file
-    ckpt_path = _download_counterfeit_single_file()
-    base = StableDiffusionPipeline.from_single_file(
-        ckpt_path,
-        torch_dtype=DTYPE,
-        safety_checker=None,
-        use_safetensors=True,
-    )
-    # 2) ControlNet
-    cn = get_cn()
-    # 3) Common kwargs (handle both image_processor/feature_extractor)
-    common = dict(
-        vae=base.vae,
-        text_encoder=base.text_encoder,
-        tokenizer=base.tokenizer,
-        unet=base.unet,
-        controlnet=cn,
-        scheduler=DPMSolverMultistepScheduler.from_config(
-            base.scheduler.config, use_karras_sigmas=True, algorithm_type="dpmsolver++"
-        ),
-        safety_checker=None,
-    )
-    if hasattr(base, "image_processor") and base.image_processor is not None:
-        common["image_processor"] = base.image_processor
-    if hasattr(base, "feature_extractor") and base.feature_extractor is not None:
-        common["feature_extractor"] = base.feature_extractor
-
-    txt = StableDiffusionControlNetPipeline(**common)
-    img = StableDiffusionControlNetImg2ImgPipeline(**common)
-    # memory-savers
-    for p in (txt, img):
-        p.enable_attention_slicing()
-        p.enable_vae_slicing()
-        p.enable_model_cpu_offload()
-    return txt, img
-
 def get_qrmon_txt2img_pipe(model_id: str):
-    global ACTIVE_MODEL_ID
-    if ACTIVE_MODEL_ID != model_id:
-        unload_all_pipes(keep_id=model_id)
-        ACTIVE_MODEL_ID = model_id
     if model_id not in _CN_TXT2IMG:
-        try:
-            pipe = StableDiffusionControlNetPipeline.from_pretrained(
-                model_id,
-                controlnet=get_cn(),
-                torch_dtype=DTYPE,
-                safety_checker=None,
-                use_safetensors=True,
-                low_cpu_mem_usage=True,
-            )
-        except EntryNotFoundError:
-            # Fallback: assemble from single file (Counterfeit)
-            txt, img = _build_controlnet_pipes_from_single_file()
-            _CN_TXT2IMG[model_id] = _base_scheduler_for(txt)
-            _CN_IMG2IMG[model_id] = _base_scheduler_for(img)
-            return _CN_TXT2IMG[model_id]
+        pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            model_id,
+            controlnet=get_cn(),
+            torch_dtype=DTYPE,
+            safety_checker=None,
+            use_safetensors=True,
+            low_cpu_mem_usage=True,
+        )
         _CN_TXT2IMG[model_id] = _base_scheduler_for(pipe)
     return _CN_TXT2IMG[model_id]
 
 def get_qrmon_img2img_pipe(model_id: str):
-    global ACTIVE_MODEL_ID
-    if ACTIVE_MODEL_ID != model_id:
-        unload_all_pipes(keep_id=model_id)
-        ACTIVE_MODEL_ID = model_id
     if model_id not in _CN_IMG2IMG:
-        try:
-            pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-                model_id,
-                controlnet=get_cn(),
-                torch_dtype=DTYPE,
-                safety_checker=None,
-                use_safetensors=True,
-                low_cpu_mem_usage=True,
-            )
-        except EntryNotFoundError:
-            # Fallback already built in get_qrmon_txt2img_pipe; reuse if exists
-            if model_id not in _CN_TXT2IMG:
-                txt, img = _build_controlnet_pipes_from_single_file()
-                _CN_TXT2IMG[model_id] = _base_scheduler_for(txt)
-                _CN_IMG2IMG[model_id] = _base_scheduler_for(img)
-                return _CN_IMG2IMG[model_id]
-            else:
-                # Build only img2img from base components again
-                txt = _CN_TXT2IMG[model_id]
-                # Pull components from txt to create img2img
-                common = dict(
-                    vae=txt.vae, text_encoder=txt.text_encoder, tokenizer=txt.tokenizer,
-                    unet=txt.unet, controlnet=txt.controlnet, scheduler=txt.scheduler,
-                    safety_checker=None,
-                )
-                if hasattr(txt, "image_processor"):
-                    common["image_processor"] = txt.image_processor
-                if hasattr(txt, "feature_extractor"):
-                    common["feature_extractor"] = txt.feature_extractor
-                pipe = StableDiffusionControlNetImg2ImgPipeline(**common)
+        pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+            model_id,
+            controlnet=get_cn(),
+            torch_dtype=DTYPE,
+            safety_checker=None,
+            use_safetensors=True,
+            low_cpu_mem_usage=True,
+        )
         _CN_IMG2IMG[model_id] = _base_scheduler_for(pipe)
     return _CN_IMG2IMG[model_id]
 
@@ -231,12 +126,16 @@ def _qr_txt2img_core(model_id: str,
                      repair_strength: float, feather: float):
 
     s = snap8(size)
+
+    # 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
     pipe = get_qrmon_txt2img_pipe(model_id)
     if torch.cuda.is_available(): torch.cuda.empty_cache()
     gc.collect()
@@ -244,8 +143,8 @@ def _qr_txt2img_core(model_id: str,
         out = pipe(
             prompt=str(style_prompt),
             negative_prompt=str(negative or ""),
-            image=qr_img,
-            controlnet_conditioning_scale=float(qr_weight),
+            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),
@@ -256,6 +155,7 @@ def _qr_txt2img_core(model_id: str,
     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)))
@@ -267,9 +167,9 @@ def _qr_txt2img_core(model_id: str,
             out2 = pipe2(
                 prompt=str(style_prompt),
                 negative_prompt=str(negative or ""),
-                image=lowres,
-                control_image=qr_img,
-                strength=float(hires_strength),
+                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),
                 control_guidance_start=0.0,
                 control_guidance_end=1.0,
@@ -283,9 +183,10 @@ def _qr_txt2img_core(model_id: 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_counterfeit(*args):
-    return _qr_txt2img_core(BASE_MODELS["counterfeit"], *args)
+def qr_txt2img_anything(*args):
+    return _qr_txt2img_core(BASE_MODELS["stable-diffusion-v1-5"], *args)
 
 @spaces.GPU(duration=120)
 def qr_txt2img_dream(*args):
@@ -293,19 +194,18 @@ def qr_txt2img_dream(*args):
 
 # ---------- UI ----------
 with gr.Blocks() as demo:
-    gr.Markdown("# ZeroGPU • Method 1: QR Control (Counterfeit V3.0 & DreamShaper 8)")
+    gr.Markdown("# ZeroGPU • Method 1: QR Control (two base models)")
 
-    with gr.Tab("Method 1 • Counterfeit V3.0"):
+    # ---- Tab 1: stable-diffusion-v1-5 (anime/illustration) ----
+    with gr.Tab("stable-diffusion-v1-5"):
         url1       = gr.Textbox(label="URL/Text", value="http://www.mybirdfire.com")
-        s_prompt1  = gr.Textbox(label="Style prompt",
-            value="ornate fantasy shrine, maple leaves, cinematic lighting, ultra detailed, anime key visual")
-        s_negative1= gr.Textbox(label="Negative prompt",
-            value="ugly, low quality, blurry, nsfw, watermark, text, low contrast, deformed, extra digits")
+        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=24, step=1, label="Steps")
+        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=6, step=1, label="QR border (quiet zone)")
-        qr_w1      = gr.Slider(0.8, 1.6, value=1.3, step=0.05, label="QR control weight")
+        border1    = gr.Slider(2, 16, value=4, step=1, label="QR border (quiet zone)")
+        qr_w1      = gr.Slider(0.6, 1.6, value=1.5, 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)")
@@ -319,22 +219,23 @@ with gr.Blocks() as demo:
         low_img1   = gr.Image(label="Low-res (Stage A) preview")
         ctrl_img1  = gr.Image(label="Control QR used")
 
-        gr.Button("Generate with Counterfeit V3.0").click(
-            qr_txt2img_counterfeit,
+        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]
         )
 
-    with gr.Tab("Method 1 • DreamShaper 8"):
+    # ---- Tab 2: DreamShaper (general art/painterly) ----
+    with gr.Tab("DreamShaper 8"):
         url2       = gr.Textbox(label="URL/Text", value="http://www.mybirdfire.com")
-        s_prompt2  = gr.Textbox(label="Style prompt", value="baroque palace interior, gilded details, chandeliers, volumetric god rays, ultra detailed")
+        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.8, 1.6, value=1.3, step=0.05, label="QR control weight")
+        qr_w2      = gr.Slider(0.6, 1.6, value=1.5, 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)")
@@ -356,5 +257,4 @@ with gr.Blocks() as demo:
         )
 
 if __name__ == "__main__":
-    gr.set_static_paths(["/tmp"])
     demo.queue(max_size=12).launch()