Upload app.py

app.py

@@ -0,0 +1,435 @@
+import os
+from dataclasses import dataclass
+from PIL import Image
+
+import cv2
+import numpy as np
+import gradio as gr
+import torch
+
+import spaces  # type: ignore
+
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
+from diffusers.models.controlnets.controlnet import ControlNetModel
+from diffusers.pipelines.controlnet.pipeline_controlnet import StableDiffusionControlNetPipeline
+from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
+from transformers import CLIPTextModel, CLIPTokenizer
+
+BIG_CSS = """
+/* Global bump */
+.gradio-container {
+  font-size: 18px !important;
+}
+
+/* Force most UI text bigger */
+.gradio-container * {
+  font-size: 18px !important;
+}
+
+/* Keep markdown headings bigger */
+.gradio-container h1 { font-size: 28px !important; }
+.gradio-container h2 { font-size: 24px !important; }
+.gradio-container h3 { font-size: 20px !important; }
+
+/* Slightly smaller helper/info text if you want */
+.gradio-container .info,
+.gradio-container .prose p,
+.gradio-container .prose li {
+  font-size: 16px !important;
+  line-height: 1.35 !important;
+}
+"""
+
+# -----------------------------
+# Pipeline builder
+# -----------------------------
+def build_controlnet_pipe(
+    base_model_name: str,
+    controlnet: ControlNetModel,
+    vae: AutoencoderKL,
+    unet: UNet2DConditionModel,
+    text_encoder: CLIPTextModel,
+    tokenizer: CLIPTokenizer,
+    device: torch.device,
+    weight_dtype: torch.dtype,
+    use_unipc: bool = True,
+) -> StableDiffusionControlNetPipeline:
+    pipe = StableDiffusionControlNetPipeline.from_pretrained(
+        base_model_name,
+        vae=vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        unet=unet,
+        controlnet=controlnet,
+        safety_checker=None,
+        torch_dtype=weight_dtype,
+    )
+    if use_unipc:
+        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+    pipe = pipe.to(device)
+    pipe.set_progress_bar_config(disable=True)
+    return pipe
+
+@dataclass
+class CannyCFG:
+    use_clahe: bool = True
+    clahe_clip: float = 2.0
+    clahe_grid: int = 8
+
+    gaussian_ksize: int = 5
+    gaussian_sigma: float = 1.2
+
+    high_pct: float = 90.0     # higher -> fewer edges (stricter)
+    low_ratio: float = 0.4     # low = low_ratio * high
+
+    aperture_size: int = 3
+    l2_gradient: bool = True
+
+
+def canny_percentile(pil_img: Image.Image, cfg: CannyCFG) -> Image.Image:
+    gray = np.array(pil_img.convert("L"), dtype=np.uint8)
+
+    if cfg.use_clahe:
+        clahe = cv2.createCLAHE(
+            clipLimit=float(cfg.clahe_clip),
+            tileGridSize=(int(cfg.clahe_grid), int(cfg.clahe_grid)),
+        )
+        gray = clahe.apply(gray)
+
+    k = int(cfg.gaussian_ksize) | 1  # ensure odd
+    blur = cv2.GaussianBlur(gray, (k, k), float(cfg.gaussian_sigma))
+
+    gx = cv2.Sobel(blur, cv2.CV_32F, 1, 0, ksize=3)
+    gy = cv2.Sobel(blur, cv2.CV_32F, 0, 1, ksize=3)
+    mag = cv2.magnitude(gx, gy)
+
+    high = float(np.percentile(mag, float(cfg.high_pct)))
+    low = float(cfg.low_ratio) * high
+    if high <= low:
+        high = low + 1.0
+
+    ap = int(cfg.aperture_size)
+    if ap not in (3, 5, 7):
+        ap = 3
+
+    edges = cv2.Canny(
+        blur,
+        threshold1=low,
+        threshold2=high,
+        apertureSize=ap,
+        L2gradient=bool(cfg.l2_gradient),
+    )
+    return Image.fromarray(edges, mode="L")
+
+
+# -----------------------------
+# Config
+# -----------------------------
+BASE_MODEL = "sd-legacy/stable-diffusion-v1-5"
+WEIGHTS_REPO = "mvp-lab/ControlNet_Weight"
+WEIGHTS_FILENAME = "diffusion_pytorch_model_1.safetensors"
+
+LOCAL_WEIGHTS = os.getenv(
+    "CONTROLNET_WEIGHTS",
+    "/home/nik/ImperialWork/GenerativeAi/sd15-controlnet-trainer/controlnet_laion/final/diffusion_pytorch_model.safetensors",
+)
+if os.path.isfile(LOCAL_WEIGHTS):
+    CONTROLNET_PATH = LOCAL_WEIGHTS
+else:
+    CONTROLNET_PATH = hf_hub_download(repo_id=WEIGHTS_REPO, filename=WEIGHTS_FILENAME, repo_type="model")
+
+DTYPE = torch.float32
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+# -----------------------------
+# Model load (once)
+# -----------------------------
+vae = AutoencoderKL.from_pretrained(BASE_MODEL, subfolder="vae", torch_dtype=DTYPE)
+unet = UNet2DConditionModel.from_pretrained(BASE_MODEL, subfolder="unet", torch_dtype=DTYPE)
+tokenizer = CLIPTokenizer.from_pretrained(BASE_MODEL, subfolder="tokenizer")
+text_encoder = CLIPTextModel.from_pretrained(BASE_MODEL, subfolder="text_encoder", torch_dtype=DTYPE)
+
+vae.requires_grad_(False)
+unet.requires_grad_(False)
+text_encoder.requires_grad_(False)
+
+controlnet = ControlNetModel.from_unet(unet, conditioning_channels=3)
+state = load_file(CONTROLNET_PATH)
+missing, unexpected = controlnet.load_state_dict(state, strict=False)
+
+pipe = build_controlnet_pipe(
+    base_model_name=BASE_MODEL,
+    controlnet=controlnet,
+    vae=vae,
+    unet=unet,
+    text_encoder=text_encoder,
+    tokenizer=tokenizer,
+    device=DEVICE,
+    weight_dtype=DTYPE,
+    use_unipc=True,
+)
+
+
+# -----------------------------
+# Helpers: fixed resize policy (longest side = 512, keep aspect, divisible by 8)
+# -----------------------------
+def round_down_to_multiple(x: int, m: int = 8) -> int:
+    return max(m, (x // m) * m)
+
+def resize_longest_side_div8(img: Image.Image, longest: int = 512) -> tuple[Image.Image, int, int]:
+    w, h = img.size
+    if w <= 0 or h <= 0:
+        raise ValueError("Invalid image size")
+
+    scale = float(longest) / float(max(w, h))
+    tw = int(round(w * scale))
+    th = int(round(h * scale))
+
+    tw = round_down_to_multiple(tw, 8)
+    th = round_down_to_multiple(th, 8)
+
+    tw = max(8, tw)
+    th = max(8, th)
+
+    resized = img.resize((tw, th), resample=Image.BICUBIC) # type: ignore
+    return resized, tw, th
+
+def compute_canny_rgb(img_rgb_resized: Image.Image, use_clahe: bool, edge_amount: float, smoothing: float) -> Image.Image:
+    high_pct = 95.0 - 20.0 * float(edge_amount)  # 0 => 95 (few), 1 => 75 (many)
+    high_pct = float(np.clip(high_pct, 70.0, 99.0))
+
+    gaussian_sigma = 0.6 + 2.2 * float(smoothing)  # 0 => 0.6, 1 => 2.8
+
+    cfg = CannyCFG(
+        use_clahe=bool(use_clahe),
+        clahe_clip=2.0,
+        clahe_grid=8,
+        gaussian_ksize=5,
+        gaussian_sigma=float(gaussian_sigma),
+        high_pct=float(high_pct),
+        low_ratio=0.4,
+        aperture_size=3,
+        l2_gradient=True,
+    )
+    edges_l = canny_percentile(img_rgb_resized, cfg)
+    return edges_l.convert("RGB")
+
+def update_canny_preview(input_image, use_clahe, edge_amount, smoothing):
+    if input_image is None:
+        return None, None, 512, 512
+
+    if not isinstance(input_image, Image.Image):
+        input_image = Image.fromarray(input_image)
+
+    img_rgb0 = input_image.convert("RGB")
+    img_rgb, width, height = resize_longest_side_div8(img_rgb0, longest=512)
+
+    canny = compute_canny_rgb(
+        img_rgb,
+        use_clahe=use_clahe,
+        edge_amount=float(edge_amount),
+        smoothing=float(smoothing),
+    )
+    return canny, canny, width, height
+
+
+@spaces.GPU
+@torch.inference_mode()
+def generate_from_canny(
+    canny: Image.Image,
+    width: int,
+    height: int,
+    prompt: str,
+    negative_prompt: str,
+    guidance_scale: float,
+    num_inference_steps: int,
+    num_images: int,
+    controlnet_conditioning_scale: float,
+):
+    if canny is None:
+        raise gr.Error("Canny conditioning image missing. Upload an image first.")
+    if int(num_images) < 1:
+        raise gr.Error("num_images must be >= 1")
+
+    gens = [torch.Generator(device=DEVICE).manual_seed(i) for i in range(int(num_images))]
+
+    imgs = pipe(
+        prompt=[prompt] * int(num_images),
+        negative_prompt=[negative_prompt] * int(num_images),
+        image=[canny] * int(num_images),
+        num_inference_steps=int(num_inference_steps),
+        guidance_scale=float(guidance_scale),
+        height=int(height),
+        width=int(width),
+        generator=gens,
+        controlnet_conditioning_scale=float(controlnet_conditioning_scale),
+    ).images # type: ignore
+
+    first = imgs[0] if imgs else None
+    return first, imgs
+
+def next_image(images, idx):
+    if not images:
+        return None, 0, "0 / 0"
+    idx = (int(idx) + 1) % len(images)
+    return images[idx], idx, f"{idx + 1} / {len(images)}"
+
+def prev_image(images, idx):
+    if not images:
+        return None, 0, "0 / 0"
+    idx = (int(idx) - 1) % len(images)
+    return images[idx], idx, f"{idx + 1} / {len(images)}"
+
+
+# -----------------------------
+# UI
+# -----------------------------
+IMG_H = 360  # uniform-ish size for both preview boxes
+
+with gr.Blocks(css=BIG_CSS) as demo:
+    gr.Markdown("# Canny-Edge ControlNet Demo")
+    gr.Markdown("**Note:** Trained on aesthetic/artistic images — best results come from similar, stylised inputs.")
+
+    # state
+    canny_state = gr.State(None)
+    width_state = gr.State(512)
+    height_state = gr.State(512)
+
+    gen_images_state = gr.State([])  # list[PIL]
+    gen_index_state = gr.State(0)
+
+    with gr.Row():
+        # ---- Left: Canny + Canny controls ----
+        with gr.Column(scale=1):
+            input_image = gr.Image(
+                label="Input Image",
+                type="pil",
+                image_mode="RGB",
+                height=IMG_H,
+            )
+
+            canny_preview = gr.Image(
+                label="Canny edges",
+                type="pil",
+                height=IMG_H,
+            )
+
+            gr.Markdown("### Edge controls")
+            use_clahe = gr.Checkbox(
+                label="Stabilise contrast (CLAHE)",
+                value=True,
+                info="Helps edges stay consistent under different lighting/contrast.",
+            )
+            edge_amount = gr.Slider(
+                label="Edge Amount",
+                minimum=0.0, maximum=1.0, value=0.6, step=0.01,
+                info="More = detect more edges (more detail). Less = cleaner outline.",
+            )
+            smoothing = gr.Slider(
+                label="Smoothing",
+                minimum=0.0, maximum=1.0, value=0.4, step=0.01,
+                info="More = reduce tiny texture/noise edges, cleaner structure.",
+            )
+
+        # ---- Right: Generated output + generation controls ----
+        with gr.Column(scale=1):
+            generated = gr.Image(
+                label="Generated image",
+                type="pil",
+                height=IMG_H,
+            )
+
+            with gr.Row():
+                prev_btn = gr.Button("◀ Prev")
+                page_label = gr.Markdown("0 / 0")
+                next_btn = gr.Button("Next ▶")
+
+            gr.Markdown("### Generation controls")
+            positive_prompt = gr.Textbox(
+                label="Positive Prompt",
+                value="",
+                lines=2,
+                info="Describe what you want. The edges guide the structure.",
+            )
+            negative_prompt = gr.Textbox(
+                label="Negative Prompt",
+                value="",
+                lines=2,
+                info="Things to avoid (e.g. blurry, deformed, low quality).",
+            )
+
+            with gr.Row():
+                guidance_scale = gr.Slider(
+                    label="Guidance Scale",
+                    minimum=1.0, maximum=15.0, value=7.5, step=0.1,
+                    info="Higher = follow text prompt more strongly (can drift from edges).",
+                )
+                controlnet_conditioning_scale = gr.Slider(
+                    label="Control Strength",
+                    minimum=0.0, maximum=2.0, value=1.0, step=0.05,
+                    info="Higher = follow edges more strongly. Too high can reduce creativity.",
+                )
+
+            with gr.Row():
+                num_inference_steps = gr.Slider(
+                    label="Steps",
+                    minimum=10, maximum=80, value=50, step=1,
+                    info="More steps can improve quality but is slower.",
+                )
+                num_images = gr.Slider(
+                    label="Samples",
+                    minimum=1, maximum=8, value=4, step=1,
+                    info="How many images to generate.",
+                )
+
+            run_btn = gr.Button("Generate", variant="primary")
+
+    # Auto-update Canny preview on changes (CPU)
+    auto_inputs = [input_image, use_clahe, edge_amount, smoothing]
+    for c in auto_inputs:
+        c.change(
+            fn=update_canny_preview,
+            inputs=auto_inputs,
+            outputs=[canny_preview, canny_state, width_state, height_state],
+        )
+
+    # Generate (GPU) -> store list -> show first -> update paging label
+    run_btn.click(
+        fn=generate_from_canny,
+        inputs=[
+            canny_state,
+            width_state,
+            height_state,
+            positive_prompt,
+            negative_prompt,
+            guidance_scale,
+            num_inference_steps,
+            num_images,
+            controlnet_conditioning_scale,
+        ],
+        outputs=[generated, gen_images_state],  # visible output first => proper "Generating..." UX
+    ).then(
+        fn=lambda imgs: (0, f"1 / {len(imgs)}") if imgs else (0, "0 / 0"),
+        inputs=[gen_images_state],
+        outputs=[gen_index_state, page_label],
+    )
+
+    # Paging buttons (CPU)
+    next_btn.click(
+        fn=next_image,
+        inputs=[gen_images_state, gen_index_state],
+        outputs=[generated, gen_index_state, page_label],
+    )
+    prev_btn.click(
+        fn=prev_image,
+        inputs=[gen_images_state, gen_index_state],
+        outputs=[generated, gen_index_state, page_label],
+    )
+
+if __name__ == "__main__":
+    demo.launch()