Update app.py

app.py

@@ -1,154 +1,316 @@
 # ---------------------------------------------------------------------------
 # Krea Realtime Video 14B — Hugging Face Space Demo
-# Attempts the real diffusers ModularPipeline path per the model card.
+# ZeroGPU compatibility version for Diffusers ModularPipeline.
 # ---------------------------------------------------------------------------
+
 import os
 
-# Redirect caches before any library import (HF Spaces operational rule)
-os.environ.setdefault("HF_HOME", "/tmp/asf-hf-cache")
+# ---------------------------------------------------------------------------
+# 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")
 
+# Compatibility mode: avoid optional hub-kernels path unless explicitly enabled.
+# The first failure you saw came from transformers -> hub_kernels -> kernels.
+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: valid on both ZeroGPU and fixed-GPU fallback hardware
+# ---------------------------------------------------------------------------
+# Safe spaces import
+# ---------------------------------------------------------------------------
+
 try:
     import spaces
+
     HAS_SPACES = True
 except Exception:
     HAS_SPACES = False
 
     class _DummySpaces:
-        def GPU(self, duration=60):
+        def GPU(self, *args, **kwargs):
             def decorator(fn):
                 return fn
+
             return decorator
+
     spaces = _DummySpaces()
 
+
+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.
+    """
+    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 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.
+_ORIG_TORCH_COMPILE = getattr(torch, "compile", None)
+
+
+def _asf_zerogpu_compile_bypass(fn=None, *args, **kwargs):
+    """
+    ZeroGPU compatibility shim.
+
+    Supports both call styles:
+      torch.compile(fn, ...)
+      @torch.compile(...)
+      def fn(...): ...
+    """
+    if fn is None:
+        def decorator(real_fn):
+            return _asf_zerogpu_compile_bypass(real_fn, *args, **kwargs)
+
+        return decorator
+
+    name = getattr(fn, "__name__", repr(fn))
+    module = getattr(fn, "__module__", "")
+
+    print(
+        f"[ASF] ZeroGPU compatibility: bypassing torch.compile for {module}.{name}",
+        flush=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
+
 import gradio as gr
 
-# ASF v198.26.20 writable Hugging Face cache guard.
-_ASF_HF_CACHE_ROOT = os.environ.get("ASF_HF_CACHE_DIR") or "/tmp/asf-hf-cache"
-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.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"))
+# ---------------------------------------------------------------------------
+# Diffusers imports
+# ---------------------------------------------------------------------------
 
 _DIFFUSERS_OK = False
 _DIFFUSERS_IMPORT_ERROR = None
 
-# Wrap diffusers imports so a build/API mismatch doesn't prevent the app from booting
 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 configuration
+# ---------------------------------------------------------------------------
+
 MODEL_ID = "krea/krea-realtime-video"
 
+_pipeline = None
+_pipeline_error = None
+_pipeline_lock = threading.Lock()
+
+
 def _log(msg):
     print(f"[KreaRealtimeVideo] {msg}", flush=True)
 
-_pipeline = None
-_pipeline_error = None
+
+def _runtime_report():
+    return {
+        "python": sys.version.replace("\n", " "),
+        "torch": getattr(torch, "__version__", "unknown"),
+        "cuda_available": bool(torch.cuda.is_available()),
+        "has_spaces": HAS_SPACES,
+        "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 _load_pipeline():
+    """
+    Lazy-load the ModularPipeline.
+
+    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.
+    """
     global _pipeline, _pipeline_error
 
-    if not _DIFFUSERS_OK:
-        _pipeline_error = _DIFFUSERS_IMPORT_ERROR or "Diffusers import failed"
-        _log(f"Pipeline load skipped: {_pipeline_error}")
-        return
+    with _pipeline_lock:
+        if _pipeline is not None:
+            return _pipeline
 
-    try:
-        _log(f"Loading ModularPipeline from {MODEL_ID} ...")
-        pipe = ModularPipeline.from_pretrained(MODEL_ID, trust_remote_code=True, token=HF_TOKEN)
-        _log("Skeleton loaded; attaching components ...")
+        if not _DIFFUSERS_OK:
+            _pipeline_error = _DIFFUSERS_IMPORT_ERROR or "Diffusers import failed"
+            _log(f"Pipeline load skipped: {_pipeline_error}")
+            return None
 
-        # Primary path: device_map="cuda" as the model card instructs.
         try:
-            pipe.load_components(
+            _log(f"Runtime report: {_runtime_report()}")
+            _log(f"Loading ModularPipeline from {MODEL_ID} ...")
+
+            pipe = ModularPipeline.from_pretrained(
+                MODEL_ID,
                 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)
-            if "Found no NVIDIA driver" in msg or "No CUDA GPUs are available" in msg or "libcudart" in msg:
-                _log(f"device_map='cuda' failed at module level ({msg}). Retrying with CPU-load + manual .to('cuda') ...")
+                token=HF_TOKEN,
+            )
+
+            _log("Skeleton loaded; attaching components ...")
+
+            # Primary path: device_map='cuda', as the Krea model card expects.
+            try:
                 pipe.load_components(
                     trust_remote_code=True,
-                    torch_dtype={"default": torch.bfloat16, "vae": torch.float16},
-                    token=HF_TOKEN)
-                # ZeroGPU emulates CUDA at module level; .to('cuda') is safe here.
-                # On fixed GPU it moves to the real device.
-                pipe = pipe.to("cuda")
-            else:
-                raise
-
-        # Optimization: fuse Q/K/V projections in attention blocks.
-        try:
-            if hasattr(pipe, "transformer") and hasattr(pipe.transformer, "blocks"):
-                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()
-                _log("Fused attention projections.")
-        except Exception as e:
-            _log(f"fuse_projections warning: {e}")
+                    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
+                )
+
+                if cuda_load_failed:
+                    _log(
+                        "device_map='cuda' failed at module level. "
+                        "Retrying with CPU-load + manual .to('cuda') ..."
+                    )
+                    _log(f"CUDA load error was: {msg}")
+
+                    pipe.load_components(
+                        trust_remote_code=True,
+                        torch_dtype={
+                            "default": torch.bfloat16,
+                            "vae": torch.float16,
+                        },
+                        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.
+            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
 
-        _pipeline = pipe
-        _pipeline_error = None
-        _log("Pipeline ready.")
-    except Exception as e:
-        _pipeline_error = f"{type(e).__name__}: {e}"
-        _log(f"Pipeline load FAILED: {_pipeline_error}")
-        traceback.print_exc()
+        except Exception as e:
+            _pipeline_error = f"{type(e).__name__}: {e}"
+            _log(f"Pipeline load FAILED: {_pipeline_error}")
+            traceback.print_exc()
+            return None
 
 
-# Attempt module-level load unless explicitly suppressed (useful for local syntax checks)
-if os.environ.get("SKIP_MODEL_LOAD") != "1":
+# 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":
     _load_pipeline()
 
 
 # ---------------------------------------------------------------------------
-# Health endpoint — cheap, no weights, no GPU work
+# Health endpoint
 # ---------------------------------------------------------------------------
+
 def health():
     return {
-        "status": "ok" if _pipeline is not None else "error",
+        "status": "ok" if _pipeline is not None else "not_loaded",
         "model_ready": _pipeline is not None,
         "pipeline_ready": _pipeline is not None,
-        "model_family": "diffusers_full_pipeline",
-        "loader_strategy": "DiffusionPipeline_or_family_specific_from_pretrained",
+        "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(),
     }
 
 
 # ---------------------------------------------------------------------------
 # Generation endpoint — real inference guarded by @spaces.GPU
 # ---------------------------------------------------------------------------
-def _gpu_duration(*args, **kwargs):
-    # inputs: prompt, num_blocks, num_inference_steps, seed
-    # Reserve enough time for load hydration + forward passes.
-    return 300
 
+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
+        steps = 4
+
+    return min(900, max(300, 180 + blocks * steps * 20))
 
-@spaces.GPU(duration=_gpu_duration)
+
+@_spaces_gpu(duration=_gpu_duration, size="xlarge")
 def generate(prompt, num_blocks, num_inference_steps, seed):
-    if _pipeline is None:
+    pipe = _load_pipeline()
+
+    if pipe is None:
         err = _pipeline_error or "Pipeline not loaded (unknown failure)"
         raise RuntimeError(f"Generation unavailable: {err}")
 
@@ -159,40 +321,68 @@ def generate(prompt, num_blocks, num_inference_steps, seed):
     num_inference_steps = int(num_inference_steps)
     seed = int(seed)
 
-    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 > 50:
-        raise ValueError("num_inference_steps must be between 1 and 50.")
+    # 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_inference_steps < 1 or num_inference_steps > 8:
+        raise ValueError(
+            "num_inference_steps must be between 1 and 8 in ZeroGPU compatibility mode."
+        )
 
     device = "cuda"
-    # If the pipeline was loaded on CPU (fallback path) ensure it is on the active CUDA device.
-    if hasattr(_pipeline, "device"):
-        current = str(_pipeline.device)
-        if "cuda" not in current:
-            _log("Moving pipeline to cuda ...")
-            _pipeline.to(device)
+
+    # Ensure 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)
+    except Exception as e:
+        _log(f"Pipeline device check/move warning: {type(e).__name__}: {e}")
 
     frames = []
     state = PipelineState()
-    generator = torch.Generator(device=device).manual_seed(seed)
+
+    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:
         for block_idx in range(num_blocks):
             _log(f"Block {block_idx + 1}/{num_blocks}")
-            state = _pipeline(
+
+            state = pipe(
                 state,
                 prompt=[prompt],
                 num_inference_steps=num_inference_steps,
                 num_blocks=num_blocks,
                 block_idx=block_idx,
-                generator=generator)
-            frames.extend(state.values["videos"][0])
+                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 {block_idx}: {e}")
-        raise RuntimeError(f"Inference error at block {block_idx}: {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.")
 
     output_path = "/tmp/krea_output.mp4"
     export_to_video(frames, output_path, fps=24)
+
     _log(f"Saved video to {output_path}")
     return output_path
 
@@ -200,14 +390,15 @@ def generate(prompt, num_blocks, num_inference_steps, seed):
 # ---------------------------------------------------------------------------
 # Gradio app
 # ---------------------------------------------------------------------------
+
 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 official Diffusers ModularPipeline path.\n\n"
-        "⚠️ **Hardware warning**: this is a 14B transformer (~28 GB) paired with a UMT5-XXL text encoder "
-        "(~12–14 GB). Standard HF Space GPUs (48 GB) are unlikely to fit the full model in bf16. "
-        "If the pipeline fails to load, `/health` and the UI will show the concrete error."
+        "This Space attempts **real local inference** for the Krea Realtime 14B "
+        "text-to-video model using the Diffusers `ModularPipeline` path.\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."
     )
 
     with gr.Row():
@@ -215,14 +406,27 @@ with gr.Blocks(title="Krea Realtime Video 14B") as demo:
             prompt = gr.Textbox(
                 label="Prompt",
                 placeholder="e.g., a cat sitting on a boat",
-                lines=2)
+                lines=2,
+            )
+
             num_blocks = gr.Slider(
-                minimum=1, maximum=9, value=3, step=1,
-                label="Number of Blocks (frames per block ≈ 3)")
+                minimum=1,
+                maximum=3,
+                value=1,
+                step=1,
+                label="Number of Blocks (ZeroGPU-safe range)",
+            )
+
             num_inference_steps = gr.Slider(
-                minimum=1, maximum=20, value=4, step=1,
-                label="Inference Steps per Block")
+                minimum=1,
+                maximum=8,
+                value=4,
+                step=1,
+                label="Inference Steps per Block",
+            )
+
             seed = gr.Number(value=42, precision=0, label="Seed")
+
             generate_btn = gr.Button("Generate Video", variant="primary")
 
         with gr.Column():
@@ -230,27 +434,35 @@ with gr.Blocks(title="Krea Realtime Video 14B") as demo:
 
     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", 2, 4, 7],
+            ["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],
         ],
         inputs=[prompt, num_blocks, num_inference_steps, seed],
         outputs=output_video,
         fn=generate,
-        cache_examples=False)
+        cache_examples=False,
+    )
 
     generate_btn.click(
         generate,
         inputs=[prompt, num_blocks, num_inference_steps, seed],
         outputs=output_video,
-        api_name="generate")
+        api_name="generate",
+    )
 
-    # Structured health endpoint
+    # Structured health endpoint.
     demo.load(
         lambda: health(),
         None,
         gr.JSON(),
-        api_name="health")
+        api_name="health",
+    )
+
 
 if __name__ == "__main__":
-    demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True)
+    demo.queue().launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+    )