indigo/LTX-2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Automated PR - 2026-05-28

Browse Source
This commit is contained in:
github-actions[bot]
2026-05-28 14:26:15 +00:00
parent 1799988521
commit 203d4842d4
45 changed files with 1969 additions and 668 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-pipelines/src/ltx_pipelines/a2vid_two_stage.py
+6 -5
View File
@@ -9,6 +9,7 @@ from ltx_core.components.schedulers import LTX2Scheduler
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import Registry
from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, AudioLatentShape, VideoPixelShape
@@ -52,7 +53,7 @@ class A2VidPipelineTwoStage:
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
@@ -71,7 +72,7 @@ class A2VidPipelineTwoStage:
loras=tuple(loras),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
@@ -82,7 +83,7 @@ class A2VidPipelineTwoStage:
loras=stage_2_loras,
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.upsampler = VideoUpsampler(
@@ -238,7 +239,7 @@ class A2VidPipelineTwoStage:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
parser = default_2_stage_arg_parser()
parser.add_argument(
"--audio-path",
@@ -266,7 +267,7 @@ def main() -> None:
gemma_root=args.gemma_root,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/distilled.py
+5 -4
View File
@@ -6,6 +6,7 @@ import torch
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import Registry
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio
@@ -53,7 +54,7 @@ class DistilledPipeline:
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
@@ -75,7 +76,7 @@ class DistilledPipeline:
loras=tuple(loras),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.upsampler = VideoUpsampler(
@@ -180,7 +181,7 @@ class DistilledPipeline:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path(distilled=True)
params = detect_params(checkpoint_path)
parser = default_2_stage_distilled_arg_parser(params=params)
@@ -191,7 +192,7 @@ def main() -> None:
gemma_root=args.gemma_root,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/hdr_ic_lora.py
+5 -2
View File
@@ -59,6 +59,7 @@ from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES, STAGE_2_DISTIL
from ltx_pipelines.utils.denoisers import SimpleDenoiser
from ltx_pipelines.utils.helpers import get_device, modality_from_latent_state
from ltx_pipelines.utils.media_io import ResizeMode, align_resolution, load_video_conditioning_hdr
from ltx_pipelines.utils.quantization_factory import QuantizationKind
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
logger = logging.getLogger(__name__)
@@ -77,7 +78,7 @@ ALIGNMENT_DIVISOR = 64
# to the pipeline constructor.
TILED_VAE_ENCODE_PIXEL_THRESHOLD = 512 * 768
_DEFAULT_QUANTIZATION = QuantizationPolicy.fp8_cast()
_DEFAULT_QUANTIZATION = QuantizationKind.FP8_CAST
# Default stage-2 configuration: one refinement phase with modest 2-way tiling
# in every dimension and a short 2-step distilled sigma schedule.
@@ -205,7 +206,7 @@ class HDRICLoraPipeline:
hdr_lora: str | Path,
text_embeddings_path: str | Path,
device: torch.device | None = None,
quantization: QuantizationPolicy = _DEFAULT_QUANTIZATION,
quantization: QuantizationPolicy | QuantizationKind | None = _DEFAULT_QUANTIZATION,
registry: Registry | None = None,
hdr_lora_config: HdrLoraConfig | None = None,
tiled_vae_encode_pixel_threshold: int = TILED_VAE_ENCODE_PIXEL_THRESHOLD,
@@ -232,6 +233,8 @@ class HDRICLoraPipeline:
"""
self.device = device or get_device()
self._tiled_vae_encode_threshold = tiled_vae_encode_pixel_threshold
if isinstance(quantization, QuantizationKind):
quantization = quantization.to_policy(checkpoint_path=distilled_checkpoint_path)
if offload_mode != OffloadMode.NONE and quantization is not None:
logger.info("Offload mode enabled — disabling quantization (not supported with layer streaming).")
quantization = None
packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py
+6 -5
View File
@@ -7,6 +7,7 @@ from ltx_core.components.noisers import GaussianNoiser
from ltx_core.conditioning import ConditioningItem
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import Registry
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, VideoPixelShape
@@ -60,7 +61,7 @@ class ICLoraPipeline:
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
@@ -82,7 +83,7 @@ class ICLoraPipeline:
loras=tuple(loras),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.stage_2 = DiffusionStage(
@@ -92,7 +93,7 @@ class ICLoraPipeline:
loras=(),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.upsampler = VideoUpsampler(
@@ -330,7 +331,7 @@ class ICLoraPipeline:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path(distilled=True)
params = detect_params(checkpoint_path)
parser = default_2_stage_distilled_arg_parser(params=params)
@@ -386,7 +387,7 @@ def main() -> None:
gemma_root=args.gemma_root,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py
+6 -5
View File
@@ -12,6 +12,7 @@ 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.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, VideoPixelShape
@@ -62,7 +63,7 @@ class KeyframeInterpolationPipeline:
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
@@ -80,7 +81,7 @@ class KeyframeInterpolationPipeline:
loras=tuple(loras),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
@@ -91,7 +92,7 @@ class KeyframeInterpolationPipeline:
loras=stage_2_loras,
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.upsampler = VideoUpsampler(
@@ -233,7 +234,7 @@ class KeyframeInterpolationPipeline:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path()
params = detect_params(checkpoint_path)
parser = default_2_stage_arg_parser(params=params)
@@ -245,7 +246,7 @@ def main() -> None:
gemma_root=args.gemma_root,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/lipdub.py
+5 -4
View File
@@ -13,6 +13,7 @@ from ltx_core.conditioning import AudioConditionByReferenceLatent
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import Registry
from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, AudioLatentShape, SpatioTemporalScaleFactors, VideoPixelShape
@@ -59,7 +60,7 @@ class LipDubPipeline:
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
) -> None:
self.device = device or get_device()
@@ -89,7 +90,7 @@ class LipDubPipeline:
loras=loras,
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.upsampler = VideoUpsampler(
@@ -289,7 +290,7 @@ def patchify_lipdub_audio_reference_latent(
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path(distilled=True)
params = detect_params(checkpoint_path)
parser = lipdub_arg_parser(params=params)
@@ -304,7 +305,7 @@ def main() -> None:
gemma_root=args.gemma_root,
ic_lora=args.lora[0],
quantization=args.quantization,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/retake.py
+5 -4
View File
@@ -11,6 +11,7 @@ from ltx_core.components.schedulers import LTX2Scheduler
from ltx_core.conditioning.types.noise_mask_cond import TemporalRegionMask
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import Registry
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import (
@@ -73,7 +74,7 @@ class RetakePipeline:
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
distilled: bool = True,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
@@ -108,7 +109,7 @@ class RetakePipeline:
loras=tuple(loras),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.video_decoder = VideoDecoder(
@@ -283,7 +284,7 @@ class RetakePipeline:
@torch.inference_mode()
def main() -> None:
"""CLI entry point for retake (regenerate a time region)."""
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
parser = video_editing_arg_parser(distilled=True)
parser.description = "Retake: regenerate a time region of a video with LTX-2."
args = parser.parse_args()
@@ -308,7 +309,7 @@ def main() -> None:
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
distilled=True,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
params = detect_params(args.distilled_checkpoint_path)
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py
+5 -4
View File
@@ -12,6 +12,7 @@ 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.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae.tiling import TilingConfig
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio
@@ -55,7 +56,7 @@ class TI2VidOneStagePipeline:
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.dtype = torch.bfloat16
@@ -82,7 +83,7 @@ class TI2VidOneStagePipeline:
loras=tuple(loras),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.video_decoder = VideoDecoder(
@@ -185,7 +186,7 @@ class TI2VidOneStagePipeline:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path()
params = detect_params(checkpoint_path)
parser = default_1_stage_arg_parser(params=params)
@@ -195,7 +196,7 @@ def main() -> None:
gemma_root=args.gemma_root,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
video, audio = pipeline(
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py
+6 -5
View File
@@ -12,6 +12,7 @@ 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.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, VideoPixelShape
@@ -61,7 +62,7 @@ class TI2VidTwoStagesPipeline:
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
@@ -85,7 +86,7 @@ class TI2VidTwoStagesPipeline:
loras=tuple(loras),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.stage_2 = DiffusionStage(
@@ -95,7 +96,7 @@ class TI2VidTwoStagesPipeline:
loras=(*tuple(loras), *distilled_lora),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
@@ -223,7 +224,7 @@ class TI2VidTwoStagesPipeline:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
checkpoint_path = detect_checkpoint_path()
params = detect_params(checkpoint_path)
parser = default_2_stage_arg_parser(params=params)
@@ -235,7 +236,7 @@ def main() -> None:
gemma_root=args.gemma_root,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_hq.py
+6 -5
View File
@@ -9,6 +9,7 @@ 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.transformer.compiling import CompilationConfig
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
@@ -60,7 +61,7 @@ class TI2VidTwoStagesHQPipeline:
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
):
self.device = device or get_device()
@@ -95,7 +96,7 @@ class TI2VidTwoStagesHQPipeline:
loras=(*loras, distilled_lora_stage_1),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
self.stage_2 = DiffusionStage(
@@ -105,7 +106,7 @@ class TI2VidTwoStagesHQPipeline:
loras=(*loras, distilled_lora_stage_2),
quantization=quantization,
registry=registry,
torch_compile=torch_compile,
compilation_config=compilation_config,
offload_mode=offload_mode,
)
@@ -242,7 +243,7 @@ class TI2VidTwoStagesHQPipeline:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
parser = hq_2_stage_arg_parser(params=LTX_2_3_HQ_PARAMS)
args = parser.parse_args()
pipeline = TI2VidTwoStagesHQPipeline(
@@ -254,7 +255,7 @@ def main() -> None:
gemma_root=args.gemma_root,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
torch_compile=args.compile,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/utils/args.py
+146 -26
View File
@@ -1,9 +1,11 @@
import argparse
import json
from collections.abc import Sequence
from pathlib import Path
from typing import NamedTuple
from typing import Any, NamedTuple
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.quantization import QuantizationPolicy
from ltx_pipelines.utils.constants import (
DEFAULT_IMAGE_CRF,
@@ -13,6 +15,7 @@ from ltx_pipelines.utils.constants import (
LTX_2_3_PARAMS,
PipelineParams,
)
from ltx_pipelines.utils.quantization_factory import QuantizationKind
from ltx_pipelines.utils.types import OffloadMode
@@ -32,7 +35,7 @@ class VideoConditioningAction(argparse.Action):
option_string: str | None = None, # noqa: ARG002
) -> None:
path, strength_str = values
resolved_path = resolve_path(path)
resolved_path = resolve_existing_path(path)
strength = float(strength_str)
current = getattr(namespace, self.dest) or []
current.append((resolved_path, strength))
@@ -58,7 +61,7 @@ class VideoMaskConditioningAction(argparse.Action):
msg = f"{option_string} requires exactly 2 arguments (MASK_PATH STRENGTH), got {len(values)}"
raise argparse.ArgumentError(self, msg)
mask_path = resolve_path(values[0])
mask_path = resolve_existing_path(values[0])
strength = float(values[1])
setattr(namespace, self.dest, (mask_path, strength))
@@ -76,7 +79,7 @@ class ImageAction(argparse.Action):
raise argparse.ArgumentError(self, msg)
conditioning = ImageConditioningInput(
path=resolve_path(values[0]),
path=resolve_existing_path(values[0]),
frame_idx=int(values[1]),
strength=float(values[2]),
crf=int(values[3]) if len(values) > 3 else DEFAULT_IMAGE_CRF,
@@ -101,7 +104,7 @@ class LoraAction(argparse.Action):
path = values[0]
strength_str = values[1] if len(values) > 1 else str(DEFAULT_LORA_STRENGTH)
resolved_path = resolve_path(path)
resolved_path = resolve_existing_path(path)
strength = float(strength_str)
current = getattr(namespace, self.dest) or []
@@ -109,11 +112,118 @@ class LoraAction(argparse.Action):
setattr(namespace, self.dest, current)
class CompileAction(argparse.Action):
"""Parse ``--compile [KEY=VALUE ...]`` into a :class:`CompilationConfig`.
The flag is absent -> ``args.compile`` stays at its default (``None``).
The flag is passed alone -> ``CompilationConfig()`` (vanilla torch defaults).
The flag is passed with args -> ``CompilationConfig`` with the given fields overridden.
Errors (unknown key, malformed value, duplicate key, empty value) raise
:class:`argparse.ArgumentError` so argparse formats them as friendly CLI
messages rather than uncaught tracebacks.
"""
_ALLOWED_KEYS = frozenset({"mode", "backend", "fullgraph", "dynamic", "inductor_config", "dynamo_config"})
def __call__(
self,
parser: argparse.ArgumentParser, # noqa: ARG002
namespace: argparse.Namespace,
values: list[str],
option_string: str | None = None, # noqa: ARG002
) -> None:
overrides: dict[str, object] = {}
for item in values:
if "=" not in item:
raise argparse.ArgumentError(self, f"expects KEY=VALUE pairs, got: {item!r}")
key, _, raw = item.partition("=")
key = key.strip()
if key not in self._ALLOWED_KEYS:
raise argparse.ArgumentError(
self,
f"{key!r} is not a CompilationConfig field; valid keys: {sorted(self._ALLOWED_KEYS)}",
)
if key in overrides:
raise argparse.ArgumentError(self, f"{key} given more than once")
if key == "mode":
overrides[key] = self._parse_mode(raw)
elif key == "backend":
overrides[key] = self._parse_non_empty(key, raw)
elif key == "fullgraph":
overrides[key] = self._parse_bool(key, raw)
elif key == "dynamic":
overrides[key] = self._parse_dynamic(raw)
elif key in ("inductor_config", "dynamo_config"):
overrides[key] = self._parse_json_dict(key, raw)
setattr(namespace, self.dest, CompilationConfig(**overrides))
def _parse_mode(self, raw: str) -> str | None:
stripped = raw.strip()
if not stripped:
raise argparse.ArgumentError(self, "mode=... value cannot be empty (use mode=none to clear)")
if stripped.lower() == "none":
return None
return stripped
def _parse_non_empty(self, key: str, raw: str) -> str:
stripped = raw.strip()
if not stripped:
raise argparse.ArgumentError(self, f"{key}=... value cannot be empty")
return stripped
def _parse_bool(self, key: str, raw: str) -> bool:
normalized = raw.strip().lower()
if normalized in ("true", "1"):
return True
if normalized in ("false", "0"):
return False
raise argparse.ArgumentError(self, f"{key}=... must be true or false; got {raw!r}")
def _parse_dynamic(self, raw: str) -> bool | None:
normalized = raw.strip().lower()
if normalized in ("auto", "none"):
return None
if normalized in ("true", "1"):
return True
if normalized in ("false", "0"):
return False
raise argparse.ArgumentError(self, f"dynamic=... must be auto/true/false; got {raw!r}")
def _parse_json_dict(self, key: str, raw: str) -> dict[str, Any]:
# Inline JSON object starts with '{'; otherwise treat the value as a path to a JSON file.
stripped = raw.strip()
if not stripped:
raise argparse.ArgumentError(self, f"{key}=... value cannot be empty")
if stripped.startswith("{"):
source = stripped
else:
path = Path(stripped).expanduser()
if not path.is_file():
raise argparse.ArgumentError(
self, f"{key}=... must be a JSON object or a path to a JSON file; got {raw!r}"
)
source = path.read_text()
try:
value = json.loads(source)
except json.JSONDecodeError as e:
raise argparse.ArgumentError(self, f"{key}=... must be a JSON object; got {raw!r} ({e.msg})") from None
if not isinstance(value, dict):
raise argparse.ArgumentError(self, f"{key}=... must decode to a JSON object; got {type(value).__name__}")
return value
def resolve_path(path: str) -> str:
return str(Path(path).expanduser().resolve().as_posix())
QUANTIZATION_POLICIES = ("fp8-cast", "fp8-scaled-mm")
def resolve_existing_path(path: str) -> str:
"""Resolve *path* and verify it exists."""
resolved = resolve_path(path)
if not Path(resolved).exists():
raise argparse.ArgumentError(None, f"Path not found: {resolved}")
return resolved
QUANTIZATION_POLICIES = tuple(k.value for k in QuantizationKind)
def _resolve_quantization(namespace: argparse.Namespace) -> None:
@@ -123,16 +233,14 @@ def _resolve_quantization(namespace: argparse.Namespace) -> None:
name = getattr(namespace, "quantization", None)
if name is None or isinstance(name, QuantizationPolicy):
return
if name == "fp8-cast":
namespace.quantization = QuantizationPolicy.fp8_cast()
try:
kind = QuantizationKind(name)
except ValueError:
return
if name == "fp8-scaled-mm":
ckpt = getattr(namespace, "checkpoint_path", None) or getattr(namespace, "distilled_checkpoint_path", None)
if ckpt is None:
raise SystemExit(
"--quantization fp8-scaled-mm requires --checkpoint-path (or --distilled-checkpoint-path)."
)
namespace.quantization = QuantizationPolicy.fp8_scaled_mm(ckpt)
ckpt = getattr(namespace, "checkpoint_path", None) or getattr(namespace, "distilled_checkpoint_path", None)
if ckpt is None:
raise SystemExit(f"--quantization {kind.value} requires --checkpoint-path (or --distilled-checkpoint-path).")
namespace.quantization = kind.to_policy(checkpoint_path=ckpt)
class _PipelineArgumentParser(argparse.ArgumentParser):
@@ -150,7 +258,7 @@ def detect_checkpoint_path(distilled: bool = False) -> str:
"""Pre-parse argv to extract the checkpoint path before building the full parser."""
pre = argparse.ArgumentParser(add_help=False)
flag = "--distilled-checkpoint-path" if distilled else "--checkpoint-path"
pre.add_argument(flag, type=resolve_path, required=True)
pre.add_argument(flag, type=resolve_existing_path, required=True)
known, _ = pre.parse_known_args()
return known.distilled_checkpoint_path if distilled else known.checkpoint_path
@@ -163,14 +271,14 @@ def basic_arg_parser(
if distilled:
parser.add_argument(
"--distilled-checkpoint-path",
type=resolve_path,
type=resolve_existing_path,
required=True,
help="Path to LTX-2 distilled model checkpoint (.safetensors file).",
)
else:
parser.add_argument(
"--checkpoint-path",
type=resolve_path,
type=resolve_existing_path,
required=True,
help="Path to LTX-2 model checkpoint (.safetensors file).",
)
@@ -185,7 +293,7 @@ def basic_arg_parser(
)
parser.add_argument(
"--gemma-root",
type=resolve_path,
type=resolve_existing_path,
required=True,
help="Path to the root directory containing the Gemma text encoder model files.",
)
@@ -276,8 +384,20 @@ def basic_arg_parser(
)
parser.add_argument(
"--compile",
action="store_true",
help="Enable torch.compile for transformer blocks to optimize performance.",
nargs="*",
action=CompileAction,
default=None,
metavar="KEY=VALUE",
help=(
"Enable torch.compile for transformer blocks. Pass alone for defaults, "
"or with KEY=VALUE overrides for any CompilationConfig field. "
"Keys: mode, backend, fullgraph, dynamic, inductor_config, dynamo_config. "
"inductor_config/dynamo_config take JSON objects (inline or a path to a .json file) "
"that fully replace the defaults. "
"Examples: --compile or --compile mode=reduce-overhead or "
"--compile mode=reduce-overhead fullgraph=true backend=eager or "
"--compile inductor_config='{\"max_autotune\": true}'"
),
)
return parser
@@ -340,7 +460,7 @@ def video_editing_arg_parser(
(no height/width/num-frames; resolution comes from input video). Default is distilled checkpoint only.
"""
parser = basic_arg_parser(distilled=distilled)
parser.add_argument("--video-path", type=resolve_path, required=True, help="Path to the source video.")
parser.add_argument("--video-path", type=resolve_existing_path, required=True, help="Path to the source video.")
parser.add_argument("--start-time", type=float, required=True, help="Start time of the region to regenerate (s).")
parser.add_argument("--end-time", type=float, required=True, help="End time of the region to regenerate (s).")
return parser
@@ -462,7 +582,7 @@ def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
default=video_guider.skip_step,
help=(
"Video skip step N controls periodic skipping during the video diffusion process: "
"only steps where step_index % (N + 1) == 0 are processed, all others are skipped "
"only steps where step_index %% (N + 1) == 0 are processed, all others are skipped "
f"(e.g., 0 = no skipping; 1 = skip every other step; 2 = skip 2 of every 3 steps; "
f"default: {video_guider.skip_step})."
),
@@ -522,7 +642,7 @@ def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
default=audio_guider.skip_step,
help=(
"Audio skip step N controls periodic skipping during the audio diffusion process: "
"only steps where step_index % (N + 1) == 0 are processed, all others are skipped "
"only steps where step_index %% (N + 1) == 0 are processed, all others are skipped "
f"(e.g., 0 = no skipping; 1 = skip every other step; 2 = skip 2 of every 3 steps; "
f"default: {audio_guider.skip_step})."
),
@@ -562,7 +682,7 @@ def default_2_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
)
parser.add_argument(
"--spatial-upsampler-path",
type=resolve_path,
type=resolve_existing_path,
required=True,
help=(
"Path to the spatial upsampler model used to increase the resolution "
@@ -605,7 +725,7 @@ def default_2_stage_distilled_arg_parser(params: PipelineParams = LTX_2_3_PARAMS
)
parser.add_argument(
"--spatial-upsampler-path",
type=resolve_path,
type=resolve_existing_path,
required=True,
help=(
"Path to the spatial upsampler model used to increase the resolution "
packages/ltx-pipelines/src/ltx_pipelines/utils/blocks.py
+122 -45
View File
@@ -6,6 +6,8 @@ removes the need for :class:`ModelLedger`.
from __future__ import annotations
import copy
import dataclasses
import logging
from collections.abc import Iterator
from contextlib import AbstractContextManager, contextmanager
@@ -21,6 +23,8 @@ from ltx_core.components.noisers import Noiser
from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.loader import SDOps
from ltx_core.loader.attention_ops import set_attention_module_op
from ltx_core.loader.fuse_loras import bf16_fuse_rule
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.primitives import BuilderProtocol, LoraPathStrengthAndSDOps, ModelBuilderProtocol
from ltx_core.loader.registry import DummyRegistry, Registry
@@ -42,7 +46,15 @@ from ltx_core.model.transformer import (
LTXModelConfigurator,
X0Model,
)
from ltx_core.model.transformer.compiling import COMPILE_TRANSFORMER, modify_sd_ops_for_compilation
from ltx_core.model.transformer.attention import (
AttentionCallable,
AttentionFunction,
)
from ltx_core.model.transformer.compiling import (
CompilationConfig,
build_compile_transformer_op,
modify_sd_ops_for_compilation,
)
from ltx_core.model.upsampler import LatentUpsamplerConfigurator, upsample_video
from ltx_core.model.video_vae import (
MEMORY_EFFICIENT_DECODE,
@@ -53,7 +65,7 @@ from ltx_core.model.video_vae import (
VideoEncoder,
VideoEncoderConfigurator,
)
from ltx_core.quantization import QuantizationPolicy
from ltx_core.quantization import QuantizationPolicy, fp8_cast_fuse_rule
from ltx_core.text_encoders.gemma import (
EMBEDDINGS_PROCESSOR_KEY_OPS,
GEMMA_LLM_KEY_OPS,
@@ -101,6 +113,28 @@ def _chain_quantization(
return chained_sd_ops, (*module_ops, *quantization.module_ops)
def _apply_compile_ops(
sd_ops: SDOps,
module_ops: tuple[ModuleOps, ...],
loras: tuple[LoraPathStrengthAndSDOps, ...],
number_of_layers: int,
compilation_config: CompilationConfig,
) -> tuple[SDOps, tuple[ModuleOps, ...], tuple[LoraPathStrengthAndSDOps, ...]]:
"""Rewrite sd_ops/module_ops/LoRAs for compiled blocks (params land under ``_orig_mod``)."""
sd_ops = modify_sd_ops_for_compilation(sd_ops, number_of_layers)
compile_op = build_compile_transformer_op(compilation_config)
module_ops = (*module_ops, compile_op)
loras = tuple(
LoraPathStrengthAndSDOps(
lora.path,
lora.strength,
modify_sd_ops_for_compilation(lora.sd_ops, number_of_layers),
)
for lora in loras
)
return sd_ops, module_ops, loras
@contextmanager
def _streaming_model(
builder: StreamingModelBuilder,
@@ -170,79 +204,110 @@ class DiffusionStage:
loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
torch_compile: bool = False,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
transformer_builder: ModelBuilderProtocol[LTXModel] | DelegatingBuilder[LTXModel] | None = None,
) -> None:
self._checkpoint_path = checkpoint_path
self._dtype = dtype
self._device = device
self._quantization = quantization
self._torch_compile = torch_compile
self._compilation_config = compilation_config
self._offload_mode = offload_mode
configurator = (
quantization.model_configurator
if quantization is not None and quantization.model_configurator is not None
else LTXModelConfigurator
)
if transformer_builder is not None:
self._transformer_builder = transformer_builder
else:
self._transformer_builder = Builder(
model_path=checkpoint_path,
model_class_configurator=LTXModelConfigurator,
model_class_configurator=configurator,
model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
loras=tuple(loras),
registry=registry or DummyRegistry(),
)
if offload_mode != OffloadMode.NONE:
if torch_compile:
if compilation_config is not None:
raise ValueError("torch.compile is not supported with layer streaming")
# WeightsProvider currently only supports plain bf16 + fp8_cast LoRA fusion
# (no companion-key emission). Quantization policies that emit
# companion keys (e.g. ``.weight_scale``) cannot be streamed yet.
if quantization is not None and quantization.fuse_rule is not fp8_cast_fuse_rule:
raise ValueError(
"Block streaming is not supported with this quantization policy "
"(only bf16 and fp8_cast are currently supported)."
)
streaming_sd_ops: SDOps = LTXV_MODEL_COMFY_RENAMING_MAP
streaming_module_ops: tuple[ModuleOps, ...] = ()
if quantization is not None:
if quantization.kind != QuantizationPolicy.Kind.FP8_CAST:
raise ValueError(
f"Layer streaming supports only QuantizationPolicy.fp8_cast(); "
f"got kind={quantization.kind!r} which produces heterogeneous block layouts."
)
streaming_sd_ops, streaming_module_ops = _chain_quantization(
streaming_sd_ops, streaming_module_ops, quantization
)
self._streaming_builder = StreamingModelBuilder(
model_class_configurator=LTXModelConfigurator,
model_class_configurator=configurator,
model_path=checkpoint_path,
model_sd_ops=streaming_sd_ops,
module_ops=streaming_module_ops,
loras=tuple(loras),
registry=registry or DummyRegistry(),
blocks_attr="velocity_model.transformer_blocks",
fuse_rule=quantization.fuse_rule if quantization is not None else bf16_fuse_rule,
blocks_attr="transformer_blocks",
blocks_prefix="transformer_blocks",
state_dict_prefix="velocity_model.",
model_wrapper=lambda m: X0Model(m).eval(),
)
def with_attention(self, attention: AttentionFunction | AttentionCallable | None) -> "DiffusionStage":
"""Return a new ``DiffusionStage`` that pins the transformer build to ``attention``.
Functional: never mutates ``self``. The returned stage shares all other
configuration with the original; only the underlying builders' ``module_ops``
gain a ``set_attention_module_op(attention)`` entry so subsequent transformer
builds use that kernel. ``attention=None`` is a no-op (returns ``self``).
"""
if attention is None:
return self
op = set_attention_module_op(attention)
new = copy.copy(self)
new._transformer_builder = self._transformer_builder.with_module_ops(
(*self._transformer_builder.module_ops, op),
)
if self._offload_mode != OffloadMode.NONE:
new._streaming_builder = dataclasses.replace(
self._streaming_builder,
module_ops=(*self._streaming_builder.module_ops, op),
)
return new
def _build_transformer(self, *, device: torch.device | None = None, **kwargs: object) -> X0Model:
target = device or self._device
sd_ops = self._transformer_builder.model_sd_ops
module_ops = self._transformer_builder.module_ops
loras = self._transformer_builder.loras
if self._torch_compile:
module_ops = (*module_ops, COMPILE_TRANSFORMER)
if self._compilation_config is not None:
number_of_layers = self._transformer_builder.model_config()["transformer"]["num_layers"]
sd_ops = modify_sd_ops_for_compilation(sd_ops, number_of_layers)
loras = tuple(
LoraPathStrengthAndSDOps(
lora.path,
lora.strength,
modify_sd_ops_for_compilation(lora.sd_ops, number_of_layers),
)
for lora in loras
sd_ops, module_ops, loras = _apply_compile_ops(
sd_ops, module_ops, loras, number_of_layers, self._compilation_config
)
if self._quantization is not None:
sd_ops, module_ops = _chain_quantization(sd_ops, module_ops, self._quantization)
builder = self._transformer_builder.with_module_ops(module_ops).with_sd_ops(sd_ops).with_loras(loras)
if self._quantization is not None:
builder = builder.with_fuse_rule(self._quantization.fuse_rule)
return X0Model(builder.build(device=target, **kwargs)).to(target).eval()
@contextmanager
def _streaming_transformer_ctx(self) -> Iterator[X0Model]:
with _streaming_model(
self._streaming_builder, self._offload_mode, self._device, self._dtype
) as streaming_wrapper:
yield X0Model(streaming_wrapper).eval()
def _transformer_ctx(self, **kwargs: object) -> AbstractContextManager:
if self._offload_mode != OffloadMode.NONE:
return _streaming_model(self._streaming_builder, self._offload_mode, self._device, self._dtype)
return self._streaming_transformer_ctx()
return gpu_model(self._build_transformer(**kwargs))
def model_context(self, **kwargs: object) -> AbstractContextManager:
@@ -348,7 +413,17 @@ class DiffusionStage:
v_shape = VideoLatentShape.from_pixel_shape(pixel_shape)
video_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape, fps)
mode = "streaming" if self._offload_mode != OffloadMode.NONE else "standard"
logger.info("Building transformer (%s) from %s", mode, self._checkpoint_path)
with self._transformer_ctx(video_tools=video_tools) as transformer:
logger.info(
"Running denoising loop (%d steps, %dx%d %d frames @ %.1f fps)",
len(sigmas) - 1,
width,
height,
frames,
fps,
)
return self.run(
transformer,
denoiser,
@@ -387,6 +462,8 @@ class PromptEncoder:
offload_mode: OffloadMode = OffloadMode.NONE,
text_encoder_builder: BuilderProtocol | None = None,
) -> None:
self._gemma_root = gemma_root
self._checkpoint_path = checkpoint_path
self._dtype = dtype
self._device = device
self._offload_mode = offload_mode
@@ -432,7 +509,7 @@ class PromptEncoder:
def _build_embeddings_processor(self) -> EmbeddingsProcessor:
"""Build the embeddings processor on the target device."""
return self._embeddings_processor_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
return self._embeddings_processor_builder.build(device=self._device, dtype=self._dtype).eval()
def _text_encoder_ctx(self) -> AbstractContextManager:
if self._offload_mode != OffloadMode.NONE:
@@ -448,6 +525,7 @@ class PromptEncoder:
enhance_prompt_seed: int = 42,
) -> list[EmbeddingsProcessorOutput]:
"""Encode *prompts* through Gemma -> embeddings processor, freeing each model after use."""
logger.info("Building text encoder from %s", self._gemma_root)
with self._text_encoder_ctx() as text_encoder:
if enhance_first_prompt:
prompts = list(prompts)
@@ -455,9 +533,12 @@ class PromptEncoder:
text_encoder, prompts[0], enhance_prompt_image, seed=enhance_prompt_seed
)
raw_outputs = [text_encoder.encode(p) for p in prompts]
logger.info("Text encoder done, building embeddings processor from %s", self._checkpoint_path)
with gpu_model(self._build_embeddings_processor()) as embeddings_processor:
return [embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs]
result = [embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs]
logger.info("Prompt encoding complete")
return result
# ---------------------------------------------------------------------------
@@ -487,7 +568,7 @@ class ImageConditioner:
)
def _build_encoder(self) -> VideoEncoder:
return self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
return self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()
def __call__(self, fn: Callable[[VideoEncoder], T]) -> T:
"""Build video encoder → call *fn(encoder)* → free encoder."""
@@ -511,6 +592,7 @@ class VideoUpsampler:
device: torch.device,
registry: Registry | None = None,
) -> None:
self._upsampler_path = upsampler_path
self._dtype = dtype
self._device = device
self._encoder_builder = Builder(
@@ -527,13 +609,10 @@ class VideoUpsampler:
def __call__(self, latent: torch.Tensor) -> torch.Tensor:
"""Upsample *latent* using video encoder + spatial upsampler, then free both."""
logger.info("Building video encoder + spatial upsampler from %s", self._upsampler_path)
with (
gpu_model(
self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
) as encoder,
gpu_model(
self._upsampler_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
) as upsampler,
gpu_model(self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()) as encoder,
gpu_model(self._upsampler_builder.build(device=self._device, dtype=self._dtype).eval()) as upsampler,
):
return upsample_video(latent=latent, video_encoder=encoder, upsampler=upsampler)
@@ -557,6 +636,7 @@ class VideoDecoder:
memory_efficient: bool = True,
decoder_builder: BuilderProtocol | None = None,
) -> None:
self._checkpoint_path = checkpoint_path
self._dtype = dtype
self._device = device
if decoder_builder is not None:
@@ -577,7 +657,8 @@ class VideoDecoder:
generator: torch.Generator | None = None,
) -> Iterator[torch.Tensor]:
"""Decode *latent* to pixel-space video chunks. Decoder freed after exhaustion."""
decoder = self._decoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
logger.info("Building video decoder from %s", self._checkpoint_path)
decoder = self._decoder_builder.build(device=self._device, dtype=self._dtype).eval()
return _cleanup_iter(decoder.decode_video(latent, tiling_config, generator), decoder)
@@ -596,6 +677,7 @@ class AudioDecoder:
device: torch.device,
registry: Registry | None = None,
) -> None:
self._checkpoint_path = checkpoint_path
self._dtype = dtype
self._device = device
self._decoder_builder = Builder(
@@ -613,13 +695,10 @@ class AudioDecoder:
def __call__(self, latent: torch.Tensor) -> Audio:
"""Decode audio *latent* through VAE decoder + vocoder, then free both."""
logger.info("Building audio decoder + vocoder from %s", self._checkpoint_path)
with (
gpu_model(
self._decoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
) as decoder,
gpu_model(
self._vocoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
) as vocoder,
gpu_model(self._decoder_builder.build(device=self._device, dtype=self._dtype).eval()) as decoder,
gpu_model(self._vocoder_builder.build(device=self._device, dtype=self._dtype).eval()) as vocoder,
):
return vae_decode_audio(latent, decoder, vocoder)
@@ -653,7 +732,5 @@ class AudioConditioner:
def __call__(self, fn: Callable[[torch.nn.Module], T]) -> T:
"""Build audio encoder → call *fn(encoder)* → free encoder."""
with gpu_model(
self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
) as encoder:
with gpu_model(self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()) as encoder:
return fn(encoder)
packages/ltx-pipelines/src/ltx_pipelines/utils/denoisers.py
+11 -1
View File
@@ -138,6 +138,8 @@ def _guided_denoise( # noqa: PLR0913,PLR0915
ptb_configs = [ptb for _, _, _, ptb in passes]
n = len(passes)
orig_b = (video_state or audio_state).latent.shape[0]
def _batched_sigma(state: LatentState) -> torch.Tensor:
"""Expand scalar sigma to (n * B,) matching the repeated state."""
return sigma.expand(state.latent.shape[0] * n)
@@ -162,8 +164,16 @@ def _guided_denoise( # noqa: PLR0913,PLR0915
enabled=not a_skip,
)
# Replicate each pass's PerturbationConfig to all `orig_b` samples it
# carries, so `BatchedPerturbationConfig.mask_like` returns a per-sample
# mask (length n*orig_b) instead of a per-pass mask (length n). Without
# this expansion the mask is broadcast against a (n*orig_b, T, D) tensor
# and the multiplication fails with a batch-dim mismatch whenever
# `orig_b > 1` (e.g. multi-prompt benchmark panels).
batched_ptb_configs = [ptb for ptb in ptb_configs for _ in range(orig_b)]
all_v, all_a = transformer(
video=batched_video, audio=batched_audio, perturbations=BatchedPerturbationConfig(ptb_configs)
video=batched_video, audio=batched_audio, perturbations=BatchedPerturbationConfig(batched_ptb_configs)
)
# Split results back and combine via guiders.
packages/ltx-pipelines/src/ltx_pipelines/utils/quantization_factory.py
+36
View File
@@ -0,0 +1,36 @@
"""User-facing quantization-policy dispatch.
``ltx-core`` exposes one ``build_policy`` factory per backend. This module
provides the user-facing string-keyed dispatch used by CLI args and pipeline
defaults — keeping the enum out of ``ltx-core`` so adding/removing backends is
a single-file change here.
"""
from enum import Enum
from typing_extensions import assert_never
from ltx_core.quantization import QuantizationPolicy
from ltx_core.quantization.fp8_cast import build_policy as _build_fp8_cast_policy
from ltx_core.quantization.fp8_scaled_mm import build_policy as _build_fp8_scaled_mm_policy
class QuantizationKind(str, Enum):
FP8_CAST = "fp8-cast"
FP8_SCALED_MM = "fp8-scaled-mm"
def to_policy(self, checkpoint_path: str | None = None) -> QuantizationPolicy:
"""Build the :class:`QuantizationPolicy` for this kind.
``checkpoint_path`` is required for both backends: ``FP8_SCALED_MM``
uses it to discover the layer set from ``.weight_scale`` tensors,
and ``FP8_CAST`` uses it to fold any prequant scales into the fp8
weight at load time.
"""
if checkpoint_path is None:
raise ValueError(f"{self.value} quantization requires checkpoint_path.")
match self:
case QuantizationKind.FP8_CAST:
return _build_fp8_cast_policy(checkpoint_path)
case QuantizationKind.FP8_SCALED_MM:
return _build_fp8_scaled_mm_policy(checkpoint_path)
case _:
assert_never(self)