import os import subprocess import sys # Disable torch.compile / dynamo before any torch import os.environ["TORCH_COMPILE_DISABLE"] = "1" os.environ["TORCHDYNAMO_DISABLE"] = "1" # Install xformers for memory-efficient attention subprocess.run([sys.executable, "-m", "pip", "install", "xformers==0.0.32.post2", "--no-build-isolation"], check=False) # Clone LTX-2 repo and install packages LTX_REPO_URL = "https://github.com/Lightricks/LTX-2.git" LTX_REPO_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "LTX-2") LTX_COMMIT_SHA = "a2c3f24078eb918171967f74b6f66b756b29ee45" if not os.path.exists(LTX_REPO_DIR): print(f"Cloning {LTX_REPO_URL}...") os.makedirs(LTX_REPO_DIR) subprocess.run(["git", "init", LTX_REPO_DIR], check=True) subprocess.run(["git", "remote", "add", "origin", LTX_REPO_URL], cwd=LTX_REPO_DIR, check=True) subprocess.run(["git", "fetch", "--depth", "1", "origin", LTX_COMMIT_SHA], cwd=LTX_REPO_DIR, check=True) subprocess.run(["git", "checkout", LTX_COMMIT_SHA], cwd=LTX_REPO_DIR, check=True) print("Installing ltx-core and ltx-pipelines from cloned repo...") subprocess.run( [sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", "-e", os.path.join(LTX_REPO_DIR, "packages", "ltx-core"), "-e", os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines")], check=True, ) sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines", "src")) sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-core", "src")) import logging import random import tempfile from pathlib import Path from collections.abc import Iterator import torch torch._dynamo.config.suppress_errors = True torch._dynamo.config.disable = True import spaces import gradio as gr import numpy as np from huggingface_hub import hf_hub_download, snapshot_download from ltx_core.components.diffusion_steps import Res2sDiffusionStep from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams from ltx_core.components.noisers import GaussianNoiser from ltx_core.components.schedulers import LTX2Scheduler from ltx_core.loader import LoraPathStrengthAndSDOps from ltx_core.loader.registry import Registry from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number from ltx_core.quantization import QuantizationPolicy from ltx_core.types import Audio, VideoLatentShape, VideoPixelShape from ltx_pipelines.utils.args import ImageConditioningInput, hq_2_stage_arg_parser from ltx_pipelines.utils.blocks import ( AudioDecoder, DiffusionStage, ImageConditioner, PromptEncoder, VideoDecoder, VideoUpsampler, ) from ltx_pipelines.utils.constants import ( LTX_2_3_HQ_PARAMS, STAGE_2_DISTILLED_SIGMAS, ) from ltx_pipelines.utils.denoisers import GuidedDenoiser, SimpleDenoiser from ltx_pipelines.utils.helpers import ( assert_resolution, combined_image_conditionings, get_device, ) from ltx_pipelines.utils.media_io import encode_video from ltx_pipelines.utils.samplers import res2s_audio_video_denoising_loop from ltx_pipelines.utils.types import ModalitySpec # Force-patch xformers attention into the LTX attention module. from ltx_core.model.transformer import attention as _attn_mod print(f"[ATTN] Before patch: memory_efficient_attention={_attn_mod.memory_efficient_attention}") try: from xformers.ops import memory_efficient_attention as _mea _attn_mod.memory_efficient_attention = _mea print(f"[ATTN] After patch: memory_efficient_attention={_attn_mod.memory_efficient_attention}") except Exception as e: print(f"[ATTN] xformers patch FAILED: {type(e).__name__}: {e}") logging.getLogger().setLevel(logging.INFO) MAX_SEED = np.iinfo(np.int32).max DEFAULT_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." ) DEFAULT_FRAME_RATE = 24.0 # Resolution presets: (width, height) RESOLUTIONS = { "high": {"16:9": (1536, 1024), "9:16": (1024, 1536), "1:1": (1024, 1024)}, "low": {"16:9": (768, 512), "9:16": (512, 768), "1:1": (768, 768)}, } class TI2VidTwoStagesHQPipeline: """ Two-stage text/image-to-video generation pipeline using the res_2s sampler. Same structure as :class:`TI2VidTwoStagesPipeline`: stage 1 generates video at half of the target resolution with CFG guidance (assuming full model is used), then Stage 2 upsamples by 2x and refines using a distilled LoRA for higher quality output. Uses the res_2s second-order sampler instead of Euler, allowing fewer steps for comparable quality. Supports optional image conditioning via the images parameter. """ def __init__( # noqa: PLR0913 self, checkpoint_path: str, distilled_lora: list[LoraPathStrengthAndSDOps], distilled_lora_strength_stage_1: float, distilled_lora_strength_stage_2: float, spatial_upsampler_path: str, gemma_root: str, loras: tuple[LoraPathStrengthAndSDOps, ...], device: torch.device | None = None, quantization: QuantizationPolicy | None = None, registry: Registry | None = None, torch_compile: bool = False, ): self.device = device or get_device() self.dtype = torch.bfloat16 self._scheduler = LTX2Scheduler() distilled_lora_stage_1 = LoraPathStrengthAndSDOps( path=distilled_lora[0].path, strength=distilled_lora_strength_stage_1, sd_ops=distilled_lora[0].sd_ops, ) distilled_lora_stage_2 = LoraPathStrengthAndSDOps( path=distilled_lora[0].path, strength=distilled_lora_strength_stage_2, sd_ops=distilled_lora[0].sd_ops, ) self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry) self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry) self.upsampler = VideoUpsampler( checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry ) self.video_decoder = VideoDecoder(checkpoint_path, self.dtype, self.device, registry=registry) self.audio_decoder = AudioDecoder(checkpoint_path, self.dtype, self.device, registry=registry) self.stage_1 = DiffusionStage( checkpoint_path, self.dtype, self.device, loras=(*loras, distilled_lora_stage_1), quantization=quantization, registry=registry, torch_compile=torch_compile, ) self.stage_2 = DiffusionStage( checkpoint_path, self.dtype, self.device, loras=(*loras, distilled_lora_stage_2), quantization=quantization, registry=registry, torch_compile=torch_compile, ) @torch.inference_mode() def __call__( # noqa: PLR0913 self, prompt: str, negative_prompt: str, seed: int, height: int, width: int, num_frames: int, frame_rate: float, num_inference_steps: int, video_guider_params: MultiModalGuiderParams, audio_guider_params: MultiModalGuiderParams, images: list[ImageConditioningInput], tiling_config: TilingConfig | None = None, enhance_prompt: bool = False, streaming_prefetch_count: int | None = None, max_batch_size: int = 1, stage_1_sigmas: torch.Tensor | None = None, stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS, ) -> tuple[Iterator[torch.Tensor], Audio]: assert_resolution(height=height, width=width, is_two_stage=True) generator = torch.Generator(device=self.device).manual_seed(seed) noiser = GaussianNoiser(generator=generator) dtype = torch.bfloat16 ctx_p, ctx_n = self.prompt_encoder( [prompt, negative_prompt], enhance_first_prompt=enhance_prompt, enhance_prompt_image=images[0][0] if len(images) > 0 else None, enhance_prompt_seed=seed, streaming_prefetch_count=streaming_prefetch_count, ) v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding # Stage 1: Generate video at half resolution with CFG guidance using res2s sampler. stage_1_output_shape = VideoPixelShape( batch=1, frames=num_frames, width=width // 2, height=height // 2, fps=frame_rate, ) stage_1_conditionings = self.image_conditioner( lambda enc: combined_image_conditionings( images=images, height=stage_1_output_shape.height, width=stage_1_output_shape.width, video_encoder=enc, dtype=dtype, device=self.device, ) ) stepper = Res2sDiffusionStep() if stage_1_sigmas is None: empty_latent = torch.empty(VideoLatentShape.from_pixel_shape(stage_1_output_shape).to_torch_shape()) stage_1_sigmas = self._scheduler.execute(latent=empty_latent, steps=num_inference_steps) sigmas = stage_1_sigmas.to(dtype=torch.float32, device=self.device) video_state, audio_state = self.stage_1( denoiser=GuidedDenoiser( v_context=v_context_p, a_context=a_context_p, video_guider=MultiModalGuider( params=video_guider_params, negative_context=v_context_n, ), audio_guider=MultiModalGuider( params=audio_guider_params, negative_context=a_context_n, ), ), sigmas=sigmas, noiser=noiser, stepper=stepper, width=stage_1_output_shape.width, height=stage_1_output_shape.height, frames=num_frames, fps=frame_rate, video=ModalitySpec(context=v_context_p, conditionings=stage_1_conditionings), audio=ModalitySpec(context=a_context_p), loop=res2s_audio_video_denoising_loop, streaming_prefetch_count=streaming_prefetch_count, max_batch_size=max_batch_size, ) # Stage 2: Upsample and refine the video at higher resolution with distilled LoRA. upscaled_video_latent = self.upsampler(video_state.latent[:1]) stage_2_sigmas = stage_2_sigmas.to(dtype=torch.float32, device=self.device) stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate) stage_2_conditionings = self.image_conditioner( lambda enc: combined_image_conditionings( images=images, height=stage_2_output_shape.height, width=stage_2_output_shape.width, video_encoder=enc, dtype=dtype, device=self.device, ) ) video_state, audio_state = self.stage_2( denoiser=SimpleDenoiser(v_context=v_context_p, a_context=a_context_p), sigmas=stage_2_sigmas, noiser=noiser, stepper=stepper, width=width, height=height, frames=num_frames, fps=frame_rate, video=ModalitySpec( context=v_context_p, conditionings=stage_2_conditionings, noise_scale=stage_2_sigmas[0].item(), initial_latent=upscaled_video_latent, ), audio=ModalitySpec( context=a_context_p, noise_scale=stage_2_sigmas[0].item(), initial_latent=audio_state.latent, ), loop=res2s_audio_video_denoising_loop, streaming_prefetch_count=streaming_prefetch_count, ) decoded_video = self.video_decoder(video_state.latent, tiling_config, generator) decoded_audio = self.audio_decoder(audio_state.latent) return decoded_video, decoded_audio @torch.inference_mode() def main() -> None: logging.getLogger().setLevel(logging.INFO) parser = hq_2_stage_arg_parser(params=LTX_2_3_HQ_PARAMS) args = parser.parse_args() pipeline = TI2VidTwoStagesHQPipeline( checkpoint_path=args.checkpoint_path, distilled_lora=args.distilled_lora, distilled_lora_strength_stage_1=args.distilled_lora_strength_stage_1, distilled_lora_strength_stage_2=args.distilled_lora_strength_stage_2, spatial_upsampler_path=args.spatial_upsampler_path, gemma_root=args.gemma_root, loras=tuple(args.lora) if args.lora else (), quantization=args.quantization, torch_compile=args.compile, ) tiling_config = TilingConfig.default() video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config) video, audio = pipeline( prompt=args.prompt, negative_prompt=args.negative_prompt, seed=args.seed, height=args.height, width=args.width, num_frames=args.num_frames, frame_rate=args.frame_rate, num_inference_steps=args.num_inference_steps, video_guider_params=MultiModalGuiderParams( cfg_scale=args.video_cfg_guidance_scale, stg_scale=args.video_stg_guidance_scale, rescale_scale=args.video_rescale_scale, modality_scale=args.a2v_guidance_scale, skip_step=args.video_skip_step, stg_blocks=args.video_stg_blocks, ), audio_guider_params=MultiModalGuiderParams( cfg_scale=args.audio_cfg_guidance_scale, stg_scale=args.audio_stg_guidance_scale, rescale_scale=args.audio_rescale_scale, modality_scale=args.v2a_guidance_scale, skip_step=args.audio_skip_step, stg_blocks=args.audio_stg_blocks, ), images=args.images, tiling_config=tiling_config, streaming_prefetch_count=args.streaming_prefetch_count, max_batch_size=args.max_batch_size, ) encode_video( video=video, fps=args.frame_rate, audio=audio, output_path=args.output_path, video_chunks_number=video_chunks_number, ) if __name__ == "__main__": main()