app.py

"""
QR Code Monster — ControlNet Space
Modell: monster-labs/control_v1p_sd15_qrcode_monster (v2)
"""

import os
import random

import gradio as gr
import qrcode
import torch
from PIL import Image
from qrcode.constants import (
    ERROR_CORRECT_L,
    ERROR_CORRECT_M,
    ERROR_CORRECT_Q,
    ERROR_CORRECT_H,
)
from diffusers import (
    ControlNetModel,
    StableDiffusionControlNetImg2ImgPipeline,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
)

# --------------------------------------------------------------------------
# ZeroGPU-Support (funktioniert auch lokal ohne das "spaces"-Paket)
# --------------------------------------------------------------------------
try:
    import spaces

    # Der erste Call schiebt ~5,7 GB Gewichte auf die Karte - deshalb großzügig.
    gpu_decorator = spaces.GPU(duration=120)
except Exception:  # lokal / eigene GPU

    def gpu_decorator(fn):
        return fn


# --------------------------------------------------------------------------
# Konfiguration
# --------------------------------------------------------------------------
BASE_MODEL = os.environ.get("BASE_MODEL", "stable-diffusion-v1-5/stable-diffusion-v1-5")
CONTROLNET_REPO = "monster-labs/control_v1p_sd15_qrcode_monster"
CONTROLNET_SUBFOLDER = "v2"  # v2 ist deutlich besser als v1
# Auf ZeroGPU ist beim Import noch keine GPU sichtbar ("Can't initialize NVML").
# Deshalb nicht auf torch.cuda.is_available() beim Start vertrauen.
IS_ZERO_GPU = os.environ.get("SPACES_ZERO_GPU", "").lower() in ("true", "1")
HAS_CUDA = IS_ZERO_GPU or torch.cuda.is_available()
DTYPE = torch.float16 if HAS_CUDA else torch.float32
GRAY = (128, 128, 128)  # #808080 – laut Model Card ideal zum "Verschmelzen"
MAX_SEED = 2**31 - 1

ERROR_LEVELS = {
    "L (7 %)": ERROR_CORRECT_L,
    "M (15 %)": ERROR_CORRECT_M,
    "Q (25 %)": ERROR_CORRECT_Q,
    "H (30 %) – empfohlen": ERROR_CORRECT_H,
}

SCHEDULERS = {
    "DPM++ 2M Karras": lambda cfg: DPMSolverMultistepScheduler.from_config(
        cfg, use_karras_sigmas=True, algorithm_type="dpmsolver++"
    ),
    "Euler a": lambda cfg: EulerAncestralDiscreteScheduler.from_config(cfg),
}

# --------------------------------------------------------------------------
# Pipeline laden (einmalig beim Start)
# --------------------------------------------------------------------------
controlnet = ControlNetModel.from_pretrained(
    CONTROLNET_REPO,
    subfolder=CONTROLNET_SUBFOLDER,
    dtype=DTYPE,  # torch_dtype ist ab diffusers 1.0 entfernt
)

pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
    BASE_MODEL,
    controlnet=controlnet,
    dtype=DTYPE,
    safety_checker=None,
    requires_safety_checker=False,
)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(
    pipe.scheduler.config, use_karras_sigmas=True, algorithm_type="dpmsolver++"
)

# Auf ZeroGPU darf CUDA erst innerhalb von @spaces.GPU angefasst werden -
# dort verschiebt generate() die Pipeline. Sonst gleich beim Start.
if HAS_CUDA and not IS_ZERO_GPU:
    pipe.to("cuda")
    pipe.enable_vae_tiling()
    pipe.enable_attention_slicing()


# --------------------------------------------------------------------------
# Hilfsfunktionen
# --------------------------------------------------------------------------
def make_qr_image(content: str, size: int, error_level: int, quiet_zone: int = 4) -> Image.Image:
    """Erzeugt ein QR-Bild mit Modulgröße ~16 px auf grauem Hintergrund."""
    qr = qrcode.QRCode(
        version=None,
        error_correction=error_level,
        box_size=16,  # Model Card: module size 16px
        border=quiet_zone,
    )
    qr.add_data(content)
    qr.make(fit=True)
    img = qr.make_image(fill_color="black", back_color="white").convert("RGB")

    # QR mittig auf graue Leinwand legen, ohne Kanten zu verwaschen (NEAREST!)
    inner = int(size * 0.9)
    img = img.resize((inner, inner), Image.NEAREST)
    canvas = Image.new("RGB", (size, size), GRAY)
    offset = (size - inner) // 2
    canvas.paste(img, (offset, offset))
    return canvas


def prepare_init_image(image: Image.Image | None, size: int) -> Image.Image:
    if image is None:
        return Image.new("RGB", (size, size), GRAY)
    return image.convert("RGB").resize((size, size), Image.LANCZOS)


# --------------------------------------------------------------------------
# Generierung
# --------------------------------------------------------------------------
@gpu_decorator
def generate(
    qr_content,
    prompt,
    negative_prompt,
    controlnet_scale,
    guidance_scale,
    steps,
    strength,
    seed,
    randomize_seed,
    size,
    error_level_name,
    scheduler_name,
    init_image,
    num_images,
    progress=gr.Progress(track_tqdm=True),
):
    if not qr_content or not qr_content.strip():
        raise gr.Error("Bitte Text oder URL für den QR-Code eingeben.")
    if not prompt or not prompt.strip():
        raise gr.Error("Bitte einen Prompt eingeben.")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)

    pipe.scheduler = SCHEDULERS[scheduler_name](pipe.scheduler.config)

    size = int(size)
    control_image = make_qr_image(qr_content, size, ERROR_LEVELS[error_level_name])
    init = prepare_init_image(init_image, size)

    device = "cuda" if HAS_CUDA else "cpu"
    if IS_ZERO_GPU:
        # Erst hier ist die ZeroGPU zugewiesen. Kein attention_slicing /
        # vae_tiling: die H200 hat reichlich VRAM, beides würde nur bremsen.
        pipe.to("cuda")

    generator = torch.Generator(device=device).manual_seed(seed)

    result = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt or None,
        image=init,
        control_image=control_image,
        width=size,
        height=size,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance_scale),
        controlnet_conditioning_scale=float(controlnet_scale),
        strength=float(strength),
        num_images_per_prompt=int(num_images),
        generator=generator,
    )

    return result.images, control_image, seed


# --------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------
DEFAULT_NEGATIVE = (
    "ugly, disfigured, low quality, blurry, jpeg artifacts, watermark, text, "
    "worst quality, lowres, deformed"
)

EXAMPLES = [
    ["https://qrcode.monster", "a japanese zen garden with raked sand, moss, soft morning light, 8k photo"],
    ["https://huggingface.co", "an ancient stone mosaic floor in a roman villa, intricate, weathered"],
    ["https://example.com", "aerial view of a snowy forest, winding paths, cinematic, highly detailed"],
]

# Gradio 6: theme/css/title gehören in launch(), nicht mehr in gr.Blocks()
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🧟 QR Code Monster
        Künstlerische, **scanbare** QR-Codes mit
        [`control_v1p_sd15_qrcode_monster`](https://huggingface.co/monster-labs/control_v1p_sd15_qrcode_monster) (v2).

        **Tipp:** Nicht jeder Code scannt beim ersten Versuch. Mehrere Seeds generieren,
        oder ControlNet-Stärke hoch + Denoising runter drehen.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            qr_content = gr.Textbox(
                label="QR-Inhalt (URL oder Text)",
                value="https://qrcode.monster",
                placeholder="https://…",
            )
            prompt = gr.Textbox(
                label="Prompt",
                lines=3,
                placeholder="z. B. a lush jungle with ancient ruins, cinematic lighting",
            )
            negative_prompt = gr.Textbox(
                label="Negativer Prompt", value=DEFAULT_NEGATIVE, lines=2
            )

            with gr.Row():
                controlnet_scale = gr.Slider(
                    0.5, 2.5, value=1.4, step=0.05,
                    label="ControlNet-Stärke (hoch = besser scanbar)",
                )
                strength = gr.Slider(
                    0.5, 1.0, value=0.9, step=0.01,
                    label="Denoising-Stärke",
                )

            with gr.Accordion("Erweiterte Einstellungen", open=False):
                with gr.Row():
                    guidance_scale = gr.Slider(1, 20, value=7.5, step=0.5, label="CFG Guidance")
                    steps = gr.Slider(10, 60, value=30, step=1, label="Steps")
                with gr.Row():
                    seed = gr.Number(value=0, label="Seed", precision=0)
                    randomize_seed = gr.Checkbox(value=True, label="Zufälliger Seed")
                with gr.Row():
                    size = gr.Radio([512, 640, 768], value=768, label="Auflösung")
                    num_images = gr.Slider(1, 4, value=1, step=1, label="Anzahl Bilder")
                error_level_name = gr.Dropdown(
                    list(ERROR_LEVELS), value="H (30 %) – empfohlen",
                    label="Fehlerkorrektur",
                )
                scheduler_name = gr.Dropdown(
                    list(SCHEDULERS), value="DPM++ 2M Karras", label="Sampler"
                )
                init_image = gr.Image(
                    label="Optionales Start-/Referenzbild (img2img)", type="pil"
                )

            run = gr.Button("QR-Code generieren", variant="primary")

        with gr.Column(scale=1):
            gallery = gr.Gallery(label="Ergebnisse", columns=2, height=520)
            control_preview = gr.Image(label="Verwendeter QR-Code (Condition)")
            used_seed = gr.Number(label="Verwendeter Seed", interactive=False)

    gr.Examples(examples=EXAMPLES, inputs=[qr_content, prompt])

    run.click(
        fn=generate,
        inputs=[
            qr_content, prompt, negative_prompt, controlnet_scale, guidance_scale,
            steps, strength, seed, randomize_seed, size, error_level_name,
            scheduler_name, init_image, num_images,
        ],
        outputs=[gallery, control_preview, used_seed],
    )

if __name__ == "__main__":
    # Queueing ist seit Gradio 5 standardmäßig aktiv.
    demo.launch(theme=gr.themes.Soft())