Update app.py

app.py

@@ -2,65 +2,80 @@ import os
 import sys
 import subprocess
 import gradio as gr
-from huggingface_hub import list_repo_files, hf_hub_download
+from huggingface_hub import hf_hub_download
 
 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")
 
-# Specific file to include from transformers/
-ALLOWED_TRANSFORMER_FILE = "ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt"
+ESSENTIAL_PATHS = [
+    # Transformers
+    #"hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt",
+    "hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt",
+    #"hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8_map.pt",
 
-def list_ckpt_files():
-    all_files = list_repo_files(repo_id=MODEL_REPO, repo_type="model", revision="main")
+    # VAE
+    "hunyuan-video-t2v-720p/vae/config.json",
+    "hunyuan-video-t2v-720p/vae/pytorch_model.pt",
 
-    filtered_files = []
-    for f in all_files:
-        # Only allow this one file from transformers/
-        if f.startswith("ckpts/hunyuan-video-t2v-720p/transformers/"):
-            if f == ALLOWED_TRANSFORMER_FILE:
-                filtered_files.append(f)
-        else:
-            filtered_files.append(f)
-    return filtered_files
+    # # llava_llama_image
+    # "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_ckpts():
-    os.makedirs(WEIGHTS_DIR, exist_ok=True)
     logs = []
+    os.makedirs(os.path.join(WEIGHTS_DIR, "ckpts"), exist_ok=True)
 
-    for filepath in list_ckpt_files():
-        relative_path = os.path.relpath(filepath, "ckpts")
-        target_path = os.path.join(WEIGHTS_DIR, "ckpts", relative_path)
-
-        if os.path.exists(target_path):
-            logs.append(f"✅ Exists: {target_path}")
+    for path in ESSENTIAL_PATHS:
+        local_path = os.path.join(WEIGHTS_DIR, "ckpts", path)
+        if os.path.exists(local_path):
+            logs.append(f"✅ Exists: {path}")
             continue
 
-        os.makedirs(os.path.dirname(target_path), exist_ok=True)
-        logs.append(f"⬇️ Downloading: {filepath}")
+        os.makedirs(os.path.dirname(local_path), exist_ok=True)
         try:
+            logs.append(f"⬇️ Downloading: {path}")
             hf_hub_download(
                 repo_id=MODEL_REPO,
-                filename=filepath,
-                local_dir=os.path.dirname(target_path),
-                local_dir_use_symlinks=False
+                filename="ckpts/" + path,
+                local_dir=os.path.dirname(local_path),
+                local_dir_use_symlinks=False,
             )
         except Exception as e:
-            logs.append(f"❌ Failed to download {filepath}: {e}")
+            logs.append(f"❌ Failed: {path} - {str(e)}")
+
     return "\n".join(logs)
 
 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):
-        return f"❌ Missing checkpoint: {checkpoint_fp8}"
+    ckpt_fp8 = os.path.join(WEIGHTS_DIR, "ckpts", "hunyuan-video-t2v-720p", "transformers", "mp_rank_00_model_states_fp8.pt")
+
+    if not os.path.isfile(ckpt_fp8):
+        return f"❌ Missing checkpoint: {ckpt_fp8}"
 
     cmd = [
         "python3", "hymm_sp/sample_gpu_poor.py",
         "--input", "assets/test.csv",
-        "--ckpt", checkpoint_fp8,
+        "--ckpt", ckpt_fp8,
         "--sample-n-frames", "129",
         "--seed", "128",
         "--image-size", "704",
@@ -92,7 +107,7 @@ def download_and_run():
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 📦 Download All Checkpoints (Except Only One File from Transformers)")
+    gr.Markdown("## 📦 Download Selective Checkpoints + Run sample_gpu_poor.py")
     output = gr.Textbox(lines=30, label="Logs")
     button = gr.Button("🚀 Download + Run")
     button.click(fn=download_and_run, outputs=output)