Update handler.py

handler.py

@@ -18,6 +18,7 @@ def get_default_args():
     parser.add_argument("--latent-channels", type=int, default=16)  # Changed from 4 to match VAE
     parser.add_argument("--precision", type=str, default="bf16", choices=["bf16", "fp32", "fp16"])
     parser.add_argument("--rope-theta", type=int, default=256)  # Changed to match original config
+    parser.add_argument("--load-key", type=str, default="module", help="Key to load the model states. 'module' for the main model, 'ema' for the EMA model.")
     
     # VAE settings
     parser.add_argument("--vae", type=str, default="884-16c-hy")
@@ -100,7 +101,8 @@ class EndpointHandler:
             raise ValueError(f"`models_root` not exists: {models_root_path}")
         
         self.model = HunyuanVideoSampler.from_pretrained(models_root_path, args=self.args)
-        
+
+    
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         """Process a single request