from dataclasses import dataclass from pathlib import Path import pathlib from typing import Dict, Any, Optional, Tuple import asyncio import base64 import io import pprint import logging import random import traceback import os import numpy as np import torch from diffusers import LTXPipeline, LTXImageToVideoPipeline from PIL import Image from varnish import Varnish # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Constraints MAX_LARGE_SIDE = 1280 MAX_SMALL_SIDE = 768 # should be 720 but it must be divisible by 32 MAX_FRAMES = (8 * 21) + 1 # visual glitches appear after about 169 frames, so we cap it # this is only a temporary solution (famous last words) def apply_dirty_hack_to_patch_file_extensions_and_bypass_filter(directory): """ Recursively rename all '.wut' files to '.pth' in the given directory Args: directory (str): Path to the directory to process """ # Convert the directory path to absolute path directory = os.path.abspath(directory) # Walk through directory and its subdirectories for root, _, files in os.walk(directory): for filename in files: if filename.endswith('.wut'): # Get full path of the file old_path = os.path.join(root, filename) # Create new filename by replacing the extension new_filename = filename.replace('.wut', '.pth') new_path = os.path.join(root, new_filename) try: os.rename(old_path, new_path) print(f"Renamed: {old_path} -> {new_path}") except OSError as e: print(f"Error renaming {old_path}: {e}") def print_directory_structure(startpath): """Print the directory structure starting from the given path.""" for root, dirs, files in os.walk(startpath): level = root.replace(startpath, '').count(os.sep) indent = ' ' * 4 * level logger.info(f"{indent}{os.path.basename(root)}/") subindent = ' ' * 4 * (level + 1) for f in files: logger.info(f"{subindent}{f}") logger.info("💡 Applying a dirty hack (patch ""/repository"" to fix file extensions):") apply_dirty_hack_to_patch_file_extensions_and_bypass_filter("/repository") #logger.info("💡 Printing directory structure of ""/repository"":") #print_directory_structure("/repository") def process_input_image(image_data: str, target_width: int, target_height: int) -> Image.Image: """ Process input image from base64, resize and crop to target dimensions Args: image_data: Base64 encoded image data target_width: Desired width target_height: Desired height Returns: Processed PIL Image """ try: # Handle data URI format if image_data.startswith('data:'): image_data = image_data.split(',', 1)[1] # Decode base64 image_bytes = base64.b64decode(image_data) image = Image.open(io.BytesIO(image_bytes)) # Convert to RGB if necessary if image.mode not in ('RGB', 'RGBA'): image = image.convert('RGB') elif image.mode == 'RGBA': # Handle transparency by compositing on white background background = Image.new('RGB', image.size, (255, 255, 255)) background.paste(image, mask=image.split()[3]) image = background # Calculate target aspect ratio target_aspect = target_width / target_height # Get current dimensions orig_width, orig_height = image.size orig_aspect = orig_width / orig_height # Calculate dimensions for resizing if orig_aspect > target_aspect: # Image is wider than target new_height = target_height new_width = int(target_height * orig_aspect) else: # Image is taller than target new_width = target_width new_height = int(target_width / orig_aspect) # Resize image image = image.resize((new_width, new_height), Image.Resampling.LANCZOS) # Center crop to target dimensions left = (new_width - target_width) // 2 top = (new_height - target_height) // 2 right = left + target_width bottom = top + target_height image = image.crop((left, top, right, bottom)) return image except Exception as e: raise ValueError(f"Failed to process input image: {str(e)}") @dataclass class GenerationConfig: """Configuration for video generation""" # general content settings prompt: str = "" negative_prompt: str = "saturated, highlight, overexposed, highlighted, overlit, shaking, too bright, worst quality, inconsistent motion, blurry, jittery, distorted, cropped, watermarked, watermark, logo, subtitle, subtitles, lowres" # video model settings (will be used during generation of the initial raw video clip) # we use small values to make things a bit faster width: int = 768 height: int = 416 # users may tend to always set this to the max, to get as much useable content as possible (which is MAX_FRAMES ie. 257). # The value must be a multiple of 8, plus 1 frame. # visual glitches appear after about 169 frames, so we don't need more actually num_frames: int = (8 * 14) + 1 # values between 3.0 and 4.0 are nice guidance_scale: float = 3.5 num_inference_steps: int = 50 # reproducible generation settings seed: int = -1 # -1 means random seed # varnish settings (will be used for post-processing after the raw video clip has been generated fps: int = 30 # FPS of the final video (only applied at the the very end, when converting to mp4) double_num_frames: bool = False # if True, the number of frames will be multiplied by 2 using RIFE super_resolution: bool = False # if True, the resolution will be multiplied by 2 using Real_ESRGAN grain_amount: float = 0.0 # be careful, adding film grian can negatively impact video compression # audio settings enable_audio: bool = False # Whether to generate audio audio_prompt: str = "" # Text prompt for audio generation audio_negative_prompt: str = "voices, voice, talking, speaking, speech" # Negative prompt for audio generation def validate_and_adjust(self) -> 'GenerationConfig': """Validate and adjust parameters to meet constraints""" # First check if it's one of our explicitly allowed resolutions if not ((self.width == MAX_LARGE_SIDE and self.height == MAX_SMALL_SIDE) or (self.width == MAX_SMALL_SIDE and self.height == MAX_LARGE_SIDE)): # For other resolutions, ensure total pixels don't exceed max MAX_TOTAL_PIXELS = MAX_SMALL_SIDE * MAX_LARGE_SIDE # or 921600 = 1280 * 720 # If total pixels exceed maximum, scale down proportionally total_pixels = self.width * self.height if total_pixels > MAX_TOTAL_PIXELS: scale = (MAX_TOTAL_PIXELS / total_pixels) ** 0.5 self.width = max(128, min(MAX_LARGE_SIDE, round(self.width * scale / 32) * 32)) self.height = max(128, min(MAX_LARGE_SIDE, round(self.height * scale / 32) * 32)) else: # Round dimensions to nearest multiple of 32 self.width = max(128, min(MAX_LARGE_SIDE, round(self.width / 32) * 32)) self.height = max(128, min(MAX_LARGE_SIDE, round(self.height / 32) * 32)) # Adjust number of frames to be in format 8k + 1 k = (self.num_frames - 1) // 8 self.num_frames = min((k * 8) + 1, MAX_FRAMES) # 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 LTX models and Varnish post-processing""" def __init__(self, model_path: str = ""): """Initialize the handler with LTX models and Varnish Args: model_path: Path to LTX model weights """ # Enable TF32 for potential speedup on Ampere GPUs #torch.backends.cuda.matmul.allow_tf32 = True # Initialize models with bfloat16 precision self.text_to_video = LTXPipeline.from_pretrained( model_path, torch_dtype=torch.bfloat16 ).to("cuda") self.image_to_video = LTXImageToVideoPipeline.from_pretrained( model_path, torch_dtype=torch.bfloat16 ).to("cuda") # Enable CPU offload for memory efficiency #self.text_to_video.enable_model_cpu_offload() #self.image_to_video.enable_model_cpu_offload() # Initialize Varnish for post-processing self.varnish = Varnish( device="cuda" if torch.cuda.is_available() else "cpu", model_base_dir="/repository/varnish", # there is currently a bug with MMAudio and/or torch and/or the weight format and/or version.. # not sure how to fix that.. :/ # # it says: # File "dist-packages/varnish.py", line 152, in __init__ # self._setup_mmaudio() # File "dist-packages/varnish/varnish.py", line 165, in _setup_mmaudio # net.load_weights(torch.load(model.model_path, map_location=self.device, weights_only=False)) # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ # File "dist-packages/torch/serialization.py", line 1384, in load # return _legacy_load( # ^^^^^^^^^^^^^ # File "dist-packages/torch/serialization.py", line 1628, in _legacy_load # magic_number = pickle_module.load(f, **pickle_load_args) # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ # _pickle.UnpicklingError: invalid load key, '<'. enable_mmaudio=False, ) async def process_frames( self, frames: torch.Tensor, config: GenerationConfig ) -> tuple[str, dict]: """Post-process generated frames using Varnish Args: frames: Generated video frames tensor config: Generation configuration Returns: Tuple of (video data URI, metadata dictionary) """ try: # Process video with Varnish result = await self.varnish( input_data=frames, # note: this might contain a certain number of frames eg. 97, which will get doubled if double_num_frames is True fps=config.fps, # this is the FPS of the final output video. This number can be used by Varnish to calculate the duration of a clip ((using frames * factor) / fps etc) double_num_frames=config.double_num_frames, # if True, the number of frames will be multiplied by 2 using RIFE super_resolution=config.super_resolution, # if True, the resolution will be multiplied by 2 using Real_ESRGAN grain_amount=config.grain_amount, enable_audio=config.enable_audio, audio_prompt=config.audio_prompt, audio_negative_prompt=config.audio_negative_prompt, ) # Convert to data URI video_uri = await result.write( type="data-uri", quality=17 ) # Collect metadata 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, } return video_uri, metadata except Exception as e: logger.error(f"Error in process_frames: {str(e)}") raise RuntimeError(f"Failed to process frames: {str(e)}") def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: """Process incoming requests for video generation Args: data: Request data containing: - inputs (dict): Dictionary containing input, which can be either "prompt" (text field) or "image" (input image) - parameters (dict): - prompt (required, string): list of concepts to keep in the video. - negative_prompt (optional, string): list of concepts to ignore in the video. - width (optional, int, default to 768): width, or horizontal size in pixels. - height (optional, int, default to 512): height, or vertical size in pixels. - num_frames (optional, int, default to 129): the numer of frames must be a multiple of 8, plus 1 frame. - guidance_scale (optional, float, default to 3.5): Guidance scale (values between 3.0 and 4.0 are nice) - num_inference_steps (optional, int, default to 50): number of inference steps - seed (optional, int, default to -1): set a random number generator seed, -1 means random seed. - fps (optional, int, default to 24): FPS of the final video (eg. 24, 25, 30, 60) - double_num_frames (optional, bool): if enabled, the number of frames will be multiplied by 2 using RIFE - super_resolution (optional, bool): if enabled, the resolution will be multiplied by 2 using Real_ESRGAN - grain_amount (optional, float): amount of film grain to add to the output video - enable_audio (optional, bool): automatically generate an audio track - audio_prompt (optional, str): prompt to use for the audio generation (concepts to add) - audio_negative_prompt (optional, str): nehative prompt to use for the audio generation (concepts to ignore) Returns: Dictionary containing: - video: Base64 encoded MP4 data URI - content-type: MIME type - metadata: Generation metadata """ inputs = data.get("inputs", dict()) input_prompt = inputs.get("prompt", "") input_image = inputs.get("image") params = data.get("parameters", dict()) if not input_image and not input_prompt: raise ValueError("Either prompt or image must be provided") if input_prompt: logger.info(f"Prompt: {input_prompt}") logger.info(f"Raw parameters:") pprint.pprint(params) # Create and validate configuration config = GenerationConfig( # general content settings prompt=input_prompt, negative_prompt=params.get("negative_prompt", GenerationConfig.negative_prompt), # video model settings (will be used during generation of the initial raw video clip) width=params.get("width", GenerationConfig.width), height=params.get("height", GenerationConfig.height), num_frames=params.get("num_frames", GenerationConfig.num_frames), guidance_scale=params.get("guidance_scale", GenerationConfig.guidance_scale), num_inference_steps=params.get("num_inference_steps", GenerationConfig.num_inference_steps), # reproducible generation settings seed=params.get("seed", GenerationConfig.seed), # varnish settings (will be used for post-processing after the raw video clip has been generated) fps=params.get("fps", GenerationConfig.fps), # FPS of the final video (only applied at the the very end, when converting to mp4) double_num_frames=params.get("double_num_frames", GenerationConfig.double_num_frames), # if True, the number of frames will be multiplied by 2 using RIFE super_resolution=params.get("super_resolution", GenerationConfig.super_resolution), # if True, the resolution will be multiplied by 2 using Real_ESRGAN grain_amount=params.get("grain_amount", GenerationConfig.grain_amount), enable_audio=params.get("enable_audio", GenerationConfig.enable_audio), audio_prompt=params.get("audio_prompt", GenerationConfig.audio_prompt), audio_negative_prompt=params.get("audio_negative_prompt", GenerationConfig.audio_negative_prompt), ).validate_and_adjust() logger.info(f"Global request settings:") pprint.pprint(config) try: with torch.no_grad(): # Set random seeds random.seed(config.seed) np.random.seed(config.seed) generator = torch.manual_seed(config.seed) # Prepare generation parameters for the video model (we omit params that are destined to Varnish, or things like the seed which is set externally) generation_kwargs = { # general content settings "prompt": config.prompt, "negative_prompt": config.negative_prompt, # video model settings (will be used during generation of the initial raw video clip) "width": config.width, "height": config.height, "num_frames": config.num_frames, "guidance_scale": config.guidance_scale, "num_inference_steps": config.num_inference_steps, # constants "output_type": "pt", "generator": generator } #logger.info(f"Video model generation settings:") #pprint.pprint(generation_kwargs) # Check if image-to-video generation is requested if input_image: processed_image = process_input_image( input_image, config.width, config.height ) generation_kwargs["image"] = processed_image frames = self.image_to_video(**generation_kwargs).frames else: frames = self.text_to_video(**generation_kwargs).frames try: loop = asyncio.get_event_loop() except RuntimeError: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) video_uri, metadata = loop.run_until_complete(self.process_frames(frames, config)) return { "video": video_uri, "content-type": "video/mp4", "metadata": metadata } except Exception as e: message = f"Error generating video ({str(e)})\n{traceback.format_exc()}" print(message) raise RuntimeError(message)