Upload handler.py

handler.py

@@ -3,87 +3,88 @@ import os
 from pathlib import Path
 import time
 from datetime import datetime
-import torch
-import base64
-from io import BytesIO
-
+import argparse
 from hyvideo.utils.file_utils import save_videos_grid
-from hyvideo.config import parse_args 
 from hyvideo.inference import HunyuanVideoSampler
+from hyvideo.config import parse_args
+from hyvideo.constants import NEGATIVE_PROMPT
+
+def get_default_args():
+    """Create default arguments instead of parsing from command line"""
+    parser = argparse.ArgumentParser()
+    
+    # Add all the arguments that were in the original parser
+    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("--precision", type=str, default="bf16", choices=["bf16", "fp32", "fp16"])
+    parser.add_argument("--batch-size", type=int, default=1)
+    parser.add_argument("--infer-steps", type=int, default=50)
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--save-path", type=str, default="outputs")
+    parser.add_argument("--video-length", type=int, default=129)  # 5 seconds
+    
+    # Parse with empty args list to avoid reading sys.argv
+    args = parser.parse_args([])
+    return args
 
 class EndpointHandler:
     def __init__(self, path: str = ""):
-        """Initialize the handler with the model path.
+        """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
         
-        Args:
-            path: Path to the model weights directory
-        """
-        self.args = parse_args()
+        # Initialize model
         models_root_path = Path(path)
         if not models_root_path.exists():
             raise ValueError(f"`models_root` not exists: {models_root_path}")
         
-        # Initialize model
         self.model = HunyuanVideoSampler.from_pretrained(models_root_path, args=self.args)
         
-        # Default parameters
-        self.default_params = {
-            "num_inference_steps": 50,
-            "guidance_scale": 1.0,
-            "flow_shift": 7.0,
-            "embedded_guidance_scale": 6.0,
-            "video_length": 129,  # 5s
-            "resolution": "1280x720"
-        }
-
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        """Process the input data and generate video.
+        """Process a single request
         
         Args:
-            data: Dictionary containing the input parameters
-                Required:
-                    - inputs (str): The prompt text
-                Optional:
-                    - resolution (str): Video resolution like "1280x720"
-                    - video_length (int): Number of frames
-                    - seed (int): Random seed (-1 for random)
-                    - num_inference_steps (int): Number of inference steps
-                    - guidance_scale (float): Guidance scale value
-                    - flow_shift (float): Flow shift value
-                    - embedded_guidance_scale (float): Embedded guidance scale value
+            data: Dictionary containing:
+                - inputs (str): The prompt text
+                - 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  
+                - embedded_guidance_scale (float, optional): Embedded guidance scale
         
         Returns:
-            Dictionary containing the base64 encoded video
+            Dictionary containing the generated video as base64 string
         """
-        # Get prompt
+        # Get inputs from request data
         prompt = data.pop("inputs", None)
         if prompt is None:
             raise ValueError("No prompt provided in the 'inputs' field")
-
-        # Get optional parameters with defaults
-        resolution = data.pop("resolution", self.default_params["resolution"])
-        video_length = int(data.pop("video_length", self.default_params["video_length"]))
-        seed = int(data.pop("seed", -1))
-        num_inference_steps = int(data.pop("num_inference_steps", self.default_params["num_inference_steps"]))
-        guidance_scale = float(data.pop("guidance_scale", self.default_params["guidance_scale"]))
-        flow_shift = float(data.pop("flow_shift", self.default_params["flow_shift"]))
-        embedded_guidance_scale = float(data.pop("embedded_guidance_scale", self.default_params["embedded_guidance_scale"]))
-
-        # Process resolution
-        width, height = resolution.split("x")
-        width, height = int(width), int(height)
-
-        # Set seed
-        seed = None if seed == -1 else seed
-
-        # Generate video
+            
+        # Parse resolution
+        resolution = data.pop("resolution", "1280x720")
+        width, height = map(int, resolution.split("x"))
+        
+        # Get other parameters with defaults
+        video_length = int(data.pop("video_length", 129))
+        seed = data.pop("seed", -1)
+        seed = None if seed == -1 else int(seed)
+        num_inference_steps = int(data.pop("num_inference_steps", 50))
+        guidance_scale = float(data.pop("guidance_scale", 1.0))
+        flow_shift = float(data.pop("flow_shift", 7.0))
+        embedded_guidance_scale = float(data.pop("embedded_guidance_scale", 6.0))
+        
+        # Run inference
         outputs = self.model.predict(
             prompt=prompt,
             height=height,
             width=width,
             video_length=video_length,
             seed=seed,
-            negative_prompt="",  # not applicable in inference
+            negative_prompt="",
             infer_steps=num_inference_steps,
             guidance_scale=guidance_scale,
             num_videos_per_prompt=1,
@@ -91,27 +92,24 @@ class EndpointHandler:
             batch_size=1,
             embedded_guidance_scale=embedded_guidance_scale
         )
-
-        # Process output video
+        
+        # Get the video tensor
         samples = outputs['samples']
         sample = samples[0].unsqueeze(0)
-
-        # Save video to temporary file
-        temp_dir = "/tmp/video_output"
-        os.makedirs(temp_dir, exist_ok=True)
         
-        time_flag = datetime.fromtimestamp(time.time()).strftime("%Y-%m-%d-%H:%M:%S")
-        video_path = f"{temp_dir}/{time_flag}_seed{outputs['seeds'][0]}.mp4"
-        save_videos_grid(sample, video_path, fps=24)
-
+        # Save to temporary file
+        temp_path = "/tmp/temp_video.mp4"
+        save_videos_grid(sample, temp_path, fps=24)
+        
         # Read video file and convert to base64
-        with open(video_path, "rb") as f:
+        with open(temp_path, "rb") as f:
             video_bytes = f.read()
+        import base64
         video_base64 = base64.b64encode(video_bytes).decode()
-
-        # Clean up
-        os.remove(video_path)
-
+        
+        # Cleanup
+        os.remove(temp_path)
+        
         return {
             "video_base64": video_base64,
             "seed": outputs['seeds'][0],