CedricPerauer/ltx-2
0
1import sys2from pathlib import Path3 4# Add packages to Python path5current_dir = Path(__file__).parent6sys.path.insert(0, str(current_dir / "packages" / "ltx-pipelines" / "src"))7sys.path.insert(0, str(current_dir / "packages" / "ltx-core" / "src"))8import numpy as np9import random10import spaces11import gradio as gr12from gradio_client import Client, handle_file13import torch14from pathlib import Path15from typing import Optional16from huggingface_hub import hf_hub_download17from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline18from ltx_core.tiling import TilingConfig19from ltx_pipelines.constants import (20 DEFAULT_SEED,21 DEFAULT_HEIGHT,22 DEFAULT_WIDTH,23 DEFAULT_NUM_FRAMES,24 DEFAULT_FRAME_RATE,25 DEFAULT_NUM_INFERENCE_STEPS,26 DEFAULT_CFG_GUIDANCE_SCALE,27 DEFAULT_LORA_STRENGTH,28)29 30MAX_SEED = np.iinfo(np.int32).max31# Custom negative prompt32DEFAULT_NEGATIVE_PROMPT = "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static"33 34# Default prompt from docstring example35DEFAULT_PROMPT = "An astronaut hatches from a fragile egg on the surface of the Moon, the shell cracking and peeling apart in gentle low-gravity motion. Fine lunar dust lifts and drifts outward with each movement, floating in slow arcs before settling back onto the ground. The astronaut pushes free in a deliberate, weightless motion, small fragments of the egg tumbling and spinning through the air. In the background, the deep darkness of space subtly shifts as stars glide with the camera's movement, emphasizing vast depth and scale. The camera performs a smooth, cinematic slow push-in, with natural parallax between the foreground dust, the astronaut, and the distant starfield. Ultra-realistic detail, physically accurate low-gravity motion, cinematic lighting, and a breath-taking, movie-like shot."36 37# HuggingFace Hub defaults38DEFAULT_REPO_ID = "Lightricks/LTX-2"39DEFAULT_CHECKPOINT_FILENAME = "ltx-2-19b-dev-fp8.safetensors"40DEFAULT_DISTILLED_LORA_FILENAME = "ltx-2-19b-distilled-lora-384.safetensors"41DEFAULT_SPATIAL_UPSAMPLER_FILENAME = "ltx-2-spatial-upscaler-x2-1.0.safetensors"42 43# Text encoder space URL44TEXT_ENCODER_SPACE = "linoyts/gemma-text-encoder"45 46def get_hub_or_local_checkpoint(repo_id: Optional[str] = None, filename: Optional[str] = None):47 """Download from HuggingFace Hub or use local checkpoint."""48 if repo_id is None and filename is None:49 raise ValueError("Please supply at least one of `repo_id` or `filename`")50 51 if repo_id is not None:52 if filename is None:53 raise ValueError("If repo_id is specified, filename must also be specified.")54 print(f"Downloading {filename} from {repo_id}...")55 ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename)56 print(f"Downloaded to {ckpt_path}")57 else:58 ckpt_path = filename59 60 return ckpt_path61 62 63# Initialize pipeline at startup64print("=" * 80)65print("Loading LTX-2 2-stage pipeline...")66print("=" * 80)67 68checkpoint_path = get_hub_or_local_checkpoint(DEFAULT_REPO_ID, DEFAULT_CHECKPOINT_FILENAME)69distilled_lora_path = get_hub_or_local_checkpoint(DEFAULT_REPO_ID, DEFAULT_DISTILLED_LORA_FILENAME)70spatial_upsampler_path = get_hub_or_local_checkpoint(DEFAULT_REPO_ID, DEFAULT_SPATIAL_UPSAMPLER_FILENAME)71 72print(f"Initializing pipeline with:")73print(f" checkpoint_path={checkpoint_path}")74print(f" distilled_lora_path={distilled_lora_path}")75print(f" spatial_upsampler_path={spatial_upsampler_path}")76print(f" text_encoder_space={TEXT_ENCODER_SPACE}")77 78# Initialize pipeline WITHOUT text encoder (gemma_root=None)79# Text encoding will be done by external space80pipeline = TI2VidTwoStagesPipeline(81 checkpoint_path=checkpoint_path,82 distilled_lora_path=distilled_lora_path,83 distilled_lora_strength=DEFAULT_LORA_STRENGTH,84 spatial_upsampler_path=spatial_upsampler_path,85 gemma_root=None,86 loras=[],87 fp8transformer=False,88 local_files_only=False89)90 91# Initialize text encoder client92print(f"Connecting to text encoder space: {TEXT_ENCODER_SPACE}")93try:94 text_encoder_client = Client(TEXT_ENCODER_SPACE)95 print("✓ Text encoder client connected!")96except Exception as e:97 print(f"⚠ Warning: Could not connect to text encoder space: {e}")98 text_encoder_client = None99 100@spaces.GPU(duration=300)101def generate_video(102 input_image,103 prompt: str,104 duration: float,105 enhance_prompt: bool = True,106 negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,107 seed: int = 42,108 randomize_seed: bool = True,109 num_inference_steps: int = 25,110 cfg_guidance_scale: float = DEFAULT_CFG_GUIDANCE_SCALE,111 height: int = DEFAULT_HEIGHT,112 width: int = DEFAULT_WIDTH,113 progress=gr.Progress(track_tqdm=True)114):115 """Generate a video based on the given parameters."""116 try:117 # Randomize seed if checkbox is enabled118 current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)119 120 # Calculate num_frames from duration (using fixed 24 fps)121 frame_rate = 24.0122 num_frames = int(duration * frame_rate) + 1 # +1 to ensure we meet the duration123 124 # Create output directory if it doesn't exist125 output_dir = Path("outputs")126 output_dir.mkdir(exist_ok=True)127 output_path = output_dir / f"video_{current_seed}.mp4"128 129 # Handle image input130 images = []131 temp_image_path = None # Initialize to None132 if input_image is not None:133 # Save uploaded image temporarily134 temp_image_path = output_dir / f"temp_input_{current_seed}.jpg"135 if hasattr(input_image, 'save'):136 input_image.save(temp_image_path)137 else:138 # If it's a file path already139 temp_image_path = Path(input_image)140 # Format: (image_path, frame_idx, strength)141 images = [(str(temp_image_path), 0, 1.0)]142 # Get embeddings from text encoder space143 print(f"Encoding prompt: {prompt}")144 145 if text_encoder_client is None:146 raise RuntimeError(147 f"Text encoder client not connected. Please ensure the text encoder space "148 f"({TEXT_ENCODER_SPACE}) is running and accessible."149 )150 151 try:152 # Prepare image for upload if it exists153 image_input = None154 if temp_image_path is not None:155 image_input = handle_file(str(temp_image_path))156 157 result = text_encoder_client.predict(158 prompt=prompt,159 enhance_prompt=enhance_prompt,160 input_image=image_input,161 seed=current_seed,162 negative_prompt=negative_prompt,163 api_name="/encode_prompt"164 )165 embedding_path = result[0] # Path to .pt file166 print(f"Embeddings received from: {embedding_path}")167 168 # Load embeddings169 embeddings = torch.load(embedding_path)170 video_context_positive = embeddings['video_context']171 audio_context_positive = embeddings['audio_context']172 173 # Load negative contexts if available174 video_context_negative = embeddings.get('video_context_negative', None)175 audio_context_negative = embeddings.get('audio_context_negative', None)176 177 print("✓ Embeddings loaded successfully")178 if video_context_negative is not None:179 print(" ✓ Negative prompt embeddings also loaded")180 except Exception as e:181 raise RuntimeError(182 f"Failed to get embeddings from text encoder space: {e}\n"183 f"Please ensure {TEXT_ENCODER_SPACE} is running properly."184 )185 186 # Run inference - progress automatically tracks tqdm from pipeline187 pipeline(188 prompt=prompt,189 negative_prompt=negative_prompt,190 output_path=str(output_path),191 seed=current_seed,192 height=height,193 width=width,194 num_frames=num_frames,195 frame_rate=frame_rate,196 num_inference_steps=num_inference_steps,197 cfg_guidance_scale=cfg_guidance_scale,198 images=images,199 tiling_config=TilingConfig.default(),200 video_context_positive=video_context_positive,201 audio_context_positive=audio_context_positive,202 video_context_negative=video_context_negative,203 audio_context_negative=audio_context_negative,204 )205 206 return str(output_path), current_seed207 208 except Exception as e:209 import traceback210 error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"211 print(error_msg)212 return None213 214 215# Create Gradio interface216with gr.Blocks(title="LTX-2 Video 🎥🔈") as demo:217 gr.Markdown("# LTX-2 🎥🔈: The First Open Source Audio-Video Model")218 gr.Markdown("State-of-the-art video & audio generation with Lightricks LTX-2 TI2V. Read more: [[model]](https://huggingface.co/Lightricks/LTX-2), [[code]](https://github.com/Lightricks/LTX-2)")219 with gr.Row():220 with gr.Column():221 input_image = gr.Image(222 label="Input Image (Optional)",223 type="pil",224 )225 226 prompt = gr.Textbox(227 label="Prompt",228 info="for best results - make it as elaborate as possible",229 value="Make this image come alive with cinematic motion, smooth animation",230 lines=3,231 placeholder="Describe the motion and animation you want..."232 )233 234 with gr.Row():235 duration = gr.Slider(236 label="Duration (seconds)",237 minimum=1.0,238 maximum=10.0,239 value=3.0,240 step=0.1241 )242 enhance_prompt = gr.Checkbox(243 label="Enhance Prompt",244 value=True245 )246 247 generate_btn = gr.Button("Generate Video", variant="primary")248 249 with gr.Accordion("Advanced Settings", open=False):250 negative_prompt = gr.Textbox(251 label="Negative Prompt",252 value=DEFAULT_NEGATIVE_PROMPT,253 lines=2254 )255 256 seed = gr.Slider(257 label="Seed",258 minimum=0,259 maximum=MAX_SEED,260 value=DEFAULT_SEED,261 step=1262 )263 264 randomize_seed = gr.Checkbox(265 label="Randomize Seed",266 value=True267 )268 269 num_inference_steps = gr.Slider(270 label="Inference Steps",271 minimum=1,272 maximum=100,273 value=25,274 step=1275 )276 277 cfg_guidance_scale = gr.Slider(278 label="CFG Guidance Scale",279 minimum=1.0,280 maximum=10.0,281 value=DEFAULT_CFG_GUIDANCE_SCALE,282 step=0.1283 )284 285 with gr.Row():286 width = gr.Number(287 label="Width",288 value=DEFAULT_WIDTH,289 precision=0290 )291 height = gr.Number(292 label="Height",293 value=DEFAULT_HEIGHT,294 precision=0295 )296 297 with gr.Column():298 output_video = gr.Video(label="Generated Video", autoplay=True)299 300 generate_btn.click(301 fn=generate_video,302 inputs=[303 input_image,304 prompt,305 duration,306 enhance_prompt,307 negative_prompt,308 seed,309 randomize_seed,310 num_inference_steps,311 cfg_guidance_scale,312 height,313 width,314 ],315 outputs=[output_video,seed]316 )317 318 # Add example319 gr.Examples(320 examples=[321 [322 "kill_bill.jpeg",323 "A low, subsonic drone pulses as Uma Thurman's character, Beatrix Kiddo, holds her razor-sharp katana blade steady in the cinematic lighting. A faint electrical hum fills the silence. Suddenly, accompanied by a deep metallic groan, the polished steel begins to soften and distort, like heated metal starting to lose its structural integrity. Discordant strings swell as the blade's perfect edge slowly warps and droops, molten steel beginning to flow downward in silvery rivulets while maintaining its metallic sheen—each drip producing a wet, viscous stretching sound. The transformation starts subtly at first—a slight bend in the blade—then accelerates as the metal becomes increasingly fluid, the groaning intensifying. The camera holds steady on her face as her piercing eyes gradually narrow, not with lethal focus, but with confusion and growing alarm as she watches her weapon dissolve before her eyes. She whispers under her breath, voice flat with disbelief: 'Wait, what?' Her heartbeat rises in the mix—thump... thump-thump—as her breathing quickens slightly while she witnesses this impossible transformation. Sharp violin stabs punctuate each breath. The melting intensifies, the katana's perfect form becoming increasingly abstract, dripping like liquid mercury from her grip. Molten droplets fall to the ground with soft, bell-like pings. Unintelligible whispers fade in and out as her expression shifts from calm readiness to bewilderment and concern, her heartbeat now pounding like a war drum, as her legendary instrument of vengeance literally liquefies in her hands, leaving her defenseless and disoriented. All sound cuts to silence—then a single devastating bass drop as the final droplet falls, leaving only her unsteady breathing in the dark.",324 5.0,325 ],326 [327 "wednesday.png",328 "A cinematic close-up of Wednesday Addams frozen mid-dance on a dark, blue-lit ballroom floor as students move indistinctly behind her, their footsteps and muffled music reduced to a distant, underwater thrum; the audio foregrounds her steady breathing and the faint rustle of fabric as she slowly raises one arm, never breaking eye contact with the camera, then after a deliberately long silence she speaks in a flat, dry, perfectly controlled voice, “I don’t dance… I vibe code,” each word crisp and unemotional, followed by an abrupt cutoff of her voice as the background sound swells slightly, reinforcing the deadpan humor, with precise lip sync, minimal facial movement, stark gothic lighting, and cinematic realism.",329 5.0,330 ],331 [332 "astronaut.jpg",333 "An astronaut hatches from a fragile egg on the surface of the Moon, the shell cracking and peeling apart in gentle low-gravity motion. Fine lunar dust lifts and drifts outward with each movement, floating in slow arcs before settling back onto the ground. The astronaut pushes free in a deliberate, weightless motion, small fragments of the egg tumbling and spinning through the air. In the background, the deep darkness of space subtly shifts as stars glide with the camera's movement, emphasizing vast depth and scale. The camera performs a smooth, cinematic slow push-in, with natural parallax between the foreground dust, the astronaut, and the distant starfield. Ultra-realistic detail, physically accurate low-gravity motion, cinematic lighting, and a breath-taking, movie-like shot.",334 3.0,335 ]336 ],337 fn=generate_video,338 inputs=[input_image, prompt, duration],339 outputs = [output_video,seed],340 label="Example",341 cache_examples=True,342 cache_mode="lazy",343 )344 345css = '''346.gradio-container .contain{max-width: 1200px !important; margin: 0 auto !important}347'''348if __name__ == "__main__":349 demo.launch(theme=gr.themes.Citrus())