"""Gradio ZeroGPU demo for VOSR image super-resolution."""

from __future__ import annotations

import spaces
import gradio as gr

from pipeline import DEFAULT_MODEL_LABEL, MODEL_CHOICES, load_pipeline, run_sr

# Eager-load default weights onto CUDA so ZeroGPU can pack them at startup.
print("Preparing default VOSR pipeline (may download weights on first boot)...")
try:
    load_pipeline(DEFAULT_MODEL_LABEL)
    print("Default pipeline ready.")
except Exception as exc:  # noqa: BLE001
    print(f"Startup load deferred: {exc}")


def _gpu_duration(image, model_label, upscale, infer_steps, *args):
    steps = int(infer_steps or 1)
    mode = MODEL_CHOICES.get(model_label, {}).get("mode", "multistep")
    pixels = 512 * 512
    if image is not None:
        w, h = image.size
        u = max(int(upscale or 1), 1)
        pixels = w * h * u * u
    mp = max(pixels / 1e6, 0.25)
    # Absolute ceiling 1800s; platform may still reject via illegal-duration / quota.
    if mode == "onestep":
        return int(min(1800, max(60, 45 + 25 * mp * steps)))
    return int(min(1800, max(90, 60 + 8 * mp * steps)))


@spaces.GPU(duration=_gpu_duration)
def upscale_image(
    image,
    model_label,
    upscale,
    infer_steps,
    cfg_scale,
    weak_cond,
    tile_size,
    tile_overlap,
    vae_tile_size,
    vae_tile_overlap,
    align_method,
    seed,
):
    return run_sr(
        image=image,
        model_label=model_label,
        upscale=upscale,
        infer_steps=infer_steps,
        cfg_scale=cfg_scale,
        weak_cond=weak_cond,
        tile_size=tile_size,
        tile_overlap=tile_overlap,
        align_method=align_method,
        seed=seed,
        vae_tile_size=vae_tile_size,
        vae_tile_overlap=vae_tile_overlap,
    )


def _on_model_change(model_label):
    spec = MODEL_CHOICES[model_label]
    is_ms = spec["mode"] == "multistep"
    return (
        gr.update(value=spec["default_steps"]),
        gr.update(interactive=is_ms),
        gr.update(interactive=is_ms),
    )


TITLE = "VOSR — Vision-Only Generative Super-Resolution"
DESCRIPTION = """
Demo of [VOSR](https://github.com/cswry/VOSR) (CVPR 2026) on **ZeroGPU**.
Upload a low-resolution image, keep the defaults (1.4B multi-step, 4×), and click **Upscale**.

For large images: set **DiT tile size** and optionally **VAE tile size** (both refer to
the upscaled result resolution; up to 8192). Prefer larger VAE tiles when VRAM allows —
tiling can still leave mild seams. Weights: [CSWRY/VOSR](https://huggingface.co/CSWRY/VOSR).
"""

with gr.Blocks(title="VOSR") as demo:
    gr.Markdown(f"# {TITLE}\n{DESCRIPTION}")
    with gr.Row():
        with gr.Column():
            inp = gr.Image(type="pil", label="Input image")
            model = gr.Dropdown(
                choices=list(MODEL_CHOICES.keys()),
                value=DEFAULT_MODEL_LABEL,
                label="Model",
            )
            upscale = gr.Slider(1, 8, value=4, step=1, label="Upscale factor")
            btn = gr.Button("Upscale", variant="primary")
            with gr.Accordion("Advanced", open=False):
                infer_steps = gr.Slider(1, 50, value=25, step=1, label="Inference steps")
                cfg_scale = gr.Slider(-1.0, 4.0, value=0.5, step=0.1, label="CFG scale (multi-step)")
                weak_cond = gr.Slider(
                    0.05, 0.25, value=0.10, step=0.01, label="Weak cond strength (multi-step)"
                )
                tile_size = gr.Slider(
                    0,
                    8192,
                    value=0,
                    step=64,
                    label="DiT tile size (0 = off; pixels on result / upscaled image)",
                )
                tile_overlap = gr.Slider(0, 1024, value=32, step=8, label="DiT tile overlap")
                vae_tile_size = gr.Slider(
                    0,
                    8192,
                    value=0,
                    step=64,
                    label="VAE tile size (0 = off; pixels on result / upscaled image)",
                )
                vae_tile_overlap = gr.Slider(
                    0, 1024, value=128, step=8, label="VAE tile overlap (pixels; ≥ tile/8 recommended)"
                )
                align_method = gr.Radio(
                    choices=["adain", "wavelet", "nofix"],
                    value="adain",
                    label="Color alignment",
                )
                seed = gr.Number(value=42, precision=0, label="Seed")
        with gr.Column():
            out = gr.Image(type="pil", label="Upscaled output")

    model.change(_on_model_change, inputs=[model], outputs=[infer_steps, cfg_scale, weak_cond])
    btn.click(
        fn=upscale_image,
        inputs=[
            inp,
            model,
            upscale,
            infer_steps,
            cfg_scale,
            weak_cond,
            tile_size,
            tile_overlap,
            vae_tile_size,
            vae_tile_overlap,
            align_method,
            seed,
        ],
        outputs=[out],
    )

if __name__ == "__main__":
    demo.queue(max_size=4).launch()