Update app.py

app.py

@@ -1,66 +1,55 @@
 import os
 import sys
 import subprocess
-from huggingface_hub import hf_hub_download, hf_hub_list
 from huggingface_hub import list_repo_files
-from pathlib import Path
-import spaces
+
 MODEL_REPO = "tencent/HunyuanVideo-Avatar"
-BASE_DIR = Path.cwd()
-WEIGHTS_DIR = BASE_DIR / "weights"
-CKPTS_DIR = WEIGHTS_DIR / "ckpts"
+REPO_ROOT_FOLDER = "ckpts"  # folder inside repo to download
 
-# The root folder in the repo you want to download files from
-REPO_ROOT_FOLDER = "ckpts"
+BASE_DIR = os.getcwd()
+WEIGHTS_DIR = os.path.join(BASE_DIR, "weights")
+OUTPUT_BASEPATH = os.path.join(BASE_DIR, "results-poor")
 
-# List all files under the folder 'ckpts' in the repo (recursively)
 def list_ckpt_files():
     print(f"🔍 Listing files under '{REPO_ROOT_FOLDER}' in repo {MODEL_REPO}...")
-    files = hf_hub_list(repo_id=MODEL_REPO, repo_type="model", revision="main", allow_regex=f"^{REPO_ROOT_FOLDER}/.*")
+    all_files = list_repo_files(repo_id=MODEL_REPO, repo_type="model", revision="main")
+    files = [f for f in all_files if f.startswith(REPO_ROOT_FOLDER + "/")]
     return files
 
-# Download all files in the list into weights/ckpts preserving folder structure
-def download_ckpts(files):
-    print(f"⬇️ Downloading {len(files)} files under '{REPO_ROOT_FOLDER}' ...")
-    for file in files:
-        # Target local path preserving folder structure after 'ckpts/'
-        relative_path = Path(file).relative_to(REPO_ROOT_FOLDER)
-        local_path = CKPTS_DIR / relative_path
-        
-        # Make sure directory exists
-        local_path.parent.mkdir(parents=True, exist_ok=True)
+def download_ckpts():
+    os.makedirs(WEIGHTS_DIR, exist_ok=True)
+    for filepath in list_ckpt_files():
+        # Target path relative to weights directory
+        relative_path = os.path.relpath(filepath, REPO_ROOT_FOLDER)
+        target_path = os.path.join(WEIGHTS_DIR, REPO_ROOT_FOLDER, relative_path)
+
+        # Skip if already exists
+        if os.path.exists(target_path):
+            print(f"✅ File exists: {target_path}")
+            continue
+
+        # Create dirs if needed
+        os.makedirs(os.path.dirname(target_path), exist_ok=True)
+
+        # Download the file blob from HF repo using huggingface_hub api (requires v0.14+)
+        print(f"⬇️ Downloading {filepath} to {target_path} ...")
+        try:
+            from huggingface_hub import hf_hub_download
+            hf_hub_download(repo_id=MODEL_REPO, filename=filepath, local_dir=os.path.dirname(target_path), local_dir_use_symlinks=False)
+        except ImportError:
+            print("⚠️ hf_hub_download not available, please upgrade huggingface_hub")
+            sys.exit(1)
 
-        # Download the file if not present
-        if not local_path.is_file():
-            print(f"Downloading {file} -> {local_path}")
-            hf_hub_download(
-                repo_id=MODEL_REPO,
-                filename=file,
-                repo_type="model",
-                cache_dir=str(WEIGHTS_DIR),
-                local_dir=str(WEIGHTS_DIR),
-                local_dir_use_symlinks=False,
-                force_filename=str(local_path.name),
-                revision="main",
-            )
-        else:
-            print(f"File {local_path} already exists, skipping download.")
-    print("✅ All checkpoint files downloaded.")
-    
-@spaces.GPU(duration=1000)
 def run_sample_gpu_poor():
-    checkpoint_fp8 = CKPTS_DIR / "hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt"
-    if not checkpoint_fp8.is_file():
-        print(f"❌ Checkpoint file {checkpoint_fp8} not found!")
+    checkpoint_fp8 = os.path.join(WEIGHTS_DIR, "ckpts", "hunyuan-video-t2v-720p", "transformers", "mp_rank_00_model_states_fp8.pt")
+    if not os.path.isfile(checkpoint_fp8):
+        print(f"❌ Checkpoint file not found at {checkpoint_fp8}. Cannot run sampling.")
         sys.exit(1)
 
-    OUTPUT_BASEPATH = BASE_DIR / "results-poor"
-    OUTPUT_BASEPATH.mkdir(parents=True, exist_ok=True)
-
     cmd = [
         "python3", "hymm_sp/sample_gpu_poor.py",
         "--input", "assets/test.csv",
-        "--ckpt", str(checkpoint_fp8),
+        "--ckpt", checkpoint_fp8,
         "--sample-n-frames", "129",
         "--seed", "128",
         "--image-size", "704",
@@ -68,7 +57,7 @@ def run_sample_gpu_poor():
         "--infer-steps", "50",
         "--use-deepcache", "1",
         "--flow-shift-eval-video", "5.0",
-        "--save-path", str(OUTPUT_BASEPATH),
+        "--save-path", OUTPUT_BASEPATH,
         "--use-fp8",
         "--cpu-offload",
         "--infer-min"
@@ -76,11 +65,11 @@ def run_sample_gpu_poor():
 
     env = os.environ.copy()
     env["PYTHONPATH"] = "./"
-    env["MODEL_BASE"] = str(WEIGHTS_DIR)
+    env["MODEL_BASE"] = WEIGHTS_DIR
     env["CPU_OFFLOAD"] = "1"
     env["CUDA_VISIBLE_DEVICES"] = "0"
 
-    print("🎬 Running sample_gpu_poor.py...")
+    print("🎬 Running sample_gpu_poor.py ...")
     proc = subprocess.run(cmd, env=env)
     if proc.returncode != 0:
         print("❌ sample_gpu_poor.py failed.")
@@ -88,12 +77,7 @@ def run_sample_gpu_poor():
     print("✅ sample_gpu_poor.py completed successfully.")
 
 def main():
-    files = list_ckpt_files()
-    if not files:
-        print("❌ No checkpoint files found in repo under ckpts folder.")
-        sys.exit(1)
-
-    download_ckpts(files)
+    download_ckpts()
     run_sample_gpu_poor()
 
 if __name__ == "__main__":