File size: 6,773 Bytes
154e1a2 b243aab 154e1a2 b243aab 154e1a2 b243aab 154e1a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | 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",
)
# A list of per-frame ndarrays rather than PIL images: consumers (including
# `diffusers-cli run`) can't distinguish a flat list of PIL frames from a batch
# of independent images, but a list of HWC arrays is unambiguously a video.
block_state.videos = list(result.frames[0])
block_state.audio = result.audio
self.set_block_state(state, block_state)
return components, state
|