Update app.py

app.py

@@ -1,76 +1,71 @@
 import os
 import sys
+from huggingface_hub import hf_hub_download, list_repo_files
 import subprocess
-import gradio as gr
-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")
 
+
+# Only download this specific file from transformers
 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",
-
-    # VAE
-    "hunyuan-video-t2v-720p/vae/config.json",
-    "hunyuan-video-t2v-720p/vae/pytorch_model.pt",
-
-    # # 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():
-    logs = []
-    os.makedirs(os.path.join(WEIGHTS_DIR, "ckpts"), exist_ok=True)
+# Download everything from these folders
+FULL_DIRS = [
+    "hunyuan-video-t2v-720p/vae",
+    #"llava_llama_image",
+    "text_encoder_2",
+    "whisper-tiny",
+    "det_align",
+]
+
+
+def list_ckpt_files():
+    """Return a list of files to download: specific files + all from some folders."""
+    try:
+        all_files = list_repo_files(MODEL_REPO)
+    except Exception as e:
+        print(f"❌ Failed to list files from {MODEL_REPO}: {e}")
+        return []
+
+    files_to_download = ESSENTIAL_PATHS.copy()
+    for path in all_files:
+        if any(path.startswith(folder + "/") for folder in FULL_DIRS):
+            files_to_download.append(path)
 
-    for path in ESSENTIAL_PATHS:
-        local_path = os.path.join(WEIGHTS_DIR, "ckpts", path)
+    return files_to_download
+
+
+def download_ckpts(files):
+    """Download selected files to local directory, preserving structure."""
+    for file_path in files:
+        local_path = os.path.join(WEIGHTS_DIR, file_path)
         if os.path.exists(local_path):
-            logs.append(f"✅ Exists: {path}")
+            print(f"✅ Already exists: {local_path}")
             continue
 
-        os.makedirs(os.path.dirname(local_path), exist_ok=True)
+        print(f"⬇️  Downloading: {file_path}")
         try:
-            logs.append(f"⬇️ Downloading: {path}")
             hf_hub_download(
                 repo_id=MODEL_REPO,
-                filename="ckpts/" + path,
+                filename=file_path,
                 local_dir=WEIGHTS_DIR,
                 local_dir_use_symlinks=False,
             )
         except Exception as e:
-            logs.append(f"❌ Failed: {path} - {str(e)}")
+            print(f"❌ Failed to download {file_path}: {e}")
 
-    return "\n".join(logs)
 
 def run_sample_gpu_poor():
     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}"
+        print(f"❌ Missing checkpoint: {ckpt_fp8}")
+        return
 
     cmd = [
         "python3", "hymm_sp/sample_gpu_poor.py",
@@ -95,22 +90,24 @@ def run_sample_gpu_poor():
     env["CPU_OFFLOAD"] = "1"
     env["CUDA_VISIBLE_DEVICES"] = "0"
 
+    print("🚀 Running sample_gpu_poor.py...")
     result = subprocess.run(cmd, env=env, capture_output=True, text=True)
+
     if result.returncode != 0:
-        return f"❌ sample_gpu_poor.py failed:\n{result.stderr}"
-    return f"✅ sample_gpu_poor.py finished:\n{result.stdout}"
+        print(f"❌ sample_gpu_poor.py failed:\n{result.stderr}")
+    else:
+        print("✅ sample_gpu_poor.py ran successfully!")
+
+
+def main():
+    files = list_ckpt_files()
+    if not files:
+        print("❌ No checkpoint files found in repo under ckpts folder.")
+        sys.exit(1)
 
-def download_and_run():
-    log1 = download_ckpts()
-    log2 = run_sample_gpu_poor()
-    return f"{log1}\n\n---\n\n{log2}"
+    download_ckpts(files)
+    run_sample_gpu_poor()
 
-# Gradio UI
-with gr.Blocks() as demo:
-    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)
 
 if __name__ == "__main__":
-    demo.launch(share=False)
+    main()