feat(ltx2): full i2v conditioning + continuation latent port
Replaces the T2V-only stub with the full internal port: * resolve_ltx2_images / _resize_and_center_crop / CRF re-encode * load_ltx2_conditioning_image + load_ltx2_conditioning_video_clip * _extract_video_latent / _insert_conditioning_latent * build_ltx2_image_conditioning composes (clean_latent, denoise_mask) from images + video clips + continuation latents (stage1 vs stage2) — returns None for plain T2V ForwardBatch + SamplingParam gain the missing fields the builder reads (ltx2_images, ltx2_image_crf, ltx2_conditioning_latent_stage1/2, ltx2_video_conditions). generate_video kwargs flow through sampling_param.update -> ForwardBatch(**shallow_asdict(...)) so Dreamverse's ltx2_image_crf=0.0 and continuation latents land on the batch correctly. 68/68 LTX-2 tests still green. Co-authored-by: Junda (David) Su <90978028+Davids048@users.noreply.github.com> Co-authored-by: Matthew Noto <99706358+RandNMR73@users.noreply.github.com> Co-authored-by: XOR-op <17672363+XOR-op@users.noreply.github.com> Co-authored-by: Zhang Peiyuan <42993249+jzhang38@users.noreply.github.com>
This commit is contained in:
SolitaryThinker
co-authored by
Junda Su
Matthew Noto
XOR-op
Zhang Peiyuan
Junda Su
Matthew Noto
XOR-op
Zhang Peiyuan
parent
040b7735c6
commit
a62b8e1e8e
@@ -117,6 +117,18 @@ class SamplingParam:
|
||||
ltx2_stg_blocks_video: list[int] = field(default_factory=lambda: [29])
|
||||
ltx2_stg_blocks_audio: list[int] = field(default_factory=lambda: [29])
|
||||
|
||||
# LTX-2 image / video / continuation conditioning. These flow from
|
||||
# generate_video(...) kwargs through ``sampling_param.update(kwargs)``
|
||||
# onto the ForwardBatch fields of the same name. ``ltx2_image_crf``
|
||||
# gates the conditioning-image H.264 re-encode; the streaming
|
||||
# session controller passes ``ltx2_image_crf=0.0`` because it
|
||||
# conditions on already-decoded VAE-quality frames.
|
||||
ltx2_images: list[tuple[str, int, float]] | None = None
|
||||
ltx2_image_crf: float = 33.0
|
||||
ltx2_conditioning_latent_stage1: Any | None = None
|
||||
ltx2_conditioning_latent_stage2: Any | None = None
|
||||
ltx2_video_conditions: list[tuple[list[str], int, float]] | None = None
|
||||
|
||||
# Stable Audio (T2A): clip start/end in seconds. Honored by
|
||||
# `StableAudioConditioningStage` + `StableAudioDecodingStage`. Other
|
||||
# families ignore them.
|
||||
|
||||
@@ -1,37 +1,53 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""LTX-2 image / continuation conditioning helpers.
|
||||
"""FastVideo-native LTX-2 image-to-video conditioning helpers.
|
||||
|
||||
This is a *minimal* port of ``FastVideo-internal/.../ltx2_i2v_conditioning.py``
|
||||
that captures just the surface needed for the spatial-refine (SR)
|
||||
pipeline's stage-2 path:
|
||||
Public-side port of ``FastVideo-internal/.../ltx2_i2v_conditioning.py``.
|
||||
The module composes a ``clean_latent`` + ``denoise_mask`` pair that the
|
||||
LTX-2 latent-prep + denoising stages mix into the noise tensor, so a
|
||||
generated segment can be anchored to:
|
||||
|
||||
* the four ``ForwardBatch.extra`` keys (clean latent, denoise mask,
|
||||
stage-1/stage-2 last-latent for continuation),
|
||||
* :func:`apply_ltx2_gaussian_noiser` — used by both denoising and refine
|
||||
to mix noise into the conditioning latent,
|
||||
* :class:`LTX2ImageConditioningState` and a T2V-only
|
||||
:func:`build_ltx2_image_conditioning` that returns ``None`` when the
|
||||
request carries no image inputs and no continuation latents.
|
||||
* one or more conditioning images at specific latent frame indices
|
||||
(``ltx2_images``),
|
||||
* a multi-frame conditioning video clip jointly VAE-encoded
|
||||
(``ltx2_video_conditions``),
|
||||
* a continuation latent carried over from the previous segment
|
||||
(``ltx2_conditioning_latent_stage1`` / ``_stage2``).
|
||||
|
||||
The full image-loading path (CRF re-encode, decode, conditioning latent
|
||||
extraction, multi-image insertion) lives in the internal module and
|
||||
will be ported separately when the public surface needs i2v / mid-roll
|
||||
continuation. SR text-to-video does not exercise it, so keeping this
|
||||
module narrow keeps the public stage graph honest.
|
||||
The streaming server's session controller populates the continuation
|
||||
latents between segments; the legacy from_pretrained path passes
|
||||
``ltx2_images`` / ``ltx2_image_crf`` through compat translation.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from io import BytesIO
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.models.vision_utils import load_image
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
|
||||
|
||||
try:
|
||||
import av
|
||||
except ImportError: # pragma: no cover - optional dependency
|
||||
av = None
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
LTX2_VIDEO_CLEAN_LATENT_KEY = "ltx2_video_clean_latent"
|
||||
LTX2_VIDEO_DENOISE_MASK_KEY = "ltx2_video_denoise_mask"
|
||||
LTX2_CONTINUATION_STAGE1_LAST_LATENT_KEY = "ltx2_continuation_stage1_last_latent"
|
||||
LTX2_CONTINUATION_STAGE2_LAST_LATENT_KEY = "ltx2_continuation_stage2_last_latent"
|
||||
|
||||
LTX2_CONTINUATION_STAGE1_LAST_LATENT_KEY = ("ltx2_continuation_stage1_last_latent")
|
||||
LTX2_CONTINUATION_STAGE2_LAST_LATENT_KEY = ("ltx2_continuation_stage2_last_latent")
|
||||
# NOTE: hard-coded for continuation quality experiments. The first
|
||||
# latent frame of the next clip is anchored to the previous clip's
|
||||
# last latent at full strength.
|
||||
LTX2_CONTINUATION_TARGET_FRAME_IDX = 0
|
||||
LTX2_CONTINUATION_STRENGTH = 1.0
|
||||
DEFAULT_LTX2_IMAGE_CRF = 33.0
|
||||
@@ -39,38 +55,244 @@ DEFAULT_LTX2_IMAGE_CRF = 33.0
|
||||
|
||||
@dataclass
|
||||
class LTX2ImageConditioningState:
|
||||
"""Result of building image / continuation conditioning.
|
||||
|
||||
Mirrors the internal shape so the refine stage can read the same
|
||||
fields. The T2V port only ever instantiates this in the conditioning
|
||||
path; the refine code's ``image_conditioning is None`` branch is
|
||||
what fires for plain text-to-video SR.
|
||||
"""
|
||||
|
||||
"""Result of building image / continuation conditioning."""
|
||||
clean_latent: torch.Tensor
|
||||
denoise_mask: torch.Tensor
|
||||
images: list[tuple[str, int, float]]
|
||||
latent_conditioned: bool
|
||||
latent_conditioned: bool = False
|
||||
|
||||
|
||||
def resolve_ltx2_images(batch: ForwardBatch) -> list[tuple[str, int, float]]:
|
||||
"""Collect any LTX-2 image conditioning inputs declared on the batch.
|
||||
"""Collect any LTX-2 image conditioning inputs from the batch.
|
||||
|
||||
The full internal version pulls from a few legacy attribute names;
|
||||
on the public batch we only support ``batch.ltx2_images`` (a list
|
||||
of ``(path, frame_index, strength)`` triples). When the field is
|
||||
missing or empty the SR path treats the request as T2V.
|
||||
Falls back to ``batch.image_path`` for the simple single-image i2v
|
||||
case (anchors the first latent frame at full strength).
|
||||
"""
|
||||
images = getattr(batch, "ltx2_images", None) or []
|
||||
if not isinstance(images, list):
|
||||
images = batch.ltx2_images
|
||||
if images is None and batch.image_path:
|
||||
images = [(batch.image_path, 0, 1.0)]
|
||||
if not images:
|
||||
return []
|
||||
out: list[tuple[str, int, float]] = []
|
||||
for entry in images:
|
||||
if not isinstance(entry, list | tuple) or len(entry) < 3:
|
||||
continue
|
||||
path, frame_idx, strength = entry[0], entry[1], entry[2]
|
||||
out.append((str(path), int(frame_idx), float(strength)))
|
||||
return out
|
||||
|
||||
resolved: list[tuple[str, int, float]] = []
|
||||
for item in images:
|
||||
if not isinstance(item, tuple | list) or len(item) != 3:
|
||||
raise ValueError("Each ltx2_images item must be a tuple/list of "
|
||||
"(path, frame_idx, strength).")
|
||||
image_path, frame_idx, strength = item
|
||||
frame_idx_int = int(frame_idx)
|
||||
strength_float = float(strength)
|
||||
if frame_idx_int < 0:
|
||||
raise ValueError(f"LTX-2 frame_idx must be >= 0, got {frame_idx_int}")
|
||||
if strength_float < 0.0 or strength_float > 1.0:
|
||||
raise ValueError(f"LTX-2 image conditioning strength must be in [0, 1], "
|
||||
f"got {strength_float}")
|
||||
resolved.append((str(image_path), frame_idx_int, strength_float))
|
||||
return resolved
|
||||
|
||||
|
||||
def _resize_and_center_crop(
|
||||
tensor: torch.Tensor,
|
||||
height: int,
|
||||
width: int,
|
||||
) -> torch.Tensor:
|
||||
if tensor.ndim != 3:
|
||||
raise ValueError(f"Expected image tensor [H, W, C], got shape {tuple(tensor.shape)}")
|
||||
|
||||
tensor = tensor.permute(2, 0, 1).unsqueeze(0)
|
||||
_, _, src_h, src_w = tensor.shape
|
||||
scale = max(height / src_h, width / src_w)
|
||||
new_h = math.ceil(src_h * scale)
|
||||
new_w = math.ceil(src_w * scale)
|
||||
tensor = F.interpolate(
|
||||
tensor,
|
||||
size=(new_h, new_w),
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
)
|
||||
|
||||
crop_top = (new_h - height) // 2
|
||||
crop_left = (new_w - width) // 2
|
||||
tensor = tensor[:, :, crop_top:crop_top + height, crop_left:crop_left + width]
|
||||
return tensor.unsqueeze(2)
|
||||
|
||||
|
||||
def _encode_single_frame(output_file: BytesIO, image_array: np.ndarray, crf: float) -> None:
|
||||
container = av.open(output_file, "w", format="mp4")
|
||||
try:
|
||||
stream = container.add_stream(
|
||||
"libx264",
|
||||
rate=1,
|
||||
options={
|
||||
"crf": str(crf),
|
||||
"preset": "veryfast",
|
||||
},
|
||||
)
|
||||
height = image_array.shape[0] // 2 * 2
|
||||
width = image_array.shape[1] // 2 * 2
|
||||
image_array = image_array[:height, :width]
|
||||
stream.height = height
|
||||
stream.width = width
|
||||
frame = av.VideoFrame.from_ndarray(image_array, format="rgb24").reformat(format="yuv420p")
|
||||
container.mux(stream.encode(frame))
|
||||
container.mux(stream.encode())
|
||||
finally:
|
||||
container.close()
|
||||
|
||||
|
||||
def _decode_single_frame(video_file: BytesIO) -> np.ndarray:
|
||||
container = av.open(video_file)
|
||||
try:
|
||||
stream = next(s for s in container.streams if s.type == "video")
|
||||
frame = next(container.decode(stream))
|
||||
finally:
|
||||
container.close()
|
||||
return frame.to_ndarray(format="rgb24")
|
||||
|
||||
|
||||
def _preprocess_conditioning_image(
|
||||
image: np.ndarray,
|
||||
image_crf: float,
|
||||
) -> np.ndarray:
|
||||
"""H.264 CRF re-encode the conditioning image to match the
|
||||
quantization the model was trained on. ``image_crf <= 0.0`` skips
|
||||
the re-encode (used by the streaming server which conditions on
|
||||
already-decoded VAE-quality frames)."""
|
||||
if image_crf <= 0.0:
|
||||
return image
|
||||
if av is None:
|
||||
logger.warning("[LTX2] PyAV is unavailable; skipping CRF "
|
||||
"conditioning preprocessing.")
|
||||
return image
|
||||
|
||||
with BytesIO() as output_file:
|
||||
_encode_single_frame(output_file, image, image_crf)
|
||||
encoded = output_file.getvalue()
|
||||
with BytesIO(encoded) as video_file:
|
||||
return _decode_single_frame(video_file)
|
||||
|
||||
|
||||
def load_ltx2_conditioning_image(
|
||||
image_path: str,
|
||||
*,
|
||||
height: int,
|
||||
width: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
image_crf: float,
|
||||
) -> torch.Tensor:
|
||||
image = load_image(image_path)
|
||||
image_np = np.array(image)[..., :3]
|
||||
image_np = _preprocess_conditioning_image(image_np, image_crf=image_crf)
|
||||
image_tensor = torch.tensor(image_np, dtype=torch.float32, device=device)
|
||||
image_tensor = _resize_and_center_crop(image_tensor, height, width)
|
||||
image_tensor = (image_tensor / 127.5 - 1.0).to(device=device, dtype=dtype)
|
||||
return image_tensor
|
||||
|
||||
|
||||
def load_ltx2_conditioning_video_clip(
|
||||
frame_paths: list[str],
|
||||
*,
|
||||
height: int,
|
||||
width: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
image_crf: float,
|
||||
) -> torch.Tensor:
|
||||
"""Load multiple frames and stack as ``[1, C, T, H, W]`` for joint
|
||||
VAE encoding so the resulting latent captures temporal/motion info.
|
||||
"""
|
||||
frame_tensors: list[torch.Tensor] = []
|
||||
for path in frame_paths:
|
||||
image = load_image(path)
|
||||
image_np = np.array(image)[..., :3]
|
||||
image_np = _preprocess_conditioning_image(image_np, image_crf=image_crf)
|
||||
t = torch.tensor(image_np, dtype=torch.float32, device=device)
|
||||
# _resize_and_center_crop returns [1, C, 1, H, W]
|
||||
t = _resize_and_center_crop(t, height, width)
|
||||
frame_tensors.append(t)
|
||||
# Concat along T dimension -> [1, C, T, H, W]
|
||||
video = torch.cat(frame_tensors, dim=2)
|
||||
return (video / 127.5 - 1.0).to(device=device, dtype=dtype)
|
||||
|
||||
|
||||
def _extract_video_latent(vae: torch.nn.Module, image: torch.Tensor) -> torch.Tensor:
|
||||
with torch.no_grad():
|
||||
encoded = vae.encode(image)
|
||||
|
||||
latent_dist = getattr(encoded, "latent_dist", None)
|
||||
if latent_dist is not None:
|
||||
encoded = latent_dist
|
||||
|
||||
latent: torch.Tensor
|
||||
if torch.is_tensor(encoded):
|
||||
latent = encoded
|
||||
elif hasattr(encoded, "mode"):
|
||||
latent = encoded.mode()
|
||||
elif hasattr(encoded, "sample"):
|
||||
latent = encoded.sample()
|
||||
elif (isinstance(encoded, tuple | list) and encoded and torch.is_tensor(encoded[0])):
|
||||
latent = encoded[0]
|
||||
else:
|
||||
raise TypeError(f"Unsupported VAE encode output type: {type(encoded)}")
|
||||
|
||||
if latent.ndim == 4:
|
||||
latent = latent.unsqueeze(2)
|
||||
if latent.ndim != 5:
|
||||
raise ValueError(f"Expected video latent with 5 dims [B,C,T,H,W], got "
|
||||
f"{tuple(latent.shape)}")
|
||||
return latent
|
||||
|
||||
|
||||
def _insert_conditioning_latent(
|
||||
*,
|
||||
conditioning_latent: torch.Tensor,
|
||||
clean_latent: torch.Tensor,
|
||||
denoise_mask: torch.Tensor,
|
||||
frame_idx: int,
|
||||
strength: float,
|
||||
source_name: str,
|
||||
) -> None:
|
||||
if conditioning_latent.ndim == 4:
|
||||
conditioning_latent = conditioning_latent.unsqueeze(2)
|
||||
if conditioning_latent.ndim != 5:
|
||||
raise ValueError(f"LTX-2 {source_name} latent must have 5 dims [B,C,T,H,W], "
|
||||
f"got {tuple(conditioning_latent.shape)}")
|
||||
|
||||
if conditioning_latent.shape[0] == 1 and clean_latent.shape[0] > 1:
|
||||
conditioned_latent = conditioning_latent.expand(
|
||||
clean_latent.shape[0],
|
||||
-1,
|
||||
-1,
|
||||
-1,
|
||||
-1,
|
||||
)
|
||||
elif conditioning_latent.shape[0] == clean_latent.shape[0]:
|
||||
conditioned_latent = conditioning_latent
|
||||
else:
|
||||
raise ValueError(f"LTX-2 {source_name} latent batch mismatch: "
|
||||
f"{conditioning_latent.shape[0]} vs {clean_latent.shape[0]}")
|
||||
|
||||
if conditioned_latent.shape[1] != clean_latent.shape[1]:
|
||||
raise ValueError(f"LTX-2 {source_name} latent channels mismatch: "
|
||||
f"{conditioned_latent.shape[1]} vs {clean_latent.shape[1]}")
|
||||
if conditioned_latent.shape[-2:] != clean_latent.shape[-2:]:
|
||||
raise ValueError(f"LTX-2 {source_name} latent spatial shape "
|
||||
f"mismatch: {tuple(conditioned_latent.shape[-2:])} "
|
||||
f"vs {tuple(clean_latent.shape[-2:])}")
|
||||
|
||||
frame_count = conditioned_latent.shape[2]
|
||||
end_idx = frame_idx + frame_count
|
||||
if frame_idx < 0 or end_idx > clean_latent.shape[2]:
|
||||
raise ValueError(f"LTX-2 {source_name} latent frame range out of bounds: "
|
||||
f"frame_idx={frame_idx}, frame_count={frame_count}, "
|
||||
f"latent_frames={clean_latent.shape[2]}")
|
||||
|
||||
clean_latent[:, :, frame_idx:end_idx] = conditioned_latent.to(
|
||||
device=clean_latent.device,
|
||||
dtype=clean_latent.dtype,
|
||||
)
|
||||
denoise_mask[:, :, frame_idx:end_idx] = 1.0 - float(strength)
|
||||
|
||||
|
||||
def build_ltx2_image_conditioning(
|
||||
@@ -83,26 +305,152 @@ def build_ltx2_image_conditioning(
|
||||
image_crf: float | None = None,
|
||||
base_clean_latent: torch.Tensor | None = None,
|
||||
) -> LTX2ImageConditioningState | None:
|
||||
"""Return image-conditioning state, or ``None`` for plain T2V.
|
||||
"""Build the (clean_latent, denoise_mask) state for the next segment.
|
||||
|
||||
The full internal builder handles three sources: image inputs,
|
||||
continuation latents, and ``ltx2_video_conditions``. The SR
|
||||
text-to-video path takes none of those, so the early-return
|
||||
``None`` branch is what fires there. We retain the parameter shape
|
||||
so future i2v / continuation upstreaming can drop in without
|
||||
changing the refine stage call site.
|
||||
Returns ``None`` for plain T2V (no images, no continuation, no
|
||||
video conditions). The denoise mask is 1 where the model should
|
||||
sample fresh, 0 where it should preserve the conditioning latent
|
||||
exactly. ``base_clean_latent is None`` corresponds to stage 1
|
||||
(fresh half-res latent); ``base_clean_latent`` set means stage 2
|
||||
(already-upsampled latent from the upsampler stage).
|
||||
"""
|
||||
del vae, height, width, image_crf, base_clean_latent # unused in T2V
|
||||
images = resolve_ltx2_images(batch)
|
||||
has_latent_conditioning = (getattr(batch, "ltx2_conditioning_latent_stage1", None) is not None
|
||||
or getattr(batch, "ltx2_conditioning_latent_stage2", None) is not None)
|
||||
has_video_conditions = bool(getattr(batch, "ltx2_video_conditions", None))
|
||||
if not images and not has_latent_conditioning and not has_video_conditions:
|
||||
conditioning_latent_stage1 = getattr(batch, "ltx2_conditioning_latent_stage1", None)
|
||||
conditioning_latent_stage2 = getattr(batch, "ltx2_conditioning_latent_stage2", None)
|
||||
is_stage1_conditioning = base_clean_latent is None
|
||||
is_stage2_conditioning = not is_stage1_conditioning
|
||||
has_latent_conditioning = False
|
||||
continuation_latent_to_insert: torch.Tensor | None = None
|
||||
if (conditioning_latent_stage1 is not None and not torch.is_tensor(conditioning_latent_stage1)):
|
||||
raise TypeError("LTX-2 stage1 continuation latent conditioning "
|
||||
"expects a torch.Tensor.")
|
||||
if (conditioning_latent_stage2 is not None and not torch.is_tensor(conditioning_latent_stage2)):
|
||||
raise TypeError("LTX-2 stage2 continuation latent conditioning "
|
||||
"expects a torch.Tensor.")
|
||||
|
||||
if (conditioning_latent_stage1 is None) != (conditioning_latent_stage2 is None):
|
||||
raise ValueError("LTX-2 continuation expects both stage1 and stage2 "
|
||||
"conditioning latents (or neither for first round).")
|
||||
if is_stage1_conditioning and conditioning_latent_stage1 is not None:
|
||||
has_latent_conditioning = True
|
||||
continuation_latent_to_insert = conditioning_latent_stage1.to(
|
||||
device=latents.device,
|
||||
dtype=latents.dtype,
|
||||
)
|
||||
elif is_stage2_conditioning and conditioning_latent_stage2 is not None:
|
||||
has_latent_conditioning = True
|
||||
continuation_latent_to_insert = conditioning_latent_stage2.to(
|
||||
device=latents.device,
|
||||
dtype=latents.dtype,
|
||||
)
|
||||
|
||||
video_conditions = getattr(batch, "ltx2_video_conditions", None) or []
|
||||
|
||||
if not images and not has_latent_conditioning and not video_conditions:
|
||||
return None
|
||||
|
||||
raise NotImplementedError("LTX-2 image / continuation conditioning is not yet ported to the "
|
||||
"public package. Open the corresponding upstream PR before "
|
||||
"exercising this code path; SR text-to-video does not need it.")
|
||||
clean_latent = (torch.zeros_like(latents) if base_clean_latent is None else base_clean_latent.clone())
|
||||
|
||||
denoise_mask = torch.ones(
|
||||
(
|
||||
latents.shape[0],
|
||||
1,
|
||||
latents.shape[2],
|
||||
latents.shape[3],
|
||||
latents.shape[4],
|
||||
),
|
||||
dtype=torch.float32,
|
||||
device=latents.device,
|
||||
)
|
||||
|
||||
if image_crf is None:
|
||||
image_crf = getattr(batch, "ltx2_image_crf", DEFAULT_LTX2_IMAGE_CRF)
|
||||
|
||||
vae_param = next(vae.parameters(), None)
|
||||
encoder_dtype = (vae_param.dtype if vae_param is not None else latents.dtype)
|
||||
encoder_device = (vae_param.device if vae_param is not None else latents.device)
|
||||
cache: dict[tuple[str, int, int, float], torch.Tensor] = {}
|
||||
latent_conditioned = False
|
||||
|
||||
if has_latent_conditioning:
|
||||
if continuation_latent_to_insert is None:
|
||||
raise RuntimeError("LTX-2 continuation latent conditioning state is invalid.")
|
||||
# NOTE: frame index and strength are intentionally hard-coded.
|
||||
# We always anchor the first frame of the next clip at full
|
||||
# strength to the previous clip's last latent.
|
||||
_insert_conditioning_latent(
|
||||
conditioning_latent=continuation_latent_to_insert,
|
||||
clean_latent=clean_latent,
|
||||
denoise_mask=denoise_mask,
|
||||
frame_idx=LTX2_CONTINUATION_TARGET_FRAME_IDX,
|
||||
strength=LTX2_CONTINUATION_STRENGTH,
|
||||
source_name="continuation",
|
||||
)
|
||||
latent_conditioned = True
|
||||
|
||||
for image_path, frame_idx, strength in images:
|
||||
cache_key = (image_path, height, width, float(image_crf))
|
||||
image_latent = cache.get(cache_key)
|
||||
if image_latent is None:
|
||||
image_tensor = load_ltx2_conditioning_image(
|
||||
image_path=image_path,
|
||||
height=height,
|
||||
width=width,
|
||||
dtype=encoder_dtype,
|
||||
device=encoder_device,
|
||||
image_crf=float(image_crf),
|
||||
)
|
||||
image_latent = _extract_video_latent(vae, image_tensor).to(
|
||||
device=latents.device,
|
||||
dtype=latents.dtype,
|
||||
)
|
||||
cache[cache_key] = image_latent
|
||||
|
||||
_insert_conditioning_latent(
|
||||
conditioning_latent=image_latent,
|
||||
clean_latent=clean_latent,
|
||||
denoise_mask=denoise_mask,
|
||||
frame_idx=frame_idx,
|
||||
strength=strength,
|
||||
source_name="image",
|
||||
)
|
||||
|
||||
for frame_paths, frame_idx, strength in video_conditions:
|
||||
video_tensor = load_ltx2_conditioning_video_clip(
|
||||
frame_paths,
|
||||
height=height,
|
||||
width=width,
|
||||
dtype=encoder_dtype,
|
||||
device=encoder_device,
|
||||
image_crf=float(image_crf),
|
||||
)
|
||||
video_latent = _extract_video_latent(vae, video_tensor).to(
|
||||
device=latents.device,
|
||||
dtype=latents.dtype,
|
||||
)
|
||||
logger.info(
|
||||
"[LTX2] Video-clip condition: %d frames -> "
|
||||
"latent T=%d at frame_idx=%d strength=%.2f",
|
||||
len(frame_paths),
|
||||
video_latent.shape[2],
|
||||
frame_idx,
|
||||
strength,
|
||||
)
|
||||
_insert_conditioning_latent(
|
||||
conditioning_latent=video_latent,
|
||||
clean_latent=clean_latent,
|
||||
denoise_mask=denoise_mask,
|
||||
frame_idx=frame_idx,
|
||||
strength=strength,
|
||||
source_name="video_clip",
|
||||
)
|
||||
|
||||
return LTX2ImageConditioningState(
|
||||
clean_latent=clean_latent,
|
||||
denoise_mask=denoise_mask,
|
||||
images=images,
|
||||
latent_conditioned=latent_conditioned,
|
||||
)
|
||||
|
||||
|
||||
def apply_ltx2_gaussian_noiser(
|
||||
@@ -114,10 +462,9 @@ def apply_ltx2_gaussian_noiser(
|
||||
) -> torch.Tensor:
|
||||
"""Mix ``noise`` into ``clean_latent`` along ``denoise_mask`` * scale.
|
||||
|
||||
Identical to the internal helper. The mask is expected to be
|
||||
broadcastable across channels and time; values close to 1 produce
|
||||
near-pure noise (used in a fresh stage-2 latent), values near 0
|
||||
leave the clean latent untouched (used in conditioning regions).
|
||||
Values close to 1 in the mask produce near-pure noise (used in a
|
||||
fresh stage-2 latent), values near 0 leave the clean latent
|
||||
untouched (used in conditioning regions).
|
||||
"""
|
||||
scaled_mask = denoise_mask * float(noise_scale)
|
||||
return (noise * scaled_mask + clean_latent * (1.0 - scaled_mask)).to(noise.dtype)
|
||||
@@ -129,9 +476,8 @@ def post_process_ltx2_denoised(
|
||||
denoise_mask: torch.Tensor,
|
||||
clean_latent: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Inverse of :func:`apply_ltx2_gaussian_noiser` — restore the
|
||||
conditioning regions of ``clean_latent`` outside the denoise mask
|
||||
after the model has filled in the masked area."""
|
||||
"""Restore the conditioning regions of ``clean_latent`` outside the
|
||||
denoise mask after the model has filled in the masked area."""
|
||||
return (denoised * denoise_mask + clean_latent.float() * (1.0 - denoise_mask)).to(denoised.dtype)
|
||||
|
||||
|
||||
@@ -146,6 +492,8 @@ __all__ = [
|
||||
"LTX2ImageConditioningState",
|
||||
"apply_ltx2_gaussian_noiser",
|
||||
"build_ltx2_image_conditioning",
|
||||
"load_ltx2_conditioning_image",
|
||||
"load_ltx2_conditioning_video_clip",
|
||||
"post_process_ltx2_denoised",
|
||||
"resolve_ltx2_images",
|
||||
]
|
||||
|
||||
@@ -79,6 +79,27 @@ class ForwardBatch:
|
||||
image_embeds: list[torch.Tensor] = field(default_factory=list)
|
||||
pil_image: torch.Tensor | PIL.Image.Image | None = None
|
||||
preprocessed_image: torch.Tensor | None = None
|
||||
# LTX-2 image conditioning. Each entry is (path, frame_idx, strength)
|
||||
# — strength 1.0 anchors the latent at frame_idx exactly to the
|
||||
# encoded image; <1.0 mixes denoised content with the conditioning
|
||||
# image. ``ltx2_image_crf`` controls the H.264 CRF re-encode applied
|
||||
# to the image before VAE encoding (matches official LTX-2 quality
|
||||
# adaptation; 0.0 disables the re-encode).
|
||||
ltx2_images: list[tuple[str, int, float]] | None = None
|
||||
ltx2_image_crf: float = 33.0
|
||||
# Stage-specific continuation latents — populated by the streaming
|
||||
# session controller between segments. ``stage1`` is the last latent
|
||||
# from the previous segment's stage-1 denoise (half-res); ``stage2``
|
||||
# is the last latent from the previous segment's stage-2 refine
|
||||
# (full-res). The image-conditioning builder reads whichever the
|
||||
# current pass needs (``base_clean_latent is None`` -> stage1).
|
||||
ltx2_conditioning_latent_stage1: torch.Tensor | None = None
|
||||
ltx2_conditioning_latent_stage2: torch.Tensor | None = None
|
||||
# Video-clip conditioning: list of (frame_paths, frame_idx, strength)
|
||||
# — multi-frame VAE-encoded conditioning so the resulting latent
|
||||
# captures temporal/motion info (used by the streaming server's
|
||||
# mid-roll continuation flow).
|
||||
ltx2_video_conditions: list[tuple[list[str], int, float]] | None = None
|
||||
|
||||
# Text inputs
|
||||
prompt: str | list[str] | None = None
|
||||
|
||||
Reference in New Issue
Block a user