Update app.py

app.py

@@ -1,3 +1,825 @@
+# ---------------------------------------------------------------------------
+# 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")
+
+# ZeroGPU compatibility mode:
+# - torch.compile is disabled below.
+# - hub kernels / torchao optimized path is intentionally not used.
+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 re
+import time
+import threading
+import traceback
+import importlib.util
+import importlib.metadata
+
+import torch
+
+_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
+
+
+def _installed_version(package_name):
+    try:
+        return importlib.metadata.version(package_name)
+    except Exception:
+        return None
+
+
+def _version_tuple(version):
+    """
+    Minimal semver-ish parser.
+    Handles strings like 0.12.0, 0.16.0, 0.16.0.dev...
+    """
+    if not version:
+        return None
+    parts = re.findall(r"\d+", version)
+    if not parts:
+        return None
+    return tuple(int(p) for p in parts[:3])
+
+
+PEFT_AVAILABLE = importlib.util.find_spec("peft") is not None
+PEFT_VERSION = _installed_version("peft")
+
+TORCHAO_VERSION = _installed_version("torchao")
+TORCHAO_INSTALLED = TORCHAO_VERSION is not None
+TORCHAO_VERSION_TUPLE = _version_tuple(TORCHAO_VERSION)
+
+# Recent PEFT rejects torchao < 0.16.0 when torchao is installed.
+# In this Space, torchao is not needed because we are not using Krea's optimized
+# torch.compile / FP8 path on ZeroGPU.
+TORCHAO_COMPATIBLE_FOR_PEFT = (
+    not TORCHAO_INSTALLED
+    or (
+        TORCHAO_VERSION_TUPLE is not None
+        and TORCHAO_VERSION_TUPLE >= (0, 16, 0)
+    )
+)
+
+LORA_BACKEND_READY = PEFT_AVAILABLE and TORCHAO_COMPATIBLE_FOR_PEFT
+
+if not PEFT_AVAILABLE:
+    LORA_BACKEND_ERROR = "PEFT is not installed. Add `peft` to requirements.txt and rebuild."
+elif not TORCHAO_COMPATIBLE_FOR_PEFT:
+    LORA_BACKEND_ERROR = (
+        f"Incompatible torchao version detected: {TORCHAO_VERSION}. "
+        "Remove `torchao==0.12.0` from requirements.txt, or upgrade torchao to >=0.16.0. "
+        "For this ZeroGPU compatibility Space, removing torchao is recommended."
+    )
+else:
+    LORA_BACKEND_ERROR = ""
+
+
+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,
+        "torch_compile_bypassed": torch.compile is _asf_zerogpu_compile_bypass,
+        "peft_available": PEFT_AVAILABLE,
+        "peft_version": PEFT_VERSION,
+        "torchao_installed": TORCHAO_INSTALLED,
+        "torchao_version": TORCHAO_VERSION,
+        "lora_backend_ready": LORA_BACKEND_READY,
+        "lora_backend_error": LORA_BACKEND_ERROR,
+        "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()),
+        "backend_ready": LORA_BACKEND_READY,
+        "backend_error": LORA_BACKEND_ERROR,
+    }
+
+
+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 LORA_BACKEND_READY:
+        raise RuntimeError(LORA_BACKEND_ERROR)
+
+    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 or incompatible, 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 style' 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 (
+            "Model is not loaded.",
+            {
+                "status": "error",
+                "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 (
+            "Base model active.",
+            {
+                "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 (
+            f"{lora_style} loaded. Trigger `{trigger}` will be added automatically.",
+            {
+                "status": "ready",
+                "message": f"LoRA loaded and activated: {lora_style}",
+                "trigger": trigger,
+                "lora": _lora_report(),
+            },
+        )
+
+    except Exception as e:
+        traceback.print_exc()
+        return (
+            f"LoRA load failed: {type(e).__name__}: {e}",
+            {
+                "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 (
+            "Model is not loaded.",
+            {
+                "status": "error",
+                "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 (
+            "Base model active.",
+            {
+                "status": "ready",
+                "message": "LoRA disabled. Base model active.",
+                "lora": _lora_report(),
+            },
+        )
+
+    except Exception as e:
+        traceback.print_exc()
+        return (
+            f"Could not switch to base model cleanly: {type(e).__name__}: {e}",
+            {
+                "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 refresh button.
+
+    Usually the model is already loaded at app startup.
+    """
+    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
+    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"
@@ -33,6 +855,13 @@ with gr.Blocks(title="Krea Realtime Video 14B") as demo:
                 load_lora_btn = gr.Button("Load style", variant="secondary")
                 disable_lora_btn = gr.Button("Base model", variant="secondary")
 
+            style_status = gr.Textbox(
+                label="Style status",
+                value="Base model active.",
+                interactive=False,
+                lines=1,
+            )
+
             with gr.Row():
                 num_blocks = gr.Slider(
                     minimum=1,
@@ -55,8 +884,8 @@ with gr.Blocks(title="Krea Realtime Video 14B") as demo:
             generate_btn = gr.Button("Generate video", variant="primary")
 
             gr.Markdown(
-                "Use **Origami** by selecting it, clicking **Load style**, then generating. "
-                "The trigger `[origami]` is added automatically."
+                "For Origami: select **Origami**, click **Load style**, then generate. "
+                "The `[origami]` trigger is added automatically."
             )
 
         with gr.Column(scale=5):
@@ -142,9 +971,6 @@ with gr.Blocks(title="Krea Realtime Video 14B") as demo:
             num_inference_steps,
             seed,
         ],
-        outputs=output_video,
-        fn=generate,
-        cache_examples=False,
     )
 
     warmup_btn.click(
@@ -157,14 +983,14 @@ with gr.Blocks(title="Krea Realtime Video 14B") as demo:
     load_lora_btn.click(
         load_selected_lora,
         inputs=[lora_style, lora_strength],
-        outputs=lora_status,
+        outputs=[style_status, lora_status],
         api_name="load_lora",
     )
 
     disable_lora_btn.click(
         disable_lora,
         inputs=None,
-        outputs=lora_status,
+        outputs=[style_status, lora_status],
         api_name="disable_lora",
     )
 
@@ -187,4 +1013,12 @@ with gr.Blocks(title="Krea Realtime Video 14B") as demo:
         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,
     )