Create handler.py

handler.py

@@ -0,0 +1,198 @@
+from dataclasses import dataclass
+from typing import Dict, Any, Optional
+import base64
+import logging
+import random
+import torch
+from diffusers import HunyuanVideoPipeline
+from varnish import Varnish
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+@dataclass
+class GenerationConfig:
+    """Configuration for video generation"""
+    # Content settings
+    prompt: str
+    negative_prompt: str = ""
+    
+    # Model settings
+    num_frames: int = 49  # Should be 4k + 1 format
+    height: int = 320
+    width: int = 576
+    num_inference_steps: int = 50
+    guidance_scale: float = 7.0
+    
+    # Reproducibility
+    seed: int = -1
+    
+    # Varnish post-processing settings
+    fps: int = 30
+    double_num_frames: bool = False
+    super_resolution: bool = False
+    grain_amount: float = 0.0
+    quality: int = 18  # CRF scale (0-51, lower is better)
+    
+    # Audio settings
+    enable_audio: bool = False
+    audio_prompt: str = ""
+    audio_negative_prompt: str = "voices, voice, talking, speaking, speech"
+
+    def validate_and_adjust(self) -> 'GenerationConfig':
+        """Validate and adjust parameters"""
+        # Ensure num_frames follows 4k + 1 format
+        k = (self.num_frames - 1) // 4
+        self.num_frames = (k * 4) + 1
+        
+        # Set random seed if not specified
+        if self.seed == -1:
+            self.seed = random.randint(0, 2**32 - 1)
+            
+        return self
+
+class EndpointHandler:
+    """Handles video generation requests using HunyuanVideo and Varnish"""
+    
+    def __init__(self, path: str = ""):
+        """Initialize handler with models
+        
+        Args:
+            path: Path to model weights
+        """
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+        # Initialize HunyuanVideo pipeline
+        self.pipeline = HunyuanVideoPipeline.from_pretrained(
+            path,
+            torch_dtype=torch.float16,
+        ).to(self.device)
+        
+        # Initialize text encoders in float16
+        self.pipeline.text_encoder = self.pipeline.text_encoder.half()
+        self.pipeline.text_encoder_2 = self.pipeline.text_encoder_2.half()
+        
+        # Initialize transformer in bfloat16 
+        self.pipeline.transformer = self.pipeline.transformer.to(torch.bfloat16)
+        
+        # Initialize VAE in float16
+        self.pipeline.vae = self.pipeline.vae.half()
+        
+        # Initialize Varnish for post-processing
+        self.varnish = Varnish(
+            device=self.device,
+            model_base_dir="/repository/varnish"
+        )
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Process video generation requests
+        
+        Args:
+            data: Request data containing:
+                - inputs (str): Prompt for video generation
+                - parameters (dict): Generation parameters
+                
+        Returns:
+            Dictionary containing:
+                - video: Base64 encoded MP4 data URI
+                - content-type: MIME type
+                - metadata: Generation metadata
+        """
+        # Extract inputs
+        inputs = data.pop("inputs", data)
+        if isinstance(inputs, dict):
+            prompt = inputs.get("prompt", "")
+        else:
+            prompt = inputs
+            
+        params = data.get("parameters", {})
+        
+        # Create and validate config
+        config = GenerationConfig(
+            prompt=prompt,
+            negative_prompt=params.get("negative_prompt", ""),
+            num_frames=params.get("num_frames", 49),
+            height=params.get("height", 320),
+            width=params.get("width", 576),
+            num_inference_steps=params.get("num_inference_steps", 50),
+            guidance_scale=params.get("guidance_scale", 7.0),
+            seed=params.get("seed", -1),
+            fps=params.get("fps", 30),
+            double_num_frames=params.get("double_num_frames", False),
+            super_resolution=params.get("super_resolution", False),
+            grain_amount=params.get("grain_amount", 0.0),
+            quality=params.get("quality", 18),
+            enable_audio=params.get("enable_audio", False),
+            audio_prompt=params.get("audio_prompt", ""),
+            audio_negative_prompt=params.get("audio_negative_prompt", "voices, voice, talking, speaking, speech"),
+        ).validate_and_adjust()
+
+        try:
+            # Set random seeds
+            if config.seed != -1:
+                torch.manual_seed(config.seed)
+                random.seed(config.seed)
+                generator = torch.Generator(device=self.device).manual_seed(config.seed)
+            else:
+                generator = None
+
+            # Generate video frames
+            with torch.inference_mode():
+                output = self.pipeline(
+                    prompt=config.prompt,
+                    negative_prompt=config.negative_prompt,
+                    num_frames=config.num_frames,
+                    height=config.height,
+                    width=config.width,
+                    num_inference_steps=config.num_inference_steps,
+                    guidance_scale=config.guidance_scale,
+                    generator=generator,
+                    output_type="pt",
+                ).frames
+
+                # Process with Varnish
+                import asyncio
+                try:
+                    loop = asyncio.get_event_loop()
+                except RuntimeError:
+                    loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(loop)
+
+                result = loop.run_until_complete(
+                    self.varnish(
+                        input_data=output,
+                        fps=config.fps,
+                        double_num_frames=config.double_num_frames,
+                        super_resolution=config.super_resolution,
+                        grain_amount=config.grain_amount,
+                        enable_audio=config.enable_audio,
+                        audio_prompt=config.audio_prompt,
+                        audio_negative_prompt=config.audio_negative_prompt,
+                    )
+                )
+
+                # Get video data URI
+                video_uri = loop.run_until_complete(
+                    result.write(
+                        type="data-uri",
+                        quality=config.quality
+                    )
+                )
+
+                return {
+                    "video": video_uri,
+                    "content-type": "video/mp4",
+                    "metadata": {
+                        "width": result.metadata.width,
+                        "height": result.metadata.height,
+                        "num_frames": result.metadata.frame_count,
+                        "fps": result.metadata.fps,
+                        "duration": result.metadata.duration,
+                        "seed": config.seed,
+                    }
+                }
+
+        except Exception as e:
+            logger.error(f"Error generating video: {str(e)}")
+            raise RuntimeError(f"Failed to generate video: {str(e)}")