import sys, gc, torch, random, numpy as np, gradio as gr, torchaudio, json from pathlib import Path from huggingface_hub import hf_hub_download from gradio_client import Client, handle_file from safetensors.torch import load_file import torchaudio.transforms as T # --- 1. THE BLOCK REGISTRY FIX --- import diffusers.models.transformers.transformer_2d from diffusers.models.modeling_utils import ModelMixin class LTX2VideoDummyBlock(ModelMixin): def __init__(self, *args, **kwargs): super().__init__() # Manually register the 3D blocks so the library doesn't crash on load for block in ["LTX2VideoDownBlock3D", "LTX2VideoUpBlock3D", "LTX2VideoMidBlock3D"]: if not hasattr(diffusers.models.transformers.transformer_2d, block): setattr(diffusers.models.transformers.transformer_2d, block, LTX2VideoDummyBlock) from diffusers import LTXPipeline, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler, LTX2VideoTransformer3DModel # --- 2. 2.4B PRUNED CONFIGURATION --- LITE_CONFIG = { "_class_name": "LTX2VideoTransformer3DModel", "activation_fn": "gelu-approximate", "attention_bias": True, "attention_head_dim": 128, "num_attention_heads": 32, "num_layers": 8, # Matches your 2.4B pruning "in_channels": 128, "out_channels": 128, "cross_attention_dim": 4096, "qk_norm": "rms_norm_across_heads", "rope_type": "split" } # Constants DEFAULT_NEGATIVE_PROMPT = "worst quality, inconsistent motion, blurry, jittery, distorted" MAX_SEED = np.iinfo(np.int32).max DEFAULT_SEED = 42 DEFAULT_CFG_GUIDANCE_SCALE = 3.0 # --- 3. MEMORY-LEAN & SHAPE-CORRECT LOADER --- def load_lean_pipeline(): # Load everything from your unified safetensors file print("Loading unified checkpoint with Transformer + VAEs...") ckpt = hf_hub_download("MihaiPopa-1/LTX-2-Lite-2.4B", "ltx-2-2.4b-pruned.safetensors") full_state_dict = load_file(ckpt) # Download VAE configs from your repo print("Loading VAE configs from Hugging Face repo...") video_vae_config_path = hf_hub_download("MihaiPopa-1/LTX-2-Lite-2.4B", "video_vae_config.json") audio_vae_config_path = hf_hub_download("MihaiPopa-1/LTX-2-Lite-2.4B", "audio_vae_config.json") with open(video_vae_config_path) as f: VIDEO_VAE_CONFIG = json.load(f) with open(audio_vae_config_path) as f: AUDIO_VAE_CONFIG = json.load(f) # Separate components by key prefix transformer_state = {} vae_state = {} audio_vae_state = {} for key, value in full_state_dict.items(): if 'audio_vae' in key: new_key = key.replace('audio_vae.', '') audio_vae_state[new_key] = value elif key.startswith('vae.'): new_key = key.replace('vae.', '') vae_state[new_key] = value else: # Everything else goes to transformer new_key = key.replace('model.', '').replace('transformer.', '') transformer_state[new_key] = value print(f"Separated: {len(transformer_state)} transformer, {len(vae_state)} vae, {len(audio_vae_state)} audio_vae keys") # Load Transformer print("Loading 2.4B Transformer...") transformer = LTX2VideoTransformer3DModel.from_config(LITE_CONFIG).to(torch.bfloat16) transformer.load_state_dict(transformer_state, strict=False) # Load Video VAE with pruned config print("Loading Video VAE with pruned config...") vae = AutoencoderKLLTX2Video.from_config(VIDEO_VAE_CONFIG).to(torch.bfloat16) missing, unexpected = vae.load_state_dict(vae_state, strict=False) if missing: print(f" Warning: {len(missing)} missing keys in Video VAE") if unexpected: print(f" Warning: {len(unexpected)} unexpected keys in Video VAE") # Load Audio VAE with pruned config print("Loading Audio VAE with pruned config...") audio_vae = AutoencoderKLLTX2Video.from_config(AUDIO_VAE_CONFIG).to(torch.bfloat16) missing, unexpected = audio_vae.load_state_dict(audio_vae_state, strict=False) if missing: print(f" Warning: {len(missing)} missing keys in Audio VAE") if unexpected: print(f" Warning: {len(unexpected)} unexpected keys in Audio VAE") scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained("Lightricks/LTX-2", subfolder="scheduler") pipe = LTXPipeline( transformer=transformer, vae=vae, scheduler=scheduler, text_encoder=None, tokenizer=None ) # Manually attach the audio VAE to the pipeline pipe.audio_vae = audio_vae # HF Spaces Memory Optimization - CPU mode # For CPU inference, keep everything on CPU pipe = pipe.to("cpu") pipe.vae.enable_tiling() pipe.vae.enable_slicing() gc.collect() return pipe pipeline = load_lean_pipeline() text_client = Client("MihaiPopa-1/gemma-text-encoder-lite") # --- 4. 48kHz STEREO VALIDATOR --- def ensure_48k_stereo(audio_path): """Ensures the generative output is properly formatted as 48kHz Stereo.""" waveform, sr = torchaudio.load(audio_path) # The LTX-2 Audio VAE generates at 48k natively, but if the container # defaults to 16k, we resample it back to the high-fidelity target. if sr != 48000: resampler = T.Resample(sr, 48000) waveform = resampler(waveform) # Expand Mono to Stereo if the model outputted a single channel if waveform.shape[0] == 1: waveform = waveform.repeat(2, 1) torchaudio.save(audio_path, waveform, 48000) return audio_path # --- 5. GENERATION --- def generate_video(img, prompt, dur, enhance_prompt, negative_prompt, seed, randomize_seed, steps, cfg, h, w): try: gc.collect() # Handle seed randomization if randomize_seed: seed = random.randint(0, MAX_SEED) torch.manual_seed(seed) # Save input image if provided if img: img.save("input_frame.jpg") # Get Embeddings from your Gemma-Lite Space res = text_client.predict( prompt=prompt, input_image=handle_file("input_frame.jpg") if img else None, api_name="/encode_prompt" ) emb = torch.load(res[0]) out_name = f"ltx2_lite_gen_{random.randint(0,999)}.mp4" # Run Inference pipeline( video_context_positive=emb['video_context'].to(torch.bfloat16), audio_context_positive=emb['audio_context'].to(torch.bfloat16), num_frames=int(dur * 24) + 1, num_inference_steps=int(steps), guidance_scale=cfg, height=int(h), width=int(w), output_path=out_name, decode_audio=True ) return out_name, seed except Exception as e: print(f"Error during generation: {str(e)}") return None, seed # --- 6. INTERFACE --- css = ''' .gradio-container .contain{max-width: 1200px !important; margin: 0 auto !important} ''' with gr.Blocks(title="LTX-2 Lite šŸŽ„šŸ”ˆ", css=css) as demo: gr.Markdown("# LTX-2 Lite šŸŽ„šŸ”ˆ: The First Open Source (And Tiny Enough) Audio-Video Model") gr.Markdown("We test LTX-2 Lite, my own pruned variant of LTX-2. It's only 2.4B parameters!") with gr.Row(): with gr.Column(): input_image = gr.Image( label="Input Image (Optional)", type="pil", ) prompt = gr.Textbox( label="Prompt", info="for best results - make it as elaborate as possible", value="Make this image come alive with cinematic motion, smooth animation", lines=3, placeholder="Describe the motion and animation you want..." ) with gr.Row(): duration = gr.Slider( label="Duration (seconds)", minimum=1.0, maximum=30.0, value=5.0, step=0.1 ) enhance_prompt = gr.Checkbox( label="Enhance Prompt", value=True ) generate_btn = gr.Button("Generate Video", variant="primary") with gr.Accordion("Advanced Settings", open=False): negative_prompt = gr.Textbox( label="Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT, lines=2 ) seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, value=DEFAULT_SEED, step=1 ) randomize_seed = gr.Checkbox( label="Randomize Seed", value=True ) num_inference_steps = gr.Slider( label="Inference Steps", minimum=1, maximum=100, value=25, step=1 ) cfg_guidance_scale = gr.Slider( label="CFG Guidance Scale", minimum=1.0, maximum=10.0, value=DEFAULT_CFG_GUIDANCE_SCALE, step=0.1 ) with gr.Row(): width = gr.Number( label="Width", value=1280, precision=0 ) height = gr.Number( label="Height", value=720, precision=0 ) with gr.Column(): output_video = gr.Video(label="Generated Video", autoplay=True) output_seed = gr.Number(label="Used Seed", precision=0) generate_btn.click( fn=generate_video, inputs=[ input_image, prompt, duration, enhance_prompt, negative_prompt, seed, randomize_seed, num_inference_steps, cfg_guidance_scale, height, width, ], outputs=[output_video, output_seed] ) # Add examples gr.Examples( examples=[ [ "kill_bill.jpeg", "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.", 5.0, ], [ "wednesday.png", "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.", 5.0, ], [ "astronaut.jpg", "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.", 3.0, ] ], fn=generate_video, inputs=[input_image, prompt, duration, enhance_prompt, negative_prompt, seed, randomize_seed, num_inference_steps, cfg_guidance_scale, height, width], outputs=[output_video, output_seed], label="Examples", cache_examples=True, cache_mode="lazy", ) if __name__ == "__main__": demo.launch(theme=gr.themes.Citrus())