Update app.py

app.py

@@ -1,58 +1,3 @@
-"""
-BFS — Best Face Swap Video  ·  Hugging Face Space
-"""
-
-from __future__ import annotations
-
-import os
-import tempfile
-
-import gradio as gr
-import numpy as np
-from PIL import Image
-
-from composer import compose_frames, crop_reserved_region
-from video_utils import (
-    compute_target_size,
-    extract_audio,
-    frames_for_duration,
-    load_video_frames,
-    resize_frames,
-    save_video,
-)
-
-# ---------------------------------------------------------------------------
-# GPU decorator — no-op locally, activates the ZeroGPU grant on HF Spaces
-# ---------------------------------------------------------------------------
-try:
-    import spaces
-    GPU = spaces.GPU
-except ImportError:
-    def GPU(fn=None, **kwargs):  # type: ignore
-        return fn if fn is not None else lambda f: f
-
-# ---------------------------------------------------------------------------
-# Global model state (loaded once per worker)
-# ---------------------------------------------------------------------------
-_pipeline_state: dict | None = None
-
-REGION_SIZE = 256
-DEFAULT_FPS = 24.0
-DEFAULT_DURATION = 5.0
-DEFAULT_RESOLUTION = 768
-
-
-def make_temp_file(suffix: str) -> str:
-    f = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
-    path = f.name
-    f.close()
-    return path
-
-
-# ---------------------------------------------------------------------------
-# Core processing function
-# ---------------------------------------------------------------------------
-
 @GPU(duration=300)
 def generate(
     guide_video_path: str | None,
@@ -65,18 +10,92 @@ def generate(
     hf_token: str = "",
     progress: gr.Progress = gr.Progress(track_tqdm=True),
 ) -> tuple[str | None, str]:
-    """
-    Full head-swap pipeline:
-      1. Load + resize guide video frames
-      2. Compose chroma face strip (ReservedRegionFrameComposer)
-      3. Run LTX-2.3 diffusion
-      4. Crop face strip from output
-      5. Mux original audio back in
-
-    Returns (output_video_path, status_message).
-    """
     global _pipeline_state
 
-    # ---- validate inputs early ----
     if not guide_video_path:
-        return None, "Please up
+        return None, "Please upload a guide video."
+    if face_image is None:
+        return None, "Please upload a reference face image."
+    if not prompt or not prompt.strip():
+        return None, "Please enter a text prompt."
+
+    if not os.path.isfile(guide_video_path):
+        return None, f"Guide video path is not a real file: {guide_video_path}"
+
+    if _pipeline_state is None:
+        from pipeline import load_pipeline
+        progress(0, desc="Loading models (first run only — ~5 min)…")
+        _pipeline_state = load_pipeline(
+            token=hf_token.strip() or None,
+            progress_cb=lambda msg: progress(0, desc=msg),
+        )
+
+    progress(0.05, desc="Loading guide video…")
+    frames, source_fps = load_video_frames(guide_video_path)
+    if len(frames) == 0:
+        return None, "Could not read frames from the guide video."
+
+    audio_tmp = make_temp_file(".wav")
+    has_audio = extract_audio(guide_video_path, audio_tmp)
+
+    progress(0.10, desc="Resizing frames…")
+    orig_h, orig_w = frames.shape[1], frames.shape[2]
+    target_w, target_h = compute_target_size(orig_w, orig_h, DEFAULT_RESOLUTION)
+    frames = resize_frames(frames, target_w, target_h)
+
+    n_frames = frames_for_duration(fps, duration)
+    if len(frames) >= n_frames:
+        frames = frames[:n_frames]
+    else:
+        pad = np.stack([frames[-1]] * (n_frames - len(frames)))
+        frames = np.concatenate([frames, pad], axis=0)
+
+    progress(0.15, desc="Compositing reference face strip…")
+    composed = compose_frames(
+        frames,
+        face_image,
+        region_position="left",
+        region_size_px=REGION_SIZE,
+    )
+
+    progress(0.20, desc="Running LTX-2.3 diffusion…")
+    from pipeline import run_inference
+    generated = run_inference(
+        _pipeline_state,
+        composed,
+        prompt=prompt,
+        fps=fps,
+        lora_strength=lora_strength,
+        seed=int(seed),
+        progress_cb=lambda msg: progress(0.20, desc=msg),
+    )
+
+    progress(0.90, desc="Cropping reserved region…")
+    cropped = crop_reserved_region(
+        generated,
+        region_position="left",
+        region_size_px=REGION_SIZE,
+        output_size=(target_w, target_h),
+    )
+
+    progress(0.95, desc="Encoding output video…")
+    out_path = make_temp_file(".mp4")
+    save_video(
+        cropped,
+        fps=fps,
+        output_path=out_path,
+        audio_path=audio_tmp if has_audio else None,
+        audio_duration=duration,
+    )
+
+    if not os.path.isfile(out_path):
+        return None, f"Output file was not created: {out_path}"
+
+    progress(1.0, desc="Done.")
+    return out_path, "Generation complete."
+
+    guide_video = gr.Video(
+    label="Guide Video",
+    sources=["upload"],
+    type="filepath",
+    )