Update app.py

app.py

@@ -19,7 +19,19 @@ from huggingface_hub import hf_hub_download, snapshot_download
 logging.basicConfig(level=logging.INFO, format="[%(asctime)s] %(levelname)s: %(message)s")
 
 ROOT = Path(__file__).resolve().parent
-OUTPUT_DIR = Path(os.getenv("SCAIL_OUTPUT_DIR", "/tmp/scail2_outputs"))
+
+
+def _default_storage_root() -> Path:
+    if os.getenv("SCAIL_STORAGE_ROOT"):
+        return Path(os.environ["SCAIL_STORAGE_ROOT"])
+    data_mount = Path("/data")
+    if data_mount.exists() and os.access(data_mount, os.W_OK):
+        return data_mount
+    return Path("/tmp")
+
+
+STORAGE_ROOT = _default_storage_root()
+OUTPUT_DIR = Path(os.getenv("SCAIL_OUTPUT_DIR", str(STORAGE_ROOT / "scail2_outputs")))
 OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
 
 MODEL_REPO_ID = os.getenv("SCAIL_MODEL_REPO_ID", "zai-org/SCAIL-2")
@@ -133,7 +145,7 @@ def _download_safetensors_if_configured() -> Path | None:
     if not SAFETENSORS_REPO_ID:
         return None
 
-    local_dir = Path(os.getenv("SCAIL_SAFETENSORS_CACHE", "/tmp/scail2_safetensors"))
+    local_dir = Path(os.getenv("SCAIL_SAFETENSORS_CACHE", str(STORAGE_ROOT / "scail2_safetensors")))
     local_dir.mkdir(parents=True, exist_ok=True)
     local_path = local_dir / SAFETENSORS_FILENAME
     if local_path.exists():
@@ -163,7 +175,7 @@ def _find_converted_safetensors(ckpt_dir: Path | None) -> Path | None:
         ROOT / "SCAIL-2.safetensors",
         ROOT / "models" / "SCAIL-2.safetensors",
         ROOT / "model.safetensors",
-        Path(os.getenv("SCAIL_CONVERTED_DIR", "/tmp/scail2_converted")) / "SCAIL-2.safetensors",
+        Path(os.getenv("SCAIL_CONVERTED_DIR", str(STORAGE_ROOT / "scail2_converted"))) / "SCAIL-2.safetensors",
     ]
     if ckpt_dir is not None:
         candidates += [
@@ -184,7 +196,7 @@ def _download_checkpoint_if_needed(include_original_dit: bool = False) -> Path:
             raise RuntimeError(f"SCAIL_CKPT_DIR does not exist: {ckpt_dir}")
         return ckpt_dir
 
-    local_dir = Path(os.getenv("SCAIL_CKPT_CACHE", "/tmp/scail2_ckpt"))
+    local_dir = Path(os.getenv("SCAIL_CKPT_CACHE", str(STORAGE_ROOT / "scail2_ckpt")))
     clip_name = "models_clip_open-clip-xlm-roberta-large-vit-huge-14-onlyvisual.pth"
     original_dit = local_dir / "model" / "1" / "fsdp2_rank_0000_checkpoint.pt"
     has_base_assets = (
@@ -267,7 +279,7 @@ def _maybe_convert_checkpoint(ckpt_dir: Path, scail_path: Path | None) -> Path:
             "because SCAIL_AUTO_CONVERT=0."
         )
 
-    save_dir = Path(os.getenv("SCAIL_CONVERTED_DIR", "/tmp/scail2_converted"))
+    save_dir = Path(os.getenv("SCAIL_CONVERTED_DIR", str(STORAGE_ROOT / "scail2_converted")))
     save_dir.mkdir(parents=True, exist_ok=True)
     save_path = save_dir / "SCAIL-2.safetensors"
     if save_path.exists():
@@ -509,16 +521,6 @@ def _duration_for_job(*args, **kwargs):
     return int(os.getenv("SCAIL_GPU_DURATION", str(GPU_DURATION_WARM)))
 
 
-@spaces.GPU(duration=_duration_for_job, size=GPU_SIZE)
-def warmup_model():
-    try:
-        _get_pipeline()
-        return "Model is loaded. The next generation should skip the cold-start model load."
-    except Exception:
-        logging.exception("Warmup failed")
-        return traceback.format_exc()
-
-
 def _run_scail_job(
     image_path,
     mask_image_path,
@@ -534,10 +536,15 @@ def _run_scail_job(
     seed,
     segment_len,
     segment_overlap,
+    progress=None,
 ):
+    if progress is not None:
+        progress(0.02, desc="Loading SCAIL-2 pipeline")
     pipeline = _get_pipeline()
     cfg = _SCAIL_CONFIGS[MODEL_NAME]
     save_file = OUTPUT_DIR / f"scail2_{uuid.uuid4().hex}.mp4"
+    if progress is not None:
+        progress(0.12, desc="Preparing inputs")
     args = SimpleNamespace(
         target_h=int(target_h),
         target_w=int(target_w),
@@ -554,6 +561,8 @@ def _run_scail_job(
         prompt=prompt or "",
     )
 
+    if progress is not None:
+        progress(0.15, desc="Generating video")
     _GENERATE_VIDEO(
         pipeline,
         prompt or "",
@@ -569,9 +578,13 @@ def _run_scail_job(
         replace_flag=bool(replace_flag),
         additional_task_input=None,
     )
+    if progress is not None:
+        progress(0.95, desc="Finalizing output")
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
+    if progress is not None:
+        progress(1.0, desc="Done")
     return str(save_file)
 
 
@@ -586,8 +599,10 @@ def generate_from_example(
     target_size,
     segment_len,
     segment_overlap,
+    progress=gr.Progress(track_tqdm=True),
 ):
     try:
+        progress(0.0, desc="Checking example inputs")
         examples = _existing_examples()
         if example_name not in examples:
             raise RuntimeError(f"Example is missing from this checkout: {example_name}")
@@ -608,6 +623,7 @@ def generate_from_example(
             seed,
             segment_len,
             segment_overlap,
+            progress=progress,
         )
         return output, "Done."
     except Exception:
@@ -630,8 +646,10 @@ def generate_from_uploads(
     target_size,
     segment_len,
     segment_overlap,
+    progress=gr.Progress(track_tqdm=True),
 ):
     try:
+        progress(0.0, desc="Checking uploaded inputs")
         required = {
             "reference image": image,
             "reference mask": mask_image,
@@ -658,6 +676,7 @@ def generate_from_uploads(
             seed,
             segment_len,
             segment_overlap,
+            progress=progress,
         )
         return output, "Done."
     except Exception:
@@ -689,6 +708,7 @@ def _startup_message():
             return "Repo layout detected, but no prepared examples were found."
         return (
             f"Ready. Found {len(examples)} prepared example(s). "
+            f"Storage root: {STORAGE_ROOT}. Output dir: {OUTPUT_DIR}.\n\n"
             f"{_ASSET_STATUS}\n\n{_RUNTIME_STATUS}\n\n"
             "Attention backend: HF Kernels flash-attn2 when available, otherwise SDPA."
         )
@@ -708,15 +728,11 @@ def build_ui():
             "V1 inference-only Space. It runs SCAIL-2 from prepared inputs: reference image, "
             "reference mask, driving/rendered video, and driving mask video. SCAIL-Pose/SAM3 "
             "preprocessing is intentionally not imported in this version.\n\n"
-            "First GPU call is a cold start: it loads T5, VAE, CLIP, and the 14B DiT before "
-            "sampling. Use the warmup button once after startup if you want to separate model "
-            "loading from generation."
+            "On ZeroGPU, each generation should be treated as a cold start: the GPU can be "
+            "released between calls, so the 14B stack may need to be loaded again before "
+            "sampling."
         )
         startup = gr.Textbox(value=_startup_message(), label="Startup status", interactive=False)
-        with gr.Row():
-            warmup_button = gr.Button("Load model / warm up GPU", variant="secondary")
-            warmup_status = gr.Textbox(label="Warmup status", lines=4, interactive=False)
-        warmup_button.click(warmup_model, outputs=[warmup_status])
 
         with gr.Tab("Prepared Examples"):
             with gr.Row():