Update handler.py

handler.py

@@ -15,52 +15,45 @@ def get_default_args():
     # Model configuration
     parser.add_argument("--model", type=str, default="HYVideo-T/2")
     parser.add_argument("--model-resolution", type=str, default="720p", choices=["540p", "720p"])
-    parser.add_argument("--latent-channels", type=int, default=4)
+    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=float, default=10000)
+    parser.add_argument("--rope-theta", type=int, default=256)  # Changed to match original config
     
     # VAE settings
     parser.add_argument("--vae", type=str, default="884-16c-hy")
-    parser.add_argument("--vae-precision", type=str, default="bf16", choices=["bf16", "fp32", "fp16"])
-    parser.add_argument("--vae-tiling", action="store_true")
+    parser.add_argument("--vae-precision", type=str, default="fp16")  # Changed to fp16 to match docs
+    parser.add_argument("--vae-tiling", action="store_true", default=True)  # Set default to True
     
     # Text encoder settings
-    parser.add_argument("--text-encoder", type=str, default="clipL", choices=["clipL", "llm"])
-    parser.add_argument("--text-encoder-precision", type=str, default="bf16", choices=["bf16", "fp32", "fp16"])
-    parser.add_argument("--text-states-dim", type=int, default=1024)
-    parser.add_argument("--text-len", type=int, default=77)
-    parser.add_argument("--tokenizer", type=str, default="clipL", choices=["clipL", "llm"])
+    parser.add_argument("--text-encoder", type=str, default="llm")  # Changed to llm
+    parser.add_argument("--text-encoder-precision", type=str, default="fp16")  # Changed to fp16
+    parser.add_argument("--text-states-dim", type=int, default=4096)  # Updated to match docs
+    parser.add_argument("--text-len", type=int, default=256)  # Updated to match docs
+    parser.add_argument("--tokenizer", type=str, default="llm")  # Changed to llm
     
     # Prompt template settings
-    parser.add_argument("--prompt-template", type=str, default="dit-llm-encode", 
-                       choices=["dit-llm-encode", "dit-llm-encode-video"])
-    parser.add_argument("--prompt-template-video", type=str, default="dit-llm-encode",
-                       choices=["dit-llm-encode", "dit-llm-encode-video"])
+    parser.add_argument("--prompt-template", type=str, default="dit-llm-encode")
+    parser.add_argument("--prompt-template-video", type=str, default="dit-llm-encode-video")
     
     # Additional text encoder settings
-    parser.add_argument("--hidden-state-skip-layer", type=int, default=0)
+    parser.add_argument("--hidden-state-skip-layer", type=int, default=2)  # Updated to match docs
     parser.add_argument("--apply-final-norm", action="store_true")
-    parser.add_argument("--text-encoder-2", type=str, default="clipL", choices=["clipL", "llm"])
-    parser.add_argument("--text-encoder-precision-2", type=str, default="bf16", choices=["bf16", "fp32", "fp16"])
-    parser.add_argument("--text-states-dim-2", type=int, default=1024)
-    parser.add_argument("--tokenizer-2", type=str, default="clipL", choices=["clipL", "llm"])
-    parser.add_argument("--text-len-2", type=int, default=77)
+    parser.add_argument("--text-encoder-2", type=str, default="clipL")
+    parser.add_argument("--text-encoder-precision-2", type=str, default="fp16")  # Changed to fp16
+    parser.add_argument("--text-states-dim-2", type=int, default=768)  # Updated to match docs
+    parser.add_argument("--tokenizer-2", type=str, default="clipL")
+    parser.add_argument("--text-len-2", type=int, default=77)  # Updated to match docs
     
     # Inference settings
-    parser.add_argument("--denoise-type", type=str, default="v-prediction")
+    parser.add_argument("--denoise-type", type=str, default="flow")  # Changed to flow
     parser.add_argument("--flow-shift", type=float, default=7.0)
-    parser.add_argument("--flow-reverse", action="store_true")
+    parser.add_argument("--flow-reverse", action="store_true", default=False)
     parser.add_argument("--flow-solver", type=str, default="euler")
     parser.add_argument("--use-linear-quadratic-schedule", action="store_true")
-    parser.add_argument("--linear-schedule-end", type=float, default=0.0)
-    
-    # Model paths and weights
-    parser.add_argument("--model-base", type=str, default=None)
-    parser.add_argument("--dit-weight", type=str, default=None)
-    parser.add_argument("--load-key", type=str, default=None)
+    parser.add_argument("--linear-schedule-end", type=int, default=25)  # Updated to match docs
     
     # Hardware settings
-    parser.add_argument("--use-cpu-offload", action="store_true")
+    parser.add_argument("--use-cpu-offload", action="store_true", default=False)
     parser.add_argument("--batch-size", type=int, default=1)
     parser.add_argument("--infer-steps", type=int, default=50)
     parser.add_argument("--disable-autocast", action="store_true")
@@ -72,22 +65,23 @@ def get_default_args():
     
     # Generation settings
     parser.add_argument("--num-videos", type=int, default=1)
-    parser.add_argument("--video-size", nargs="+", type=int, default=None)
+    parser.add_argument("--video-size", nargs="+", type=int, default=[720, 1280])
     parser.add_argument("--video-length", type=int, default=129)
     parser.add_argument("--prompt", type=str, default=None)
-    parser.add_argument("--seed-type", type=str, default="random", choices=["file", "random", "fixed", "auto"])
-    parser.add_argument("--seed", type=int, default=-1)
+    parser.add_argument("--seed-type", type=str, default="auto", choices=["file", "random", "fixed", "auto"])
+    parser.add_argument("--seed", type=int, default=None)
     parser.add_argument("--neg-prompt", type=str, default="")
     parser.add_argument("--cfg-scale", type=float, default=1.0)
     parser.add_argument("--embedded-cfg-scale", type=float, default=6.0)
     parser.add_argument("--reproduce", action="store_true")
     
-    # Additional degrees
-    parser.add_argument("--ulysses-degree", type=float, default=1.0)
-    parser.add_argument("--ring-degree", type=float, default=1.0)
+    # Parallel settings
+    parser.add_argument("--ulysses-degree", type=int, default=1)
+    parser.add_argument("--ring-degree", type=int, default=1)
     
     # Parse with empty args list to avoid reading sys.argv
     args = parser.parse_args([])
+    
     return args
 
 class EndpointHandler:
@@ -95,7 +89,10 @@ class EndpointHandler:
         """Initialize the handler with model path and default config."""
         # Use default args instead of parsing from command line
         self.args = get_default_args()
-        self.args.model_base = path  # Use the provided model path
+        
+        # Set up model paths
+        self.args.model_base = path
+        self.args.dit_weight = str(Path(path) / "hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt")
         
         # Initialize model
         models_root_path = Path(path)
@@ -110,12 +107,12 @@ class EndpointHandler:
         Args:
             data: Dictionary containing:
                 - inputs (str): The prompt text
-                - resolution (str, optional): Video resolution like "1280x720"
+                - resolution (str, optional): Video resolution like "1280x720"  
                 - video_length (int, optional): Number of frames
                 - num_inference_steps (int, optional): Number of inference steps
                 - seed (int, optional): Random seed (-1 for random)
                 - guidance_scale (float, optional): Guidance scale value
-                - flow_shift (float, optional): Flow shift value  
+                - flow_shift (float, optional): Flow shift value
                 - embedded_guidance_scale (float, optional): Embedded guidance scale
         
         Returns:
@@ -126,7 +123,7 @@ class EndpointHandler:
         if prompt is None:
             raise ValueError("No prompt provided in the 'inputs' field")
             
-        # Parse resolution
+        # Parse resolution 
         resolution = data.pop("resolution", "1280x720")
         width, height = map(int, resolution.split("x"))
         
@@ -144,7 +141,7 @@ class EndpointHandler:
             prompt=prompt,
             height=height,
             width=width,
-            video_length=video_length,
+            video_length=video_length, 
             seed=seed,
             negative_prompt="",
             infer_steps=num_inference_steps,
@@ -176,4 +173,4 @@ class EndpointHandler:
             "video_base64": video_base64,
             "seed": outputs['seeds'][0],
             "prompt": outputs['prompts'][0]
-        }
+        }