app.py

# ---------------------------------------------------------------------------
# Krea Realtime Video 14B — Hugging Face Space Demo
# ZeroGPU compatibility version for Diffusers ModularPipeline.
# ---------------------------------------------------------------------------

import os

# ---------------------------------------------------------------------------
# HF Spaces / cache configuration — must happen before HF imports
# ---------------------------------------------------------------------------

_ASF_HF_CACHE_ROOT = os.environ.get("ASF_HF_CACHE_DIR") or "/tmp/asf-hf-cache"

os.environ.setdefault("HF_HOME", _ASF_HF_CACHE_ROOT)
os.environ.setdefault("HF_HUB_CACHE", os.path.join(_ASF_HF_CACHE_ROOT, "hub"))
os.environ.setdefault("HUGGINGFACE_HUB_CACHE", os.path.join(_ASF_HF_CACHE_ROOT, "hub"))
os.environ.setdefault("TRANSFORMERS_CACHE", os.path.join(_ASF_HF_CACHE_ROOT, "transformers"))
os.environ.setdefault("DIFFUSERS_CACHE", os.path.join(_ASF_HF_CACHE_ROOT, "diffusers"))
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

# Keep hub kernels disabled in this ZeroGPU compatibility mode.
os.environ.setdefault("DIFFUSERS_ENABLE_HUB_KERNELS", "0")
os.environ.setdefault("USE_HUB_KERNELS", "NO")

os.makedirs(_ASF_HF_CACHE_ROOT, exist_ok=True)
os.makedirs(os.path.join(_ASF_HF_CACHE_ROOT, "hub"), exist_ok=True)
os.makedirs(os.path.join(_ASF_HF_CACHE_ROOT, "transformers"), exist_ok=True)
os.makedirs(os.path.join(_ASF_HF_CACHE_ROOT, "diffusers"), exist_ok=True)
os.makedirs(os.environ["HF_MODULES_CACHE"], exist_ok=True)
os.makedirs(os.environ["MPLCONFIGDIR"], exist_ok=True)

HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")

# ---------------------------------------------------------------------------
# Safe spaces import
# ---------------------------------------------------------------------------

try:
    import spaces

    HAS_SPACES = True
except Exception:
    HAS_SPACES = False

    class _DummySpaces:
        def GPU(self, *args, **kwargs):
            def decorator(fn):
                return fn

            return decorator

    spaces = _DummySpaces()


def _spaces_gpu(*args, **kwargs):
    """
    Wrapper around spaces.GPU.

    Some versions of the spaces package may not support size=...
    In that case, we fall back to the same decorator without size.
    """
    try:
        return spaces.GPU(*args, **kwargs)
    except TypeError:
        kwargs.pop("size", None)
        return spaces.GPU(*args, **kwargs)


# ---------------------------------------------------------------------------
# Imports + ZeroGPU torch.compile bypass
# ---------------------------------------------------------------------------

import sys
import time
import threading
import traceback
import importlib.util

import torch

# ZeroGPU does not support torch.compile.
# Krea remote code compiles torch.nn.attention.flex_attention.
# We no-op torch.compile globally for this compatibility Space.
_ORIG_TORCH_COMPILE = getattr(torch, "compile", None)
_ASF_COMPILE_BYPASS_ANNOUNCED = False


def _asf_zerogpu_compile_bypass(fn=None, *args, **kwargs):
    """
    ZeroGPU compatibility shim.

    Supports both call styles:
      torch.compile(fn, ...)
      @torch.compile(...)
      def fn(...): ...

    Default behavior is quiet because FlexAttention may call this repeatedly
    during generation.
    """
    global _ASF_COMPILE_BYPASS_ANNOUNCED

    if fn is None:
        def decorator(real_fn):
            return _asf_zerogpu_compile_bypass(real_fn, *args, **kwargs)

        return decorator

    if os.environ.get("ASF_VERBOSE_COMPILE_BYPASS", "0") == "1":
        name = getattr(fn, "__name__", repr(fn))
        module = getattr(fn, "__module__", "")
        print(
            f"[ASF] ZeroGPU compatibility: bypassing torch.compile for {module}.{name}",
            flush=True,
        )
    elif not _ASF_COMPILE_BYPASS_ANNOUNCED:
        print(
            "[ASF] ZeroGPU compatibility: torch.compile bypass is active.",
            flush=True,
        )
        _ASF_COMPILE_BYPASS_ANNOUNCED = True

    return fn


if _ORIG_TORCH_COMPILE is not None and os.environ.get("ASF_ENABLE_TORCH_COMPILE", "0") != "1":
    torch.compile = _asf_zerogpu_compile_bypass

try:
    import torch._dynamo

    torch._dynamo.config.suppress_errors = True
except Exception:
    pass

PEFT_AVAILABLE = importlib.util.find_spec("peft") is not None

import gradio as gr

# ---------------------------------------------------------------------------
# Diffusers imports
# ---------------------------------------------------------------------------

_DIFFUSERS_OK = False
_DIFFUSERS_IMPORT_ERROR = None

try:
    from diffusers import ModularPipeline
    from diffusers.modular_pipelines import PipelineState
    from diffusers.utils import export_to_video

    _DIFFUSERS_OK = True
except Exception as e:
    _DIFFUSERS_IMPORT_ERROR = f"{type(e).__name__}: {e}"
    traceback.print_exc()


# ---------------------------------------------------------------------------
# Model / LoRA configuration
# ---------------------------------------------------------------------------

MODEL_ID = "krea/krea-realtime-video"

KNOWN_LORAS = {
    "Base model": None,
    "Origami": {
        "repo_id": "shauray/Origami_WanLora",
        "prefix": "diffusion_model",
        "weight_name": "origami_000000500.safetensors",
        "adapter_name": "origami",
        "trigger": "[origami]",
    },
}

_pipeline = None
_pipeline_error = None
_pipeline_lock = threading.Lock()

_loaded_loras = set()
_active_lora = None
_active_lora_label = "Base model"
_active_lora_strength = 1.0
_lora_lock = threading.Lock()


def _log(msg):
    print(f"[KreaRealtimeVideo] {msg}", flush=True)


def _runtime_report():
    return {
        "python": sys.version.replace("\n", " "),
        "torch": getattr(torch, "__version__", "unknown"),
        "cuda_available": bool(torch.cuda.is_available()),
        "cuda_device_count": int(torch.cuda.device_count()) if torch.cuda.is_available() else 0,
        "has_spaces": HAS_SPACES,
        "peft_available": PEFT_AVAILABLE,
        "torch_compile_bypassed": torch.compile is _asf_zerogpu_compile_bypass,
        "hf_home": os.environ.get("HF_HOME", ""),
        "hf_modules_cache": os.environ.get("HF_MODULES_CACHE", ""),
    }


def _lora_report():
    return {
        "active_lora": _active_lora_label,
        "active_adapter": _active_lora,
        "active_strength": _active_lora_strength,
        "loaded_loras": sorted(list(_loaded_loras)),
        "available_loras": list(KNOWN_LORAS.keys()),
        "peft_available": PEFT_AVAILABLE,
    }


def _call_from_pretrained_compat(*args, **kwargs):
    """
    Compatibility wrapper because some diffusers/HF Hub combinations
    may use token= while older ones expect use_auth_token= or no token.
    """
    try:
        return ModularPipeline.from_pretrained(*args, **kwargs)
    except TypeError as e:
        if "token" in str(e):
            kwargs.pop("token", None)
            if HF_TOKEN:
                kwargs["use_auth_token"] = HF_TOKEN
            return ModularPipeline.from_pretrained(*args, **kwargs)
        raise


def _load_components_compat(pipe, **kwargs):
    """
    Compatibility wrapper around pipe.load_components().
    """
    try:
        return pipe.load_components(**kwargs)
    except TypeError as e:
        msg = str(e)
        if "token" in msg:
            kwargs.pop("token", None)
            if HF_TOKEN:
                kwargs["use_auth_token"] = HF_TOKEN
            return pipe.load_components(**kwargs)
        raise


def _load_pipeline():
    """
    Load the ModularPipeline once at app startup.

    For ZeroGPU, this gives the best UX:
    - model warms up when the app starts;
    - generation remains protected by @spaces.GPU;
    - LoRAs can be loaded manually before generation.
    """
    global _pipeline, _pipeline_error

    with _pipeline_lock:
        if _pipeline is not None:
            return _pipeline

        if not _DIFFUSERS_OK:
            _pipeline_error = _DIFFUSERS_IMPORT_ERROR or "Diffusers import failed"
            _log(f"Pipeline load skipped: {_pipeline_error}")
            return None

        try:
            _log(f"Runtime report: {_runtime_report()}")
            _log(f"Loading ModularPipeline from {MODEL_ID} ...")

            pipe = _call_from_pretrained_compat(
                MODEL_ID,
                trust_remote_code=True,
                token=HF_TOKEN,
            )

            _log("Skeleton loaded; attaching components ...")

            try:
                _load_components_compat(
                    pipe,
                    trust_remote_code=True,
                    device_map="cuda",
                    torch_dtype={
                        "default": torch.bfloat16,
                        "vae": torch.float16,
                    },
                    token=HF_TOKEN,
                )
            except RuntimeError as err:
                msg = str(err)
                cuda_load_failed = (
                    "Found no NVIDIA driver" in msg
                    or "No CUDA GPUs are available" in msg
                    or "libcudart" in msg
                    or "CUDA error" in msg
                    or "CUDA driver" in msg
                )

                if cuda_load_failed:
                    _log(
                        "device_map='cuda' failed during startup. "
                        "Retrying CPU-load + manual .to('cuda') ..."
                    )
                    _log(f"CUDA load error was: {msg}")

                    _load_components_compat(
                        pipe,
                        trust_remote_code=True,
                        torch_dtype={
                            "default": torch.bfloat16,
                            "vae": torch.float16,
                        },
                        token=HF_TOKEN,
                    )

                    pipe = pipe.to("cuda")
                else:
                    raise

            # Krea model-card optimization: fuse projections.
            # This is safe; it is not torch.compile.
            try:
                if hasattr(pipe, "transformer") and hasattr(pipe.transformer, "blocks"):
                    fused = 0
                    for block in pipe.transformer.blocks:
                        self_attn = getattr(block, "self_attn", None)
                        if self_attn is not None and hasattr(self_attn, "fuse_projections"):
                            self_attn.fuse_projections()
                            fused += 1
                    _log(f"Fused attention projections on {fused} blocks.")
            except Exception as e:
                _log(f"fuse_projections warning: {type(e).__name__}: {e}")

            _pipeline = pipe
            _pipeline_error = None
            _log("Pipeline ready.")
            return _pipeline

        except Exception as e:
            _pipeline_error = f"{type(e).__name__}: {e}"
            _log(f"Pipeline load FAILED: {_pipeline_error}")
            traceback.print_exc()
            return None


# ---------------------------------------------------------------------------
# LoRA helpers
# ---------------------------------------------------------------------------

def _load_lora_if_needed(pipe, lora_label):
    """
    Load a known LoRA adapter once.

    This is intentionally not decorated with @spaces.GPU when called through
    the UI load button, so it does not reserve ZeroGPU generation time.
    """
    global _loaded_loras

    cfg = KNOWN_LORAS.get(lora_label)
    if not cfg:
        return None

    if not PEFT_AVAILABLE:
        raise RuntimeError(
            "PEFT is required for LoRA support. Add `peft` to requirements.txt "
            "and rebuild the Space."
        )

    adapter_name = cfg["adapter_name"]

    if adapter_name in _loaded_loras:
        return adapter_name

    transformer = getattr(pipe, "transformer", None)
    if transformer is None or not hasattr(transformer, "load_lora_adapter"):
        raise RuntimeError("This pipeline transformer does not expose load_lora_adapter().")

    _log(f"Loading LoRA adapter: {lora_label} ({adapter_name})")

    transformer.load_lora_adapter(
        cfg["repo_id"],
        prefix=cfg["prefix"],
        weight_name=cfg["weight_name"],
        adapter_name=adapter_name,
    )

    _loaded_loras.add(adapter_name)
    _log(f"LoRA loaded: {adapter_name}")

    return adapter_name


def _safe_disable_lora(transformer):
    """
    Disable PEFT LoRA if available.

    Some diffusers methods raise if PEFT is not installed, so this is defensive.
    """
    if transformer is None:
        return

    if hasattr(transformer, "disable_lora"):
        try:
            transformer.disable_lora()
            return
        except Exception as e:
            _log(f"disable_lora warning: {type(e).__name__}: {e}")

    if hasattr(transformer, "set_adapters"):
        try:
            transformer.set_adapters([], adapter_weights=[])
            return
        except Exception as e:
            _log(f"set_adapters([]) warning: {type(e).__name__}: {e}")


def _set_lora(pipe, lora_label, lora_strength, allow_load=True):
    """
    Activate the selected LoRA, or disable LoRA for base model.

    If allow_load=False, this function will not download/load a missing adapter.
    This keeps generate() fast and avoids hidden loading inside @spaces.GPU.
    """
    global _active_lora, _active_lora_label, _active_lora_strength

    transformer = getattr(pipe, "transformer", None)
    if transformer is None:
        raise RuntimeError("Pipeline has no transformer.")

    cfg = KNOWN_LORAS.get(lora_label)

    if not cfg:
        _safe_disable_lora(transformer)

        _active_lora = None
        _active_lora_label = "Base model"
        _active_lora_strength = 1.0
        return ""

    adapter_name = cfg["adapter_name"]

    if adapter_name not in _loaded_loras:
        if not allow_load:
            raise RuntimeError(
                f"LoRA '{lora_label}' is selected but not loaded. "
                "Click 'Load selected LoRA' before generating."
            )
        adapter_name = _load_lora_if_needed(pipe, lora_label)

    if hasattr(transformer, "enable_lora"):
        transformer.enable_lora()

    if hasattr(transformer, "set_adapters"):
        transformer.set_adapters(
            [adapter_name],
            adapter_weights=[float(lora_strength)],
        )
    elif hasattr(transformer, "set_adapter"):
        transformer.set_adapter(adapter_name)

    _active_lora = adapter_name
    _active_lora_label = lora_label
    _active_lora_strength = float(lora_strength)

    return cfg.get("trigger", "").strip()


def load_selected_lora(lora_style, lora_strength):
    """
    Manual LoRA loading button.

    Not decorated with @spaces.GPU on purpose:
    loading the adapter should happen before generation and not consume
    the generation reservation window.
    """
    pipe = _load_pipeline()
    if pipe is None:
        return {
            "status": "error",
            "message": "Pipeline failed to load.",
            "error": _pipeline_error or "Pipeline failed to load",
            "lora": _lora_report(),
        }

    if lora_style == "Base model":
        with _lora_lock:
            _set_lora(pipe, "Base model", 1.0, allow_load=False)

        return {
            "status": "ready",
            "message": "Base model active. No LoRA selected.",
            "lora": _lora_report(),
        }

    try:
        with _lora_lock:
            trigger = _set_lora(
                pipe,
                lora_style,
                float(lora_strength),
                allow_load=True,
            )

        return {
            "status": "ready",
            "message": f"LoRA loaded and activated: {lora_style}",
            "trigger": trigger,
            "lora": _lora_report(),
        }

    except Exception as e:
        traceback.print_exc()
        return {
            "status": "error",
            "message": "LoRA load failed.",
            "error": f"{type(e).__name__}: {e}",
            "lora": _lora_report(),
        }


def disable_lora():
    """
    Disable LoRA and return to base model.

    This does not necessarily remove the adapter from memory; it only disables it.
    Keeping the adapter cached makes switching back faster.
    """
    pipe = _load_pipeline()
    if pipe is None:
        return {
            "status": "error",
            "message": "Pipeline failed to load.",
            "error": _pipeline_error or "Pipeline failed to load",
            "lora": _lora_report(),
        }

    try:
        with _lora_lock:
            _set_lora(pipe, "Base model", 1.0, allow_load=False)

        return {
            "status": "ready",
            "message": "LoRA disabled. Base model active.",
            "lora": _lora_report(),
        }

    except Exception as e:
        traceback.print_exc()
        return {
            "status": "error",
            "message": "Could not disable LoRA cleanly.",
            "error": f"{type(e).__name__}: {e}",
            "lora": _lora_report(),
        }


# ---------------------------------------------------------------------------
# Eager app runtime warm-up
# ---------------------------------------------------------------------------

if os.environ.get("SKIP_MODEL_LOAD") != "1":
    _load_pipeline()


# ---------------------------------------------------------------------------
# Health / warm-up endpoints
# ---------------------------------------------------------------------------

def health():
    return {
        "status": "ready" if _pipeline is not None else "not_loaded",
        "model_ready": _pipeline is not None,
        "pipeline_ready": _pipeline is not None,
        "model_id": MODEL_ID,
        "runtime_mode": "zerogpu_compatibility_compile_bypass",
        "last_error": _pipeline_error or "",
        "runtime": _runtime_report(),
        "lora": _lora_report(),
    }


def warmup_model():
    """
    Manual warm-up button.

    Usually the model is already loaded at app startup.
    This remains useful if startup load failed and we want to retry from the UI.
    """
    pipe = _load_pipeline()
    if pipe is None:
        return {
            "status": "error",
            "message": "Pipeline failed to load.",
            "error": _pipeline_error or "Pipeline failed to load",
            "runtime": _runtime_report(),
            "lora": _lora_report(),
        }

    return {
        "status": "ready",
        "message": "Model loaded and cached in this Space process.",
        "model_id": MODEL_ID,
        "runtime": _runtime_report(),
        "lora": _lora_report(),
    }


# ---------------------------------------------------------------------------
# Generation endpoint — real inference guarded by @spaces.GPU
# ---------------------------------------------------------------------------

def _gpu_duration(
    prompt,
    lora_style,
    lora_strength,
    num_blocks,
    num_inference_steps,
    seed,
    *args,
    **kwargs,
):
    try:
        blocks = int(num_blocks)
        steps = int(num_inference_steps)
    except Exception:
        blocks = 9
        steps = 6

    # Model and LoRA are expected to be loaded before generation.
    # Observed on this Space:
    #   9 blocks × 4 steps  < 75s
    #   9 blocks × 8 steps  < 80s
    #
    # Keep a safety buffer without over-reserving ZeroGPU.
    return min(120, max(30, int(35 + blocks * steps * 1.2)))


@_spaces_gpu(duration=_gpu_duration, size="xlarge")
def generate(
    prompt,
    lora_style,
    lora_strength,
    num_blocks,
    num_inference_steps,
    seed,
    progress=gr.Progress(track_tqdm=True),
):
    pipe = _load_pipeline()

    if pipe is None:
        err = _pipeline_error or "Pipeline not loaded (unknown failure)"
        raise RuntimeError(f"Generation unavailable: {err}")

    if not isinstance(prompt, str) or not prompt.strip():
        raise ValueError("Prompt must be a non-empty string.")

    num_blocks = int(num_blocks)
    num_inference_steps = int(num_inference_steps)
    seed = int(seed)
    lora_strength = float(lora_strength)

    if num_blocks < 1 or num_blocks > 12:
        raise ValueError("num_blocks must be between 1 and 12.")

    if num_inference_steps < 1 or num_inference_steps > 8:
        raise ValueError("num_inference_steps must be between 1 and 8.")

    device = "cuda"

    try:
        pipe = pipe.to(device)
    except Exception as e:
        _log(f"Pipeline .to('cuda') warning: {type(e).__name__}: {e}")

    with _lora_lock:
        trigger = _set_lora(
            pipe,
            lora_style,
            lora_strength,
            allow_load=False,
        )

    final_prompt = prompt.strip()
    if trigger and not final_prompt.startswith(trigger):
        final_prompt = f"{trigger} {final_prompt}"

    frames = []
    state = PipelineState()

    try:
        generator = torch.Generator(device=device).manual_seed(seed)
    except Exception as e:
        _log(f"CUDA generator failed, falling back to CPU generator: {type(e).__name__}: {e}")
        generator = torch.Generator(device="cpu").manual_seed(seed)

    try:
        progress(0, desc="Preparing generation")

        for block_idx in progress.tqdm(
            range(num_blocks),
            desc="Generating video blocks",
        ):
            _log(f"Block {block_idx + 1}/{num_blocks}")

            progress(
                block_idx / max(1, num_blocks),
                desc=f"Generating block {block_idx + 1}/{num_blocks}",
            )

            state = pipe(
                state,
                prompt=[final_prompt],
                num_inference_steps=num_inference_steps,
                num_blocks=num_blocks,
                block_idx=block_idx,
                generator=generator,
            )

            videos = state.values.get("videos")
            if not videos:
                raise RuntimeError("Pipeline state did not contain `videos` after inference.")

            frames.extend(videos[0])

    except Exception as e:
        _log(f"Inference failed at block {locals().get('block_idx', 'unknown')}: {e}")
        traceback.print_exc()
        raise RuntimeError(
            f"Inference error at block {locals().get('block_idx', 'unknown')}: {e}"
        )

    if not frames:
        raise RuntimeError("No frames were generated.")

    progress(0.95, desc="Exporting video")

    output_path = f"/tmp/krea_output_{int(time.time())}.mp4"
    export_to_video(frames, output_path, fps=24)

    progress(1.0, desc="Done")
    _log(f"Saved video to {output_path}")

    return output_path


# ---------------------------------------------------------------------------
# Gradio app
# ---------------------------------------------------------------------------

with gr.Blocks(title="Krea Realtime Video 14B") as demo:
    gr.Markdown(
        "# Krea Realtime Video 14B\n\n"
        "Real local inference with Diffusers `ModularPipeline` on ZeroGPU.\n\n"
        "The model loads at app startup. Generation uses ZeroGPU. "
        "LoRAs can be loaded manually before generation."
    )

    with gr.Row():
        with gr.Column(scale=4):
            with gr.Group():
                gr.Markdown("## 1. Prompt")
                prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="Describe the video you want to generate...",
                    lines=4,
                )

            with gr.Group():
                gr.Markdown("## 2. Optional Style / LoRA")

                with gr.Row():
                    lora_style = gr.Dropdown(
                        choices=list(KNOWN_LORAS.keys()),
                        value="Base model",
                        label="Style / LoRA",
                        scale=2,
                    )

                    lora_strength = gr.Slider(
                        minimum=0.0,
                        maximum=1.5,
                        value=1.0,
                        step=0.05,
                        label="Strength",
                        scale=3,
                    )

                with gr.Row():
                    load_lora_btn = gr.Button("Load selected LoRA", variant="secondary")
                    disable_lora_btn = gr.Button("Use Base Model", variant="secondary")

                gr.Markdown(
                    "For **Origami**, the app automatically adds the `[origami]` trigger "
                    "when generating."
                )

            with gr.Group():
                gr.Markdown("## 3. Generation Settings")

                with gr.Row():
                    num_blocks = gr.Slider(
                        minimum=1,
                        maximum=12,
                        value=9,
                        step=1,
                        label="Video Length / Blocks",
                    )

                    num_inference_steps = gr.Slider(
                        minimum=1,
                        maximum=8,
                        value=6,
                        step=1,
                        label="Steps per Block",
                    )

                seed = gr.Number(value=42, precision=0, label="Seed")

                generate_btn = gr.Button("Generate Video", variant="primary")

        with gr.Column(scale=5):
            output_video = gr.Video(label="Generated Video")

            with gr.Accordion("Runtime status", open=False):
                model_status = gr.JSON(label="Model Status")
                lora_status = gr.JSON(label="LoRA Status")
                warmup_btn = gr.Button("Refresh model status", variant="secondary")

    gr.Examples(
        examples=[
            [
                "Astronaut in a jungle, cold color palette, muted colors, detailed, cinematic, 8k",
                "Base model",
                1.0,
                9,
                6,
                42,
            ],
            [
                "A tiny wooden boat drifting through a misty lake at sunrise, a curious cat sitting at the front, soft cinematic lighting, calm water reflections",
                "Base model",
                1.0,
                9,
                6,
                123,
            ],
            [
                "A futuristic city at sunset, flying vehicles between glass towers, neon reflections, cinematic camera movement, atmospheric haze",
                "Base model",
                1.0,
                9,
                6,
                7,
            ],
            [
                "A car racing down a snowy mountain road, dramatic chase shot, powder snow flying behind the wheels, cold blue lighting, high speed motion",
                "Base model",
                1.0,
                9,
                6,
                99,
            ],
            [
                "A surreal underwater library, glowing jellyfish floating between bookshelves, slow cinematic dolly shot, dreamlike atmosphere",
                "Base model",
                1.0,
                9,
                6,
                314,
            ],
            [
                "a cat sitting on a boat",
                "Origami",
                1.0,
                9,
                6,
                2026,
            ],
            [
                "a dragon flying over a mountain village at sunrise, paper-folded geometry, delicate handmade texture, soft shadows",
                "Origami",
                1.0,
                9,
                6,
                777,
            ],
            [
                "a small fox walking through a paper forest, handcrafted origami style, warm lantern light, cinematic close-up",
                "Origami",
                0.9,
                9,
                6,
                888,
            ],
        ],
        inputs=[
            prompt,
            lora_style,
            lora_strength,
            num_blocks,
            num_inference_steps,
            seed,
        ],
        outputs=output_video,
        fn=generate,
        cache_examples=False,
    )

    warmup_btn.click(
        warmup_model,
        inputs=None,
        outputs=model_status,
        api_name="warmup",
    )

    load_lora_btn.click(
        load_selected_lora,
        inputs=[lora_style, lora_strength],
        outputs=lora_status,
        api_name="load_lora",
    )

    disable_lora_btn.click(
        disable_lora,
        inputs=None,
        outputs=lora_status,
        api_name="disable_lora",
    )

    generate_btn.click(
        generate,
        inputs=[
            prompt,
            lora_style,
            lora_strength,
            num_blocks,
            num_inference_steps,
            seed,
        ],
        outputs=output_video,
        api_name="generate",
    )

    demo.load(
        lambda: health(),
        inputs=None,
        outputs=model_status,
        api_name="health",
    )


if __name__ == "__main__":
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
    )