import copy import json import math import os import re import types from typing import Callable, Iterator import torch import torchaudio from accelerate import init_empty_weights from safetensors.torch import load_file from shared.utils import files_locator as fl from shared.utils.hdr import VIDEO_PROMPT_HDR_OUTPUT_FLAG, hdr_linear_to_vae_range from .ltx_core.conditioning import AudioConditionByLatent, AudioConditionByLatentPrefix, AudioConditionByReferenceLatent from .ltx_core.model.audio_vae import ( VOCODER_COMFY_KEYS_FILTER, AudioDecoderConfigurator, AudioEncoderConfigurator, AudioProcessor, VocoderConfigurator, ) from .ltx_core.model.transformer import ( LTXV_MODEL_COMFY_RENAMING_MAP, LTXModelConfigurator, X0Model, ) from .ltx_core.model.upsampler import LatentUpsamplerConfigurator from .ltx_core.model.video_vae import VideoDecoderConfigurator, VideoEncoderConfigurator from .ltx_core.text_encoders.gemma import ( GemmaTextEmbeddingsConnectorModelConfigurator, TEXT_EMBEDDING_PROJECTION_KEY_OPS, TEXT_EMBEDDINGS_CONNECTOR_KEY_OPS, build_gemma_text_encoder, ) from .ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear from .ltx_core.model.video_vae import SpatialTilingConfig, TemporalTilingConfig, TilingConfig from .ltx_core.types import AudioLatentShape, VideoPixelShape from .lora_utils import is_ic_lora_filename, phase2_ic_lora_name from .ltx_pipelines.distilled import DistilledPipeline from .ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline from .ltx_pipelines.utils.constants import AUDIO_SAMPLE_RATE, DEFAULT_NEGATIVE_PROMPT _GEMMA_FOLDER = "gemma-3-12b-it-qat-q4_0-unquantized" _SPATIAL_UPSCALER_FILENAME = "ltx-2-spatial-upscaler-x2-1.0.safetensors" LTX2_USE_FP32_ROPE_FREQS = True LTX2_ID_LORA_GUIDANCE_SCALE = 3.0 LTX2_ID_LORA_AUDIO_CFG_SCALE = 7.0 LTX2_ID_LORA_MAX_REFERENCE_SECONDS = 121.0 / 25.0 LTX2_OUTPAINT_GAMMA = 2.0 LTX2_HDR_TRANSFORM = "logc3" LTX2_DISABLE_STAGE2_WITH_CONTROL_VIDEO = True LTX2_ENABLE_EMBEDDING_LORAS = False LTX2_EMBEDDING_LORA_PREFIXES = ( "text_embedding_projection.", "feature_extractor_linear.", "text_embeddings_connector.", "embeddings_connector.", "video_embeddings_connector.", "audio_embeddings_connector.", ) def _normalize_config(config_value): if isinstance(config_value, dict): return config_value if isinstance(config_value, (bytes, bytearray, memoryview)): try: config_value = bytes(config_value).decode("utf-8") except Exception: return {} if isinstance(config_value, str): try: return json.loads(config_value) except json.JSONDecodeError: return {} return {} def _is_editanything_model(model_def) -> bool: return bool((model_def or {}).get("ltx2_edit_anything", False)) def _load_config_from_checkpoint(path, fallback_config_path: str | None = None): from mmgp import quant_router if isinstance(path, (list, tuple)): if not path: return {} path = path[0] if not path: return {} def _read_config_metadata(one_path: str) -> dict: if not one_path: return {} _, metadata = quant_router.load_metadata_state_dict(one_path) if not metadata: return {} return _normalize_config(metadata.get("config")) config = _read_config_metadata(path) if config: return config if not fallback_config_path: return {} try: with open(fallback_config_path, "r", encoding="utf-8") as reader: return _normalize_config(json.load(reader)) except Exception: return {} def _strip_model_prefix(key: str) -> str: for prefix in ("model.", "velocity_model."): if key.startswith(prefix): return _strip_model_prefix(key[len(prefix) :]) return key def _apply_sd_ops(state_dict: dict, quantization_map: dict | None, sd_ops): if sd_ops is not None: has_match = False for key in state_dict.keys(): key = _strip_model_prefix(key) if sd_ops.apply_to_key(key) is not None: has_match = True break if not has_match: new_sd = {_strip_model_prefix(k): v for k, v in state_dict.items()} new_qm = {} if quantization_map: new_qm = {_strip_model_prefix(k): v for k, v in quantization_map.items()} return new_sd, new_qm new_sd = {} for key, value in state_dict.items(): key = _strip_model_prefix(key) if sd_ops is None: new_sd[key] = value continue else: new_key = sd_ops.apply_to_key(key) if new_key is None: continue new_pairs = sd_ops.apply_to_key_value(new_key, value) for pair in new_pairs: new_sd[pair.new_key] = pair.new_value new_qm = {} if quantization_map: for key, value in quantization_map.items(): key = _strip_model_prefix(key) if sd_ops is None: new_key = key else: new_key = sd_ops.apply_to_key(key) if new_key is None: continue new_qm[new_key] = value return new_sd, new_qm def _make_sd_postprocess(sd_ops): def postprocess(state_dict, quantization_map): return _apply_sd_ops(state_dict, quantization_map, sd_ops) return postprocess def _split_vae_state_dict(state_dict: dict, prefix: str): new_sd = {} for key, value in state_dict.items(): key = _strip_model_prefix(key) if key.startswith(prefix): key = key[len(prefix) :] elif key.startswith(("encoder.", "decoder.", "per_channel_statistics.")): key = key else: continue if key.startswith("per_channel_statistics."): suffix = key[len("per_channel_statistics.") :] new_sd[f"encoder.per_channel_statistics.{suffix}"] = value.clone() new_sd[f"decoder.per_channel_statistics.{suffix}"] = value.clone() else: new_sd[key] = value return new_sd, {} def _make_vae_postprocess(prefix: str): def postprocess(state_dict, quantization_map): return _split_vae_state_dict(state_dict, prefix) return postprocess class _AudioVAEWrapper(torch.nn.Module): def __init__(self, decoder: torch.nn.Module) -> None: super().__init__() per_stats = getattr(decoder, "per_channel_statistics", None) if per_stats is not None: self.per_channel_statistics = per_stats self.decoder = decoder class _VAEContainer(torch.nn.Module): def __init__(self, encoder: torch.nn.Module, decoder: torch.nn.Module) -> None: super().__init__() self.encoder = encoder self.decoder = decoder class _ExternalConnectorWrapper: def __init__(self, module: torch.nn.Module) -> None: self._module = module def __call__(self, *args, **kwargs): return self._module(*args, **kwargs) class LTX2SuperModel(torch.nn.Module): def __init__(self, ltx2_model: "LTX2") -> None: super().__init__() object.__setattr__(self, "_ltx2", ltx2_model) transformer = ltx2_model.model velocity_model = getattr(transformer, "velocity_model", transformer) self.velocity_model = velocity_model split_map = getattr(transformer, "split_linear_modules_map", None) if split_map is not None: self.split_linear_modules_map = split_map self.text_embedding_projection = ltx2_model.text_embedding_projection self.text_embeddings_connector = ltx2_model.text_embeddings_connector @property def _interrupt(self) -> bool: return self._ltx2._interrupt @_interrupt.setter def _interrupt(self, value: bool) -> None: self._ltx2._interrupt = value def forward(self, *args, **kwargs): return self._ltx2.model(*args, **kwargs) def generate(self, *args, **kwargs): return self._ltx2.generate(*args, **kwargs) def get_trans_lora(self): return self, None def __getattr__(self, name): try: return super().__getattr__(name) except AttributeError: return getattr(self._ltx2, name) class _LTX2VAEHelper: def __init__(self, block_size: int = 64) -> None: self.block_size = block_size def get_VAE_tile_size( self, vae_config: int, device_mem_capacity: float, mixed_precision: bool, output_height: int | None = None, output_width: int | None = None, ) -> int | tuple[int, int]: if vae_config >= 4: vae_config = 0 if vae_config == 0: if mixed_precision: device_mem_capacity = device_mem_capacity / 1.5 if device_mem_capacity >= 24000: use_vae_config = 1 elif device_mem_capacity >= 8000: use_vae_config = 2 else: use_vae_config = 3 else: use_vae_config = vae_config ref_size = output_height if output_height is not None else output_width if ref_size is not None and ref_size > 480: use_vae_config += 1 spatial_tile_size = 128 if use_vae_config <= 1: spatial_tile_size = 0 elif use_vae_config == 2: spatial_tile_size = 512 elif use_vae_config == 3: spatial_tile_size = 256 return spatial_tile_size def _attach_lora_preprocessor(transformer: torch.nn.Module) -> None: def preprocess_loras(self: torch.nn.Module, model_type: str, sd: dict) -> dict: if not sd: return sd module_names = getattr(self, "_lora_module_names", None) if module_names is None: module_names = {name for name, _ in self.named_modules()} self._lora_module_names = module_names def split_lora_key(lora_key: str) -> tuple[str | None, str]: if lora_key.endswith(".alpha"): return lora_key[: -len(".alpha")], ".alpha" if lora_key.endswith(".diff"): return lora_key[: -len(".diff")], ".diff" if lora_key.endswith(".diff_b"): return lora_key[: -len(".diff_b")], ".diff_b" if lora_key.endswith(".dora_scale"): return lora_key[: -len(".dora_scale")], ".dora_scale" pos = lora_key.rfind(".lora_") if pos > 0: return lora_key[:pos], lora_key[pos:] return None, "" new_sd = {} dropped_keys = [] for key, value in sd.items(): original_key = key if key.startswith("model."): key = key[len("model.") :] if key.startswith("diffusion_model."): key = key[len("diffusion_model.") :] if key.startswith("transformer."): key = key[len("transformer.") :] if not LTX2_ENABLE_EMBEDDING_LORAS and key.startswith(LTX2_EMBEDDING_LORA_PREFIXES): continue if key.startswith("embeddings_connector."): key = f"text_embeddings_connector.video_embeddings_connector.{key[len('embeddings_connector.'):]}" if key.startswith("video_embeddings_connector."): key = f"text_embeddings_connector.{key}" if key.startswith("audio_embeddings_connector."): key = f"text_embeddings_connector.{key}" if key.startswith("feature_extractor_linear."): key = f"text_embedding_projection.{key[len('feature_extractor_linear.'):]}" module_name, suffix = split_lora_key(key) if not module_name: dropped_keys.append(original_key) continue if module_name not in module_names: prefixed_name = f"velocity_model.{module_name}" if prefixed_name in module_names: module_name = prefixed_name else: dropped_keys.append(original_key) continue new_sd[f"{module_name}{suffix}"] = value if dropped_keys: sample = ", ".join(dropped_keys[:8]) if len(dropped_keys) > 8: sample += ", ..." raise ValueError( f"LTX2 LoRA preprocessing dropped {len(dropped_keys)} unmatched keys for model '{model_type}': {sample}" ) return new_sd transformer.preprocess_loras = types.MethodType(preprocess_loras, transformer) def _coerce_image_list(image_value): if isinstance(image_value, list): return image_value[0] if image_value else None return image_value def _duplicate_ref_image_as_video(ref_image, frame_count: int = 9): if ref_image is None: return None frame_count = max(1, int(frame_count)) if isinstance(ref_image, (list, tuple)): ref_image = ref_image[0] if ref_image else None if ref_image is None: return None if torch.is_tensor(ref_image): image = ref_image.detach() if image.ndim == 3: if image.shape[0] in (1, 3, 4): return image.unsqueeze(1).repeat(1, frame_count, 1, 1) return image.unsqueeze(0).repeat(frame_count, 1, 1, 1) if image.ndim == 4: if image.shape[0] in (1, 3, 4): return image[:, :1].repeat(1, frame_count, 1, 1) if image.shape[-1] in (1, 3, 4): return image[:1].repeat(frame_count, 1, 1, 1) return image import numpy as np from PIL import Image if isinstance(ref_image, str): with Image.open(ref_image) as image: frame = np.array(image.convert("RGB")) else: frame = np.array(ref_image)[..., :3] return np.repeat(frame[None, ...], frame_count, axis=0) def _to_latent_index(frame_idx: int, stride: int) -> int: frame_idx = int(frame_idx) stride = int(stride) if frame_idx <= 0: return 0 # Causal LTX VAEs keep pixel frame 0 in its own latent slot. return (frame_idx - 1) // stride + 1 def _normalize_tiling_size(tile_size: int) -> int: tile_size = int(tile_size) if tile_size <= 0: return 0 tile_size = max(64, tile_size) if tile_size % 32 != 0: tile_size = int(math.ceil(tile_size / 32) * 32) return tile_size def _normalize_temporal_tiling_size(tile_frames: int) -> int: tile_frames = int(tile_frames) if tile_frames <= 0: return 0 tile_frames = max(16, tile_frames) if tile_frames % 8 != 0: tile_frames = int(math.ceil(tile_frames / 8) * 8) return tile_frames def _normalize_temporal_overlap(overlap_frames: int, tile_frames: int) -> int: overlap_frames = max(0, int(overlap_frames)) if overlap_frames % 8 != 0: overlap_frames = int(round(overlap_frames / 8) * 8) overlap_frames = max(0, min(overlap_frames, max(0, tile_frames - 8))) return overlap_frames def _build_tiling_config(tile_size: int | tuple | list | None, num_frames: int | None) -> TilingConfig | None: temporal_tiling_divisor = 1 spatial_config = None if isinstance(tile_size, (tuple, list)): if len(tile_size) == 0: tile_size = None else: if len(tile_size) > 1: temporal_tiling_divisor = max(1, int(tile_size[0] or 1)) tile_size = tile_size[-1] if tile_size is not None: tile_size = _normalize_tiling_size(tile_size) if tile_size > 0: overlap = max(0, tile_size // 4) overlap = int(math.floor(overlap / 32) * 32) if overlap >= tile_size: overlap = max(0, tile_size - 32) spatial_config = SpatialTilingConfig(tile_size_in_pixels=tile_size, tile_overlap_in_pixels=overlap) temporal_config = None if num_frames is not None and num_frames > 241: temporal_tiling_divisor = max(1, temporal_tiling_divisor) tile_frames = _normalize_temporal_tiling_size(int(math.ceil(232 / temporal_tiling_divisor))) if tile_frames > 0: overlap_frames = int(round(tile_frames * 3 / 8)) overlap_frames = _normalize_temporal_overlap(overlap_frames, tile_frames) temporal_config = TemporalTilingConfig( tile_size_in_frames=tile_frames, tile_overlap_in_frames=overlap_frames, ) if spatial_config is None and temporal_config is None: return None return TilingConfig(spatial_config=spatial_config, temporal_config=temporal_config) def _infer_ic_lora_downscale_factor(loras_selected) -> int | None: factors = [] for lora_path in loras_selected or []: name = os.path.basename(str(lora_path)).lower() if not is_ic_lora_filename(name): continue match = re.search(r"-ref([0-9]+(?:\.[0-9]+)?)", name) if not match: factors.append(1) continue ref_ratio = float(match.group(1)) if ref_ratio <= 0: factors.append(1) continue factors.append(max(1, int(round(1.0 / ref_ratio)))) if not factors: return None return min(factors) def _collect_video_chunks( video: Iterator[torch.Tensor] | torch.Tensor, interrupt_check: Callable[[], bool] | None = None, expected_frames: int | None = None, expected_height: int | None = None, expected_width: int | None = None, ) -> torch.Tensor | None: iterator = None if video is None: return None try: if torch.is_tensor(video): frames = video if expected_height is not None or expected_width is not None: frames = frames[:, :expected_height, :expected_width] return frames.permute(3, 0, 1, 2) else: iterator = iter(video) video_tensor = None write_pos = 0 for chunk in iterator: if interrupt_check is not None and interrupt_check(): return None if chunk is None: continue chunk = chunk if torch.is_tensor(chunk) else torch.tensor(chunk) if expected_height is not None or expected_width is not None: chunk = chunk[:, :expected_height, :expected_width] if video_tensor is None: channels = int(chunk.shape[-1]) frame_capacity = int(expected_frames) if expected_frames is not None and expected_frames > 0 else int(chunk.shape[0]) video_tensor = torch.empty( (channels, frame_capacity, chunk.shape[1], chunk.shape[2]), dtype=chunk.dtype, device=chunk.device, ) frame_count = min(int(chunk.shape[0]), int(video_tensor.shape[1] - write_pos)) if frame_count <= 0: break video_tensor[:, write_pos : write_pos + frame_count].copy_(chunk[:frame_count].permute(3, 0, 1, 2)) write_pos += frame_count if video_tensor is None: return None return video_tensor[:, :write_pos] finally: if iterator is not None: close = getattr(iterator, "close", None) if close is not None: close() # frames = frames.to(dtype=torch.float32).div_(127.5).sub_(1.0) # return frames.permute(3, 0, 1, 2).contiguous() def _build_frozen_control_video( input_frames: torch.Tensor | None, input_video: torch.Tensor | None, frame_num: int, prefix_frames_count: int, latent_stride: int = 8, ) -> torch.Tensor: if input_frames is None: raise ValueError("LTX2 audio-from-control-video mode requires a raw Control Video.") requested_frames = int(frame_num) prefix_frames = 0 if input_video is not None and prefix_frames_count > 0: prefix_frames = min(int(prefix_frames_count), int(input_video.shape[1])) target_frames = min(requested_frames, prefix_frames + int(input_frames.shape[1])) target_frames = ((target_frames - 1) // int(latent_stride)) * int(latent_stride) + 1 pieces = [] remaining_frames = target_frames if prefix_frames > 0: prefix = input_video[:, : min(prefix_frames, target_frames)] pieces.append(prefix) remaining_frames -= int(prefix.shape[1]) if remaining_frames > 0: tail = input_frames if tail.shape[1] > remaining_frames: tail = tail[:, -remaining_frames:] if pieces else tail[:, :remaining_frames] pieces.append(tail) if not pieces: raise ValueError("LTX2 audio-from-control-video mode received no Control Video frames.") frozen_video = torch.cat(pieces, dim=1) if len(pieces) > 1 else pieces[0] return frozen_video[:, :target_frames] def _normalize_outpainting_dims(outpainting_dims) -> list[float] | None: if outpainting_dims is None: return None if isinstance(outpainting_dims, str): outpainting_dims = outpainting_dims.strip() if not outpainting_dims or outpainting_dims.startswith("#"): return None outpainting_dims = outpainting_dims.split() if not isinstance(outpainting_dims, (list, tuple)) or len(outpainting_dims) != 4: return None dims = [max(0.0, float(v)) for v in outpainting_dims] return dims if any(dims) else None def _get_outpainting_inner_rect(height: int, width: int, outpainting_dims) -> tuple[int, int, int, int] | None: dims = _normalize_outpainting_dims(outpainting_dims) if dims is None or height <= 0 or width <= 0: return None from shared.utils.utils import get_outpainting_frame_location inner_height, inner_width, margin_top, margin_left = get_outpainting_frame_location(int(height), int(width), dims, 1) top = max(0, min(int(margin_top), int(height))) left = max(0, min(int(margin_left), int(width))) bottom = max(top, min(top + int(inner_height), int(height))) right = max(left, min(left + int(inner_width), int(width))) return (top, bottom, left, right) if bottom > top and right > left else None def _apply_gamma_to_media(media_tensor: torch.Tensor | None, gamma: float) -> bool: if media_tensor is None or not torch.is_tensor(media_tensor) or media_tensor.dim() < 2 or gamma <= 0 or media_tensor.numel() == 0: return False exponent = 1.0 / float(gamma) if media_tensor.dtype == torch.uint8: corrected = media_tensor.to(dtype=torch.float32).div_(255.0).clamp_(0.0, 1.0).pow_(exponent) media_tensor.copy_(corrected.mul_(255.0).round_().clamp_(0.0, 255.0).to(dtype=torch.uint8)) return True corrected = media_tensor.to(dtype=torch.float32).add_(1.0).mul_(0.5).clamp_(0.0, 1.0).pow_(exponent) media_tensor.copy_(corrected.mul_(2.0).sub_(1.0).to(dtype=media_tensor.dtype)) return True def _apply_gamma_to_video_rect(video_tensor: torch.Tensor | None, rect: tuple[int, int, int, int] | None, gamma: float) -> bool: if video_tensor is None or not torch.is_tensor(video_tensor) or rect is None or video_tensor.dim() < 4: return False top, bottom, left, right = rect region = video_tensor[..., top:bottom, left:right] return _apply_gamma_to_media(region, gamma) class LTX2: def __init__( self, model_filename, model_type: str, base_model_type: str, model_def: dict, dtype: torch.dtype = torch.bfloat16, VAE_dtype: torch.dtype = torch.float32, text_encoder_filename: str | None = None, text_encoder_filepath = None, checkpoint_paths: dict | None = None, ) -> None: self.device = torch.device("cuda") self.dtype = dtype self.VAE_dtype = VAE_dtype self.base_model_type = base_model_type self.model_def = model_def self._interrupt = False self._hdr_scene_context = None self.vae = _LTX2VAEHelper() from .ltx_core.model.transformer import rope as rope_utils self.use_fp32_rope_freqs = bool(model_def.get("ltx2_rope_freqs_fp32", LTX2_USE_FP32_ROPE_FREQS)) rope_utils.set_use_fp32_rope_freqs(self.use_fp32_rope_freqs) if isinstance(model_filename, (list, tuple)): if not model_filename: raise ValueError("Missing LTX-2 checkpoint path.") transformer_path = list(model_filename) else: transformer_path = model_filename component_paths = checkpoint_paths or {} if component_paths: transformer_path = component_paths.get("transformer") if not transformer_path: raise ValueError("Missing transformer path in checkpoint_paths.") gemma_root = text_encoder_filepath if text_encoder_filename is None else text_encoder_filename if not gemma_root: raise ValueError("Missing Gemma text encoder path.") if component_paths: spatial_upsampler_path = component_paths.get("spatial_upsampler") else: spatial_upsampler_path = None if not spatial_upsampler_path: spatial_upsampler_name = model_def.get("ltx2_spatial_upscaler_file", _SPATIAL_UPSCALER_FILENAME) spatial_upsampler_path = fl.locate_file(spatial_upsampler_name) # Internal FP8 handling is disabled; mmgp manages quantization/dtypes. pipeline_kind = model_def.get("ltx2_pipeline", "two_stage") pipeline_models = self._init_models( transformer_path=transformer_path, component_paths=component_paths, gemma_root=gemma_root, spatial_upsampler_path=spatial_upsampler_path, ) if pipeline_kind == "distilled": self.pipeline = DistilledPipeline( device=self.device, models=pipeline_models, ) else: self.pipeline = TI2VidTwoStagesPipeline( device=self.device, stage_1_models=pipeline_models, stage_2_models=pipeline_models, ) self._build_diffuser_model() def _init_models( self, transformer_path, component_paths: dict, gemma_root: str, spatial_upsampler_path: str, ): from mmgp import offload as mmgp_offload fallback_config_path = component_paths.get("model_config") if component_paths else None base_config = _load_config_from_checkpoint(transformer_path, fallback_config_path=fallback_config_path) if not base_config: raise ValueError("Missing config in transformer checkpoint.") def _component_path(key: str): if component_paths: path = component_paths.get(key) if not path: raise ValueError(f"Missing '{key}' path in checkpoint_paths.") return path return transformer_path def _component_config(path): config = _load_config_from_checkpoint(path, fallback_config_path=fallback_config_path) return config or base_config def _load_component(model, path, sd_ops=None, postprocess=None, ignore_unused_weights=False): if postprocess is None and sd_ops is not None: postprocess = _make_sd_postprocess(sd_ops) mmgp_offload.load_model_data( model, path, postprocess_sd=postprocess, default_dtype=self.dtype, writable_tensors=False, ignore_missing_keys=False, ignore_unused_weights=ignore_unused_weights, ) model.eval().requires_grad_(False) return model transformer_sd_ops = LTXV_MODEL_COMFY_RENAMING_MAP with init_empty_weights(): velocity_model = LTXModelConfigurator.from_config(base_config) velocity_model = _load_component(velocity_model, transformer_path, transformer_sd_ops, ignore_unused_weights=True) transformer_modules = component_paths.get("transformer_modules") if component_paths else None if transformer_modules: from .editanything import install_editanything_modules install_editanything_modules(velocity_model, transformer_modules, self.model_def) transformer = X0Model(velocity_model) transformer.eval().requires_grad_(False) VAE_URLs = self.model_def.get("VAE_URLs", None) video_vae_path = fl.locate_file(VAE_URLs[0]) if VAE_URLs is not None and len(VAE_URLs) else _component_path("video_vae") video_config = copy.deepcopy(_component_config(video_vae_path)) video_config_vae = video_config.setdefault("vae", {}) video_config_vae["spatial_padding_mode"] = "reflect" video_config_vae["encoder_spatial_padding_mode"] = "reflect" video_config_vae["decoder_spatial_padding_mode"] = "reflect" # print("[LTX2 VAE Config] forcing encoder/decoder spatial_padding_mode=reflect") with init_empty_weights(): video_encoder = VideoEncoderConfigurator.from_config(video_config) video_decoder = VideoDecoderConfigurator.from_config(video_config) video_vae = _VAEContainer(video_encoder, video_decoder) video_vae = _load_component(video_vae, video_vae_path, postprocess=_make_vae_postprocess("vae."), ignore_unused_weights=True) video_encoder = video_vae.encoder video_decoder = video_vae.decoder audio_vae_path = _component_path("audio_vae") audio_config = _component_config(audio_vae_path) with init_empty_weights(): audio_encoder = AudioEncoderConfigurator.from_config(audio_config) audio_decoder = AudioDecoderConfigurator.from_config(audio_config) audio_vae = _VAEContainer(audio_encoder, audio_decoder) audio_vae = _load_component(audio_vae, audio_vae_path, postprocess=_make_vae_postprocess("audio_vae.")) audio_encoder = audio_vae.encoder audio_decoder = audio_vae.decoder vocoder_path = _component_path("vocoder") vocoder_config = _component_config(vocoder_path) with init_empty_weights(): vocoder = VocoderConfigurator.from_config(vocoder_config) vocoder = _load_component(vocoder, vocoder_path, VOCODER_COMFY_KEYS_FILTER) text_projection_path = _component_path("text_embedding_projection") text_projection_config = _component_config(text_projection_path) with init_empty_weights(): text_embedding_projection = GemmaFeaturesExtractorProjLinear.from_config(text_projection_config) text_embedding_projection = _load_component( text_embedding_projection, text_projection_path, TEXT_EMBEDDING_PROJECTION_KEY_OPS ) text_connector_path = _component_path("text_embeddings_connector") text_connector_config = _component_config(text_connector_path) with init_empty_weights(): text_embeddings_connector = GemmaTextEmbeddingsConnectorModelConfigurator.from_config(text_connector_config) text_embeddings_connector = _load_component( text_embeddings_connector, text_connector_path, TEXT_EMBEDDINGS_CONNECTOR_KEY_OPS ) text_encoder = build_gemma_text_encoder(gemma_root, default_dtype=self.dtype) text_encoder.eval().requires_grad_(False) upsampler_config = _load_config_from_checkpoint(spatial_upsampler_path) with init_empty_weights(): spatial_upsampler = LatentUpsamplerConfigurator.from_config(upsampler_config) spatial_upsampler = _load_component(spatial_upsampler, spatial_upsampler_path, None) self.text_encoder = text_encoder self.text_embedding_projection = text_embedding_projection self.text_embeddings_connector = text_embeddings_connector self.video_embeddings_connector = text_embeddings_connector.video_embeddings_connector self.audio_embeddings_connector = text_embeddings_connector.audio_embeddings_connector self.video_encoder = video_encoder self.video_decoder = video_decoder self.audio_encoder = audio_encoder self.audio_decoder = audio_decoder self.vocoder = vocoder self.spatial_upsampler = spatial_upsampler self.model = transformer self.model2 = None return types.SimpleNamespace( text_encoder=self.text_encoder, text_embedding_projection=self.text_embedding_projection, text_embeddings_connector=self.text_embeddings_connector, video_encoder=self.video_encoder, video_decoder=self.video_decoder, audio_encoder=self.audio_encoder, audio_decoder=self.audio_decoder, vocoder=self.vocoder, spatial_upsampler=self.spatial_upsampler, transformer=self.model, ) def _load_hdr_scene_context(self, lora_dir: str | None = None) -> tuple[torch.Tensor, torch.Tensor]: cached = self._hdr_scene_context if cached is not None: return cached path = fl.locate_file(self.model_def.get("ltx2_hdr_scene_embeddings_file", ""), error_if_none=False) tensors = load_file(path, device="cpu") self._hdr_scene_context = (tensors["video_context"].detach().cpu(), tensors["audio_context"].detach().cpu()) return self._hdr_scene_context def _detach_text_encoder_connectors(self) -> None: text_encoder = getattr(self, "text_encoder", None) if text_encoder is None: return connectors = {} feature_extractor = getattr(self, "text_embedding_projection", None) video_connector = getattr(self, "video_embeddings_connector", None) audio_connector = getattr(self, "audio_embeddings_connector", None) if feature_extractor is not None: connectors["feature_extractor_linear"] = feature_extractor if video_connector is not None: connectors["embeddings_connector"] = video_connector if audio_connector is not None: connectors["audio_embeddings_connector"] = audio_connector if not connectors: return for name, module in connectors.items(): if name in text_encoder._modules: del text_encoder._modules[name] setattr(text_encoder, name, _ExternalConnectorWrapper(module)) self._text_connectors = connectors def _build_diffuser_model(self) -> None: self._detach_text_encoder_connectors() self.diffuser_model = LTX2SuperModel(self) _attach_lora_preprocessor(self.diffuser_model) def get_trans_lora(self): trans = getattr(self, "diffuser_model", None) if trans is None: trans = self.model return trans, None def get_loras_transformer(self, get_model_recursive_prop, model_type, video_prompt_type, base_model_type=None, model_def = None, lora_dir = None, sample_solver = None, **kwargs): control_map = { "O": "pose_align", "P": "pose", "D": "depth", "E": "canny", } from shared.utils.utils import get_outpainting_dims loras = [] loras_mult = [] guidance_phases = max(1, int(kwargs["guidance_phases"])) audio_prompt_type = kwargs["audio_prompt_type"] outpainting_ratio = kwargs["video_guide_outpainting_ratio"].strip() outpainting_setting = str(kwargs["video_guide_outpainting"]) pipeline_kind = model_def.get("ltx2_pipeline", "two_stage") resolved_base_model_type = base_model_type sample_solver = (sample_solver or "").lower() selected_loras = {os.path.basename(lora).lower() for lora in kwargs.get("activated_loras", [])} preload_urls = get_model_recursive_prop(model_type, "preload_URLs", return_list=True) if isinstance(preload_urls, str): preload_urls = [preload_urls] def _get_preload_lora_url(signature): matched_url = None for entry in preload_urls: if isinstance(entry, str) and entry.endswith("|%lora_dir"): source_url = entry.split("|", 1)[0] if signature in os.path.basename(source_url).lower(): matched_url = source_url return matched_url def _append_system_lora(name, multiplier, signature): signature = signature.lower() url = _get_preload_lora_url(signature) or model_def.get(f"ltx2_lora_{name}", "") if not url: return if any(signature in os.path.basename(lora).lower() for lora in loras): return for lora in selected_loras: if signature in lora: print(f"Default system '{signature}' lora and corresponding multiplier will be ignored as User has provided its own lora ({lora})") return loras.append(url) loras_mult.append(multiplier) if pipeline_kind != "distilled" and (guidance_phases > 1 or sample_solver in {"distilled_8_steps", "res2s"}): use_hq_sampler = sample_solver == "res2s" use_distilled_8_steps = sample_solver == "distilled_8_steps" use_id_lora = "1" in audio_prompt_type if guidance_phases == 1 and use_hq_sampler: mult = 0.2 elif guidance_phases == 1 and use_distilled_8_steps: mult = 0.5 elif use_hq_sampler: mult = "0.25;0.5" elif use_id_lora: mult = "0;0.8" elif use_distilled_8_steps: mult = "0.5;0.5" else: mult = "0;1" _append_system_lora("distilled", mult, "distilled-lora") if resolved_base_model_type == "ltx2_22B" and VIDEO_PROMPT_HDR_OUTPUT_FLAG in video_prompt_type: _append_system_lora("hdr", 1.0, "ic-lora-hdr") if any(letter in video_prompt_type for letter in control_map): _append_system_lora("union_control", 1.0, "union-control") if resolved_base_model_type == "ltx2_22B" and get_outpainting_dims(outpainting_setting, outpainting_ratio) is not None: _append_system_lora("outpaint", 1.0, "outpaint") if "1" in audio_prompt_type: _append_system_lora("id", 1.0 if guidance_phases == 1 else "1;0", "id-lora-celebvhq") return loras, loras_mult def generate( self, input_prompt: str, n_prompt: str | None = None, image_start=None, image_end=None, sampling_steps: int = 40, guide_scale: float = 4.0, alt_guide_scale: float = 1.0, input_video=None, prefix_frames_count: int = 0, conditioning_latents_size: int = 0, window_no: int = 1, input_frames=None, input_frames2=None, frames_to_inject = None, input_masks=None, input_masks2=None, frames_relative_positions_list=None, masking_strength: float | None = None, input_video_strength: float | None = None, return_latent_slice=None, video_prompt_type: str = "", audio_prompt_type: str = "", denoising_strength: float | None = None, cfg_star_switch: int = 0, apg_switch: int = 0, perturbation_switch: int = 0, perturbation_layers: list[int] | None = None, perturbation_start: float = 0.0, perturbation_end: float = 1.0, audio_cfg_scale: float | None = None, alt_scale: float = 0.0, sample_solver: str = "", NAG_scale: float = 1.0, NAG_tau: float = 3.5, NAG_alpha: float = 0.5, self_refiner_setting: int = 0, self_refiner_plan: str = "", self_refiner_f_uncertainty: float = 0.1, self_refiner_certain_percentage: float = 0.999, loras_slists=None, loras_selected=None, text_connectors=None, input_ref_images=None, input_ref_masks=None, input_waveform=None, input_waveform_sample_rate=None, audio_scale: float | None = None, masking_source: dict | None = None, outpainting_dims: list[int] | None = None, frame_num: int = 121, height: int = 1024, width: int = 1536, fps: float = 24.0, seed: int = 0, callback=None, set_progress_status=None, VAE_tile_size=None, guide_phases= 1, **kwargs, ): if self._interrupt: return None distill = self.model_def.get("ltx2_pipeline", "two_stage") == "distilled" editanything = _is_editanything_model(self.model_def) hdr_enabled = self.base_model_type == "ltx2_22B" and VIDEO_PROMPT_HDR_OUTPUT_FLAG in video_prompt_type input_video_is_hdr = bool(kwargs.get("input_video_is_hdr", False)) hdr_scene_context = self._load_hdr_scene_context(kwargs.get("lora_dir")) if hdr_enabled else None if hdr_enabled: NAG_scale = 1.0 audio_prompt_type = "" input_waveform = None audio_from_control_video = "2" in audio_prompt_type image_start = _coerce_image_list(image_start) image_end = _coerce_image_list(image_end) if frames_to_inject is None: frames_to_inject = [] if frames_relative_positions_list is None: frames_relative_positions_list = [] elif isinstance(frames_relative_positions_list, (list, tuple)): frames_relative_positions_list = list(frames_relative_positions_list) else: frames_relative_positions_list = [frames_relative_positions_list] if image_start is None: new_frames_to_inject = [] new_frames_relative_positions_list = [] for frame_to_inject, frame_relative_position in zip(frames_to_inject,frames_relative_positions_list): if frame_relative_position == 0: image_start = frame_to_inject else: new_frames_to_inject.append(frame_to_inject) new_frames_relative_positions_list.append(frame_relative_position) frames_to_inject = new_frames_to_inject frames_relative_positions_list = new_frames_relative_positions_list outpainting_dims = _normalize_outpainting_dims(outpainting_dims) any_outpainting = outpainting_dims is not None and "V" in video_prompt_type self_refiner_max_plans = self.model_def.get("self_refiner_max_plans", 1) requested_outpaint_gamma_roundtrip = self.base_model_type == "ltx2_22B" and any_outpainting if hdr_enabled: requested_outpaint_gamma_roundtrip = False if any_outpainting: guide_phases = 1 use_outpaint_gamma_roundtrip = False latent_stride = 8 if hasattr(self.pipeline, "pipeline_components"): scale_factors = getattr(self.pipeline.pipeline_components, "video_scale_factors", None) if scale_factors is not None: latent_stride = int(getattr(scale_factors, "time", scale_factors[0])) input_video_strength = max(0.0, min(1.0, input_video_strength)) if requested_outpaint_gamma_roundtrip: conditioning_gamma_applied = _apply_gamma_to_media(image_start, LTX2_OUTPAINT_GAMMA) conditioning_gamma_applied = _apply_gamma_to_media(image_end, LTX2_OUTPAINT_GAMMA) or conditioning_gamma_applied if torch.is_tensor(input_video) and prefix_frames_count > 0: conditioning_gamma_applied = _apply_gamma_to_media(input_video[:, :prefix_frames_count], LTX2_OUTPAINT_GAMMA) or conditioning_gamma_applied for ref_image in frames_to_inject: conditioning_gamma_applied = _apply_gamma_to_media(ref_image, LTX2_OUTPAINT_GAMMA) or conditioning_gamma_applied if conditioning_gamma_applied: print("[WAN2GP][LTX2] Applying full-frame gamma preprocessing for outpainting IC-LoRA conditioning images.") use_outpaint_gamma_roundtrip = True if "G" not in video_prompt_type: denoising_strength = 1.0 masking_strength = 0.0 if hdr_enabled and input_video_is_hdr and torch.is_tensor(input_video): input_video = hdr_linear_to_vae_range(input_video, transform=LTX2_HDR_TRANSFORM).to(dtype=input_video.dtype) control_strength = denoising_strength ic_lora_downscale_factor = None ic_lora_downscale_factor = _infer_ic_lora_downscale_factor(loras_selected) video_conditioning_downscale_factor = ic_lora_downscale_factor or 1 # merge_conditioning_and_guide = False has_prefix_frames = input_video is not None is_start_image_only = image_start is not None and (not has_prefix_frames or prefix_frames_count <= 1) merge_conditioning_and_guide = continuous_conditioning_and_guide = False video_conditioning = None frozen_control_video = None masking_source = None if input_frames is not None or input_frames2 is not None: if audio_from_control_video: frozen_control_video = _build_frozen_control_video(input_frames, input_video, frame_num, prefix_frames_count, latent_stride) frame_num = int(frozen_control_video.shape[1]) else: # continuous_conditioning_and_guide = has_prefix_frames and (ic_lora_downscale_factor or 1) == 1 and not is_start_image_only # merge_conditioning_and_guide = has_prefix_frames and any_outpainting continuous_conditioning_and_guide = has_prefix_frames and any_outpainting skip_first_guide_latent = has_prefix_frames and (not is_start_image_only) and not (merge_conditioning_and_guide or continuous_conditioning_and_guide) if requested_outpaint_gamma_roundtrip: control_tensor = input_frames if input_frames is not None else input_frames2 control_rect = None if control_tensor is None else _get_outpainting_inner_rect(control_tensor.shape[-2], control_tensor.shape[-1], outpainting_dims) if control_rect is not None and _apply_gamma_to_video_rect(control_tensor, control_rect, LTX2_OUTPAINT_GAMMA): print("[WAN2GP][LTX2] Applying preserved-area gamma preprocessing for outpainting IC-LoRA control video.") use_outpaint_gamma_roundtrip = True control_start_frame = prefix_frames_count if merge_conditioning_and_guide or continuous_conditioning_and_guide: if prefix_frames_count == 1: input_frames[:, 0] = input_video[:, 0] else: input_frames = torch.concat( [input_video[:, :prefix_frames_count], input_frames[:, 1:]], dim=1) if continuous_conditioning_and_guide: control_start_frame = -prefix_frames_count else: prefix_frames_count = 0 control_start_frame = 0 input_video = None elif skip_first_guide_latent: control_start_frame = -prefix_frames_count conditioning_entries = [] if input_frames is not None: conditioning_entries.append((input_frames, control_start_frame, control_strength)) if input_frames2 is not None: conditioning_entries.append((input_frames2, control_start_frame, control_strength)) if conditioning_entries: video_conditioning = conditioning_entries if masking_strength > 0.0: if input_masks is not None and input_frames is not None: masking_source = { "video": input_frames, "mask": input_masks, "start_frame": control_start_frame, } elif input_masks2 is not None and input_frames2 is not None: masking_source = { "video": input_frames2, "mask": input_masks2, "start_frame": control_start_frame, } if not editanything and "I" in video_prompt_type and "F" not in video_prompt_type and "K" not in video_prompt_type and input_ref_images is not None: ref_frame_count = self.model_def.get("ltx2_ic_lora_ref_video_frames", 1) ref_video = _duplicate_ref_image_as_video(input_ref_images, ref_frame_count) if ref_video is not None: if video_conditioning is None: video_conditioning = [] video_conditioning.append((ref_video, 0, control_strength)) latent_conditioning_stage2 = None images = [] guiding_images = [] guiding_images_stage2 = [] images_stage2 = [] stage2_override = False def _append_prefix_entries(target_list, extra_list=None): if input_video is None or is_start_image_only: return frame_count = min(prefix_frames_count, input_video.shape[1]) if frame_count <= 0: return entry = (input_video[:, :frame_count].permute(1, 2, 3, 0), 0, input_video_strength) target_list.append(entry) if extra_list is not None: extra_list.append(entry) def _append_injected_ref_entries(target_list, extra_list=None): injected_ref_count = min(len(frames_to_inject), len(frames_relative_positions_list)) for ref_image, frame_idx in zip(frames_to_inject[:injected_ref_count], frames_relative_positions_list[:injected_ref_count]): entry = (ref_image, int(frame_idx), input_video_strength, "lanczos") target_list.append(entry) if extra_list is not None: extra_list.append(entry) if image_start is None: _append_prefix_entries(images, images_stage2) else: entry = (image_start, _to_latent_index(0, latent_stride), input_video_strength, "lanczos") images.append(entry) images_stage2.append(entry) if image_end is not None: entry = (image_end, int(frame_num - 1), input_video_strength) guiding_images.append(entry) guiding_images_stage2.append(entry) _append_injected_ref_entries(guiding_images, guiding_images_stage2) tiling_config = _build_tiling_config(VAE_tile_size, frame_num) interrupt_check = lambda: self._interrupt text_connectors = text_connectors or getattr(self, "_text_connectors", None) editanything_ref_images = input_ref_images if editanything else None audio_conditionings = None audio_conditionings_stage2 = None audio_identity_guidance_scale = 0.0 if input_waveform is not None: if audio_scale is None: audio_scale = 1.0 audio_strength = max(0.0, min(1.0, float(audio_scale))) if audio_strength > 0.0: if self._interrupt: return None waveform, waveform_sample_rate = torch.from_numpy(input_waveform), input_waveform_sample_rate if self._interrupt: return None if waveform.ndim == 1: waveform = waveform.unsqueeze(0).unsqueeze(0) elif waveform.ndim == 2: waveform = waveform.T.unsqueeze(0) target_channels = int(getattr(self.audio_encoder, "in_channels", waveform.shape[1])) if target_channels <= 0: target_channels = waveform.shape[1] if waveform.shape[1] != target_channels: if waveform.shape[1] == 1 and target_channels > 1: waveform = waveform.repeat(1, target_channels, 1) elif target_channels == 1: waveform = waveform.mean(dim=1, keepdim=True) else: waveform = waveform[:, :target_channels, :] if waveform.shape[1] < target_channels: pad_channels = target_channels - waveform.shape[1] pad = torch.zeros( (waveform.shape[0], pad_channels, waveform.shape[2]), dtype=waveform.dtype, ) waveform = torch.cat([waveform, pad], dim=1) waveform = waveform.to(device="cpu", dtype=torch.float32) if "1" in audio_prompt_type: max_samples = int(round(float(waveform_sample_rate) * LTX2_ID_LORA_MAX_REFERENCE_SECONDS)) waveform = waveform[:, :, :max_samples] audio_processor = AudioProcessor( sample_rate=self.audio_encoder.sample_rate, mel_bins=self.audio_encoder.mel_bins, mel_hop_length=self.audio_encoder.mel_hop_length, n_fft=self.audio_encoder.n_fft, ) skip_audio_conditioning = False waveform_sample_rate = int(waveform_sample_rate or 0) input_samples = int(waveform.shape[-1]) if "1" not in audio_prompt_type and audio_processor.waveform_too_short_for_mel(waveform, waveform_sample_rate): print(f"[WAN2GP][LTX2] Audio conditioning is too short for mel encoding ({input_samples} samples at {waveform_sample_rate} Hz); disabling it so audio frames are denoised.") skip_audio_conditioning = True if not skip_audio_conditioning: audio_processor = audio_processor.to(waveform.device) mel = audio_processor.waveform_to_mel(waveform, waveform_sample_rate) if self._interrupt: return None audio_params = next(self.audio_encoder.parameters(), None) audio_device = audio_params.device if audio_params is not None else self.device audio_dtype = audio_params.dtype if audio_params is not None else self.dtype mel = mel.to(device=audio_device, dtype=audio_dtype) with torch.inference_mode(): audio_latent = self.audio_encoder(mel) if self._interrupt: return None audio_downsample = getattr( getattr(self.audio_encoder, "patchifier", None), "audio_latent_downsample_factor", 4, ) audio_latent = audio_latent.to(device=self.device, dtype=self.dtype) if "1" in audio_prompt_type: audio_conditionings = [AudioConditionByReferenceLatent(audio_latent)] audio_conditionings_stage2 = [] audio_identity_guidance_scale = LTX2_ID_LORA_GUIDANCE_SCALE else: target_shape = AudioLatentShape.from_video_pixel_shape( VideoPixelShape( batch=audio_latent.shape[0], frames=int(frame_num), width=1, height=1, fps=float(fps), ), channels=audio_latent.shape[1], mel_bins=audio_latent.shape[3], sample_rate=self.audio_encoder.sample_rate, hop_length=self.audio_encoder.mel_hop_length, audio_latent_downsample_factor=audio_downsample, ) target_frames = target_shape.frames if audio_latent.shape[2] < target_frames: audio_conditionings = [AudioConditionByLatentPrefix(audio_latent)] else: if audio_latent.shape[2] > target_frames: audio_latent = audio_latent[:, :, :target_frames, :] audio_conditionings = [AudioConditionByLatent(audio_latent, audio_strength)] target_height = int(height) target_width = int(width) resolution_divisor = 64 if target_height % resolution_divisor != 0: target_height = int(math.ceil(target_height / resolution_divisor) * resolution_divisor) if target_width % resolution_divisor != 0: target_width = int(math.ceil(target_width / resolution_divisor) * resolution_divisor) if latent_conditioning_stage2 is not None: expected_lat_h = target_height // 32 expected_lat_w = target_width // 32 if ( latent_conditioning_stage2.shape[3] != expected_lat_h or latent_conditioning_stage2.shape[4] != expected_lat_w ): latent_conditioning_stage2 = None else: latent_conditioning_stage2 = latent_conditioning_stage2.to(device=self.device, dtype=self.dtype) video_conditioning_stage2 = None negative_prompt = n_prompt if n_prompt else DEFAULT_NEGATIVE_PROMPT skip_stage_2 = guide_phases <= 1 phase2_ic_lora = phase2_ic_lora_name(loras_selected, loras_slists, force_phase2_control=editanything, force_name="EditAnything") if video_conditioning else None if video_conditioning and phase2_ic_lora is not None: video_conditioning_stage2 = video_conditioning if audio_cfg_scale is None: effective_audio_cfg_scale = LTX2_ID_LORA_AUDIO_CFG_SCALE if "1" in audio_prompt_type else float(guide_scale) else: effective_audio_cfg_scale = float(audio_cfg_scale) if "1" in audio_prompt_type and effective_audio_cfg_scale <= 1.0: effective_audio_cfg_scale = LTX2_ID_LORA_AUDIO_CFG_SCALE sample_solver = sample_solver.lower() prompt_relay_frame_offset = 0 if int(window_no or 1) > 1 or (input_video is not None and not is_start_image_only): prompt_relay_frame_offset = max(0, int(prefix_frames_count or 0)) ltx2_22B_class = self.model_def.get("ltx2_22B_class", False) if isinstance(self.pipeline, TI2VidTwoStagesPipeline): pipeline_output = self.pipeline( prompt=input_prompt, negative_prompt=negative_prompt, seed=int(seed), height=target_height, width=target_width, num_frames=int(frame_num), frame_rate=float(fps), prompt_relay_frame_offset=prompt_relay_frame_offset, num_inference_steps=int(sampling_steps), cfg_guidance_scale=float(guide_scale), audio_cfg_guidance_scale=effective_audio_cfg_scale, cfg_star_switch=cfg_star_switch, apg_switch=apg_switch, perturbation_switch=perturbation_switch, perturbation_layers=perturbation_layers, perturbation_start=perturbation_start, perturbation_end=perturbation_end, alt_guidance_scale=float(alt_guide_scale), alt_scale=float(alt_scale), sample_solver=sample_solver, images=images, guiding_images=guiding_images or None, guiding_images_stage2=guiding_images_stage2 or None, images_stage2=images_stage2 if stage2_override else None, video_conditioning=video_conditioning, video_conditioning_downscale_factor=video_conditioning_downscale_factor, video_conditioning_stage2=video_conditioning_stage2, latent_conditioning_stage2=latent_conditioning_stage2, tiling_config=tiling_config, enhance_prompt=False, audio_conditionings=audio_conditionings, audio_conditionings_stage2=audio_conditionings_stage2, audio_identity_guidance_scale=audio_identity_guidance_scale, callback=callback, set_progress_status=set_progress_status, interrupt_check=interrupt_check, loras_slists=loras_slists, text_connectors=text_connectors, masking_source=masking_source, masking_strength=masking_strength, return_latent_slice=return_latent_slice, continuous_conditioning_and_guide=continuous_conditioning_and_guide, skip_stage_2=skip_stage_2, frozen_video_conditioning=frozen_control_video, frozen_output_video=frozen_control_video, self_refiner_setting=self_refiner_setting, self_refiner_plan=self_refiner_plan, self_refiner_f_uncertainty=self_refiner_f_uncertainty, self_refiner_certain_percentage=self_refiner_certain_percentage, self_refiner_max_plans=self_refiner_max_plans, editanything_ref_images=editanything_ref_images, ltx2_22B_class=ltx2_22B_class, ) else: distilled_kwargs = {} if distill: distilled_kwargs.update( { "NAG_scale": float(NAG_scale), "NAG_tau": float(NAG_tau), "NAG_alpha": float(NAG_alpha), } ) pipeline_output = self.pipeline( prompt=input_prompt, negative_prompt=negative_prompt, seed=int(seed), height=target_height, width=target_width, num_frames=int(frame_num), frame_rate=float(fps), prompt_relay_frame_offset=prompt_relay_frame_offset, images=images, guiding_images=guiding_images or None, guiding_images_stage2=guiding_images_stage2 or None, images_stage2=images_stage2 if stage2_override else None, alt_guidance_scale=float(alt_guide_scale), audio_cfg_guidance_scale=effective_audio_cfg_scale, video_conditioning=video_conditioning, video_conditioning_downscale_factor=video_conditioning_downscale_factor, video_conditioning_stage2=video_conditioning_stage2, latent_conditioning_stage2=latent_conditioning_stage2, tiling_config=tiling_config, enhance_prompt=False, audio_conditionings=audio_conditionings, audio_conditionings_stage2=audio_conditionings_stage2, audio_identity_guidance_scale=audio_identity_guidance_scale, callback=callback, set_progress_status=set_progress_status, interrupt_check=interrupt_check, loras_slists=loras_slists, text_connectors=text_connectors, masking_source=masking_source, masking_strength=masking_strength, return_latent_slice=return_latent_slice, hdr_transform=LTX2_HDR_TRANSFORM if hdr_enabled else None, precomputed_contexts=hdr_scene_context, skip_audio=hdr_enabled, continuous_conditioning_and_guide=continuous_conditioning_and_guide, skip_stage_2=skip_stage_2, frozen_video_conditioning=frozen_control_video, frozen_output_video=frozen_control_video, self_refiner_setting=self_refiner_setting, self_refiner_plan=self_refiner_plan, self_refiner_f_uncertainty=self_refiner_f_uncertainty, self_refiner_certain_percentage=self_refiner_certain_percentage, self_refiner_max_plans=self_refiner_max_plans, editanything_ref_images=editanything_ref_images, ltx2_22B_class=ltx2_22B_class, **distilled_kwargs, ) latent_slice = None if isinstance(pipeline_output, tuple) and len(pipeline_output) == 3: video, audio, latent_slice = pipeline_output else: video, audio = pipeline_output if video is None or (audio is None and not hdr_enabled): return None if self._interrupt: return None video_tensor = _collect_video_chunks( video, interrupt_check=interrupt_check, expected_frames=int(frame_num), expected_height=int(height), expected_width=int(width), ) if video_tensor is None: return None video_tensor = video_tensor[:, :frame_num, :height, :width] if use_outpaint_gamma_roundtrip: if torch.is_inference(video_tensor): raise RuntimeError("LTX2 decoded video output is still an inference tensor; decode_video_to_tensor must allocate the output buffer outside inference mode.") exponent = float(LTX2_OUTPAINT_GAMMA) if video_tensor.dtype == torch.uint8: corrected = video_tensor.to(dtype=torch.float32).div_(255.0).clamp_(0.0, 1.0).pow_(exponent) video_tensor.copy_(corrected.mul_(255.0).round_().clamp_(0.0, 255.0).to(dtype=torch.uint8)) else: corrected = video_tensor.to(dtype=torch.float32).add_(1.0).mul_(0.5).clamp_(0.0, 1.0).pow_(exponent) video_tensor.copy_(corrected.mul_(2.0).sub_(1.0).to(dtype=video_tensor.dtype)) audio_np = None if hdr_enabled else audio.detach().float().cpu().numpy() if audio is not None else None if audio_np is not None and audio_np.ndim == 2: if audio_np.shape[0] in (1, 2) and audio_np.shape[1] > audio_np.shape[0]: audio_np = audio_np.T output_audio_sampling_rate = int(getattr(self.vocoder, "output_sampling_rate", AUDIO_SAMPLE_RATE)) result = { "x": video_tensor, "audio": audio_np, "audio_sampling_rate": output_audio_sampling_rate, } if hdr_enabled: result["hdr"] = True result["hdr_format"] = "linear_srgb" result["hdr_transform"] = LTX2_HDR_TRANSFORM if latent_slice is not None: result["latent_slice"] = latent_slice return result