Update app.py

app.py

@@ -2,63 +2,57 @@ import os
 import sys
 import subprocess
 import time
-import shutil
-import psutil
+from pathlib import Path
 from huggingface_hub import snapshot_download
 
+# --------------------
+# CONFIGURATION
+# --------------------
 MODEL_REPO = "tencent/HunyuanVideo-Avatar"
-BASE_DIR = os.getcwd()
-WEIGHTS_DIR = os.path.join(BASE_DIR, "weights")
-OUTPUT_BASEPATH = os.path.join(BASE_DIR, "results-poor")
-CUSTOM_CACHE_DIR = "/data/hf_cache"  # ✅ change if different
-HF_HOME_ENV = {"HF_HOME": CUSTOM_CACHE_DIR}
+HF_CACHE_DIR = Path("/home/user/.cache/huggingface/hf_cache/hunyuan_avatar")
+HF_CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
-# Paths to required checkpoints
-CHECKPOINT_FILE = os.path.join(WEIGHTS_DIR, "ckpts", "hunyuan-video-t2v-720p", "transformers", "mp_rank_00_model_states.pt")
-CHECKPOINT_FP8_FILE = os.path.join(WEIGHTS_DIR, "ckpts", "hunyuan-video-t2v-720p", "transformers", "mp_rank_00_model_states_fp8.pt")
+CHECKPOINT_FILE = HF_CACHE_DIR / "ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt"
+CHECKPOINT_FP8_FILE = HF_CACHE_DIR / "ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt"
 
-def check_disk_space(min_free_gb=10):
-    total, used, free = shutil.disk_usage("/")
-    free_gb = free // (2**30)
-    print(f"💾 Disk space: {free_gb}GB free")
-    if free_gb < min_free_gb:
-        print(f"❌ Not enough disk space. {free_gb}GB available, {min_free_gb}GB required.")
-        sys.exit(1)
-
-def clear_hf_cache():
-    hf_cache = os.path.expanduser("~/.cache/huggingface")
-    if os.path.exists(hf_cache):
-        print("🧹 Cleaning Hugging Face cache...")
-        shutil.rmtree(hf_cache)
+ASSETS_CSV = "assets/test.csv"
+OUTPUT_DIR = Path("results-poor")
+OUTPUT_DIR.mkdir(exist_ok=True)
 
+# --------------------
+# Download the model (if needed)
+# --------------------
 def download_model():
-    print("⬇️  Downloading model to weights directory...")
-    os.makedirs(WEIGHTS_DIR, exist_ok=True)
+    if CHECKPOINT_FILE.exists() and CHECKPOINT_FP8_FILE.exists():
+        print("✅ Model checkpoint already exists. Skipping download.")
+        return
 
+    print("⬇️ Downloading model into HF Space cache...")
     snapshot_download(
         repo_id=MODEL_REPO,
-        local_dir=WEIGHTS_DIR,
-        local_dir_use_symlinks=False,
-        **HF_HOME_ENV
+        local_dir=HF_CACHE_DIR,
+        local_dir_use_symlinks=False
     )
 
-    if not os.path.isfile(CHECKPOINT_FILE):
-        print(f"❌ Checkpoint missing at {CHECKPOINT_FILE}")
+    if not CHECKPOINT_FILE.exists():
+        print(f"❌ Missing checkpoint: {CHECKPOINT_FILE}")
         sys.exit(1)
-    if not os.path.isfile(CHECKPOINT_FP8_FILE):
-        print(f"❌ FP8 Checkpoint missing at {CHECKPOINT_FP8_FILE}")
+
+    if not CHECKPOINT_FP8_FILE.exists():
+        print(f"❌ Missing FP8 checkpoint: {CHECKPOINT_FP8_FILE}")
         sys.exit(1)
 
-    print("✅ Model downloaded successfully.")
-    clear_hf_cache()
+    print("✅ Model download complete.")
 
+# --------------------
+# Run sample_gpu_poor.py
+# --------------------
 def run_sample_gpu_poor():
     print("🎬 Running sample_gpu_poor.py...")
-
     cmd = [
         "python3", "hymm_sp/sample_gpu_poor.py",
-        "--input", "assets/test.csv",
-        "--ckpt", CHECKPOINT_FP8_FILE,
+        "--input", ASSETS_CSV,
+        "--ckpt", str(CHECKPOINT_FP8_FILE),
         "--sample-n-frames", "129",
         "--seed", "128",
         "--image-size", "704",
@@ -66,74 +60,38 @@ def run_sample_gpu_poor():
         "--infer-steps", "50",
         "--use-deepcache", "1",
         "--flow-shift-eval-video", "5.0",
-        "--save-path", OUTPUT_BASEPATH,
+        "--save-path", str(OUTPUT_DIR),
         "--use-fp8",
         "--cpu-offload",
         "--infer-min"
     ]
 
     env = os.environ.copy()
-    env.update(HF_HOME_ENV)
     env["PYTHONPATH"] = "./"
-    env["MODEL_BASE"] = WEIGHTS_DIR
+    env["MODEL_BASE"] = str(HF_CACHE_DIR)
     env["CPU_OFFLOAD"] = "1"
     env["CUDA_VISIBLE_DEVICES"] = "0"
 
-    result = subprocess.run(cmd, env=env)
-    if result.returncode != 0:
+    proc = subprocess.run(cmd, env=env)
+    if proc.returncode != 0:
         print("❌ sample_gpu_poor.py failed.")
         sys.exit(1)
+    print("✅ sample_gpu_poor.py completed successfully.")
 
-    print("✅ sample_gpu_poor.py completed.")
-
-def run_flask_audio():
-    print("🚀 Starting flask_audio.py...")
-    cmd = [
-        "torchrun",
-        "--nnodes=1",
-        "--nproc_per_node=8",
-        "--master_port=29605",
-        "hymm_gradio/flask_audio.py",
-        "--input", "assets/test.csv",
-        "--ckpt", CHECKPOINT_FILE,
-        "--sample-n-frames", "129",
-        "--seed", "128",
-        "--image-size", "704",
-        "--cfg-scale", "7.5",
-        "--infer-steps", "50",
-        "--use-deepcache", "1",
-        "--flow-shift-eval-video", "5.0"
-    ]
-    subprocess.Popen(cmd)
-
+# --------------------
+# Optional: Start UI
+# --------------------
 def run_gradio_ui():
-    print("🟢 Launching Gradio UI...")
+    print("🟢 Launching Gradio interface...")
     cmd = ["python3", "hymm_gradio/gradio_audio.py"]
     subprocess.Popen(cmd)
 
-def create_gitignore():
-    gitignore_path = os.path.join(BASE_DIR, ".gitignore")
-    lines_to_add = ["weights/", "results-*/", ".cache/"]
-    with open(gitignore_path, "a+") as f:
-        f.seek(0)
-        existing = f.read().splitlines()
-        for line in lines_to_add:
-            if line not in existing:
-                f.write(f"{line}\n")
-
+# --------------------
+# Entry point
+# --------------------
 def main():
-    create_gitignore()
-    check_disk_space(min_free_gb=10)
-
-    if os.path.isfile(CHECKPOINT_FILE) and os.path.isfile(CHECKPOINT_FP8_FILE):
-        print("✅ Checkpoints exist. Skipping download.")
-    else:
-        download_model()
-
+    download_model()
     run_sample_gpu_poor()
-
-    # Optional: Launch UIs
-    run_flask_audio()
     time.sleep(5)
     run_gradio_ui()