| import PIL.Image |
| import torch |
|
|
| from diffusers import ( |
| AutoencoderKLLTX2Audio, |
| AutoencoderKLLTX2Video, |
| FlowMatchEulerDiscreteScheduler, |
| LTX2ConditionPipeline, |
| LTX2VideoTransformer3DModel, |
| ) |
| from diffusers.modular_pipelines import ComponentSpec, InputParam, ModularPipelineBlocks, OutputParam, PipelineState |
| from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors |
| from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition |
| from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE |
| from diffusers.utils import load_image |
| from transformers import Gemma3ForConditionalGeneration, GemmaTokenizerFast |
|
|
|
|
| _LTX2_REPO = "diffusers/LTX-2.3-Diffusers" |
|
|
|
|
| class LTX2FirstLastFrameBlock(ModularPipelineBlocks): |
| model_name = "ltx2" |
|
|
| @property |
| def description(self) -> str: |
| return ( |
| "Generates a video with LTX-2.3 conditioned on a first frame (`image`) and an optional " |
| "last frame (`image_2`). Images can be PIL images, local paths, or URLs — path/URL strings " |
| "are loaded automatically, so plain image inputs work without constructing " |
| "`LTX2VideoCondition` objects." |
| ) |
|
|
| @property |
| def expected_components(self) -> list[ComponentSpec]: |
| return [ |
| ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler, pretrained_model_name_or_path=_LTX2_REPO, subfolder="scheduler"), |
| ComponentSpec("vae", AutoencoderKLLTX2Video, pretrained_model_name_or_path=_LTX2_REPO, subfolder="vae"), |
| ComponentSpec("audio_vae", AutoencoderKLLTX2Audio, pretrained_model_name_or_path=_LTX2_REPO, subfolder="audio_vae"), |
| ComponentSpec("text_encoder", Gemma3ForConditionalGeneration, pretrained_model_name_or_path=_LTX2_REPO, subfolder="text_encoder"), |
| ComponentSpec("tokenizer", GemmaTokenizerFast, pretrained_model_name_or_path=_LTX2_REPO, subfolder="tokenizer"), |
| ComponentSpec("connectors", LTX2TextConnectors, pretrained_model_name_or_path=_LTX2_REPO, subfolder="connectors"), |
| ComponentSpec("transformer", LTX2VideoTransformer3DModel, pretrained_model_name_or_path=_LTX2_REPO, subfolder="transformer"), |
| ComponentSpec("vocoder", LTX2VocoderWithBWE, pretrained_model_name_or_path=_LTX2_REPO, subfolder="vocoder"), |
| ] |
|
|
| @property |
| def inputs(self) -> list[InputParam]: |
| return [ |
| InputParam("prompt", type_hint=str, required=True, description="Text prompt describing the video."), |
| InputParam( |
| "image", |
| type_hint="PIL.Image.Image | str", |
| required=True, |
| description="First frame to condition on. PIL image, local path, or URL.", |
| ), |
| InputParam( |
| "image_2", |
| type_hint="PIL.Image.Image | str | None", |
| description="Optional last frame to condition on. PIL image, local path, or URL.", |
| ), |
| InputParam("negative_prompt", type_hint="str | None", description="Negative text prompt."), |
| InputParam("height", type_hint=int, default=512, description="Output video height in pixels."), |
| InputParam("width", type_hint=int, default=768, description="Output video width in pixels."), |
| InputParam("num_frames", type_hint=int, default=121, description="Number of frames to generate."), |
| InputParam("frame_rate", type_hint=float, default=24.0, description="Frames per second."), |
| InputParam("num_inference_steps", type_hint=int, default=40, description="Number of denoising steps."), |
| InputParam("guidance_scale", type_hint=float, default=4.0, description="Classifier-free guidance scale."), |
| InputParam( |
| "first_frame_strength", type_hint=float, default=1.0, description="Conditioning strength of `image`." |
| ), |
| InputParam( |
| "last_frame_strength", type_hint=float, default=1.0, description="Conditioning strength of `image_2`." |
| ), |
| InputParam( |
| "offload", |
| type_hint=bool, |
| default=True, |
| description="Enable model CPU offload so the ~100GB of components fit a single 80GB GPU.", |
| ), |
| InputParam("generator", description="torch.Generator for deterministic sampling."), |
| ] |
|
|
| @property |
| def intermediate_outputs(self) -> list[OutputParam]: |
| return [ |
| OutputParam("videos", type_hint="list[np.ndarray]", description="Generated video frames."), |
| OutputParam("audio", type_hint=torch.Tensor, description="Generated audio waveform."), |
| ] |
|
|
| @torch.no_grad() |
| def __call__(self, components, state: PipelineState) -> PipelineState: |
| block_state = self.get_block_state(state) |
|
|
| first = block_state.image |
| if isinstance(first, str): |
| first = load_image(first) |
| conditions = [LTX2VideoCondition(frames=first, index=0, strength=block_state.first_frame_strength)] |
|
|
| last = block_state.image_2 |
| if last is not None: |
| if isinstance(last, str): |
| last = load_image(last) |
| conditions.append(LTX2VideoCondition(frames=last, index=-1, strength=block_state.last_frame_strength)) |
|
|
| pipeline = LTX2ConditionPipeline( |
| scheduler=components.scheduler, |
| vae=components.vae, |
| audio_vae=components.audio_vae, |
| text_encoder=components.text_encoder, |
| tokenizer=components.tokenizer, |
| connectors=components.connectors, |
| transformer=components.transformer, |
| vocoder=components.vocoder, |
| ) |
| if block_state.offload: |
| pipeline.enable_model_cpu_offload() |
|
|
| result = pipeline( |
| conditions=conditions, |
| prompt=block_state.prompt, |
| negative_prompt=block_state.negative_prompt, |
| height=block_state.height, |
| width=block_state.width, |
| num_frames=block_state.num_frames, |
| frame_rate=block_state.frame_rate, |
| num_inference_steps=block_state.num_inference_steps, |
| guidance_scale=block_state.guidance_scale, |
| generator=block_state.generator, |
| output_type="np", |
| ) |
| |
| |
| |
| block_state.videos = list(result.frames[0]) |
| block_state.audio = result.audio |
|
|
| self.set_block_state(state, block_state) |
| return components, state |
|
|
