Update app.py

app.py

@@ -1,63 +1,101 @@
 import os
-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.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} ...")
+import sys
+import subprocess
+import gradio as gr
+from huggingface_hub import list_repo_files, 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"
+
+def list_ckpt_files():
+    all_files = list_repo_files(repo_id=MODEL_REPO, repo_type="model", revision="main")
+
+    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
+
+def download_ckpts():
+    os.makedirs(WEIGHTS_DIR, exist_ok=True)
+    logs = []
+
+    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}")
+            continue
+
+        os.makedirs(os.path.dirname(target_path), exist_ok=True)
+        logs.append(f"⬇️ Downloading: {filepath}")
         try:
             hf_hub_download(
-                repo_id=REPO_ID,
-                filename=source_path,
-                repo_type="model",
-                local_dir=local_dir,
+                repo_id=MODEL_REPO,
+                filename=filepath,
+                local_dir=os.path.dirname(target_path),
                 local_dir_use_symlinks=False
             )
         except Exception as e:
-            print(f"❌ Failed to download {source_path}: {e}")
+            logs.append(f"❌ Failed to download {filepath}: {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}"
+
+    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"
+
+    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}"
+
+def download_and_run():
+    log1 = download_ckpts()
+    log2 = run_sample_gpu_poor()
+    return f"{log1}\n\n---\n\n{log2}"
+
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## 📦 Download All Checkpoints (Except Only One File from Transformers)")
+    output = gr.Textbox(lines=30, label="Logs")
+    button = gr.Button("🚀 Download + Run")
+    button.click(fn=download_and_run, outputs=output)
 
 if __name__ == "__main__":
-    download_files()
-    print("\n✅ All selected model weights downloaded to:")
-    print(f"{os.path.abspath(LOCAL_BASE)}")
+    demo.launch(share=False)