Update app.py

app.py

@@ -1,4 +1,4 @@
-import os, json, tempfile, subprocess, shutil, uuid
+import os, sys, json, tempfile, subprocess, shutil, uuid
 from pathlib import Path
 from typing import Optional, Tuple, List
 
@@ -6,6 +6,9 @@ import gradio as gr
 import spaces
 from huggingface_hub import snapshot_download
 
+from loguru import logger
+import torch, torchaudio
+
 # ========= Paths & Config =========
 ROOT = Path(__file__).parent.resolve()
 REPO_DIR = ROOT / "HunyuanVideo-Foley"
@@ -15,14 +18,15 @@ OUT_DIR = ROOT / "outputs"
 ASSETS = ROOT / "assets"
 ASSETS.mkdir(exist_ok=True)
 
-# You can keep these env vars silently; we just won't mention them in the UI
 APP_TITLE   = os.environ.get("APP_TITLE", "Foley Studio · ZeroGPU")
 APP_TAGLINE = os.environ.get("APP_TAGLINE", "Generate scene-true foley for short clips (ZeroGPU-ready).")
-PRIMARY_COLOR = os.environ.get("PRIMARY_COLOR", "#6B5BFF")
+PRIMARY_COLOR = os.environ.get("PRIMARY_COLOR", "#6B5BFF")  # UI accent only
 
-MAX_SECS = int(os.environ.get("MAX_SECS", "22"))   # ZeroGPU-friendly clip length
-TARGET_H = int(os.environ.get("TARGET_H", "480"))  # downscale target height
-SR       = int(os.environ.get("TARGET_SR", "48000"))  # WAV sample rate
+# ZeroGPU-safe defaults (tweak in Space Secrets if needed)
+MAX_SECS = int(os.environ.get("MAX_SECS", "15"))        # keep clips short for ZeroGPU window
+TARGET_H = int(os.environ.get("TARGET_H", "480"))       # downscale target height
+SR       = int(os.environ.get("TARGET_SR", "48000"))    # WAV sample rate
+ZEROGPU_DURATION = int(os.environ.get("ZEROGPU_DURATION", "110"))  # <= platform maximum
 
 def sh(cmd: str):
     print(">>", cmd)
@@ -71,7 +75,7 @@ def _clone_without_lfs():
     sparse_file = REPO_DIR / ".git" / "info" / "sparse-checkout"
     sparse_file.parent.mkdir(parents=True, exist_ok=True)
     sparse_file.write_text("\n".join([
-        "infer.py",
+        "hunyuanvideo_foley/",
         "configs/",
         "gradio_app.py",
         "requirements.txt",
@@ -89,19 +93,78 @@ def _clone_without_lfs():
 def prepare_once():
     """Clone code (skip LFS), download weights, set env, prepare dirs."""
     _clone_without_lfs()
+
+    # Ensure we can import their package
+    if str(REPO_DIR) not in sys.path:
+        sys.path.insert(0, str(REPO_DIR))
+
     WEIGHTS_DIR.mkdir(parents=True, exist_ok=True)
     snapshot_download(
         repo_id="tencent/HunyuanVideo-Foley",
         local_dir=str(WEIGHTS_DIR),
         local_dir_use_symlinks=False,
         repo_type="model",
+        resume_download=True,
     )
     os.environ["HIFI_FOLEY_MODEL_PATH"] = str(WEIGHTS_DIR)
+
     CACHE_DIR.mkdir(exist_ok=True)
     OUT_DIR.mkdir(exist_ok=True)
 
 prepare_once()
 
+# Now safe to import their internals
+from hunyuanvideo_foley.utils.model_utils import load_model, denoise_process
+from hunyuanvideo_foley.utils.feature_utils import feature_process
+from hunyuanvideo_foley.utils.media_utils import merge_audio_video
+
+# ========= Native Model Setup =========
+MODEL_PATH = os.environ.get("HIFI_FOLEY_MODEL_PATH", str(WEIGHTS_DIR))
+CONFIG_PATH = str(REPO_DIR / "configs" / "hunyuanvideo-foley-xxl.yaml")
+
+_model_dict = None
+_cfg = None
+_device = None
+
+def _setup_device(device_str: str = "auto", gpu_id: int = 0) -> torch.device:
+    if device_str == "auto":
+        if torch.cuda.is_available():
+            d = torch.device(f"cuda:{gpu_id}")
+            logger.info(f"Using CUDA {d}")
+        elif torch.backends.mps.is_available():
+            d = torch.device("mps")
+            logger.info("Using MPS")
+        else:
+            d = torch.device("cpu")
+            logger.info("Using CPU")
+    else:
+        d = torch.device(device_str if device_str != "cuda" else f"cuda:{gpu_id}")
+        logger.info(f"Using specified device: {d}")
+    return d
+
+def auto_load_models() -> str:
+    """Download weights if needed + load model natively."""
+    global _model_dict, _cfg, _device
+
+    if not os.path.exists(MODEL_PATH):
+        os.makedirs(MODEL_PATH, exist_ok=True)
+
+    if not os.path.exists(CONFIG_PATH):
+        return f"❌ Config file not found: {CONFIG_PATH}"
+
+    _device = _setup_device("auto", 0)
+    logger.info("Loading HunyuanVideo-Foley model...")
+    logger.info(f"MODEL_PATH:  {MODEL_PATH}")
+    logger.info(f"CONFIG_PATH: {CONFIG_PATH}")
+    _model_dict, _cfg = load_model(MODEL_PATH, CONFIG_PATH, _device)
+    logger.info("✅ Model loaded")
+    return "✅ Model loaded"
+
+# Init logger and load model once
+logger.remove()
+logger.add(lambda msg: print(msg, end=''), level="INFO")
+logger.info(auto_load_models())
+
 # ========= Preprocessing =========
 def preprocess_video(in_path: str) -> Tuple[str, float]:
     """
@@ -145,42 +208,44 @@ def preprocess_video(in_path: str) -> Tuple[str, float]:
     return str(processed), final_dur
 
 # ========= Inference (ZeroGPU) =========
-@spaces.GPU(duration=240)  # ~4 minutes per call (ZeroGPU window)
-def run_model(video_path: str, prompt_text: str) -> str:
+@spaces.GPU(duration=ZEROGPU_DURATION)  # tune via env if needed
+@torch.inference_mode()
+def run_model(video_path: str, prompt_text: str,
+              guidance_scale: float = 4.5,
+              num_inference_steps: int = 50,
+              sample_nums: int = 1) -> Tuple[List[str], int]:
     """
-    Call Tencent's infer.py on GPU and return a 48 kHz WAV path.
+    Native inference (no shell). Returns ([wav_paths], sample_rate).
     """
-    job_id = uuid.uuid4().hex[:8]
-    work_out = OUT_DIR / f"job_{job_id}"
-    work_out.mkdir(parents=True, exist_ok=True)
-
-    cmd = [
-        "python", str(REPO_DIR / "infer.py"),
-        "--model_path", str(WEIGHTS_DIR),
-        "--config_path", str(REPO_DIR / "configs" / "hunyuanvideo-foley-xxl.yaml"),
-        "--single_video", video_path,
-        "--single_prompt", json.dumps(prompt_text or ""),
-        "--output_dir", str(work_out),
-        "--device", "cuda",
-    ]
-    sh(" ".join(cmd))
-
-    # Find produced wav
-    wav = None
-    for p in work_out.rglob("*.wav"):
-        wav = p
-        break
-    if not wav:
-        raise RuntimeError("No audio produced by the model.")
-
-    # Normalize / resample to SR stereo
-    fixed = work_out / "foley_48k.wav"
-    sh(" ".join([
-        "ffmpeg", "-y", "-i", f"\"{str(wav)}\"",
-        "-ar", str(SR), "-ac", "2",
-        f"\"{str(fixed)}\""
-    ]))
-    return str(fixed)
+    if _model_dict is None or _cfg is None:
+        raise RuntimeError("Model not loaded yet.")
+
+    text_prompt = (prompt_text or "").strip()
+
+    # Extract features
+    visual_feats, text_feats, audio_len_s = feature_process(
+        video_path, text_prompt, _model_dict, _cfg
+    )
+
+    # Generate audio (B x C x T)
+    logger.info(f"Generating {sample_nums} sample(s)...")
+    audio_batch, sr = denoise_process(
+        visual_feats, text_feats, audio_len_s, _model_dict, _cfg,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        batch_size=sample_nums
+    )
+
+    # Save each sample as WAV
+    out_dir = OUT_DIR / f"job_{uuid.uuid4().hex[:8]}"
+    out_dir.mkdir(parents=True, exist_ok=True)
+    wav_paths = []
+    for i in range(sample_nums):
+        wav_p = out_dir / f"generated_audio_{i+1}.wav"
+        torchaudio.save(str(wav_p), audio_batch[i], sr)
+        wav_paths.append(str(wav_p))
+
+    return wav_paths, sr
 
 # ========= Optional: Mux Foley back to video =========
 def mux_audio_with_video(video_path: str, audio_path: str) -> str:
@@ -197,26 +262,34 @@ def mux_audio_with_video(video_path: str, audio_path: str) -> str:
     return str(out_path)
 
 # ========= UI Handlers =========
-def single_generate(video: str, prompt: str, want_mux: bool, project_name: str) -> Tuple[Optional[str], Optional[str], str, list]:
+def single_generate(video: str, prompt: str, want_mux: bool, project_name: str):
     history = []
     try:
         if not video:
             return None, None, "⚠️ Please upload a video.", history
         history.append(["Preprocess", "Downscaling & trimming"])
         pre_path, final_dur = preprocess_video(video)
-        history.append(["Inference", "Running on ZeroGPU"])
-        wav = run_model(pre_path, prompt or "")
+
+        history.append(["Inference", "ZeroGPU native pipeline"])
+        wav_list, sr = run_model(
+            pre_path, prompt or "", guidance_scale=4.5, num_inference_steps=50, sample_nums=1
+        )
+        if not wav_list:
+            raise RuntimeError("No audio produced.")
+        wav = wav_list[0]
+
         muxed = None
         if want_mux:
             history.append(["Mux", "Merging foley with video"])
             muxed = mux_audio_with_video(pre_path, wav)
+
         history.append(["Done", f"OK · ~{final_dur:.1f}s"])
         return wav, muxed, f"✅ Completed (~{final_dur:.1f}s)", history
     except Exception as e:
         history.append(["Error", str(e)])
         return None, None, f"❌ {type(e).__name__}: {e}", history
 
-def batch_lite_generate(files: List[str], prompt: str, want_mux: bool) -> Tuple[str, list]:
+def batch_lite_generate(files: List[str], prompt: str, want_mux: bool):
     log = []
     if not files:
         return "⚠️ Please upload 1–3 videos.", log
@@ -230,7 +303,10 @@ def batch_lite_generate(files: List[str], prompt: str, want_mux: bool) -> Tuple[
             log.append([f"Preprocess {i}", Path(f).name])
             pre, final_dur = preprocess_video(f)
             log.append([f"Run {i}", f"ZeroGPU ~{final_dur:.1f}s"])
-            wav = run_model(pre, prompt or "")
+            wav_list, sr = run_model(pre, prompt or "", sample_nums=1)
+            if not wav_list:
+                raise RuntimeError("No audio produced.")
+            wav = wav_list[0]
             muxed = mux_audio_with_video(pre, wav) if want_mux else None
             outputs.append((wav, muxed))
             log.append([f"Done {i}", "OK"])
@@ -359,7 +435,7 @@ with gr.Blocks(css=THEME_CSS, title=APP_TITLE, analytics_enabled=False) as demo:
 **Usage guidelines**
 - Keep clips short (the tool trims to **≤ {MAX_SECS}s** automatically).
 - The video is downscaled to **{TARGET_H}p** to fit the ZeroGPU time window.
-- If you see a quota message, just try again later (ZeroGPU limits GPU minutes per visitor).
+- If you see a quota message, try again later (ZeroGPU limits GPU minutes per visitor).
 
 **Outputs**
 - WAV is **{SR//1000} kHz** stereo.