Update app.py

app.py

@@ -1,81 +1,48 @@
 #!/usr/bin/env python
 
-import spaces
-import os
-import random
-
+import os, random, numpy as np, cv2, torch
 import gradio as gr
-import numpy as np
-import PIL.Image
-import torch
-import torchvision.transforms.functional as TF
-
-from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL
-from diffusers import DDIMScheduler, EulerAncestralDiscreteScheduler
-from controlnet_aux import PidiNetDetector, HEDdetector
-from diffusers.utils import load_image
-from huggingface_hub import HfApi
 from pathlib import Path
 from PIL import Image, ImageOps
-import torch
-import numpy as np
-import cv2
-import os
-import random
+import PIL.Image
+
+import spaces
+from diffusers import (
+    ControlNetModel,
+    StableDiffusionXLControlNetPipeline,
+    AutoencoderKL,
+    EulerAncestralDiscreteScheduler,
+)
+from controlnet_aux import HEDdetector
 from gradio_imageslider import ImageSlider
 
+# ──────────────────────────────────────────────────────────────────────────────
+# Small JS helper to force dark theme (kept from your version)
+# ──────────────────────────────────────────────────────────────────────────────
 js_func = """
 function refresh() {
     const url = new URL(window.location);
-
     if (url.searchParams.get('__theme') !== 'dark') {
         url.searchParams.set('__theme', 'dark');
         window.location.href = url.href;
     }
 }
 """
-def nms(x, t, s):
-    x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
-
-    f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
-    f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
-    f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
-    f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
-
-    y = np.zeros_like(x)
-
-    for f in [f1, f2, f3, f4]:
-        np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
-
-    z = np.zeros_like(y, dtype=np.uint8)
-    z[y > t] = 255
-    return z
 
-def HWC3(x):
-    assert x.dtype == np.uint8
-    if x.ndim == 2:
-        x = x[:, :, None]
-    assert x.ndim == 3
-    H, W, C = x.shape
-    assert C == 1 or C == 3 or C == 4
-    if C == 3:
-        return x
-    if C == 1:
-        return np.concatenate([x, x, x], axis=2)
-    if C == 4:
-        color = x[:, :, 0:3].astype(np.float32)
-        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
-        y = color * alpha + 255.0 * (1.0 - alpha)
-        y = y.clip(0, 255).astype(np.uint8)
-        return y
-
-DESCRIPTION = '''# Scribble SDXL 🖋️🌄
-sketch to image with SDXL, using [@xinsir](https://huggingface.co/xinsir) [scribble sdxl controlnet](https://huggingface.co/xinsir/controlnet-scribble-sdxl-1.0), [sdxl controlnet canny](https://huggingface.co/xinsir/controlnet-canny-sdxl-1.0)
+# ──────────────────────────────────────────────────────────────────────────────
+# UI text
+# ──────────────────────────────────────────────────────────────────────────────
+DESCRIPTION = '''# Scribble SDXL 🖋️🌄 — live updates
+Sketch → image with SDXL ControlNet (scribble/canny). Now with **auto re-inference** when you draw or tweak settings (debounced).
+Models: [xinsir/controlnet-scribble-sdxl-1.0], [xinsir/controlnet-canny-sdxl-1.0], base [stabilityai/stable-diffusion-xl-base-1.0]
 '''
 
 if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
+    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo is intended for GPU Spaces for good latency.</p>"
 
+# ──────────────────────────────────────────────────────────────────────────────
+# Styles (unchanged, but refactored into a compact mapping)
+# ──────────────────────────────────────────────────────────────────────────────
 style_list = [
     {
         "name": "(No style)",
@@ -114,7 +81,7 @@ style_list = [
     },
     {
         "name": "Fantasy art",
-        "prompt": "ethereal fantasy concept art of  {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
+        "prompt": "ethereal fantasy concept art of {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
         "negative_prompt": "photographic, realistic, realism, 35mm film, dslr, cropped, frame, text, deformed, glitch, noise, noisy, off-center, deformed, cross-eyed, closed eyes, bad anatomy, ugly, disfigured, sloppy, duplicate, mutated, black and white",
     },
     {
@@ -128,204 +95,243 @@ style_list = [
         "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, realism, photorealistic, Western comic style",
     },
 ]
-
-styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
+styles = {s["name"]: (s["prompt"], s["negative_prompt"]) for s in style_list}
 STYLE_NAMES = list(styles.keys())
 DEFAULT_STYLE_NAME = "(No style)"
 
-
 def apply_style(style_name: str, positive: str, negative: str = "") -> tuple[str, str]:
     p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
-    return p.replace("{prompt}", positive), n + negative
+    return p.replace("{prompt}", positive), (n + " " + (negative or "")).strip()
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Utilities
+# ──────────────────────────────────────────────────────────────────────────────
+def HWC3(x: np.ndarray) -> np.ndarray:
+    assert x.dtype == np.uint8
+    if x.ndim == 2:
+        x = x[:, :, None]
+    H, W, C = x.shape
+    assert C in (1, 3, 4)
+    if C == 3:
+        return x
+    if C == 1:
+        return np.concatenate([x, x, x], axis=2)
+    # C == 4
+    color = x[:, :, 0:3].astype(np.float32)
+    alpha = x[:, :, 3:4].astype(np.float32) / 255.0
+    y = color * alpha + 255.0 * (1.0 - alpha)
+    y = y.clip(0, 255).astype(np.uint8)
+    return y
 
+def nms(x, t, s):
+    x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
+    f1 = np.array([[0,0,0],[1,1,1],[0,0,0]], dtype=np.uint8)
+    f2 = np.array([[0,1,0],[0,1,0],[0,1,0]], dtype=np.uint8)
+    f3 = np.array([[1,0,0],[0,1,0],[0,0,1]], dtype=np.uint8)
+    f4 = np.array([[0,0,1],[0,1,0],[1,0,0]], dtype=np.uint8)
+    y = np.zeros_like(x)
+    for f in [f1,f2,f3,f4]:
+        np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
+    z = np.zeros_like(y, dtype=np.uint8)
+    z[y > t] = 255
+    return z
 
+def clamp_size_to_megapixels(w: int, h: int, max_mpx: float = 1.0) -> tuple[int, int]:
+    """Scale so that w*h ≈ max_mpx*1e6 (default ~1024x1024 area)."""
+    area = w * h
+    target = max_mpx * 1_000_000.0
+    if area <= target:
+        return w, h
+    r = (target / area) ** 0.5
+    return max(64, int(w * r)) // 8 * 8, max(64, int(h * r)) // 8 * 8  # SDXL likes multiples of 8
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Load models once
+# ──────────────────────────────────────────────────────────────────────────────
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-eulera_scheduler = EulerAncestralDiscreteScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler")
-
+scheduler = EulerAncestralDiscreteScheduler.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler"
+)
 
-controlnet = ControlNetModel.from_pretrained(
-    "xinsir/controlnet-scribble-sdxl-1.0",
-    torch_dtype=torch.float16
+controlnet_scribble = ControlNetModel.from_pretrained(
+    "xinsir/controlnet-scribble-sdxl-1.0", torch_dtype=torch.float16 if device.type=="cuda" else torch.float32
 )
 controlnet_canny = ControlNetModel.from_pretrained(
-    "xinsir/controlnet-canny-sdxl-1.0",
-    torch_dtype=torch.float16
+    "xinsir/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16 if device.type=="cuda" else torch.float32
+)
+vae = AutoencoderKL.from_pretrained(
+    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16 if device.type=="cuda" else torch.float32
 )
-# when test with other base model, you need to change the vae also.
-vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 
-pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+pipe_scribble = StableDiffusionXLControlNetPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
-    controlnet=controlnet,
+    controlnet=controlnet_scribble,
     vae=vae,
-    torch_dtype=torch.float16,
-    scheduler=eulera_scheduler,
+    torch_dtype=torch.float16 if device.type=="cuda" else torch.float32,
+    scheduler=scheduler,
 )
-pipe.to(device)
-# Load model.
 pipe_canny = StableDiffusionXLControlNetPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
     controlnet=controlnet_canny,
     vae=vae,
-    safety_checker=None,
-    torch_dtype=torch.float16,
-    scheduler=eulera_scheduler,
+    torch_dtype=torch.float16 if device.type=="cuda" else torch.float32,
+    scheduler=scheduler,
 )
 
-pipe_canny.to(device)
+for p in (pipe_scribble, pipe_canny):
+    if device.type == "cuda":
+        try:
+            p.enable_xformers_memory_efficient_attention()
+        except Exception:
+            pass
+        p.enable_attention_slicing()
+    p.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
-processor = HEDdetector.from_pretrained('lllyasviel/Annotators')
-def nms(x, t, s):
-    x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
+hed = HEDdetector.from_pretrained("lllyasviel/Annotators")
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Core inference
+# ──────────────────────────────────────────────────────────────────────────────
+def _prepare_control_image(image_editor_value, use_hed: bool, use_canny: bool) -> Image.Image:
+    """
+    Accepts the dict from gr.ImageEditor (contains 'composite'), or a PIL.Image.
+    Returns a PIL.Image with control map (scribble/canny/hed result).
+    """
+    if image_editor_value is None:
+        return None
+
+    if isinstance(image_editor_value, dict) and "composite" in image_editor_value:
+        img = image_editor_value["composite"]
+    elif isinstance(image_editor_value, PIL.Image.Image):
+        img = image_editor_value
+    else:
+        return None
 
-    f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
-    f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
-    f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
-    f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
+    # Convert to RGB for detectors
+    if img.mode != "RGB":
+        img = img.convert("RGB")
 
-    y = np.zeros_like(x)
+    if use_canny:
+        arr = np.array(img)
+        edge = cv2.Canny(arr, 100, 200)
+        edge = HWC3(edge)
+        return Image.fromarray(edge)
+
+    if use_hed:
+        control = hed(img, scribble=False)
+        control = np.array(control)
+        control = nms(control, 127, 3)
+        control = cv2.GaussianBlur(control, (0, 0), 3)
+
+        # Simulate human sketch width with a soft random threshold
+        thr = int(round(random.uniform(0.01, 0.10), 2) * 255)
+        control[control > thr] = 255
+        control[control < 255] = 0
+        return Image.fromarray(control)
+
+    # Default: use the editor composite as "scribble"
+    return img
+
+def _image_size_from_editor(image_editor_value, target_mpx=1.0) -> tuple[int, int]:
+    if image_editor_value is None:
+        return 1024, 1024
+    if isinstance(image_editor_value, dict) and "composite" in image_editor_value:
+        w, h = image_editor_value["composite"].size
+    elif isinstance(image_editor_value, PIL.Image.Image):
+        w, h = image_editor_value.size
+    else:
+        w, h = 1024, 1024
+    return clamp_size_to_megapixels(w, h, max_mpx=target_mpx)
 
-    for f in [f1, f2, f3, f4]:
-        np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
+def _pick_pipe(use_canny: bool):
+    return pipe_canny if use_canny else pipe_scribble
 
-    z = np.zeros_like(y, dtype=np.uint8)
-    z[y > t] = 255
-    return z
+def _maybe_seed(seed: int):
+    if seed is None or seed < 0:
+        return None
+    return torch.Generator(device=device).manual_seed(int(seed))
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
+    return random.randint(0, MAX_SEED) if randomize_seed else int(seed)
 
 @spaces.GPU
 def run(
-    image: PIL.Image.Image,
+    image,                        # dict from ImageEditor or PIL.Image
     prompt: str,
     negative_prompt: str,
     style_name: str = DEFAULT_STYLE_NAME,
-    num_steps: int = 25,
-    guidance_scale: float = 5,
+    num_steps: int = 12,
+    guidance_scale: float = 5.0,
     controlnet_conditioning_scale: float = 1.0,
     seed: int = 0,
     use_hed: bool = False,
     use_canny: bool = False,
     progress=gr.Progress(track_tqdm=True),
-) -> PIL.Image.Image:
-    width, height  = image['composite'].size
-    ratio = np.sqrt(1024. * 1024. / (width * height))
-    new_width, new_height = int(width * ratio), int(height * ratio)
-    image = image['composite'].resize((new_width, new_height))
-    
-    if use_canny:
-        controlnet_img = np.array(image)
-        controlnet_img = cv2.Canny(controlnet_img, 100, 200)
-        controlnet_img = HWC3(controlnet_img)
-        image = Image.fromarray(controlnet_img)
-    
-    elif not use_hed:
-          controlnet_img = image
-    else:
-        controlnet_img = processor(image, scribble=False)
-      # following is some processing to simulate human sketch draw, different threshold can generate different width of lines
-        controlnet_img = np.array(controlnet_img)
-        controlnet_img = nms(controlnet_img, 127, 3)
-        controlnet_img = cv2.GaussianBlur(controlnet_img, (0, 0), 3)
-
-        # higher threshold, thiner line
-        random_val = int(round(random.uniform(0.01, 0.10), 2) * 255)
-        controlnet_img[controlnet_img > random_val] = 255
-        controlnet_img[controlnet_img < 255] = 0
-        image = Image.fromarray(controlnet_img)
-    
-    
-    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)
-
-    generator = torch.Generator(device=device).manual_seed(seed)
-    if use_canny:
-        out = pipe_canny(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        image=image,
-        num_inference_steps=num_steps,
-        generator=generator,
-        controlnet_conditioning_scale=controlnet_conditioning_scale,
-        guidance_scale=guidance_scale,
-        width=new_width,
-        height=new_height,
+):
+    if image is None or (isinstance(prompt, str) and prompt.strip() == ""):
+        return (None, None)
+
+    # Prepare control image + target size (≈1MP for speed)
+    ctrl_img = _prepare_control_image(image, use_hed=use_hed, use_canny=use_canny)
+    w, h = _image_size_from_editor(image, target_mpx=1.0)
+
+    # Style injection
+    prompt_styled, neg_styled = apply_style(style_name, prompt, negative_prompt or "")
+
+    g = _maybe_seed(seed)
+    pipe = _pick_pipe(use_canny)
+
+    out = pipe(
+        prompt=prompt_styled,
+        negative_prompt=neg_styled,
+        image=ctrl_img,
+        num_inference_steps=int(num_steps),
+        controlnet_conditioning_scale=float(controlnet_conditioning_scale),
+        guidance_scale=float(guidance_scale),
+        generator=g,
+        width=w, height=h,
     ).images[0]
-    else:
-        out = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        image=image,
-        num_inference_steps=num_steps,
-        generator=generator,
-        controlnet_conditioning_scale=controlnet_conditioning_scale,
-        guidance_scale=guidance_scale,
-        width=new_width,
-        height=new_height,).images[0]
-
-    return (controlnet_img, out)
 
+    # Return (control, output) for ImageSlider
+    if isinstance(ctrl_img, Image.Image):
+        ci = ctrl_img
+    else:
+        ci = Image.fromarray(ctrl_img) if ctrl_img is not None else None
+    return (ci, out)
 
-with gr.Blocks(css="style.css", js=js_func) as demo:
+# ──────────────────────────────────────────────────────────────────────────────
+# UI (with live updates wired via .change on inputs)
+# ──────────────────────────────────────────────────────────────────────────────
+with gr.Blocks(css="style.css", js=js_func, title="Scribble SDXL — Live") as demo:
     gr.Markdown(DESCRIPTION, elem_id="description")
-    gr.DuplicateButton(
-        value="Duplicate Space for private use",
-        elem_id="duplicate-button",
-        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
-    )
 
     with gr.Row():
         with gr.Column():
             with gr.Group():
-                image = gr.ImageEditor(type="pil", image_mode="L", crop_size=(512, 512))
-                prompt = gr.Textbox(label="Prompt")
+                image = gr.ImageEditor(type="pil", image_mode="L", crop_size=(512, 512), label="Draw / Edit")
+                prompt = gr.Textbox(label="Prompt", value="a detailed robot mascot, studio lighting, clean lines")
                 style = gr.Dropdown(label="Style", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME)
-                use_hed = gr.Checkbox(label="use HED detector", value=False, info="check this box if you upload an image and want to turn it to a sketch")
-                use_canny = gr.Checkbox(label="use Canny", value=False, info="check this to use ControlNet canny instead of scribble")
+                use_hed = gr.Checkbox(label="Use HED detector (turn photo → sketch)", value=False)
+                use_canny = gr.Checkbox(label="Use Canny (ControlNet Canny)", value=False)
                 run_button = gr.Button("Run")
             with gr.Accordion("Advanced options", open=False):
                 negative_prompt = gr.Textbox(
                     label="Negative prompt",
                     value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
                 )
-                num_steps = gr.Slider(
-                    label="Number of steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=25,
-                )
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.1,
-                    maximum=10.0,
-                    step=0.1,
-                    value=5,
-                )
+                num_steps = gr.Slider(label="Steps (lower = faster)", minimum=4, maximum=40, step=1, value=12)
+                guidance_scale = gr.Slider(label="Guidance", minimum=0.1, maximum=12.0, step=0.1, value=5.0)
                 controlnet_conditioning_scale = gr.Slider(
-                    label="controlnet conditioning scale",
-                    minimum=0.5,
-                    maximum=5.0,
-                    step=0.1,
-                    value=0.9,
+                    label="Control strength", minimum=0.5, maximum=2.0, step=0.05, value=1.0
                 )
-                seed = gr.Slider(
-                    label="Seed",
-                    minimum=0,
-                    maximum=MAX_SEED,
-                    step=1,
-                    value=0,
-                )
-                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-               
+                seed = gr.Slider(label="Seed (-1 random)", minimum=-1, maximum=MAX_SEED, step=1, value=-1)
+                randomize_seed = gr.Checkbox(label="Randomize seed on Run", value=True)
+
         with gr.Column():
             with gr.Group():
-                image_slider = ImageSlider(position=0.5)
-
+                image_slider = ImageSlider(position=0.5, label="Control ↔ Output")
 
     inputs = [
         image,
@@ -337,19 +343,19 @@ with gr.Blocks(css="style.css", js=js_func) as demo:
         controlnet_conditioning_scale,
         seed,
         use_hed,
-        use_canny
+        use_canny,
     ]
     outputs = [image_slider]
+
+    # Manual "Run" flow (seed randomization, clear slider, then infer)
     run_button.click(
-        fn=randomize_seed_fn,
-        inputs=[seed, randomize_seed],
-        outputs=seed,
-        queue=False,
-        api_name=False,
-    ).then(lambda x: None, inputs=None, outputs=image_slider).then(
-        fn=run, inputs=inputs, outputs=outputs
-    )
-    
-    
-
-demo.queue().launch()
+        fn=randomize_seed_fn, inputs=[seed, randomize_seed], outputs=seed, queue=False, api_name=False
+    ).then(lambda: None, inputs=None, outputs=image_slider).then(fn=run, inputs=inputs, outputs=outputs)
+
+    # ── Live re-inference hooks (debounced) ───────────────────────────────────
+    # Fire when drawing or tweaking settings. 'every' = debounce seconds.
+    for comp in [image, prompt, negative_prompt, style, num_steps, guidance_scale,
+                 controlnet_conditioning_scale, seed, use_hed, use_canny]:
+        comp.change(fn=run, inputs=inputs, outputs=outputs, every=0.5, queue=True)
+
+demo.queue(concurrency_count=2, max_size=20).launch()