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 diffusers.hooks import apply_enhance_a_video, EnhanceAVideoConfig from PIL import Image from teacache import TeaCacheConfig, enable_teacache, disable_teacache from varnish import Varnish from varnish.utils import is_truthy, process_input_image # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Get token from environment hf_token = os.getenv("HF_API_TOKEN") # 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 # Check environment variable for pipeline support support_image_prompt = is_truthy(os.getenv("SUPPORT_INPUT_IMAGE_PROMPT")) @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 # this is a hack to fool LTX-Video into believing our input image is an actual video frame with poor encoding quality # after a quick benchmark using the value 70 seems like a sweet spot input_image_quality: int = 70 # 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 # The range of the CRF scale is 0–51, where: # 0 is lossless (for 8 bit only, for 10 bit use -qp 0) # 23 is the default # 51 is worst quality possible # A lower value generally leads to higher quality, and a subjectively sane range is 17–28. # Consider 17 or 18 to be visually lossless or nearly so; # it should look the same or nearly the same as the input but it isn't technically lossless. # The range is exponential, so increasing the CRF value +6 results in roughly half the bitrate / file size, while -6 leads to roughly twice the bitrate. quality: int = 18 # TeaCache settings enable_teacache: bool = True teacache_threshold: float = 0.05 # values: 0 (original), 0.03 (1.6x speedup), 0.05 (2.1x speedup). # Enhance-A-Video settings enable_enhance_a_video: bool = True enhance_a_video_weight: float = 5.0 # LoRA settings lora_model_name: str = "" # HuggingFace repo ID or path to LoRA model lora_model_weight_file: str = "" # Specific weight file to load from the LoRA model lora_model_trigger: str = "" # Optional trigger word to prepend to the prompt 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 if support_image_prompt: self.image_to_video = LTXImageToVideoPipeline.from_pretrained( model_path, torch_dtype=torch.bfloat16 ).to("cuda") else: # Initialize models with bfloat16 precision self.text_to_video = LTXPipeline.from_pretrained( model_path, torch_dtype=torch.bfloat16 ).to("cuda") # Initialize LoRA tracking self._current_lora_model = None #if support_image_prompt: # # Enable CPU offload for memory efficiency # self.image_to_video.enable_model_cpu_offload() # # Inject enhance-a-video functionality # inject_enhance_for_ltx(self.image_to_video.transformer) #else: # # Enable CPU offload for memory efficiency # self.text_to_video.enable_model_cpu_offload() # # Inject enhance-a-video functionality # inject_enhance_for_ltx(self.text_to_video.transformer) # Initialize Varnish for post-processing self.varnish = Varnish( device="cuda", 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=True, ) # Store TeaCache config for each model self.text_to_video_teacache = None self.image_to_video_teacache = None def _configure_teacache(self, model, config: GenerationConfig): """Configure TeaCache for a model based on generation config Args: model: The model to configure TeaCache for config: Generation configuration """ if config.enable_teacache: # Create and enable TeaCache if it should be enabled teacache_config = TeaCacheConfig( enabled=True, rel_l1_thresh=config.teacache_threshold, num_inference_steps=config.num_inference_steps ) enable_teacache(model.transformer.__class__, teacache_config) logger.info(f"TeaCache enabled with threshold {config.teacache_threshold}") else: # Disable TeaCache if it was previously enabled if hasattr(model.transformer.__class__, 'teacache_config'): disable_teacache(model.transformer.__class__) logger.info("TeaCache disabled") 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=config.quality) # 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. - input_image_quality (optional, int, default to 100): this is a trick we use to convert a "pristine" image into a "dirty" video frame. This helps fooling LTX-Video into turning the image into an animated one. - 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) - quality (optional, str, default to 18): The range of the CRF scale is 0–51, where 0 is lossless (for 8 bit only, for 10 bit use -qp 0), 23 is the default, and 51 is worst quality possible. - enable_teacache (optional, bool, default to True): Generate faster at the cost of a slight quality loss - teacache_threshold (optional, float, default to 0.05): Amount of cache, 0 (original), 0.03 (1.6x speedup), 0.05 (Default, 2.1x speedup). - enable_enhance_a_video (optional, bool, default to True): enable the enhance_a_video optimization - enhance_a_video_weight(optional, float, default to 5.0): amount of video enhancement to apply - lora_model_name(optional, str, default to ""): HuggingFace repo ID or path to LoRA model - lora_model_weight_file(optional, str, default to ""): Specific weight file to load from the LoRA model - lora_model_trigger(optional, str, default to ""): Optional trigger word to prepend to the prompt 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") #logger.debug(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), input_image_quality=params.get("input_image_quality", GenerationConfig.input_image_quality), 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), quality=params.get("quality", GenerationConfig.quality), # TeaCache settings enable_teacache=params.get("enable_teacache", True), # values: 0 (original), 0.03 (1.6x speedup), 0.05 (2.1x speedup). teacache_threshold=params.get("teacache_threshold", 0.05), # Add enhance-a-video settings enable_enhance_a_video=params.get("enable_enhance_a_video", True), enhance_a_video_weight=params.get("enhance_a_video_weight", 5.0), # LoRA settings lora_model_name=params.get("lora_model_name", ""), lora_model_weight_file=params.get("lora_model_weight_file", ""), lora_model_trigger=params.get("lora_model_trigger", ""), ).validate_and_adjust() #logger.debug(f"Global request settings:") #pprint.pprint(config) try: with torch.inference_mode(): # Set random seeds random.seed(config.seed) np.random.seed(config.seed) torch.manual_seed(config.seed) generator = torch.Generator(device='cuda') generator = generator.manual_seed(config.seed) # Configure enhance-a-video #if config.enable_enhance_a_video: # enable_enhance() # set_enhance_weight(config.enhance_a_video_weight) # 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, # Timestep for decoding VAE noise: the timestep at which generated video is decoded "decode_timestep": 0.05, # Noise level for decoding VAE noise: the interpolation factor between random noise and denoised latents at the decode timestep "decode_noise_scale": 0.025, } #logger.info(f"Video model generation settings:") #pprint.pprint(generation_kwargs) # Handle LoRA loading/unloading if hasattr(self, '_current_lora_model'): if self._current_lora_model != (config.lora_model_name, config.lora_model_weight_file): # Unload previous LoRA if it exists and is different if hasattr(self.text_to_video, 'unload_lora_weights'): self.text_to_video.unload_lora_weights() if support_image_prompt and hasattr(self.image_to_video, 'unload_lora_weights'): self.image_to_video.unload_lora_weights() if config.lora_model_name: # Load new LoRA if hasattr(self.text_to_video, 'load_lora_weights'): self.text_to_video.load_lora_weights( config.lora_model_name, weight_name=config.lora_model_weight_file if config.lora_model_weight_file else None, token=hf_token, ) if support_image_prompt and hasattr(self.image_to_video, 'load_lora_weights'): self.image_to_video.load_lora_weights( config.lora_model_name, weight_name=config.lora_model_weight_file if config.lora_model_weight_file else None, token=hf_token, ) self._current_lora_model = (config.lora_model_name, config.lora_model_weight_file) # Modify prompt if trigger word is provided if config.lora_model_trigger: generation_kwargs["prompt"] = f"{config.lora_model_trigger} {generation_kwargs['prompt']}" enhance_a_video_config = EnhanceAVideoConfig( weight=config.enhance_a_video_weight if config.enable_enhance_a_video else 0.0, # doing some testing num_frames_callback=lambda: (8 + 1), # num_frames_callback=lambda: config.num_frames, # num_frames_callback=lambda: (config.num_frames - 1), _attention_type=1 ) # Check if image-to-video generation is requested if support_image_prompt and input_image: self._configure_teacache(self.image_to_video, config) processed_image = process_input_image( input_image, config.width, config.height, config.input_image_quality, ) generation_kwargs["image"] = processed_image # disabled (we cannot install the hook multiple times, we would have to uninstall it first or find another way to dynamically enable it, eg. using the weight only) # apply_enhance_a_video(self.image_to_video.transformer, enhance_a_video_config) frames = self.image_to_video(**generation_kwargs).frames else: self._configure_teacache(self.text_to_video, config) # disabled (we cannot install the hook multiple times, we would have to uninstall it first or find another way to dynamically enable it, eg. using the weight only) # apply_enhance_a_video(self.text_to_video.transformer, enhance_a_video_config) 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)