Update app.py

app.py

@@ -1,84 +1,63 @@
 import os
-import sys
-import subprocess
-from huggingface_hub import list_repo_files
-
-MODEL_REPO = "tencent/HunyuanVideo-Avatar"
-REPO_ROOT_FOLDER = "ckpts"  # folder inside repo to download
-
-BASE_DIR = os.getcwd()
-WEIGHTS_DIR = os.path.join(BASE_DIR, "weights")
-OUTPUT_BASEPATH = os.path.join(BASE_DIR, "results-poor")
-
-def list_ckpt_files():
-    print(f"🔍 Listing files under '{REPO_ROOT_FOLDER}' in repo {MODEL_REPO}...")
-    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
-
-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} ...")
+from huggingface_hub import hf_hub_download
+
+REPO_ID = "tencent/HunyuanVideo-Avatar"
+BASE_PATH = "ckpts"
+LOCAL_BASE = os.path.join(os.getcwd(), "weights", "ckpts")
+
+# List of essential files/folders to download (you can expand this if needed)
+ESSENTIAL_PATHS = [
+    # Transformers checkpoints
+    "hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt",
+    "hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt",
+    "hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8_map.pt",
+
+    # VAE
+    "hunyuan-video-t2v-720p/vae/config.json",
+    "hunyuan-video-t2v-720p/vae/pytorch_model.pt",
+
+    # llava_llama_image shard files (adjust count if needed)
+    "llava_llama_image/model-00001-of-00004.safetensors",
+    "llava_llama_image/model-00002-of-00004.safetensors",
+    "llava_llama_image/model-00003-of-00004.safetensors",
+    "llava_llama_image/model-00004-of-00004.safetensors",
+    "llava_llama_image/config.json",
+
+    # text_encoder_2
+    "text_encoder_2/config.json",
+    "text_encoder_2/pytorch_model.bin",
+
+    # whisper-tiny
+    "whisper-tiny/config.json",
+    "whisper-tiny/pytorch_model.bin",
+    "whisper-tiny/tokenizer.json",
+    "whisper-tiny/tokenizer_config.json",
+    "whisper-tiny/vocab.json",
+
+    # det_align
+    "det_align/config.json",
+    "det_align/pytorch_model.bin",
+]
+
+def download_files():
+    for relative_path in ESSENTIAL_PATHS:
+        source_path = f"{BASE_PATH}/{relative_path}"
+        local_dir = os.path.join(LOCAL_BASE, os.path.dirname(relative_path))
+        os.makedirs(local_dir, exist_ok=True)
+
+        print(f"⬇️ Downloading {source_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)
-
-def run_sample_gpu_poor():
-    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)
-
-    cmd = [
-        "python3", "hymm_sp/sample_gpu_poor.py",
-        "--input", "assets/test.csv",
-        "--ckpt", checkpoint_fp8,
-        "--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",
-        "--save-path", OUTPUT_BASEPATH,
-        "--use-fp8",
-        "--cpu-offload",
-        "--infer-min"
-    ]
-
-    env = os.environ.copy()
-    env["PYTHONPATH"] = "./"
-    env["MODEL_BASE"] = WEIGHTS_DIR
-    env["CPU_OFFLOAD"] = "1"
-    env["CUDA_VISIBLE_DEVICES"] = "0"
-
-    print("🎬 Running sample_gpu_poor.py ...")
-    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.")
-
-def main():
-    download_ckpts()
-    run_sample_gpu_poor()
+            hf_hub_download(
+                repo_id=REPO_ID,
+                filename=source_path,
+                repo_type="model",
+                local_dir=local_dir,
+                local_dir_use_symlinks=False
+            )
+        except Exception as e:
+            print(f"❌ Failed to download {source_path}: {e}")
 
 if __name__ == "__main__":
-    main()
+    download_files()
+    print("\n✅ All selected model weights downloaded to:")
+    print(f"{os.path.abspath(LOCAL_BASE)}")