Update app.py

app.py

@@ -1,8 +1,6 @@
 import os
 import sys
 import subprocess
-import argparse
-from pathlib import Path
 import torch
 import datetime
 import numpy as np
@@ -13,15 +11,14 @@ import spaces
 # --- Part 1: Auto-Setup (Clone Repo & Download Weights) ---
 
 REPO_URL = "https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5.git"
-REPO_DIR = "HunyuanVideo-1.5"
-MODEL_DIR = "ckpts"
+REPO_DIR = os.path.abspath("HunyuanVideo-1.5")
+# Use Absolute Path to ensure the loader finds the folder
+MODEL_DIR = os.path.abspath("ckpts") 
 HF_REPO_ID = "tencent/HunyuanVideo"
 
 # Configuration
 TRANSFORMER_VERSION = "480p_i2v_distilled"
 DTYPE = torch.bfloat16
-# Set to False if you have >40GB VRAM and want everything on GPU constantly.
-# Set to True (Default) to allow running on 16GB-24GB cards via CPU offloading.
 ENABLE_OFFLOADING = True 
 
 def setup_environment():
@@ -31,19 +28,20 @@ def setup_environment():
     
     # 1. Clone Repository
     if not os.path.exists(REPO_DIR):
-        print(f"Cloning repository from {REPO_URL}...")
-        subprocess.run(["git", "clone", REPO_URL], check=True)
+        print(f"Cloning repository to {REPO_DIR}...")
+        subprocess.run(["git", "clone", REPO_URL, REPO_DIR], check=True)
     else:
-        print(f"Repository {REPO_DIR} exists.")
+        print(f"Repository exists at {REPO_DIR}")
 
     # 2. Add Repo to Python Path
-    repo_path = os.path.abspath(REPO_DIR)
-    if repo_path not in sys.path:
-        sys.path.insert(0, repo_path)
+    if REPO_DIR not in sys.path:
+        sys.path.insert(0, REPO_DIR)
 
     # 3. Download Weights
-    if not os.path.exists(MODEL_DIR) or not os.listdir(MODEL_DIR):
-        print(f"Downloading weights from {HF_REPO_ID} to {MODEL_DIR}...")
+    # Check if key folders exist to verify download
+    transformer_path = os.path.join(MODEL_DIR, "transformer", TRANSFORMER_VERSION)
+    if not os.path.exists(transformer_path):
+        print(f"Downloading weights to {MODEL_DIR}...")
         try:
             from huggingface_hub import snapshot_download
             allow_patterns = [
@@ -54,11 +52,18 @@ def setup_environment():
                 "scheduler/*",
                 "tokenizer/*"
             ]
-            snapshot_download(repo_id=HF_REPO_ID, local_dir=MODEL_DIR, allow_patterns=allow_patterns)
+            snapshot_download(
+                repo_id=HF_REPO_ID, 
+                local_dir=MODEL_DIR, 
+                allow_patterns=allow_patterns
+            )
             print("Download complete.")
         except Exception as e:
             print(f"Error downloading weights: {e}")
             sys.exit(1)
+    else:
+        print(f"Weights found in {MODEL_DIR}")
+        
     print("Environment Ready.")
     print("=" * 50)
 
@@ -70,12 +75,12 @@ setup_environment()
 # Set Env Vars for HyVideo
 if 'PYTORCH_CUDA_ALLOC_CONF' not in os.environ:
     os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
+# Even for single GPU, HyVideo code expects these env vars to be set
 os.environ['RANK'] = '0'
 os.environ['WORLD_SIZE'] = '1'
 
 try:
     from hyvideo.pipelines.hunyuan_video_pipeline import HunyuanVideo_1_5_Pipeline
-    from hyvideo.commons.parallel_states import initialize_parallel_state
     from hyvideo.commons.infer_state import initialize_infer_state
 except ImportError as e:
     print(f"CRITICAL ERROR: Could not import hyvideo modules. {e}")
@@ -85,17 +90,14 @@ import gradio as gr
 
 # --- Part 3: Model Initialization (Pre-Load) ---
 
-# Initialize Distributed/Infer States
-#parallel_dims = initialize_parallel_state(sp=1)
-#if torch.cuda.is_available():
-#    torch.cuda.set_device(0)
-
+# Mock args for inference configuration (required by internal logic)
 class ArgsNamespace:
     def __init__(self):
         self.use_sageattn = False
         self.sage_blocks_range = "0-53"
         self.enable_torch_compile = False
 
+# Initialize internal state mock
 initialize_infer_state(ArgsNamespace())
 
 # Global Pipeline Variable
@@ -104,7 +106,13 @@ pipe = None
 def pre_load_model():
     """Loads the model into memory/GPU before UI launch."""
     global pipe
-    print(f"⏳ Initializing Pipeline ({TRANSFORMER_VERSION})... this may take a moment...")
+    
+    # Double check path exists
+    if not os.path.isdir(MODEL_DIR):
+        print(f"❌ Error: Model directory not found at {MODEL_DIR}")
+        sys.exit(1)
+
+    print(f"⏳ Initializing Pipeline ({TRANSFORMER_VERSION}) from {MODEL_DIR}...")
     
     try:
         pipe = HunyuanVideo_1_5_Pipeline.create_pipeline(
@@ -115,12 +123,10 @@ def pre_load_model():
             transformer_dtype=DTYPE,
         )
         print("✅ Model loaded successfully!")
-        if not ENABLE_OFFLOADING:
-            print("   Model is fully resident on GPU.")
-        else:
-            print("   Model loaded with CPU Offloading enabled (optimizes VRAM usage).")
     except Exception as e:
         print(f"❌ Failed to load model: {e}")
+        import traceback
+        traceback.print_exc()
         sys.exit(1)
 
 def save_video_tensor(video_tensor, path, fps=24):
@@ -176,7 +182,6 @@ def generate(input_image, prompt, length, steps, shift, seed, guidance):
 def create_ui():
     with gr.Blocks(title="HunyuanVideo 1.5 I2V") as demo:
         gr.Markdown(f"### 🎬 HunyuanVideo 1.5 I2V ({TRANSFORMER_VERSION})")
-        gr.Markdown("Model is pre-loaded. Ready to generate.")
         
         with gr.Row():
             with gr.Column():