Update app.py

app.py

@@ -75,6 +75,7 @@ import sys
 import time
 import threading
 import traceback
+import importlib.util
 
 import torch
 
@@ -132,6 +133,8 @@ try:
 except Exception:
     pass
 
+PEFT_AVAILABLE = importlib.util.find_spec("peft") is not None
+
 import gradio as gr
 
 # ---------------------------------------------------------------------------
@@ -191,6 +194,7 @@ def _runtime_report():
         "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", ""),
@@ -204,6 +208,7 @@ def _lora_report():
         "active_strength": _active_lora_strength,
         "loaded_loras": sorted(list(_loaded_loras)),
         "available_loras": list(KNOWN_LORAS.keys()),
+        "peft_available": PEFT_AVAILABLE,
     }
 
 
@@ -356,6 +361,12 @@ def _load_lora_if_needed(pipe, 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:
@@ -380,6 +391,30 @@ def _load_lora_if_needed(pipe, lora_label):
     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.
@@ -396,13 +431,7 @@ def _set_lora(pipe, lora_label, lora_strength, allow_load=True):
     cfg = KNOWN_LORAS.get(lora_label)
 
     if not cfg:
-        if hasattr(transformer, "disable_lora"):
-            transformer.disable_lora()
-        elif hasattr(transformer, "set_adapters"):
-            try:
-                transformer.set_adapters([], adapter_weights=[])
-            except Exception:
-                pass
+        _safe_disable_lora(transformer)
 
         _active_lora = None
         _active_lora_label = "Base model"
@@ -449,8 +478,8 @@ def load_selected_lora(lora_style, lora_strength):
     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(),
         }
 
@@ -460,8 +489,7 @@ def load_selected_lora(lora_style, lora_strength):
 
         return {
             "status": "ready",
-            "message": "Base model active. No LoRA loaded.",
-            "runtime": _runtime_report(),
+            "message": "Base model active. No LoRA selected.",
             "lora": _lora_report(),
         }
 
@@ -478,7 +506,6 @@ def load_selected_lora(lora_style, lora_strength):
             "status": "ready",
             "message": f"LoRA loaded and activated: {lora_style}",
             "trigger": trigger,
-            "runtime": _runtime_report(),
             "lora": _lora_report(),
         }
 
@@ -486,8 +513,8 @@ def load_selected_lora(lora_style, lora_strength):
         traceback.print_exc()
         return {
             "status": "error",
+            "message": "LoRA load failed.",
             "error": f"{type(e).__name__}: {e}",
-            "runtime": _runtime_report(),
             "lora": _lora_report(),
         }
 
@@ -503,8 +530,8 @@ def disable_lora():
     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(),
         }
 
@@ -515,7 +542,6 @@ def disable_lora():
         return {
             "status": "ready",
             "message": "LoRA disabled. Base model active.",
-            "runtime": _runtime_report(),
             "lora": _lora_report(),
         }
 
@@ -523,8 +549,8 @@ def disable_lora():
         traceback.print_exc()
         return {
             "status": "error",
+            "message": "Could not disable LoRA cleanly.",
             "error": f"{type(e).__name__}: {e}",
-            "runtime": _runtime_report(),
             "lora": _lora_report(),
         }
 
@@ -547,11 +573,8 @@ def health():
         "model_ready": _pipeline is not None,
         "pipeline_ready": _pipeline is not None,
         "model_id": MODEL_ID,
-        "model_family": "diffusers_modular_pipeline",
-        "loader_strategy": "ModularPipeline_from_pretrained_trust_remote_code",
         "runtime_mode": "zerogpu_compatibility_compile_bypass",
         "last_error": _pipeline_error or "",
-        "expected_output_type": "video",
         "runtime": _runtime_report(),
         "lora": _lora_report(),
     }
@@ -568,6 +591,7 @@ def warmup_model():
     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(),
@@ -575,9 +599,8 @@ def warmup_model():
 
     return {
         "status": "ready",
-        "model_id": MODEL_ID,
-        "runtime_mode": "zerogpu_compatibility_compile_bypass",
         "message": "Model loaded and cached in this Space process.",
+        "model_id": MODEL_ID,
         "runtime": _runtime_report(),
         "lora": _lora_report(),
     }
@@ -645,14 +668,11 @@ def generate(
 
     device = "cuda"
 
-    # Make sure the pipeline is on CUDA inside the ZeroGPU-decorated function.
     try:
         pipe = pipe.to(device)
     except Exception as e:
         _log(f"Pipeline .to('cuda') warning: {type(e).__name__}: {e}")
 
-    # Activate selected LoRA without hidden loading during generation.
-    # If user selected a LoRA but did not load it, fail with a clear message.
     with _lora_lock:
         trigger = _set_lora(
             pipe,
@@ -731,79 +751,82 @@ def generate(
 with gr.Blocks(title="Krea Realtime Video 14B") as demo:
     gr.Markdown(
         "# Krea Realtime Video 14B\n\n"
-        "This Space runs **real local inference** for the Krea Realtime 14B "
-        "text-to-video model using Diffusers `ModularPipeline`.\n\n"
-        "⚠️ **ZeroGPU compatibility mode**: `torch.compile` is disabled because "
-        "ZeroGPU does not support it.\n\n"
-        "**Video length** is controlled by the number of blocks. "
-        "Roughly: 1 block ≈ ~1 second, 9 blocks ≈ ~9 seconds. "
-        "Values above 9 are experimental.\n\n"
-        "**LoRA support**: select a style, click **Load selected LoRA**, then generate. "
-        "The Origami preset automatically prefixes the prompt with `[origami]`."
+        "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():
-            model_status = gr.JSON(label="Model Status")
-
-            with gr.Row():
-                warmup_btn = gr.Button("Warm up / Check model", variant="secondary")
-
-            gr.Markdown("## Prompt")
-
-            prompt = gr.Textbox(
-                label="Prompt",
-                placeholder="e.g., a cat sitting on a boat",
-                lines=3,
-            )
+        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,
+                )
 
-            gr.Markdown("## Style / LoRA")
+            with gr.Group():
+                gr.Markdown("## 2. Optional Style / LoRA")
 
-            lora_style = gr.Dropdown(
-                choices=list(KNOWN_LORAS.keys()),
-                value="Base model",
-                label="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="LoRA Strength",
-            )
+                    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("Disable LoRA / Use Base Model", variant="secondary")
+                with gr.Row():
+                    load_lora_btn = gr.Button("Load selected LoRA", variant="secondary")
+                    disable_lora_btn = gr.Button("Use Base Model", variant="secondary")
 
-            lora_status = gr.JSON(label="LoRA Status")
+                gr.Markdown(
+                    "For **Origami**, the app automatically adds the `[origami]` trigger "
+                    "when generating."
+                )
 
-            gr.Markdown("## Generation Settings")
+            with gr.Group():
+                gr.Markdown("## 3. Generation Settings")
 
-            num_blocks = gr.Slider(
-                minimum=1,
-                maximum=12,
-                value=9,
-                step=1,
-                label="Video Length / Number of Blocks",
-            )
+                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="Inference Steps per Block",
-            )
+                    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")
+                seed = gr.Number(value=42, precision=0, label="Seed")
 
-            generate_btn = gr.Button("Generate Video", variant="primary")
+                generate_btn = gr.Button("Generate Video", variant="primary")
 
-        with gr.Column():
+        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=[
             [