from typing import Dict, Any, Union, Optional, Tuple import torch from diffusers import LTXPipeline, LTXImageToVideoPipeline from PIL import Image import base64 import io import tempfile import random import numpy as np from moviepy.editor import ImageSequenceClip import os import logging # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) ENABLE_CPU_OFFLOAD = True EXPERIMENTAL_STUFF = False random.seed(0) np.random.seed(0) generator = torch.manual_seed(0) # you can notice we don't use device=cuda, for more info see: # https://huggingface.co/docs/diffusers/v0.16.0/en/using-diffusers/reproducibility#gpu class EndpointHandler: # Default configuration DEFAULT_FPS = 24 DEFAULT_DURATION = 4 # seconds DEFAULT_NUM_FRAMES = (DEFAULT_DURATION * DEFAULT_FPS) + 1 # 97 frames DEFAULT_NUM_STEPS = 25 DEFAULT_WIDTH = 768 DEFAULT_HEIGHT = 512 # Constraints MAX_WIDTH = 1280 MAX_HEIGHT = 720 MAX_FRAMES = 257 def __init__(self, path: str = ""): """Initialize the LTX Video handler with both text-to-video and image-to-video pipelines. Args: path (str): Path to the model weights directory """ if EXPERIMENTAL_STUFF: torch.backends.cuda.matmul.allow_tf32 = True # Load both pipelines with bfloat16 precision as recommended in docs self.text_to_video = LTXPipeline.from_pretrained( path, torch_dtype=torch.bfloat16 ).to("cuda") self.image_to_video = LTXImageToVideoPipeline.from_pretrained( path, torch_dtype=torch.bfloat16 ).to("cuda") if ENABLE_CPU_OFFLOAD: self.text_to_video.enable_model_cpu_offload() self.image_to_video.enable_model_cpu_offload() def _validate_and_adjust_resolution(self, width: int, height: int) -> Tuple[int, int]: """Validate and adjust resolution to meet constraints. Args: width (int): Requested width height (int): Requested height Returns: Tuple[int, int]: Adjusted (width, height) """ # Round to nearest multiple of 32 width = round(width / 32) * 32 height = round(height / 32) * 32 # Enforce maximum dimensions width = min(width, self.MAX_WIDTH) height = min(height, self.MAX_HEIGHT) # Enforce minimum dimensions width = max(width, 32) height = max(height, 32) return width, height def _validate_and_adjust_frames(self, num_frames: Optional[int] = None, fps: Optional[int] = None) -> Tuple[int, int]: """Validate and adjust frame count and FPS to meet constraints. Args: num_frames (Optional[int]): Requested number of frames fps (Optional[int]): Requested frames per second Returns: Tuple[int, int]: Adjusted (num_frames, fps) """ # Use defaults if not provided fps = fps or self.DEFAULT_FPS num_frames = num_frames or self.DEFAULT_NUM_FRAMES # Adjust frames to be in format 8k + 1 k = (num_frames - 1) // 8 num_frames = (k * 8) + 1 # Enforce maximum frame count num_frames = min(num_frames, self.MAX_FRAMES) return num_frames, fps def _create_video_file(self, frames: torch.Tensor, fps: int = DEFAULT_FPS) -> bytes: """Convert frames to an MP4 video file. Args: frames (torch.Tensor): Generated frames tensor fps (int): Frames per second for the output video Returns: bytes: MP4 video file content """ # Log frame information num_frames = frames.shape[1] duration = num_frames / fps logger.info(f"Creating video with {num_frames} frames at {fps} FPS (duration: {duration:.2f} seconds)") # Convert tensor to numpy array video_np = frames.squeeze(0).permute(0, 2, 3, 1).cpu().float().numpy() video_np = (video_np * 255).astype(np.uint8) # Get dimensions _, height, width, _ = video_np.shape logger.info(f"Video dimensions: {width}x{height}") # Create temporary file output_path = tempfile.mktemp(suffix=".mp4") try: # Create video clip and write to file clip = ImageSequenceClip(list(video_np), fps=fps) # potential speed optimizations: # there is a preset= field, to trade encoding speed with file size (but not quality) # values are: ultrafast, superfast, veryfast, faster, fast, medium, slow, slower, veryslow, placebo # # there is a threads= field, by default None, which can be set to 2, 3, 4 etc.. clip.write_videofile(output_path, codec="libx264", audio=False) # Read the video file with open(output_path, "rb") as f: video_content = f.read() return video_content finally: # Cleanup if os.path.exists(output_path): os.remove(output_path) # Clear memory del video_np torch.cuda.empty_cache() def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: """Process the input data and generate video using LTX. Args: data (Dict[str, Any]): Input data containing: - prompt (str): Text description for video generation - image (Optional[str]): Base64 encoded image for image-to-video generation - width (Optional[int]): Video width (default: 768) - height (Optional[int]): Video height (default: 512) - num_frames (Optional[int]): Number of frames (default: 97) - fps (Optional[int]): Frames per second (default: 24) - num_inference_steps (Optional[int]): Number of inference steps (default: 25) - guidance_scale (Optional[float]): Guidance scale (default: 7.5) Returns: Dict[str, Any]: Dictionary containing: - video: video encoded in Base64 (h.264 MP4 video). This is a data-uri (prefixed with "data:"). - content-type: MIME type of the video (right now always "video/mp4") - metadata: Dictionary with actual values used for generation """ # Get inputs from request data prompt = data.get("inputs", None) if not prompt: raise ValueError("No prompt provided in the 'inputs' field") # Get and validate resolution width = data.get("width", self.DEFAULT_WIDTH) height = data.get("height", self.DEFAULT_HEIGHT) width, height = self._validate_and_adjust_resolution(width, height) # Get and validate frames and FPS num_frames = data.get("num_frames", self.DEFAULT_NUM_FRAMES) fps = data.get("fps", self.DEFAULT_FPS) num_frames, fps = self._validate_and_adjust_frames(num_frames, fps) # Get other parameters with defaults guidance_scale = data.get("guidance_scale", 7.5) num_inference_steps = data.get("num_inference_steps", self.DEFAULT_NUM_STEPS) seed = data.get("seed", -1) seed = random.randint(0, 2**32 - 1) if seed == -1 else int(seed) logger.info(f"Generating video with prompt: '{prompt}'") logger.info(f"Video params: size={width}x{height}, num_frames={num_frames}, fps={fps}") logger.info(f"Generation params: seed={seed}, guidance_scale={guidance_scale}, num_inference_steps={num_inference_steps}") try: with torch.no_grad(): random.seed(seed) np.random.seed(seed) generator.manual_seed(seed) generation_kwargs = { "prompt": prompt, "height": height, "width": width, "num_frames": num_frames, "guidance_scale": guidance_scale, "num_inference_steps": num_inference_steps, "output_type": "pt", "generator": generator } # Check if image is provided for image-to-video generation image_data = data.get("image") if image_data: if image_data.startswith('data:'): image_data = image_data.split(',', 1)[1] image_bytes = base64.b64decode(image_data) image = Image.open(io.BytesIO(image_bytes)).convert("RGB") logger.info("Using image-to-video generation mode") generation_kwargs["image"] = image output = self.image_to_video(**generation_kwargs).frames else: logger.info("Using text-to-video generation mode") output = self.text_to_video(**generation_kwargs).frames # Convert frames to video file video_content = self._create_video_file(output, fps=fps) # Encode video to base64 video_base64 = base64.b64encode(video_content).decode('utf-8') content_type = "video/mp4" # Add MP4 data URI prefix video_data_uri = f"data:{content_type};base64,{video_base64}" return { "video": video_data_uri, "content-type": content_type, "metadata": { "width": width, "height": height, "num_frames": num_frames, "fps": fps, "duration": num_frames / fps, "num_inference_steps": num_inference_steps, "seed": seed } } except Exception as e: logger.error(f"Error generating video: {str(e)}") raise RuntimeError(f"Error generating video: {str(e)}")