Update handler.py

handler.py

@@ -1,5 +1,6 @@
 from typing import Dict, Any
 import os
+import shutil
 from pathlib import Path
 import time
 from datetime import datetime
@@ -12,6 +13,30 @@ from hyvideo.constants import NEGATIVE_PROMPT
 # Configure logger
 logger.add("handler_debug.log", rotation="500 MB")
 
+def setup_vae_path(vae_path: Path) -> Path:
+    """Create a temporary directory with correctly named VAE config file"""
+    tmp_vae_dir = Path("/tmp/vae")
+    if tmp_vae_dir.exists():
+        shutil.rmtree(tmp_vae_dir)
+    tmp_vae_dir.mkdir(parents=True)
+    
+    # Copy files to temp directory
+    logger.info(f"Setting up VAE in temporary directory: {tmp_vae_dir}")
+    
+    # Copy and rename config file
+    original_config = vae_path / "hunyuan-video-t2v-720p_vae_config.json"
+    new_config = tmp_vae_dir / "config.json"
+    shutil.copy2(original_config, new_config)
+    logger.info(f"Copied VAE config from {original_config} to {new_config}")
+    
+    # Copy model file
+    original_model = vae_path / "pytorch_model.pt"
+    new_model = tmp_vae_dir / "pytorch_model.pt"
+    shutil.copy2(original_model, new_model)
+    logger.info(f"Copied VAE model from {original_model} to {new_model}")
+    
+    return tmp_vae_dir
+
 def get_default_args():
     """Create default arguments instead of parsing from command line"""
     parser = argparse.ArgumentParser()
@@ -99,7 +124,6 @@ def get_default_args():
 class EndpointHandler:
     def __init__(self, path: str = ""):
         """Initialize the handler with model path and default config."""
-        # Log the initial path
         logger.info(f"Initializing EndpointHandler with path: {path}")
         
         # Use default args instead of parsing from command line
@@ -114,19 +138,28 @@ class EndpointHandler:
         
         # Set paths for model components
         dit_weight_path = Path(path) / "hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt"
-        vae_path = Path(path) / "hunyuan-video-t2v-720p/vae"
+        original_vae_path = Path(path) / "hunyuan-video-t2v-720p/vae"
         
         # Log all critical paths
         logger.info(f"Model base path: {self.args.model_base}")
         logger.info(f"DiT weight path: {dit_weight_path}")
-        logger.info(f"VAE path: {vae_path}")
+        logger.info(f"Original VAE path: {original_vae_path}")
         
         # Verify paths exist
         logger.info("Checking if paths exist:")
         logger.info(f"DiT weight exists: {dit_weight_path.exists()}")
-        logger.info(f"VAE path exists: {vae_path.exists()}")
-        if vae_path.exists():
-            logger.info(f"VAE path contents: {list(vae_path.glob('*'))}")
+        logger.info(f"VAE path exists: {original_vae_path.exists()}")
+        
+        if original_vae_path.exists():
+            logger.info(f"VAE path contents: {list(original_vae_path.glob('*'))}")
+            
+            # Set up VAE in temporary directory with correct file names
+            tmp_vae_path = setup_vae_path(original_vae_path)
+            
+            # Override the VAE path in constants to use our temporary directory
+            from hyvideo.constants import VAE_PATH
+            VAE_PATH["884-16c-hy"] = str(tmp_vae_path)
+            logger.info(f"Updated VAE_PATH to: {VAE_PATH['884-16c-hy']}")
         
         self.args.dit_weight = str(dit_weight_path)
         
@@ -135,11 +168,6 @@ class EndpointHandler:
         if not models_root_path.exists():
             raise ValueError(f"models_root_path does not exist: {models_root_path}")
         
-        # Log directory contents for debugging
-        logger.info("Directory contents:")
-        for item in models_root_path.glob("**/*"):
-            logger.info(f"  {item}")
-        
         try:
             logger.info("Attempting to initialize HunyuanVideoSampler...")
             self.model = HunyuanVideoSampler.from_pretrained(models_root_path, args=self.args)