Update handler.py

handler.py

@@ -13,6 +13,13 @@ from hyvideo.constants import NEGATIVE_PROMPT
 # Configure logger
 logger.add("handler_debug.log", rotation="500 MB")
 
+DEFAULT_RESOLUTION = "720p"
+DEFAULT_WIDTH = 1280
+DEFAULT_HEIGHT = 720
+DEFAULT_NB_FRAMES = (4 * 30) + 1 # or 129 (note: hunyan requires an extra +1 frame)
+DEFAULT_NB_STEPS = 22 # or 50
+DEFAULT_FPS = 24
+
 def setup_vae_path(vae_path: Path) -> Path:
     """Create a temporary directory with correctly named VAE config file"""
     tmp_vae_dir = Path("/tmp/vae")
@@ -43,7 +50,7 @@ def get_default_args():
     
     # Model configuration
     parser.add_argument("--model", type=str, default="HYVideo-T/2-cfgdistill")
-    parser.add_argument("--model-resolution", type=str, default="720p", choices=["540p", "720p"])
+    parser.add_argument("--model-resolution", type=str, default=DEFAULT_RESOLUTION, choices=["540p", "720p"])
     parser.add_argument("--latent-channels", type=int, default=16)
     parser.add_argument("--precision", type=str, default="bf16", choices=["bf16", "fp32", "fp16"])
     parser.add_argument("--rope-theta", type=int, default=256)
@@ -93,7 +100,7 @@ def get_default_args():
     # Hardware settings
     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("--infer-steps", type=int, default=DEFAULT_NB_STEPS)
     parser.add_argument("--disable-autocast", action="store_true")
     
     # Output settings
@@ -103,8 +110,8 @@ def get_default_args():
     
     # Generation settings
     parser.add_argument("--num-videos", type=int, default=1)
-    parser.add_argument("--video-size", nargs="+", type=int, default=[720, 1280])
-    parser.add_argument("--video-length", type=int, default=129)
+    parser.add_argument("--video-size", nargs="+", type=int, default=[DEFAULT_HEIGHT, DEFAULT_WIDTH])
+    parser.add_argument("--video-length", type=int, default=DEFAULT_NB_FRAMES)
     parser.add_argument("--prompt", type=str, default=None)
     parser.add_argument("--seed-type", type=str, default="auto", choices=["file", "random", "fixed", "auto"])
     parser.add_argument("--seed", type=int, default=None)
@@ -214,14 +221,14 @@ class EndpointHandler:
             raise ValueError("No prompt provided in the 'inputs' field")
             
         # Parse resolution 
-        resolution = data.pop("resolution", "1280x720")
+        resolution = data.pop("resolution", f"{DEFAULT_WIDTH}x{DEFAULT_HEIGHT}")
         width, height = map(int, resolution.split("x"))
         
         # Get other parameters with defaults
-        video_length = int(data.pop("video_length", 129))
+        video_length = int(data.pop("video_length", DEFAULT_NB_FRAMES))
         seed = data.pop("seed", -1)
         seed = None if seed == -1 else int(seed)
-        num_inference_steps = int(data.pop("num_inference_steps", 50))
+        num_inference_steps = int(data.pop("num_inference_steps", DEFAULT_NB_STEPS))
         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))
@@ -253,7 +260,7 @@ class EndpointHandler:
             
             # Save to temporary file
             temp_path = "/tmp/temp_video.mp4"
-            save_videos_grid(sample, temp_path, fps=24)
+            save_videos_grid(sample, temp_path, fps=DEFAULT_FPS)
             
             # Read video file and convert to base64
             with open(temp_path, "rb") as f:
@@ -261,16 +268,15 @@ class EndpointHandler:
             import base64
             video_base64 = base64.b64encode(video_bytes).decode()
             
+            # Add MP4 data URI prefix
+            video_data_uri = f"data:video/mp4;base64,{video_base64}"
+            
             # Cleanup
             os.remove(temp_path)
             
             logger.info("Successfully generated and encoded video")
             
-            return {
-                "video_base64": video_base64,
-                "seed": outputs['seeds'][0],
-                "prompt": outputs['prompts'][0]
-            }
+            return video_data_uri
             
         except Exception as e:
             logger.error(f"Error during video generation: {str(e)}")