Automated PR - 2026-05-28

packages/ltx-pipelines/src/ltx_pipelines/utils/args.py

@@ -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

@@ -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

@@ -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

@@ -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)