app.py

"""
Anima RDBT — Gradio front-end for ZeroGPU Hugging Face Spaces, with an embedded
ComfyUI backend for native RDBT KSampler execution.
"""
from __future__ import annotations

import functools
import os
import traceback
from typing import Any

import gradio as gr

# ZeroGPU: https://huggingface.co/docs/hub/main/spaces-zerogpu
try:
    import spaces
except ImportError:  # local / tests without the runtime wheel
    class _Spaces:
        @staticmethod
        def GPU(*args: Any, **kwargs: Any):
            if args and callable(args[0]) and not kwargs:
                return args[0]

            def _deco(fn: Any) -> Any:
                return fn

            return _deco

    spaces = _Spaces()  # type: ignore[misc, assignment]

from src.comfy_backend import run_generation
from src.errors import UserFacingError
from src.validation import validate_and_clamp

from src import config

LICENSE_BANNER = (
    "**License:** Anima is distributed under the **CircleStone Labs Non-Commercial License** "
    "and is a derivative of **Cosmos-Predict2-2B-Text2Image** (NVIDIA terms where applicable). "
    "Use **non-commercially** unless you have a separate license. See the "
    "[Anima model card](https://huggingface.co/circlestone-labs/Anima)."
)

FALLBACK_NOTE = (
    "Generation uses ComfyUI's native nodes and KSampler under the Gradio UI. "
    "Sampler, scheduler, cfg, denoise, seed, size, and batch are passed through after validation."
)


def _compute_gpu_duration(
    _prompt: str,
    _neg: str,
    _w: int,
    _h: int,
    steps: float,
    _cfg: float,
    batch: float,
    _s: str,
    _sc: str,
    _d: float,
    _seed: float,
    _randomize_seed: bool,
    _progress: Any = None,
) -> int:
    """Must accept the same positional args as _generate_task (including Gradio-injected progress)."""
    try:
        s = int(steps)
    except (TypeError, ValueError):
        s = 16
    try:
        b = int(batch)
    except (TypeError, ValueError):
        b = 1
    s = max(1, min(50, s))
    b = max(1, min(config.MAX_BATCH, b))
    # Queue-friendly: high-res + batch → more time; cap 300s (ZeroGPU allowlist)
    est = 25 + s * 4 * b
    w = int(os.environ.get("ANIMA_MAX_GPU_DURATION", "300"))
    return min(int(w), max(30, int(est)))


@spaces.GPU(duration=_compute_gpu_duration)  # type: ignore[misc]
def _generate_task(
    prompt: str,
    negative_prompt: str,
    width: float,
    height: float,
    steps: float,
    cfg: float,
    batch_size: float,
    sampler_name: str,
    scheduler: str,
    denoise: float,
    seed: float,
    randomize_seed: bool,
    progress: gr.Progress = gr.Progress(),
) -> tuple[Any, str, str, str]:
    """Returns (gallery or None, status markdown, actual positive str, actual negative str)."""
    try:
        params = validate_and_clamp(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            steps=steps,
            cfg=cfg,
            batch_size=batch_size,
            sampler_name=sampler_name,
            scheduler=scheduler,
            denoise=denoise,
            seed=seed,
            randomize_seed=randomize_seed,
        )
    except UserFacingError as e:
        lines = [f"**Error:** {e.user_message}"]
        if e.details:
            lines.append(f"`{e.details}`")
        return (
            None,
            "\n\n".join(lines),
            "— (nothing sent — validation failed before the pipeline)",
            "— (nothing sent — validation failed before the pipeline)",
        )
    try:
        images, det, pos_sent, neg_sent = run_generation(params, progress=progress)
    except UserFacingError as e:
        return (
            None,
            f"**Error:** {e.user_message}",
            params.prompt,
            params.negative_prompt,
        )
    except Exception as e:  # noqa: BLE001
        tb = traceback.format_exc()
        return (
            None,
            f"**Error:** Generation failed: `{e!s}`\n\n```\n{tb[-4000:]}\n```",
            params.prompt,
            params.negative_prompt,
        )
    status_parts = [
        "**Done.** " + det,
    ]
    if params.warnings:
        status_parts.append("**Notices:**\n" + "\n".join(f"- {w}" for w in params.warnings))
    status_text = "\n\n".join(status_parts)
    if images:
        return images, status_text, pos_sent, neg_sent
    return None, status_text, pos_sent, neg_sent


def build_ui() -> gr.Blocks:
    # Gradio 6.0: theme and related UI options belong on launch(), not Blocks().
    with gr.Blocks(
        title="Anima RDBT (Gradio)",
    ) as demo:
        gr.Markdown(
            f"# Anima RDBT — {config.RDBT_UNET_NAME}\n\n{LICENSE_BANNER}\n\n{FALLBACK_NOTE}"
        )
        with gr.Row():
            with gr.Column(scale=1):
                prompt = gr.Textbox(
                    label="prompt",
                    lines=3,
                    value="digital anime illustration, 1girl, smile",
                )
                neg = gr.Textbox(
                    label="negative_prompt",
                    lines=2,
                    value="",
                )
                with gr.Row():
                    w = gr.Slider(
                        config.MIN_WH,
                        config.MAX_WH,
                        value=config.DEFAULT_WIDTH,
                        step=64,
                        label="width",
                    )
                    h = gr.Slider(
                        config.MIN_WH,
                        config.MAX_WH,
                        value=config.DEFAULT_HEIGHT,
                        step=64,
                        label="height",
                    )
                st = gr.Slider(
                    config.MIN_STEPS,
                    config.MAX_STEPS,
                    value=config.DEFAULT_STEPS,
                    step=1,
                    label="steps",
                )
                cfg = gr.Slider(
                    config.MIN_CFG,
                    config.MAX_CFG,
                    value=config.DEFAULT_CFG,
                    step=config.CFG_STEP,
                    label="cfg",
                )
            with gr.Column(scale=1):
                with gr.Accordion("Advanced", open=True):
                    batch = gr.Slider(
                        config.MIN_BATCH,
                        config.MAX_BATCH,
                        value=config.DEFAULT_BATCH_SIZE,
                        step=1,
                        label="batch_size",
                    )
                    with gr.Row():
                        seed = gr.Number(
                            value=config.DEFAULT_SEED,
                            label="seed",
                            precision=0,
                            minimum=config.MIN_SEED,
                            maximum=config.MAX_SEED,
                            step=1,
                            info="Used when randomize_seed is off.",
                        )
                        randomize_seed = gr.Checkbox(
                            value=config.DEFAULT_RANDOMIZE_SEED,
                            label="randomize_seed",
                            info="Generate a fresh seed for each run.",
                        )
                    sampler = gr.Dropdown(
                        choices=list(config.SAMPLER_CHOICES),
                        value=config.DEFAULT_SAMPLER,
                        label="sampler_name",
                        info="RDBT default: euler_ancestral. Unknown values fall back with a notice.",
                    )
                    sched = gr.Dropdown(
                        choices=list(config.SCHEDULER_CHOICES),
                        value=config.DEFAULT_SCHEDULER,
                        label="scheduler",
                        info="RDBT default: simple.",
                    )
                    denoise = gr.Slider(
                        config.MIN_DENOISE,
                        config.MAX_DENOISE,
                        value=config.DEFAULT_DENOISE,
                        step=config.DENOISE_STEP,
                        label="denoise",
                    )
                with gr.Accordion(
                    "Actual strings sent to ComfyUI (post-validation)",
                    open=True,
                ):
                    gr.Markdown(
                        "These are the exact `prompt` and `negative_prompt` strings passed to "
                        "ComfyUI's `CLIPTextEncode` nodes for the last run (not raw textbox state if validation changed things)."
                    )
                    actual_prompt_out = gr.Textbox(
                    label="prompt → ComfyUI",
                        lines=4,
                        max_lines=30,
                        interactive=False,
                        value="",
                    )
                    actual_neg_out = gr.Textbox(
                    label="negative_prompt → ComfyUI",
                        lines=3,
                        max_lines=20,
                        interactive=False,
                        value="",
                    )
                gallery = gr.Gallery(
                    label="output",
                    columns=2,
                    object_fit="contain",
                    height=600,
                )
                status = gr.Markdown(
                    "**While the Space starts:** the RDBT checkpoint downloads (and optional hub assets from config; no GPU needed). "
                    "**On first Generate:** ZeroGPU assigns a GPU, then ComfyUI loads UNET/CLIP/VAE and runs the native graph — "
                    "that first click can take extra time to build and load the model objects."
                )
                go = gr.Button("Generate", variant="primary")
        go.click(
            fn=_generate_task,
            inputs=[
                prompt,
                neg,
                w,
                h,
                st,
                cfg,
                batch,
                sampler,
                sched,
                denoise,
                seed,
                randomize_seed,
            ],
            outputs=[gallery, status, actual_prompt_out, actual_neg_out],
            show_progress="full",
        )
    return demo


# Hugging Face imports this module as `import app` (not __main__), so the block
# below must NOT be the only place that calls `queue()`. Otherwise the Space
# serves `demo` without a queue and shutdown can get noisy async teardown.
demo = build_ui()
demo.queue()

# Prepare ComfyUI source + weights at container import time so the first HTTP
# request does not pay download cost; the Space "Starting..." state covers this.
if not config.skip_startup_bootstrap():
    from src.comfy_backend import run_at_container_startup

    run_at_container_startup()

# Gradio 6: `theme` is a launch() argument. ZeroGPU calls `demo.launch()` without
# running our `if __name__` block, so wrap to always default the Soft theme.
_base_launch = demo.launch

@functools.wraps(_base_launch)  # type: ignore[misc]
def _launch_with_default_theme(*args: Any, **kwargs: Any):
    kwargs.setdefault("theme", gr.themes.Soft())
    return _base_launch(*args, **kwargs)


demo.launch = _launch_with_default_theme  # type: ignore[method-assign, assignment]

if __name__ == "__main__":
    port = int(os.environ.get("PORT", "7860"))
    demo.launch(
        server_name="0.0.0.0",
        server_port=port,
        share=False,
    )