Update app.py

app.py

@@ -153,15 +153,32 @@ except Exception as e:
 
 
 # ---------------------------------------------------------------------------
-# Model configuration
+# 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)
@@ -180,6 +197,16 @@ def _runtime_report():
     }
 
 
+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()),
+    }
+
+
 def _call_from_pretrained_compat(*args, **kwargs):
     """
     Compatibility wrapper because some diffusers/HF Hub combinations
@@ -216,11 +243,10 @@ def _load_pipeline():
     """
     Load the ModularPipeline once at app startup.
 
-    For ZeroGPU, this is the preferred UX path: the app warms up at runtime,
-    while actual generation remains protected by @spaces.GPU.
-
-    If startup loading fails because CUDA is not fully available yet,
-    generate() will retry loading under @spaces.GPU.
+    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
 
@@ -245,7 +271,6 @@ def _load_pipeline():
 
             _log("Skeleton loaded; attaching components ...")
 
-            # Primary path: Krea model-card style.
             try:
                 _load_components_compat(
                     pipe,
@@ -314,6 +339,196 @@ def _load_pipeline():
             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
+
+    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 _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:
+        if hasattr(transformer, "disable_lora"):
+            transformer.disable_lora()
+        elif hasattr(transformer, "set_adapters"):
+            try:
+                transformer.set_adapters([], adapter_weights=[])
+            except Exception:
+                pass
+
+        _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",
+            "error": _pipeline_error or "Pipeline failed to load",
+            "runtime": _runtime_report(),
+            "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 loaded.",
+            "runtime": _runtime_report(),
+            "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,
+            "runtime": _runtime_report(),
+            "lora": _lora_report(),
+        }
+
+    except Exception as e:
+        traceback.print_exc()
+        return {
+            "status": "error",
+            "error": f"{type(e).__name__}: {e}",
+            "runtime": _runtime_report(),
+            "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",
+            "error": _pipeline_error or "Pipeline failed to load",
+            "runtime": _runtime_report(),
+            "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.",
+            "runtime": _runtime_report(),
+            "lora": _lora_report(),
+        }
+
+    except Exception as e:
+        traceback.print_exc()
+        return {
+            "status": "error",
+            "error": f"{type(e).__name__}: {e}",
+            "runtime": _runtime_report(),
+            "lora": _lora_report(),
+        }
+
+
 # ---------------------------------------------------------------------------
 # Eager app runtime warm-up
 # ---------------------------------------------------------------------------
@@ -323,7 +538,7 @@ if os.environ.get("SKIP_MODEL_LOAD") != "1":
 
 
 # ---------------------------------------------------------------------------
-# Health endpoint
+# Health / warm-up endpoints
 # ---------------------------------------------------------------------------
 
 def health():
@@ -338,10 +553,10 @@ def health():
         "last_error": _pipeline_error or "",
         "expected_output_type": "video",
         "runtime": _runtime_report(),
+        "lora": _lora_report(),
     }
 
 
-@_spaces_gpu(duration=120, size="xlarge")
 def warmup_model():
     """
     Manual warm-up button.
@@ -355,6 +570,7 @@ def warmup_model():
             "status": "error",
             "error": _pipeline_error or "Pipeline failed to load",
             "runtime": _runtime_report(),
+            "lora": _lora_report(),
         }
 
     return {
@@ -363,6 +579,7 @@ def warmup_model():
         "runtime_mode": "zerogpu_compatibility_compile_bypass",
         "message": "Model loaded and cached in this Space process.",
         "runtime": _runtime_report(),
+        "lora": _lora_report(),
     }
 
 
@@ -370,29 +587,37 @@ def warmup_model():
 # Generation endpoint — real inference guarded by @spaces.GPU
 # ---------------------------------------------------------------------------
 
-def _gpu_duration(prompt, num_blocks, num_inference_steps, seed, *args, **kwargs):
+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 = 3
-        steps = 4
+        blocks = 9
+        steps = 6
 
-    # Model is loaded at app startup.
-    # Observed: 9 blocks × 4 steps ≈ <75s.
-    # Keep a small safety buffer without over-reserving ZeroGPU.
+    # Model and LoRA are expected to be loaded before generation.
+    # Observed on this Space:
+    #   9 blocks × 4 steps  < 75s
+    #   9 blocks �� 8 steps  < 80s
     #
-    # 1×4  -> 30s
-    # 3×4  -> 44s
-    # 6×4  -> 68s
-    # 9×4  -> 92s
-    # 9×8  -> 150s cap
-    return min(150, max(30, int(20 + blocks * steps * 2)))
+    # 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,
@@ -410,9 +635,10 @@ def generate(
     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 > 9:
-        raise ValueError("num_blocks must be between 1 and 9.")
+    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.")
@@ -425,6 +651,20 @@ def generate(
     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,
+            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()
 
@@ -450,7 +690,7 @@ def generate(
 
             state = pipe(
                 state,
-                prompt=[prompt],
+                prompt=[final_prompt],
                 num_inference_steps=num_inference_steps,
                 num_blocks=num_blocks,
                 block_idx=block_idx,
@@ -494,23 +734,57 @@ with gr.Blocks(title="Krea Realtime Video 14B") as demo:
         "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. This may be slower than the optimized Krea runtime.\n\n"
+        "ZeroGPU does not support it.\n\n"
         "**Video length** is controlled by the number of blocks. "
-        "Roughly: 1 block ≈ ~1 second, 3 blocks ≈ ~3 seconds, 9 blocks ≈ ~9 seconds."
+        "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]`."
     )
 
     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=2,
+                lines=3,
+            )
+
+            gr.Markdown("## Style / LoRA")
+
+            lora_style = gr.Dropdown(
+                choices=list(KNOWN_LORAS.keys()),
+                value="Base model",
+                label="Style / LoRA",
             )
 
+            lora_strength = gr.Slider(
+                minimum=0.0,
+                maximum=1.5,
+                value=1.0,
+                step=0.05,
+                label="LoRA Strength",
+            )
+
+            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")
+
+            lora_status = gr.JSON(label="LoRA Status")
+
+            gr.Markdown("## Generation Settings")
+
             num_blocks = gr.Slider(
                 minimum=1,
-                maximum=9,
-                value=3,
+                maximum=12,
+                value=9,
                 step=1,
                 label="Video Length / Number of Blocks",
             )
@@ -518,29 +792,93 @@ with gr.Blocks(title="Krea Realtime Video 14B") as demo:
             num_inference_steps = gr.Slider(
                 minimum=1,
                 maximum=8,
-                value=4,
+                value=6,
                 step=1,
                 label="Inference Steps per Block",
             )
 
             seed = gr.Number(value=42, precision=0, label="Seed")
 
-            with gr.Row():
-                warmup_btn = gr.Button("Load / Warm up model", variant="secondary")
-                generate_btn = gr.Button("Generate Video", variant="primary")
-
-            warmup_status = gr.JSON(label="Model Status")
+            generate_btn = gr.Button("Generate Video", variant="primary")
 
         with gr.Column():
             output_video = gr.Video(label="Generated Video")
 
     gr.Examples(
         examples=[
-            ["a cat sitting on a boat", 3, 4, 42],
-            ["a futuristic city at sunset", 3, 4, 123],
-            ["a panda playing guitar in a forest", 3, 4, 7],
+            [
+                "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,
         ],
-        inputs=[prompt, num_blocks, num_inference_steps, seed],
         outputs=output_video,
         fn=generate,
         cache_examples=False,
@@ -549,13 +887,34 @@ with gr.Blocks(title="Krea Realtime Video 14B") as demo:
     warmup_btn.click(
         warmup_model,
         inputs=None,
-        outputs=warmup_status,
+        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, num_blocks, num_inference_steps, seed],
+        inputs=[
+            prompt,
+            lora_style,
+            lora_strength,
+            num_blocks,
+            num_inference_steps,
+            seed,
+        ],
         outputs=output_video,
         api_name="generate",
     )
@@ -563,7 +922,7 @@ with gr.Blocks(title="Krea Realtime Video 14B") as demo:
     demo.load(
         lambda: health(),
         inputs=None,
-        outputs=warmup_status,
+        outputs=model_status,
         api_name="health",
     )