Update app.py

app.py

@@ -19,8 +19,7 @@ os.environ.setdefault("DIFFUSERS_CACHE", os.path.join(_ASF_HF_CACHE_ROOT, "diffu
 os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
 os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
 
-# Compatibility mode: avoid optional hub-kernels path unless explicitly enabled.
-# The first failure you saw came from transformers -> hub_kernels -> kernels.
+# Keep hub kernels disabled in this ZeroGPU compatibility mode.
 os.environ.setdefault("DIFFUSERS_ENABLE_HUB_KERNELS", "0")
 os.environ.setdefault("USE_HUB_KERNELS", "NO")
 
@@ -58,8 +57,8 @@ def _spaces_gpu(*args, **kwargs):
     """
     Wrapper around spaces.GPU.
 
-    Some spaces versions may not support size=...
-    In that case, fall back to the same decorator without size.
+    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)
@@ -73,14 +72,15 @@ def _spaces_gpu(*args, **kwargs):
 # ---------------------------------------------------------------------------
 
 import sys
+import time
 import threading
 import traceback
 
 import torch
 
 # ZeroGPU does not support torch.compile.
-# Krea remote code compiles FlexAttention at import/load time.
-# Therefore we no-op torch.compile globally for this Space.
+# 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)
 
 
@@ -159,6 +159,7 @@ def _runtime_report():
         "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,
         "hf_home": os.environ.get("HF_HOME", ""),
@@ -166,13 +167,47 @@ def _runtime_report():
     }
 
 
+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():
     """
-    Lazy-load the ModularPipeline.
+    Load the ModularPipeline once at app startup.
 
-    Important for ZeroGPU:
-    - The GPU is only actually allocated inside @spaces.GPU-decorated functions.
-    - Therefore loading inside generate() is safer than loading at module import.
+    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.
     """
     global _pipeline, _pipeline_error
 
@@ -189,7 +224,7 @@ def _load_pipeline():
             _log(f"Runtime report: {_runtime_report()}")
             _log(f"Loading ModularPipeline from {MODEL_ID} ...")
 
-            pipe = ModularPipeline.from_pretrained(
+            pipe = _call_from_pretrained_compat(
                 MODEL_ID,
                 trust_remote_code=True,
                 token=HF_TOKEN,
@@ -197,9 +232,10 @@ def _load_pipeline():
 
             _log("Skeleton loaded; attaching components ...")
 
-            # Primary path: device_map='cuda', as the Krea model card expects.
+            # Primary path: Krea model-card style.
             try:
-                pipe.load_components(
+                _load_components_compat(
+                    pipe,
                     trust_remote_code=True,
                     device_map="cuda",
                     torch_dtype={
@@ -215,16 +251,18 @@ def _load_pipeline():
                     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 at module level. "
-                        "Retrying with CPU-load + manual .to('cuda') ..."
+                        "device_map='cuda' failed during startup. "
+                        "Retrying CPU-load + manual .to('cuda') ..."
                     )
                     _log(f"CUDA load error was: {msg}")
 
-                    pipe.load_components(
+                    _load_components_compat(
+                        pipe,
                         trust_remote_code=True,
                         torch_dtype={
                             "default": torch.bfloat16,
@@ -233,13 +271,12 @@ def _load_pipeline():
                         token=HF_TOKEN,
                     )
 
-                    # On ZeroGPU, CUDA should be active inside @spaces.GPU.
                     pipe = pipe.to("cuda")
                 else:
                     raise
 
-            # Krea model-card optimization: fuse Q/K/V projections.
-            # This is not torch.compile and should be safe.
+            # 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
@@ -264,9 +301,11 @@ def _load_pipeline():
             return None
 
 
-# Optional eager load for non-ZeroGPU debugging only.
-# Keep default lazy for ZeroGPU.
-if os.environ.get("ASF_LOAD_AT_STARTUP") == "1" and os.environ.get("SKIP_MODEL_LOAD") != "1":
+# ---------------------------------------------------------------------------
+# Eager app runtime warm-up
+# ---------------------------------------------------------------------------
+
+if os.environ.get("SKIP_MODEL_LOAD") != "1":
     _load_pipeline()
 
 
@@ -276,7 +315,7 @@ if os.environ.get("ASF_LOAD_AT_STARTUP") == "1" and os.environ.get("SKIP_MODEL_L
 
 def health():
     return {
-        "status": "ok" if _pipeline is not None else "not_loaded",
+        "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,
@@ -289,21 +328,51 @@ def health():
     }
 
 
+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",
+            "error": _pipeline_error or "Pipeline failed to load",
+            "runtime": _runtime_report(),
+        }
+
+    return {
+        "status": "ready",
+        "model_id": MODEL_ID,
+        "runtime_mode": "zerogpu_compatibility_compile_bypass",
+        "message": "Model loaded and cached in this Space process.",
+        "runtime": _runtime_report(),
+    }
+
+
 # ---------------------------------------------------------------------------
 # Generation endpoint — real inference guarded by @spaces.GPU
 # ---------------------------------------------------------------------------
 
 def _gpu_duration(prompt, num_blocks, num_inference_steps, seed):
-    # Reserve enough time for lazy load + forward passes.
-    # Keep conservative for ZeroGPU quota/queue behavior.
     try:
         blocks = int(num_blocks)
         steps = int(num_inference_steps)
     except Exception:
-        blocks = 1
+        blocks = 3
         steps = 4
 
-    return min(900, max(300, 180 + blocks * steps * 20))
+    # Model is loaded at app startup.
+    # Duration only covers generation.
+    # Aggressive ZeroGPU reservation:
+    # 1x4  -> 30s
+    # 3x4  -> 51s
+    # 6x4  -> 87s
+    # 9x4  -> 123s
+    # 9x8  -> 180s cap
+    return min(180, max(30, 15 + blocks * steps * 3))
 
 
 @_spaces_gpu(duration=_gpu_duration, size="xlarge")
@@ -321,26 +390,19 @@ def generate(prompt, num_blocks, num_inference_steps, seed):
     num_inference_steps = int(num_inference_steps)
     seed = int(seed)
 
-    # Conservative bounds for ZeroGPU compatibility mode.
-    if num_blocks < 1 or num_blocks > 3:
-        raise ValueError("num_blocks must be between 1 and 3 in ZeroGPU compatibility mode.")
+    if num_blocks < 1 or num_blocks > 9:
+        raise ValueError("num_blocks must be between 1 and 9.")
 
     if num_inference_steps < 1 or num_inference_steps > 8:
-        raise ValueError(
-            "num_inference_steps must be between 1 and 8 in ZeroGPU compatibility mode."
-        )
+        raise ValueError("num_inference_steps must be between 1 and 8.")
 
     device = "cuda"
 
-    # Ensure the pipeline is on CUDA inside the ZeroGPU-decorated function.
+    # Make sure the pipeline is on CUDA inside the ZeroGPU-decorated function.
     try:
-        if hasattr(pipe, "device"):
-            current = str(pipe.device)
-            if "cuda" not in current:
-                _log("Moving pipeline to cuda ...")
-                pipe.to(device)
+        pipe = pipe.to(device)
     except Exception as e:
-        _log(f"Pipeline device check/move warning: {type(e).__name__}: {e}")
+        _log(f"Pipeline .to('cuda') warning: {type(e).__name__}: {e}")
 
     frames = []
     state = PipelineState()
@@ -380,7 +442,7 @@ def generate(prompt, num_blocks, num_inference_steps, seed):
     if not frames:
         raise RuntimeError("No frames were generated.")
 
-    output_path = "/tmp/krea_output.mp4"
+    output_path = f"/tmp/krea_output_{int(time.time())}.mp4"
     export_to_video(frames, output_path, fps=24)
 
     _log(f"Saved video to {output_path}")
@@ -394,11 +456,12 @@ def generate(prompt, num_blocks, num_inference_steps, seed):
 with gr.Blocks(title="Krea Realtime Video 14B") as demo:
     gr.Markdown(
         "# Krea Realtime Video 14B\n\n"
-        "This Space attempts **real local inference** for the Krea Realtime 14B "
-        "text-to-video model using the Diffusers `ModularPipeline` path.\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. This may be slower than the optimized Krea runtime.\n\n"
-        "Start with 1 block and 4 steps to validate the runtime before increasing settings."
+        "**Video length** is controlled by the number of blocks. "
+        "Roughly: 1 block ≈ ~1 second, 3 blocks ≈ ~3 seconds, 9 blocks ≈ ~9 seconds."
     )
 
     with gr.Row():
@@ -411,10 +474,10 @@ with gr.Blocks(title="Krea Realtime Video 14B") as demo:
 
             num_blocks = gr.Slider(
                 minimum=1,
-                maximum=3,
-                value=1,
+                maximum=9,
+                value=3,
                 step=1,
-                label="Number of Blocks (ZeroGPU-safe range)",
+                label="Video Length / Number of Blocks",
             )
 
             num_inference_steps = gr.Slider(
@@ -427,16 +490,20 @@ with gr.Blocks(title="Krea Realtime Video 14B") as demo:
 
             seed = gr.Number(value=42, precision=0, label="Seed")
 
-            generate_btn = gr.Button("Generate Video", variant="primary")
+            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")
 
         with gr.Column():
             output_video = gr.Video(label="Generated Video")
 
     gr.Examples(
         examples=[
-            ["a cat sitting on a boat", 1, 4, 42],
-            ["a futuristic city at sunset", 1, 4, 123],
-            ["a panda playing guitar in a forest", 1, 4, 7],
+            ["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],
         ],
         inputs=[prompt, num_blocks, num_inference_steps, seed],
         outputs=output_video,
@@ -444,6 +511,13 @@ with gr.Blocks(title="Krea Realtime Video 14B") as demo:
         cache_examples=False,
     )
 
+    warmup_btn.click(
+        warmup_model,
+        inputs=None,
+        outputs=warmup_status,
+        api_name="warmup",
+    )
+
     generate_btn.click(
         generate,
         inputs=[prompt, num_blocks, num_inference_steps, seed],
@@ -451,11 +525,10 @@ with gr.Blocks(title="Krea Realtime Video 14B") as demo:
         api_name="generate",
     )
 
-    # Structured health endpoint.
     demo.load(
         lambda: health(),
-        None,
-        gr.JSON(),
+        inputs=None,
+        outputs=warmup_status,
         api_name="health",
     )