app.py

import os, re, gc, random
import numpy as np
from contextlib import nullcontext
from typing import Tuple

import gradio as gr
from PIL import Image, ImageFilter
import qrcode
from qrcode.constants import ERROR_CORRECT_H

import torch
from diffusers import (
    StableDiffusionPipeline,
    StableDiffusionControlNetImg2ImgPipeline,
    ControlNetModel,
    DPMSolverMultistepScheduler,
)
import spaces  # ZeroGPU decorator

# =========================================================
# Auth (optional for private models)
# =========================================================
hf_token = os.getenv("HF_TOKEN")
AUTH_KW = {"token": hf_token} if hf_token else {}

# =========================================================
# Helpers (untouched logic)
# =========================================================
def normalize_color(c):
    if c is None: return "white"
    if isinstance(c, (tuple, list)):
        r, g, b = (int(max(0, min(255, round(float(x))))) for x in c[:3]); return (r, g, b)
    if isinstance(c, str):
        s = c.strip()
        if s.startswith("#"): return s
        m = re.match(r"rgba?\(\s*([0-9.]+)\s*,\s*([0-9.]+)\s*,\s*([0-9.]+)", s, re.IGNORECASE)
        if m:
            r = int(max(0, min(255, round(float(m.group(1))))))
            g = int(max(0, min(255, round(float(m.group(2))))))
            b = int(max(0, min(255, round(float(m.group(3))))))
            return (r, g, b)
        return s
    return "white"

def strengthen_qr_prompts(pos: str, neg: str) -> Tuple[str, str]:
    # DON’T say “QR code” here – let ControlNet impose it
    pos = (pos or "").strip()
    neg = (neg or "").strip()
    pos2 = f"{pos}, high contrast lighting, clean details, cohesive composition".strip(", ")
    add_neg = "frame, border, ornate frame, watermark, text, numbers, checkerboard, mosaic, halftone, repeated pattern, glitch"
    neg2 = (neg + (", " if neg else "") + add_neg).strip(", ").strip()
    return pos2, neg2

def enforce_qr_contrast(stylized: Image.Image, qr_img: Image.Image, strength: float = 0.6, feather: float = 1.0) -> Image.Image:
    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)))
    black = np.asarray(black_mask, dtype=np.float32) / 255.0
    white = 1.0 - black
    s = np.asarray(stylized.convert("RGB"), dtype=np.float32) / 255.0
    s = s * (1.0 - float(strength) * black[..., None])
    s = s + (1.0 - s) * (float(strength) * 0.85 * white[..., None])
    s = np.clip(s, 0.0, 1.0)
    return Image.fromarray((s * 255.0).astype(np.uint8), mode="RGB")

# =========================================================
# Models & loading (ZeroGPU-friendly lazy load)
# =========================================================
BASE_15        = "runwayml/stable-diffusion-v1-5"
QR_MONSTER_15  = "monster-labs/control_v1p_sd15_qrcode_monster"  # v2 subfolder is handled by authors; base path is fine
BRIGHTNESS_15  = "latentcat/control_v1p_sd15_brightness"          # optional helper

_sd = {"pipe": None}
_cn = {"pipe": None}

def _setup_scheduler(pipe):
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(
        pipe.scheduler.config,
        use_karras_sigmas=True,
        algorithm_type="dpmsolver++"
    )

def _enable_memory_savers(pipe):
    # Good defaults for Spaces/ZeroGPU
    pipe.enable_attention_slicing()
    pipe.enable_vae_slicing()
    pipe.enable_vae_tiling()
    pipe.enable_model_cpu_offload()

def _load_sd_txt2img():
    if _sd["pipe"] is None:
        pipe = StableDiffusionPipeline.from_pretrained(
            BASE_15,
            torch_dtype=torch.float16,
            safety_checker=None,
            use_safetensors=True,
            low_cpu_mem_usage=True,
            **AUTH_KW
        )
        _setup_scheduler(pipe)
        _enable_memory_savers(pipe)
        _sd["pipe"] = pipe
    return _sd["pipe"]

def _load_cn_img2img():
    if _cn["pipe"] is None:
        qrnet  = ControlNetModel.from_pretrained(
            QR_MONSTER_15, torch_dtype=torch.float16, use_safetensors=True, **AUTH_KW
        )
        bright = ControlNetModel.from_pretrained(
            BRIGHTNESS_15, torch_dtype=torch.float16, use_safetensors=True, **AUTH_KW
        )
        pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
            BASE_15,
            controlnet=[qrnet, bright],
            torch_dtype=torch.float16,
            safety_checker=None,
            use_safetensors=True,
            low_cpu_mem_usage=True,
            **AUTH_KW
        )
        _setup_scheduler(pipe)
        _enable_memory_savers(pipe)
        _cn["pipe"] = pipe
    return _cn["pipe"]

# =========================================================
# Generation utilities (use inside @spaces.GPU)
# =========================================================
def sd_generate(prompt, negative, steps, guidance, seed, size=512):
    pipe = _load_sd_txt2img()
    # Reproducible generator — on GPU if available
    gen = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
    if int(seed) != 0:
        gen = gen.manual_seed(int(seed))
    else:
        gen = gen.manual_seed(random.randint(0, 2**31 - 1))

    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    gc.collect()

    out = pipe(
        prompt=prompt,
        negative_prompt=negative or "",
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        width=int(size), height=int(size),
        generator=gen
    )
    return out.images[0]

def make_qr(url="http://www.mybirdfire.com", size=512, border=10, back_color="#808080", blur_radius=0.0):
    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)
    bg  = normalize_color(back_color)
    img = qr.make_image(fill_color="black", back_color=bg).convert("RGB").resize((size, size), Image.NEAREST)
    if blur_radius and blur_radius > 0:
        img = img.filter(ImageFilter.GaussianBlur(radius=float(blur_radius)))
    return img

NEG_DEFAULT = "lowres, low contrast, blurry, jpeg artifacts, worst quality, bad anatomy, extra digits"

# =========================================================
# Main two-stage generator (ZeroGPU-guarded)
# =========================================================
@spaces.GPU(duration=120)  # allocate GPU only while generating
def qr_art_two_stage(
    prompt, negative,
    base_steps, base_cfg, base_seed,
    stylize_steps, stylize_cfg, stylize_seed,
    size, url, border, back_color,
    denoise, qr_weight, bright_weight,
    qr_start, qr_end, bright_start, bright_end,
    control_blur, repair_strength, feather_px
):
    size = max(384, int(size) // 8 * 8)

    # Stage A: base art (txt2img)
    p_pos, p_neg = strengthen_qr_prompts(prompt, negative)
    base_img = sd_generate(p_pos, p_neg, base_steps, base_cfg, base_seed, size=size)

    # Stage B: img2img + ControlNet
    qr_img = make_qr(url=url, size=size, border=border, back_color=back_color, blur_radius=control_blur)
    pipe   = _load_cn_img2img()

    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    gc.collect()

    gen = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
    if int(stylize_seed) != 0:
        gen = gen.manual_seed(int(stylize_seed))
    else:
        gen = gen.manual_seed(random.randint(0, 2**31 - 1))

    kwargs = dict(
        prompt=p_pos,
        negative_prompt=p_neg or NEG_DEFAULT,
        image=base_img,                              # init image for img2img
        control_image=[qr_img, qr_img],              # Monster + Brightness
        strength=float(denoise),                     # how much we allow change
        num_inference_steps=int(stylize_steps),
        guidance_scale=float(stylize_cfg),
        generator=gen,
        controlnet_conditioning_scale=[float(qr_weight), float(bright_weight)],
        width=size, height=size,  # (diffusers uses init image size; harmless here)
    )

    try:
        out = pipe(
            **kwargs,
            control_guidance_start=[float(qr_start), float(bright_start)],
            control_guidance_end=[float(qr_end), float(bright_end)],
        )
    except TypeError:
        out = pipe(
            **kwargs,
            controlnet_start=[float(qr_start), float(bright_start)],
            controlnet_end=[float(qr_end), float(bright_end)],
        )

    img = out.images[0]

    # Optional post repair to push blacks/whites where modules demand
    img = enforce_qr_contrast(img, qr_img, strength=float(repair_strength), feather=float(feather_px))
    return img, base_img, qr_img

# =========================================================
# UI (Gradio Space)
# =========================================================
with gr.Blocks() as demo:
    gr.Markdown("## 🧩 QR-Code Monster — Two-Stage (txt2img → img2img + ControlNet) — ZeroGPU")

    with gr.Tab("Two-Stage QR Art"):
        with gr.Row():
            with gr.Column():
                url      = gr.Textbox(label="URL/Text", value="http://www.mybirdfire.com")
                prompt   = gr.Textbox(
                    label="Style prompt (no 'QR code' here)",
                    value="baroque palace interior with intricate roots, cinematic, dramatic lighting, ultra detailed"
                )
                negative = gr.Textbox(label="Negative", value="")
                size     = gr.Slider(512, 1024, value=768, step=64, label="Canvas (px)")

                gr.Markdown("**Stage A — Base art (txt2img)**")
                base_steps = gr.Slider(10, 60, value=26, step=1, label="Base steps")
                base_cfg   = gr.Slider(1.0, 12.0, value=6.0, step=0.1, label="Base CFG")
                base_seed  = gr.Number(value=0, precision=0, label="Base seed (0=random)")

                gr.Markdown("**Stage B — ControlNet img2img**")
                stylize_steps = gr.Slider(10, 60, value=28, step=1, label="Stylize steps")
                stylize_cfg   = gr.Slider(1.0, 12.0, value=6.0, step=0.1, label="Stylize CFG")
                stylize_seed  = gr.Number(value=0, precision=0, label="Stylize seed (0=random)")
                denoise       = gr.Slider(0.1, 0.8, value=0.48, step=0.01, label="Denoising strength (keep composition lower)")

                qr_weight     = gr.Slider(0.5, 1.7, value=1.2, step=0.05, label="QR Monster weight")
                bright_weight = gr.Slider(0.0, 1.0, value=0.20, step=0.05, label="Brightness weight")

                qr_start      = gr.Slider(0.0, 1.0, value=0.05, step=0.01, label="QR start")
                qr_end        = gr.Slider(0.0, 1.0, value=0.95, step=0.01, label="QR end")
                bright_start  = gr.Slider(0.0, 1.0, value=0.40, step=0.01, label="Brightness start")
                bright_end    = gr.Slider(0.0, 1.0, value=0.85, step=0.01, label="Brightness end")

                border        = gr.Slider(4, 20, value=12, step=1, label="QR border (quiet zone)")
                back_color    = gr.ColorPicker(value="#808080", label="QR background (mid-gray blends better)")
                control_blur  = gr.Slider(0.0, 3.0, value=1.2, step=0.1, label="Soften control (Gaussian blur radius)")
                repair_strength = gr.Slider(0.0, 1.0, value=0.65, step=0.05, label="Post repair strength")
                feather_px      = gr.Slider(0.0, 3.0, value=1.0, step=0.1, label="Repair feather (px)")

                go = gr.Button("Generate QR Art", variant="primary")

            with gr.Column():
                final_img = gr.Image(label="Final stylized QR")
                base_img  = gr.Image(label="Base art (Stage A)")
                ctrl_img  = gr.Image(label="Control image (QR used)")

        go.click(
            qr_art_two_stage,
            inputs=[prompt, negative,
                    base_steps, base_cfg, base_seed,
                    stylize_steps, stylize_cfg, stylize_seed,
                    size, url, border, back_color,
                    denoise, qr_weight, bright_weight,
                    qr_start, qr_end, bright_start, bright_end,
                    control_blur, repair_strength, feather_px],
            outputs=[final_img, base_img, ctrl_img]
        )

if __name__ == "__main__":
    demo.launch()