Update app.py

app.py

@@ -82,6 +82,7 @@ import torch
 # 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)
+_ASF_COMPILE_BYPASS_ANNOUNCED = False
 
 
 def _asf_zerogpu_compile_bypass(fn=None, *args, **kwargs):
@@ -92,20 +93,32 @@ def _asf_zerogpu_compile_bypass(fn=None, *args, **kwargs):
       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
 
-    name = getattr(fn, "__name__", repr(fn))
-    module = getattr(fn, "__module__", "")
+    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
 
-    print(
-        f"[ASF] ZeroGPU compatibility: bypassing torch.compile for {module}.{name}",
-        flush=True,
-    )
     return fn
 
 
@@ -328,6 +341,7 @@ def health():
     }
 
 
+@_spaces_gpu(duration=120, size="xlarge")
 def warmup_model():
     """
     Manual warm-up button.
@@ -365,18 +379,25 @@ def _gpu_duration(prompt, num_blocks, num_inference_steps, seed):
         steps = 4
 
     # 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))
+    # Observed: 9 blocks × 4 steps ≈ <75s.
+    # Keep a small safety buffer without over-reserving ZeroGPU.
+    #
+    # 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)))
 
 
 @_spaces_gpu(duration=_gpu_duration, size="xlarge")
-def generate(prompt, num_blocks, num_inference_steps, seed):
+def generate(
+    prompt,
+    num_blocks,
+    num_inference_steps,
+    seed,
+    progress=gr.Progress(track_tqdm=True),
+):
     pipe = _load_pipeline()
 
     if pipe is None:
@@ -414,9 +435,19 @@ def generate(prompt, num_blocks, num_inference_steps, seed):
         generator = torch.Generator(device="cpu").manual_seed(seed)
 
     try:
-        for block_idx in range(num_blocks):
+        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=[prompt],
@@ -442,10 +473,14 @@ def generate(prompt, num_blocks, num_inference_steps, seed):
     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