Merge pull request #111 from Lightricks/pr-2026-02-09

Open a PR to sync - 2026-02-09

LICENSE

@@ -90,20 +90,20 @@ to be bound by this Agreement.
       subject to the terms and provisions of a different license (the
       "Commercial Use Agreement"), as will be provided by the Licensor.
       Commercial Entities interested in such a commercial license are
-      required to contact Licensor. Any commercial use of LTX-2 or
-      Derivatives of LTX-2 by the Commercial Entities not in accordance
-      with this Agreement and/or the Commercial Use Agreement is strictly
-      prohibited and shall be deemed a material breach of this Agreement.
-      Such material breach will be subject, in addition to any license
-      fees owed to Licensor for the period such Commercial Entity used
-      LTX-2 (as will be determined by Licensor), to liquidated damages,
-      which will be paid to Licensor immediately upon demand, in an
-      amount equal to double the amount that would otherwise have been
-      paid by you for the relevant period of time. Such amount reflects
-      a reasonable estimation of the losses and administrative costs
-      incurred due to such breach. You agree and understand that this
-      remedy does not limit the Licensor's right to pursue other remedies
-      available at law or equity.
+      required to [contact Licensor](https://ltx.io/model/licensing).
+      Any commercial use of LTX-2 or Derivatives of LTX-2 by the
+      Commercial Entities not in accordance with this Agreement and/or
+      the Commercial Use Agreement is strictly prohibited and shall be
+      deemed a material breach of this Agreement. Such material breach
+      will be subject, in addition to any license fees owed to Licensor
+      for the period such Commercial Entity used LTX-2 (as will be
+      determined by Licensor), to liquidated damages, which will be paid
+      to Licensor immediately upon demand, in an amount equal to double
+      the amount that would otherwise have been paid by you for the
+      relevant period of time. Such amount reflects a reasonable estimation
+      of the losses and administrative costs incurred due to such breach.
+      You agree and understand that this remedy does not limit the Licensor's
+      right to pursue other remedies available at law or equity.
 
    3. Distribution and Redistribution. You may host for third parties
       remote access purposes (e.g. software-as-a-service), reproduce

packages/ltx-core/README.md

@@ -10,6 +10,7 @@ The foundational library for the LTX-2 Audio-Video generation model. This packag
 - **`loader/`**: Utilities for loading weights from `.safetensors`, fusing LoRAs, and managing memory
 - **`model/`**: PyTorch implementations of the LTX-2 Transformer, Video VAE, Audio VAE, Vocoder and Upscaler
 - **`text_encoders/gemma`**: Gemma text encoder implementation with tokenizers, feature extractors, and separate encoders for audio-video and video-only generation
+- **`quantization/`**: FP8 quantization backends (FP8-TensorRT-LLM scaled MM, FP8 cast) for reduced memory footprint.
 
 ## 🚀 Quick Start
 
@@ -53,6 +54,63 @@ pip install -e packages/ltx-core
 ### Utilities
 
 - **Loader** ([`loader/`](src/ltx_core/loader/)): Model loading from `.safetensors`, LoRA fusion, weight remapping, and memory management
+- **Quantization** ([`quantization/`](src/ltx_core/quantization/)): FP8 quantization backends for reduced memory footprint and faster inference
+
+### Quantization
+
+The `quantization/` module provides FP8 quantization support for the LTX-2 transformer, significantly reducing memory usage while maintaining quality. Two backends are available:
+
+#### FP8 Scaled MM (TensorRT-LLM)
+
+Uses NVIDIA TensorRT-LLM's `cublas_scaled_mm` for efficient FP8 matrix multiplication. Weights are stored in FP8 format with per-tensor scaling, and inputs are quantized dynamically (or statically with calibration data).
+
+**Requirements**: `uv sync --frozen --extra fp8-trtllm`
+
+**Usage with QuantizationPolicy:**
+
+```python
+from ltx_core.quantization import QuantizationPolicy
+
+# Dynamic input quantization (no calibration needed)
+policy = QuantizationPolicy.fp8_scaled_mm()
+
+# Static input quantization with calibration file
+policy = QuantizationPolicy.fp8_scaled_mm(calibration_amax_path="/path/to/amax.json")
+```
+
+The policy provides `sd_ops` and `module_ops` that can be passed to the model builder:
+
+```python
+from ltx_core.loader import SingleGPUModelBuilder
+
+builder = SingleGPUModelBuilder(
+    model=model,
+    device=device,
+    sd_ops=policy.sd_ops,
+    module_ops=policy.module_ops,
+)
+builder.load(checkpoint_path)
+```
+
+**Calibration File Format** (for static input quantization):
+
+```json
+{
+  "amax_values": {
+    "transformer_blocks.0.attn.to_q.input_quantizer": 12.5,
+    "transformer_blocks.0.attn.to_k.input_quantizer": 8.3,
+    ...
+  }
+}
+```
+
+#### FP8 Cast
+
+A simpler approach that casts weights to FP8 for storage and upcasts during inference:
+
+```python
+policy = QuantizationPolicy.fp8_cast()
+```
 
 For complete, production-ready pipeline implementations that combine these building blocks, see the [`ltx-pipelines`](../ltx-pipelines/) package.
 

packages/ltx-core/pyproject.toml

@@ -9,7 +9,7 @@ dependencies = [
     "torchaudio",
     "einops",
     "numpy",
-    "transformers~=4.57.0",
+    "transformers>=4.52",
     "safetensors",
     "accelerate",
     "scipy>=1.14",
@@ -17,16 +17,34 @@ dependencies = [
 
 [project.optional-dependencies]
 xformers = ["xformers"]
+fp8-trtllm = [
+    "tensorrt-llm==1.0.0",
+    "onnx>=1.16.0,<1.20.0",
+    "openmpi",
+]
 
+[tool.uv]
+conflicts = [
+  [
+    { extra = "xformers" },
+    { extra = "fp8-trtllm" },
+  ],
+]
 
 [tool.uv.sources]
 xformers = { index = "pytorch" }
+tensorrt-llm = { index = "nvidia" }
 
 [[tool.uv.index]]
 name = "pytorch"
 url = "https://download.pytorch.org/whl/cu129"
 explicit = true
 
+[[tool.uv.index]]
+name = "nvidia"
+url = "https://pypi.nvidia.com/"
+explicit = true
+
 [build-system]
 requires = ["uv_build>=0.9.8,<0.10.0"]
 build-backend = "uv_build"

packages/ltx-core/src/ltx_core/loader/fuse_loras.py

@@ -1,44 +1,46 @@
 import torch
-import triton
 
-from ltx_core.loader.kernels import fused_add_round_kernel
 from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict
-
-BLOCK_SIZE = 1024
+from ltx_core.quantization.fp8_cast import calculate_weight_float8
+from ltx_core.quantization.fp8_scaled_mm import quantize_weight_to_fp8_per_tensor
 
 
-def fused_add_round_launch(target_weight: torch.Tensor, original_weight: torch.Tensor, seed: int) -> torch.Tensor:
-    if original_weight.dtype == torch.float8_e4m3fn:
-        exponent_bits, mantissa_bits, exponent_bias = 4, 3, 7
-    elif original_weight.dtype == torch.float8_e5m2:
-        exponent_bits, mantissa_bits, exponent_bias = 5, 2, 15  # noqa: F841
-    else:
-        raise ValueError("Unsupported dtype")
+def apply_loras(
+    model_sd: StateDict,
+    lora_sd_and_strengths: list[LoraStateDictWithStrength],
+    dtype: torch.dtype | None = None,
+    destination_sd: StateDict | None = None,
+) -> StateDict:
+    sd = {}
+    if destination_sd is not None:
+        sd = destination_sd.sd
+    size = 0
+    device = torch.device("meta")
+    inner_dtypes = set()
+    for key, weight in model_sd.sd.items():
+        if weight is None:
+            continue
+        # Skip scale keys - they are handled together with their weight keys
+        if key.endswith(".weight_scale"):
+            continue
+        device = weight.device
+        target_dtype = dtype if dtype is not None else weight.dtype
+        deltas_dtype = target_dtype if target_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2] else torch.bfloat16
 
-    if target_weight.dtype != torch.bfloat16:
-        raise ValueError("target_weight dtype must be bfloat16")
+        scale_key = key.replace(".weight", ".weight_scale") if key.endswith(".weight") else None
+        is_scaled_fp8 = scale_key is not None and scale_key in model_sd.sd
 
-    # Calculate grid and block sizes
-    n_elements = original_weight.numel()
-    grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
+        deltas = _prepare_deltas(lora_sd_and_strengths, key, deltas_dtype, device)
+        fused = _fuse_deltas(deltas, weight, key, sd, target_dtype, device, is_scaled_fp8, scale_key, model_sd)
 
-    # Launch kernel
-    fused_add_round_kernel[grid](
-        original_weight,
-        target_weight,
-        seed,
-        n_elements,
-        exponent_bias,
-        mantissa_bits,
-        BLOCK_SIZE,
-    )
-    return target_weight
+        sd.update(fused)
+        for tensor in fused.values():
+            inner_dtypes.add(tensor.dtype)
+            size += tensor.nbytes
 
-
-def calculate_weight_float8_(target_weights: torch.Tensor, original_weights: torch.Tensor) -> torch.Tensor:
-    result = fused_add_round_launch(target_weights, original_weights, seed=0).to(target_weights.dtype)
-    target_weights.copy_(result, non_blocking=True)
-    return target_weights
+    if destination_sd is not None:
+        return destination_sd
+    return StateDict(sd, device, size, inner_dtypes)
 
 
 def _prepare_deltas(
@@ -60,41 +62,89 @@ def _prepare_deltas(
     return torch.sum(torch.stack(deltas, dim=0), dim=0)
 
 
-def apply_loras(
+def _fuse_deltas(
+    deltas: torch.Tensor | None,
+    weight: torch.Tensor,
+    key: str,
+    sd: dict[str, torch.Tensor],
+    target_dtype: torch.dtype,
+    device: torch.device,
+    is_scaled_fp8: bool,
+    scale_key: str | None,
     model_sd: StateDict,
-    lora_sd_and_strengths: list[LoraStateDictWithStrength],
-    dtype: torch.dtype,
-    destination_sd: StateDict | None = None,
-) -> StateDict:
-    sd = {}
-    if destination_sd is not None:
-        sd = destination_sd.sd
-    size = 0
-    device = torch.device("meta")
-    inner_dtypes = set()
-    for key, weight in model_sd.sd.items():
-        if weight is None:
-            continue
-        device = weight.device
-        target_dtype = dtype if dtype is not None else weight.dtype
-        deltas_dtype = target_dtype if target_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2] else torch.bfloat16
-        deltas = _prepare_deltas(lora_sd_and_strengths, key, deltas_dtype, device)
-        if deltas is None:
-            if key in sd:
-                continue
-            deltas = weight.clone().to(dtype=target_dtype, device=device)
-        elif weight.dtype == torch.float8_e4m3fn:
-            if str(device).startswith("cuda"):
-                deltas = calculate_weight_float8_(deltas, weight)
-            else:
-                deltas.add_(weight.to(dtype=deltas.dtype, device=device))
-        elif weight.dtype == torch.bfloat16:
-            deltas.add_(weight)
+) -> dict[str, torch.Tensor]:
+    if deltas is None:
+        if key in sd:
+            return {}
+        fused = _copy_weight_without_lora(weight, key, target_dtype, device, is_scaled_fp8, scale_key, model_sd)
+    elif weight.dtype == torch.float8_e4m3fn:
+        if is_scaled_fp8:
+            fused = _fuse_delta_with_scaled_fp8(deltas, weight, key, scale_key, model_sd)
         else:
-            raise ValueError(f"Unsupported dtype: {weight.dtype}")
-        sd[key] = deltas.to(dtype=target_dtype)
-        inner_dtypes.add(target_dtype)
-        size += deltas.nbytes
-    if destination_sd is not None:
-        return destination_sd
-    return StateDict(sd, device, size, inner_dtypes)
+            fused = _fuse_delta_with_cast_fp8(deltas, weight, key, target_dtype, device)
+    elif weight.dtype == torch.bfloat16:
+        fused = _fuse_delta_with_bfloat16(deltas, weight, key, target_dtype)
+    else:
+        raise ValueError(f"Unsupported dtype: {weight.dtype}")
+
+    return fused
+
+
+def _copy_weight_without_lora(
+    weight: torch.Tensor,
+    key: str,
+    target_dtype: torch.dtype,
+    device: torch.device,
+    is_scaled_fp8: bool,
+    scale_key: str | None,
+    model_sd: StateDict,
+) -> dict[str, torch.Tensor]:
+    """Copy original weight (and scale if applicable) when no LoRA affects this key."""
+    result = {key: weight.clone().to(dtype=target_dtype, device=device)}
+    if is_scaled_fp8:
+        result[scale_key] = model_sd.sd[scale_key].clone()
+    return result
+
+
+def _fuse_delta_with_scaled_fp8(
+    deltas: torch.Tensor,
+    weight: torch.Tensor,
+    key: str,
+    scale_key: str,
+    model_sd: StateDict,
+) -> dict[str, torch.Tensor]:
+    """Dequantize scaled FP8 weight, add LoRA delta, and re-quantize."""
+    weight_scale = model_sd.sd[scale_key]
+
+    original_weight = weight.t().to(torch.float32) * weight_scale
+
+    new_weight = original_weight + deltas.to(torch.float32)
+
+    new_fp8_weight, new_weight_scale = quantize_weight_to_fp8_per_tensor(new_weight)
+    return {key: new_fp8_weight, scale_key: new_weight_scale}
+
+
+def _fuse_delta_with_cast_fp8(
+    deltas: torch.Tensor,
+    weight: torch.Tensor,
+    key: str,
+    target_dtype: torch.dtype,
+    device: torch.device,
+) -> dict[str, torch.Tensor]:
+    """Fuse LoRA delta with cast-only FP8 weight (no scale factor)."""
+    if str(device).startswith("cuda"):
+        deltas = calculate_weight_float8(deltas, weight)
+    else:
+        deltas.add_(weight.to(dtype=deltas.dtype, device=device))
+    return {key: deltas.to(dtype=target_dtype)}
+
+
+def _fuse_delta_with_bfloat16(
+    deltas: torch.Tensor,
+    weight: torch.Tensor,
+    key: str,
+    target_dtype: torch.dtype,
+) -> dict[str, torch.Tensor]:
+    """Fuse LoRA delta with bfloat16 weight."""
+    deltas.add_(weight)
+    return {key: deltas.to(dtype=target_dtype)}

packages/ltx-core/src/ltx_core/model/transformer/__init__.py

@@ -4,21 +4,15 @@ from ltx_core.model.transformer.modality import Modality
 from ltx_core.model.transformer.model import LTXModel, X0Model
 from ltx_core.model.transformer.model_configurator import (
     LTXV_MODEL_COMFY_RENAMING_MAP,
-    LTXV_MODEL_COMFY_RENAMING_WITH_TRANSFORMER_LINEAR_DOWNCAST_MAP,
-    UPCAST_DURING_INFERENCE,
     LTXModelConfigurator,
     LTXVideoOnlyModelConfigurator,
-    UpcastWithStochasticRounding,
 )
 
 __all__ = [
     "LTXV_MODEL_COMFY_RENAMING_MAP",
-    "LTXV_MODEL_COMFY_RENAMING_WITH_TRANSFORMER_LINEAR_DOWNCAST_MAP",
-    "UPCAST_DURING_INFERENCE",
     "LTXModel",
     "LTXModelConfigurator",
     "LTXVideoOnlyModelConfigurator",
     "Modality",
-    "UpcastWithStochasticRounding",
     "X0Model",
 ]

packages/ltx-core/src/ltx_core/model/transformer/attention.py

@@ -150,6 +150,7 @@ class Attention(torch.nn.Module):
         norm_eps: float = 1e-6,
         rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
         attention_function: AttentionCallable | AttentionFunction = AttentionFunction.DEFAULT,
+        apply_gated_attention: bool = False,
     ) -> None:
         super().__init__()
         self.rope_type = rope_type
@@ -168,6 +169,12 @@ class Attention(torch.nn.Module):
         self.to_k = torch.nn.Linear(context_dim, inner_dim, bias=True)
         self.to_v = torch.nn.Linear(context_dim, inner_dim, bias=True)
 
+        # Optional per-head gating
+        if apply_gated_attention:
+            self.to_gate_logits = torch.nn.Linear(query_dim, heads, bias=True)
+        else:
+            self.to_gate_logits = None
+
         self.to_out = torch.nn.Sequential(torch.nn.Linear(inner_dim, query_dim, bias=True), torch.nn.Identity())
 
     def forward(
@@ -191,5 +198,18 @@ class Attention(torch.nn.Module):
             k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type)
 
         # attention_function can be an enum *or* a custom callable
-        out = self.attention_function(q, k, v, self.heads, mask)
+        out = self.attention_function(q, k, v, self.heads, mask)  # (B, T, H*D)
+
+        # Apply per-head gating if enabled
+        if self.to_gate_logits is not None:
+            gate_logits = self.to_gate_logits(x)  # (B, T, H)
+            b, t, _ = out.shape
+            # Reshape to (B, T, H, D) for per-head gating
+            out = out.view(b, t, self.heads, self.dim_head)
+            # Apply gating: 2 * sigmoid(x) so that zero-init gives identity (2 * 0.5 = 1.0)
+            gates = 2.0 * torch.sigmoid(gate_logits)  # (B, T, H)
+            out = out * gates.unsqueeze(-1)  # (B, T, H, D) * (B, T, H, 1)
+            # Reshape back to (B, T, H*D)
+            out = out.view(b, t, self.heads * self.dim_head)
+
         return self.to_out(out)

packages/ltx-core/src/ltx_core/model/transformer/model.py

@@ -61,6 +61,7 @@ class LTXModel(torch.nn.Module):
         av_ca_timestep_scale_multiplier: int = 1,
         rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
         double_precision_rope: bool = False,
+        apply_gated_attention: bool = False,
     ):
         super().__init__()
         self._enable_gradient_checkpointing = False
@@ -113,6 +114,7 @@ class LTXModel(torch.nn.Module):
             audio_cross_attention_dim=audio_cross_attention_dim,
             norm_eps=norm_eps,
             attention_type=attention_type,
+            apply_gated_attention=apply_gated_attention,
         )
 
     def _init_video(
@@ -272,6 +274,7 @@ class LTXModel(torch.nn.Module):
         audio_cross_attention_dim: int,
         norm_eps: float,
         attention_type: AttentionFunction | AttentionCallable,
+        apply_gated_attention: bool,
     ) -> None:
         """Initialize transformer blocks for LTX."""
         video_config = (
@@ -280,6 +283,7 @@ class LTXModel(torch.nn.Module):
                 heads=self.num_attention_heads,
                 d_head=attention_head_dim,
                 context_dim=cross_attention_dim,
+                apply_gated_attention=apply_gated_attention,
             )
             if self.model_type.is_video_enabled()
             else None
@@ -290,6 +294,7 @@ class LTXModel(torch.nn.Module):
                 heads=self.audio_num_attention_heads,
                 d_head=audio_attention_head_dim,
                 context_dim=audio_cross_attention_dim,
+                apply_gated_attention=apply_gated_attention,
             )
             if self.model_type.is_audio_enabled()
             else None

packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py

@@ -1,8 +1,4 @@
-import torch
-
-from ltx_core.loader.fuse_loras import fused_add_round_launch
-from ltx_core.loader.module_ops import ModuleOps
-from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
+from ltx_core.loader.sd_ops import SDOps
 from ltx_core.model.model_protocol import ModelConfigurator
 from ltx_core.model.transformer.attention import AttentionFunction
 from ltx_core.model.transformer.model import LTXModel, LTXModelType
@@ -63,6 +59,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
             av_ca_timestep_scale_multiplier=config.get("av_ca_timestep_scale_multiplier", 1),
             rope_type=LTXRopeType(config.get("rope_type", "interleaved")),
             double_precision_rope=config.get("frequencies_precision", False) == "float64",
+            apply_gated_attention=config.get("apply_gated_attention", False),
         )
 
 
@@ -109,129 +106,12 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
             use_middle_indices_grid=config.get("use_middle_indices_grid", True),
             rope_type=LTXRopeType(config.get("rope_type", "interleaved")),
             double_precision_rope=config.get("frequencies_precision", False) == "float64",
+            apply_gated_attention=config.get("apply_gated_attention", False),
         )
 
 
-def _naive_weight_or_bias_downcast(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
-    """
-    Downcast the weight or bias to the float8_e4m3fn dtype.
-    """
-    return [KeyValueOperationResult(key, value.to(dtype=torch.float8_e4m3fn))]
-
-
-def _upcast_and_round(
-    weight: torch.Tensor, dtype: torch.dtype, with_stochastic_rounding: bool = False, seed: int = 0
-) -> torch.Tensor:
-    """
-    Upcast the weight to the given dtype and optionally apply stochastic rounding.
-    Input weight needs to have float8_e4m3fn or float8_e5m2 dtype.
-    """
-    if not with_stochastic_rounding:
-        return weight.to(dtype)
-    return fused_add_round_launch(torch.zeros_like(weight, dtype=dtype), weight, seed)
-
-
-def replace_fwd_with_upcast(layer: torch.nn.Linear, with_stochastic_rounding: bool = False, seed: int = 0) -> None:
-    """
-    Replace linear.forward and rms_norm.forward with a version that:
-      - upcasts weight and bias to input's dtype
-      - returns F.linear or F.rms_norm calculated in that dtype
-    """
-
-    layer.original_forward = layer.forward
-
-    def new_linear_forward(*args, **_kwargs) -> torch.Tensor:
-        # assume first arg is the input tensor
-        x = args[0]
-        w_up = _upcast_and_round(layer.weight, x.dtype, with_stochastic_rounding, seed)
-        b_up = None
-
-        if layer.bias is not None:
-            b_up = _upcast_and_round(layer.bias, x.dtype, with_stochastic_rounding, seed)
-
-        return torch.nn.functional.linear(x, w_up, b_up)
-
-    layer.forward = new_linear_forward
-
-
-def amend_forward_with_upcast(
-    model: torch.nn.Module, with_stochastic_rounding: bool = False, seed: int = 0
-) -> torch.nn.Module:
-    """
-    Replace the forward method of the model's Linear and RMSNorm layers to forward
-    with upcast and optional stochastic rounding.
-    """
-    for m in model.modules():
-        if isinstance(m, (torch.nn.Linear)):
-            replace_fwd_with_upcast(m, with_stochastic_rounding, seed)
-    return model
-
-
 LTXV_MODEL_COMFY_RENAMING_MAP = (
     SDOps("LTXV_MODEL_COMFY_PREFIX_MAP")
     .with_matching(prefix="model.diffusion_model.")
     .with_replacement("model.diffusion_model.", "")
 )
-
-LTXV_MODEL_COMFY_RENAMING_WITH_TRANSFORMER_LINEAR_DOWNCAST_MAP = (
-    SDOps("LTXV_MODEL_COMFY_PREFIX_MAP")
-    .with_matching(prefix="model.diffusion_model.")
-    .with_replacement("model.diffusion_model.", "")
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_q.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_q.bias", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_k.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_k.bias", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_v.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_v.bias", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_out.0.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_out.0.bias", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".ff.net.0.proj.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".ff.net.0.proj.bias", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".ff.net.2.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".ff.net.2.bias", operation=_naive_weight_or_bias_downcast
-    )
-)
-
-UPCAST_DURING_INFERENCE = ModuleOps(
-    name="upcast_fp8_during_linear_forward",
-    matcher=lambda model: isinstance(model, LTXModel),
-    mutator=lambda model: amend_forward_with_upcast(model, False),
-)
-
-
-class UpcastWithStochasticRounding(ModuleOps):
-    """
-    ModuleOps for upcasting the model's float8_e4m3fn weights and biases to the bfloat16 dtype
-    and applying stochastic rounding during linear forward.
-    """
-
-    def __new__(cls, seed: int = 0):
-        return super().__new__(
-            cls,
-            name="upcast_fp8_during_linear_forward_with_stochastic_rounding",
-            matcher=lambda model: isinstance(model, LTXModel),
-            mutator=lambda model: amend_forward_with_upcast(model, True, seed),
-        )

packages/ltx-core/src/ltx_core/model/transformer/transformer.py

@@ -16,6 +16,7 @@ class TransformerConfig:
     heads: int
     d_head: int
     context_dim: int
+    apply_gated_attention: bool = False
 
 
 class BasicAVTransformerBlock(torch.nn.Module):
@@ -40,6 +41,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 rope_type=rope_type,
                 norm_eps=norm_eps,
                 attention_function=attention_function,
+                apply_gated_attention=video.apply_gated_attention,
             )
             self.attn2 = Attention(
                 query_dim=video.dim,
@@ -49,6 +51,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 rope_type=rope_type,
                 norm_eps=norm_eps,
                 attention_function=attention_function,
+                apply_gated_attention=video.apply_gated_attention,
             )
             self.ff = FeedForward(video.dim, dim_out=video.dim)
             self.scale_shift_table = torch.nn.Parameter(torch.empty(6, video.dim))
@@ -62,6 +65,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 rope_type=rope_type,
                 norm_eps=norm_eps,
                 attention_function=attention_function,
+                apply_gated_attention=audio.apply_gated_attention,
             )
             self.audio_attn2 = Attention(
                 query_dim=audio.dim,
@@ -71,6 +75,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 rope_type=rope_type,
                 norm_eps=norm_eps,
                 attention_function=attention_function,
+                apply_gated_attention=audio.apply_gated_attention,
             )
             self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim)
             self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(6, audio.dim))
@@ -85,6 +90,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 rope_type=rope_type,
                 norm_eps=norm_eps,
                 attention_function=attention_function,
+                apply_gated_attention=video.apply_gated_attention,
             )
 
             # Q: Audio, K,V: Video
@@ -96,6 +102,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 rope_type=rope_type,
                 norm_eps=norm_eps,
                 attention_function=attention_function,
+                apply_gated_attention=audio.apply_gated_attention,
             )
 
             self.scale_shift_table_a2v_ca_audio = torch.nn.Parameter(torch.empty(5, audio.dim))

packages/ltx-core/src/ltx_core/model/video_vae/tiling.py

@@ -44,6 +44,31 @@ def compute_trapezoidal_mask_1d(
     return mask.clamp_(0, 1)
 
 
+def compute_rectangular_mask_1d(
+    length: int,
+    left_ramp: int,
+    right_ramp: int,
+) -> torch.Tensor:
+    """
+    Generate a 1D rectangular (pulse) mask.
+    Args:
+        length: Output length of the mask.
+        left_ramp: Number of elements at the start of the mask to set to 0.
+        right_ramp: Number of elements at the end of the mask to set to 0.
+    Returns:
+        A 1D tensor of shape `(length,)` with values 0 or 1.
+    """
+    if length <= 0:
+        raise ValueError("Mask length must be positive.")
+
+    mask = torch.ones(length)
+    if left_ramp > 0:
+        mask[:left_ramp] = 0
+    if right_ramp > 0:
+        mask[-right_ramp:] = 0
+    return mask
+
+
 @dataclass(frozen=True)
 class SpatialTilingConfig:
     """Configuration for dividing each frame into spatial tiles with optional overlap.
@@ -114,12 +139,16 @@ class TilingConfig:
 
 @dataclass(frozen=True)
 class DimensionIntervals:
-    """Intervals which a single dimension of the latent space is split into.
-    Each interval is defined by its start, end, left ramp, and right ramp.
-    The start and end are the indices of the first and last element (exclusive) in the interval.
-    Ramps are regions of the interval where the value of the mask tensor is
-    interpolated between 0 and 1 for blending with neighboring intervals.
-    The left ramp and right ramp values are the lengths of the left and right ramps.
+    """Defines how a single dimension is split into overlapping intervals (tiles).
+    Each list has length N where N is the number of intervals. The i-th element
+    of each list describes the i-th interval.
+    Attributes:
+        starts: Start index of each interval (inclusive).
+        ends: End index of each interval (exclusive).
+        left_ramps: Length of the left blend ramp for each interval.
+            Used to create masks that fade in from 0 to 1.
+        right_ramps: Length of the right blend ramp for each interval.
+            Used to create masks that fade out from 1 to 0.
     """
 
     starts: List[int]
@@ -129,9 +158,11 @@ class DimensionIntervals:
 
 
 @dataclass(frozen=True)
-class LatentIntervals:
-    """Intervals which the latent tensor of given shape is split into.
-    Each dimension of the latent space is split into intervals based on the length along said dimension.
+class TensorTilingSpec:
+    """Specifies how a tensor of a given shape is split into intervals (tiles) along each dimension.
+    Attributes:
+        original_shape: Shape of the tensor being tiled.
+        dimension_intervals: Per-dimension intervals (starts, ends, ramps) for each axis.
     """
 
     original_shape: torch.Size
@@ -209,7 +240,7 @@ class Tile(NamedTuple):
 
 
 def create_tiles_from_intervals_and_mappers(
-    intervals: LatentIntervals,
+    intervals: TensorTilingSpec,
     mappers: List[MappingOperation],
 ) -> List[Tile]:
     full_dim_input_slices = []
@@ -241,20 +272,20 @@ def create_tiles_from_intervals_and_mappers(
 
 
 def create_tiles(
-    latent_shape: torch.Size,
+    tensor_shape: torch.Size,
     splitters: List[SplitOperation],
     mappers: List[MappingOperation],
 ) -> List[Tile]:
-    if len(splitters) != len(latent_shape):
+    if len(splitters) != len(tensor_shape):
         raise ValueError(
-            f"Number of splitters must be equal to number of dimensions in latent shape, "
-            f"got {len(splitters)} and {len(latent_shape)}"
+            f"Number of splitters must be equal to number of dimensions in tensor shape, "
+            f"got {len(splitters)} and {len(tensor_shape)}"
         )
-    if len(mappers) != len(latent_shape):
+    if len(mappers) != len(tensor_shape):
         raise ValueError(
-            f"Number of mappers must be equal to number of dimensions in latent shape, "
-            f"got {len(mappers)} and {len(latent_shape)}"
+            f"Number of mappers must be equal to number of dimensions in tensor shape, "
+            f"got {len(mappers)} and {len(tensor_shape)}"
         )
-    intervals = [splitter(length) for splitter, length in zip(splitters, latent_shape, strict=True)]
-    latent_intervals = LatentIntervals(original_shape=latent_shape, dimension_intervals=tuple(intervals))
-    return create_tiles_from_intervals_and_mappers(latent_intervals, mappers)
+    intervals = [splitter(length) for splitter, length in zip(splitters, tensor_shape, strict=True)]
+    tiling_spec = TensorTilingSpec(original_shape=tensor_shape, dimension_intervals=tuple(intervals))
+    return create_tiles_from_intervals_and_mappers(tiling_spec, mappers)

packages/ltx-core/src/ltx_core/model/video_vae/video_vae.py

@@ -1,3 +1,4 @@
+import logging
 from dataclasses import replace
 from typing import Any, Callable, Iterator, List, Tuple
 
@@ -20,10 +21,13 @@ from ltx_core.model.video_vae.tiling import (
     SplitOperation,
     Tile,
     TilingConfig,
+    compute_rectangular_mask_1d,
     compute_trapezoidal_mask_1d,
     create_tiles,
 )
-from ltx_core.types import SpatioTemporalScaleFactors, VideoLatentShape
+from ltx_core.types import VIDEO_SCALE_FACTORS, SpatioTemporalScaleFactors, VideoLatentShape
+
+logger: logging.Logger = logging.getLogger(__name__)
 
 
 def _make_encoder_block(
@@ -248,18 +252,22 @@ class VideoEncoder(nn.Module):
         r"""
         Encode video frames into normalized latent representation.
         Args:
-            sample: Input video (B, C, F, H, W). F must be 1 + 8*k (e.g., 1, 9, 17, 25, 33...).
+            sample: Input video (B, C, F, H, W). F should be 1 + 8*k (e.g., 1, 9, 17, 25, 33...).
+                If not, the encoder crops the last frames to the nearest valid length.
         Returns:
             Normalized latent means (B, 128, F', H', W') where F' = 1+(F-1)/8, H' = H/32, W' = W/32.
             Example: (B, 3, 33, 512, 512) -> (B, 128, 5, 16, 16).
         """
-        # Validate frame count
+        # Validate frame count (crop to nearest valid length if needed)
         frames_count = sample.shape[2]
         if ((frames_count - 1) % 8) != 0:
-            raise ValueError(
-                "Invalid number of frames: Encode input must have 1 + 8 * x frames "
-                "(e.g., 1, 9, 17, ...). Please check your input."
+            frames_to_crop = (frames_count - 1) % 8
+            logger.warning(
+                "Invalid number of frames %s for encode; cropping last %s frames to satisfy 1 + 8*k.",
+                frames_count,
+                frames_to_crop,
             )
+            sample = sample[:, :, :-frames_to_crop, ...]
 
         # Initial spatial compression: trade spatial resolution for channel depth
         # This reduces H,W by patch_size and increases channels, making convolutions more efficient
@@ -311,6 +319,152 @@ class VideoEncoder(nn.Module):
         means, _ = torch.chunk(sample, 2, dim=1)
         return self.per_channel_statistics.normalize(means)
 
+    def tiled_encode(
+        self,
+        video: torch.Tensor,
+        tiling_config: TilingConfig | None = None,
+    ) -> torch.Tensor:
+        """Encode video to latent using tiled processing of the given video tensor.
+        Device Handling:
+            - Input video can be on CPU or GPU
+            - Accumulation buffers are created on model's device
+            - Each tile is automatically moved to model's device before encoding
+            - Output latent is returned on model's device
+        Args:
+            video: Input video tensor (B, 3, F, H, W) in range [-1, 1]
+            tiling_config: Tiling configuration for the video tensor
+        Returns:
+            Latent tensor (B, 128, F', H', W') on model's device
+            where F' = 1 + (F-1)/8, H' = H/32, W' = W/32
+        """
+        # Detect model device and dtype
+        model_device = next(self.parameters()).device
+        model_dtype = next(self.parameters()).dtype
+
+        # Extract shape components
+        batch, _, frames, height, width = video.shape
+
+        # Check frame count and crop if needed
+        if (frames - 1) % VIDEO_SCALE_FACTORS.time != 0:
+            frames_to_crop = (frames - 1) % VIDEO_SCALE_FACTORS.time
+            logger.warning(
+                f"Number of frames {frames} of input video is not ({VIDEO_SCALE_FACTORS.time} * k + 1), "
+                f"last {frames_to_crop} frames will be cropped"
+            )
+            video = video[:, :, :-frames_to_crop, ...]
+            # Update frames after cropping
+            frames = video.shape[2]
+
+        # Calculate output latent shape (inverse of upscale)
+        latent_shape = VideoLatentShape(
+            batch=batch,
+            channels=self.latent_channels,  # 128 for standard VAE
+            frames=(frames - 1) // VIDEO_SCALE_FACTORS.time + 1,
+            height=height // VIDEO_SCALE_FACTORS.height,
+            width=width // VIDEO_SCALE_FACTORS.width,
+        )
+
+        # Prepare tiles (operates on VIDEO dimensions)
+        tiles = prepare_tiles_for_encoding(video, tiling_config)
+
+        # Initialize accumulation buffers on model device
+        latent_buffer = torch.zeros(
+            latent_shape.to_torch_shape(),
+            device=model_device,
+            dtype=model_dtype,
+        )
+        weights_buffer = torch.zeros_like(latent_buffer)
+
+        # Process each tile
+        for tile in tiles:
+            # Extract video tile from input (may be on CPU)
+            video_tile = video[tile.in_coords]
+
+            # Move tile to model device if needed
+            if video_tile.device != model_device or video_tile.dtype != model_dtype:
+                video_tile = video_tile.to(device=model_device, dtype=model_dtype)
+
+            # Encode tile to latent (output on model device)
+            latent_tile = self.forward(video_tile)
+
+            # Move blend mask to model device
+            mask = tile.blend_mask.to(
+                device=model_device,
+                dtype=model_dtype,
+            )
+
+            # Weighted accumulation in latent space
+            latent_buffer[tile.out_coords] += latent_tile * mask
+            weights_buffer[tile.out_coords] += mask
+
+            del latent_tile, mask, video_tile
+
+        # Normalize by accumulated weights
+        weights_buffer = weights_buffer.clamp(min=1e-8)
+        return latent_buffer / weights_buffer
+
+
+def prepare_tiles_for_encoding(
+    video: torch.Tensor,
+    tiling_config: TilingConfig | None = None,
+) -> List[Tile]:
+    """Prepare tiles for VAE encoding.
+    Args:
+        video: Input video tensor (B, 3, F, H, W) in range [-1, 1]
+        tiling_config: Tiling configuration for the video tensor
+    Returns:
+        List of tiles for the video tensor
+    """
+
+    splitters = [DEFAULT_SPLIT_OPERATION] * len(video.shape)
+    mappers = [DEFAULT_MAPPING_OPERATION] * len(video.shape)
+    minimum_spatial_overlap_px = 64
+    minimum_temporal_overlap_frames = 16
+
+    if tiling_config is not None and tiling_config.spatial_config is not None:
+        cfg = tiling_config.spatial_config
+
+        tile_size_px = cfg.tile_size_in_pixels
+        overlap_px = cfg.tile_overlap_in_pixels
+
+        # Set minimum spatial overlap to 64 pixels in order to allow cutting padding from
+        # the front and back of the tiles and concatenate tiles without artifacts.
+        # The encoder uses symmetric padding (pad=1) in H and W at each conv layer. At tile
+        # boundaries, convs see padding (zeros/reflect) instead of real neighbor pixels, causing
+        # incorrect context near edges.
+        # For each overlap we discard 1 latent per edge (32px at scale 32) and concatenate tiles at a
+        # shared region with the next tile.
+        if overlap_px < minimum_spatial_overlap_px:
+            logger.warning(
+                f"Overlap pixels {overlap_px} in spatial tiling is less than \
+            {minimum_spatial_overlap_px}, setting to minimum required {minimum_spatial_overlap_px}"
+            )
+            overlap_px = minimum_spatial_overlap_px
+
+        # Define split and map operations for the spatial dimensions
+
+        # Height axis (H)
+        splitters[3] = split_with_symmetric_overlaps(tile_size_px, overlap_px)
+        mappers[3] = make_mapping_operation(map_spatial_interval_to_latent, scale=VIDEO_SCALE_FACTORS.height)
+
+        # Width axis (W)
+        splitters[4] = split_with_symmetric_overlaps(tile_size_px, overlap_px)
+        mappers[4] = make_mapping_operation(map_spatial_interval_to_latent, scale=VIDEO_SCALE_FACTORS.width)
+
+    if tiling_config is not None and tiling_config.temporal_config is not None:
+        cfg = tiling_config.temporal_config
+        tile_size_frames = cfg.tile_size_in_frames
+        overlap_frames = cfg.tile_overlap_in_frames
+
+        if overlap_frames < minimum_temporal_overlap_frames:
+            logger.warning(f"Overlap frames {overlap_frames} is less than 16, setting to minimum required 16")
+            overlap_frames = minimum_temporal_overlap_frames
+
+        splitters[2] = split_temporal_frames(tile_size_frames, overlap_frames)
+        mappers[2] = make_mapping_operation(map_temporal_interval_to_latent, scale=VIDEO_SCALE_FACTORS.time)
+
+    return create_tiles(video.shape, splitters, mappers)
+
 
 def _make_decoder_block(
     block_name: str,
@@ -631,8 +785,8 @@ class VideoDecoder(nn.Module):
                 axis_length = latent.shape[axis_idx]
                 lower_threshold = max(2, overlap + 1)
                 tile_size = max(lower_threshold, round(size * axis_length / long_side))
-                splitters[axis_idx] = split_in_spatial(tile_size, overlap)
-                mappers[axis_idx] = to_mapping_operation(map_spatial_slice, factor)
+                splitters[axis_idx] = split_with_symmetric_overlaps(tile_size, overlap)
+                mappers[axis_idx] = make_mapping_operation(map_spatial_interval_to_pixel, scale=factor)
 
             enable_on_axis(3, self.video_downscale_factors.height)
             enable_on_axis(4, self.video_downscale_factors.width)
@@ -641,8 +795,8 @@ class VideoDecoder(nn.Module):
             cfg = tiling_config.temporal_config
             tile_size = cfg.tile_size_in_frames // self.video_downscale_factors.time
             overlap = cfg.tile_overlap_in_frames // self.video_downscale_factors.time
-            splitters[2] = split_in_temporal(tile_size, overlap)
-            mappers[2] = to_mapping_operation(map_temporal_slice, self.video_downscale_factors.time)
+            splitters[2] = split_temporal_latents(tile_size, overlap)
+            mappers[2] = make_mapping_operation(map_temporal_interval_to_frame, scale=self.video_downscale_factors.time)
 
         return create_tiles(latent.shape, splitters, mappers)
 
@@ -856,7 +1010,7 @@ def get_video_chunks_number(num_frames: int, tiling_config: TilingConfig | None
     return (num_frames - 1 + frame_stride - 1) // frame_stride
 
 
-def split_in_spatial(size: int, overlap: int) -> SplitOperation:
+def split_with_symmetric_overlaps(size: int, overlap: int) -> SplitOperation:
     def split(dimension_size: int) -> DimensionIntervals:
         if dimension_size <= size:
             return DEFAULT_SPLIT_OPERATION(dimension_size)
@@ -871,26 +1025,81 @@ def split_in_spatial(size: int, overlap: int) -> SplitOperation:
     return split
 
 
-def split_in_temporal(size: int, overlap: int) -> SplitOperation:
-    non_causal_split = split_in_spatial(size, overlap)
+def split_temporal_latents(size: int, overlap: int) -> SplitOperation:
+    """Split a temporal axis into overlapping tiles with causal handling.
+    Example with size=24, overlap=8 (units are whatever axis you split):
+        Non-causal split would produce:
+            Tile 0: [0, 24), left_ramp=0,  right_ramp=8
+            Tile 1: [16, 40), left_ramp=8, right_ramp=8
+            Tile 2: [32, 56), left_ramp=8, right_ramp=0
+        Causal split produces:
+            Tile 0: [0, 24), left_ramp=0,  right_ramp=8  (unchanged - starts at anchor)
+            Tile 1: [15, 40), left_ramp=9, right_ramp=8  (shifted back 1, ramp +1)
+            Tile 2: [31, 56), left_ramp=9, right_ramp=0  (shifted back 1, ramp +1)
+    This ensures each tile can causally depend on frames from previous tiles while maintaining
+    proper temporal continuity through the blend ramps.
+    Args:
+        size: Tile size in *axis units* (latent steps for LTX time tiling)
+        overlap: Overlap between tiles in the same units
+    Returns:
+        Split operation that divides temporal dimension with causal handling
+    """
+    non_causal_split = split_with_symmetric_overlaps(size, overlap)
 
     def split(dimension_size: int) -> DimensionIntervals:
         if dimension_size <= size:
             return DEFAULT_SPLIT_OPERATION(dimension_size)
         intervals = non_causal_split(dimension_size)
+
         starts = intervals.starts
         starts[1:] = [s - 1 for s in starts[1:]]
+
+        # Extend blend ramps by 1 for non-first tiles to blend over the extra frame
         left_ramps = intervals.left_ramps
         left_ramps[1:] = [r + 1 for r in left_ramps[1:]]
+
         return replace(intervals, starts=starts, left_ramps=left_ramps)
 
     return split
 
 
-def to_mapping_operation(
-    map_func: Callable[[int, int, int, int, int], Tuple[slice, torch.Tensor]],
+def split_temporal_frames(tile_size_frames: int, overlap_frames: int) -> SplitOperation:
+    """Split a temporal axis in video frame space into overlapping tiles.
+    Args:
+        tile_size_frames: Tile length in frames.
+        overlap_frames: Overlap between consecutive tiles in frames.
+    Returns:
+        Split operation that takes frame count and returns DimensionIntervals in frame indices.
+    """
+    non_causal_split = split_with_symmetric_overlaps(tile_size_frames, overlap_frames)
+
+    def split(dimension_size: int) -> DimensionIntervals:
+        if dimension_size <= tile_size_frames:
+            return DEFAULT_SPLIT_OPERATION(dimension_size)
+        intervals = non_causal_split(dimension_size)
+        ends = intervals.ends
+        ends[:-1] = [e + 1 for e in ends[:-1]]
+        right_ramps = [0] * len(intervals.right_ramps)
+        return replace(intervals, ends=ends, right_ramps=right_ramps)
+
+    return split
+
+
+def make_mapping_operation(
+    map_func: Callable[[int, int, int, int, int], Tuple[slice, torch.Tensor | None]],
     scale: int,
 ) -> MappingOperation:
+    """Create a mapping operation over a set of tiling intervals.
+    The given mapping function is applied to each interval in the input dimension. The result function is used for
+    creating tiles in the output dimension.
+    Args:
+        map_func: Mapping function to create the mapping operation from
+        scale: Scale factor for the transformation, used as an argument for the mapping function
+    Returns:
+        Mapping operation that takes a set of tiling intervals and returns a set of slices and masks in the output
+        dimension.
+    """
+
     def map_op(intervals: DimensionIntervals) -> tuple[list[slice], list[torch.Tensor | None]]:
         output_slices: list[slice] = []
         masks_1d: list[torch.Tensor | None] = []
@@ -908,19 +1117,104 @@ def to_mapping_operation(
     return map_op
 
 
-def map_temporal_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]:
+def map_temporal_interval_to_frame(
+    begin: int,
+    end: int,
+    left_ramp: int,
+    right_ramp: int,
+    scale: int,
+) -> Tuple[slice, torch.Tensor]:
+    """Map temporal interval in latent space to video frame space.
+    Args:
+        begin: Start position in latent space
+        end: End position in latent space
+        left_ramp: Left ramp size in latent space
+        right_ramp: Right ramp size in latent space
+        scale: Scale factor for transformation
+    Returns:
+        Tuple of (output_slice, blend_mask)
+    """
     start = begin * scale
     stop = 1 + (end - 1) * scale
-    left_ramp = 1 + (left_ramp - 1) * scale
-    right_ramp = right_ramp * scale
 
-    return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp, right_ramp, True)
+    left_ramp_frames = 0 if left_ramp == 0 else 1 + (left_ramp - 1) * scale
+    right_ramp_frames = right_ramp * scale
+
+    mask_1d = compute_trapezoidal_mask_1d(stop - start, left_ramp_frames, right_ramp_frames, True)
+    return slice(start, stop), mask_1d
 
 
-def map_spatial_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]:
+def map_temporal_interval_to_latent(
+    begin: int, end: int, left_ramp: int, right_ramp: int | None = None, scale: int = 1
+) -> Tuple[slice, torch.Tensor]:
+    """
+    Map temporal interval in video frame space to latent space.
+    Args:
+        begin: Start position in video frame space
+        end: End position in video frame space
+        left_ramp: Left ramp size in video frame space
+        right_ramp: Right ramp size in video frame space
+        scale: Scale factor for transformation
+    Returns:
+        Tuple of (output_slice, blend_mask)
+    """
+    start = begin // scale
+    stop = (end - 1) // scale + 1
+
+    left_ramp_latents = 0 if left_ramp == 0 else 1 + (left_ramp - 1) // scale
+    right_ramp_latents = right_ramp // scale
+
+    if right_ramp_latents != 0:
+        raise ValueError("For tiled encoding, temporal tiles are expected to have a right ramp equal to 0")
+
+    mask_1d = compute_rectangular_mask_1d(stop - start, left_ramp_latents, right_ramp_latents)
+
+    return slice(start, stop), mask_1d
+
+
+def map_spatial_interval_to_pixel(
+    begin: int,
+    end: int,
+    left_ramp: int,
+    right_ramp: int,
+    scale: int,
+) -> Tuple[slice, torch.Tensor]:
+    """Map spatial interval in latent space to pixel space.
+    Args:
+        begin: Start position in latent space
+        end: End position in latent space
+        left_ramp: Left ramp size in latent space
+        right_ramp: Right ramp size in latent space
+        scale: Scale factor for transformation
+    """
     start = begin * scale
     stop = end * scale
-    left_ramp = left_ramp * scale
-    right_ramp = right_ramp * scale
+    mask_1d = compute_trapezoidal_mask_1d(stop - start, left_ramp * scale, right_ramp * scale, False)
+    return slice(start, stop), mask_1d
 
-    return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp, right_ramp, False)
+
+def map_spatial_interval_to_latent(
+    begin: int,
+    end: int,
+    left_ramp: int,
+    right_ramp: int,
+    scale: int,
+) -> Tuple[slice, torch.Tensor]:
+    """Map spatial interval in pixel space to latent space.
+       Args:
+        begin: Start position in pixel space
+        end: End position in pixel space
+        left_ramp: Left ramp size in pixel space
+        right_ramp: Right ramp size in pixel space
+        scale: Scale factor for transformation
+    Returns:
+        Tuple of (output_slice, blend_mask)
+    """
+    start = begin // scale
+    stop = end // scale
+    left_ramp = max(0, left_ramp // scale - 1)
+
+    right_ramp = 0 if right_ramp == 0 else 1
+
+    mask_1d = compute_rectangular_mask_1d(stop - start, left_ramp, right_ramp)
+    return slice(start, stop), mask_1d

packages/ltx-core/src/ltx_core/quantization/__init__.py

@@ -0,0 +1,16 @@
+from ltx_core.quantization.fp8_cast import (
+    TRANSFORMER_LINEAR_DOWNCAST_MAP,
+    UPCAST_DURING_INFERENCE,
+    UpcastWithStochasticRounding,
+)
+from ltx_core.quantization.fp8_scaled_mm import FP8_PREPARE_MODULE_OPS, FP8_TRANSPOSE_SD_OPS
+from ltx_core.quantization.policy import QuantizationPolicy
+
+__all__ = [
+    "FP8_PREPARE_MODULE_OPS",
+    "FP8_TRANSPOSE_SD_OPS",
+    "TRANSFORMER_LINEAR_DOWNCAST_MAP",
+    "UPCAST_DURING_INFERENCE",
+    "QuantizationPolicy",
+    "UpcastWithStochasticRounding",
+]

packages/ltx-core/src/ltx_core/quantization/fp8_cast.py

@@ -0,0 +1,163 @@
+import torch
+
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
+from ltx_core.model.transformer.model import LTXModel
+
+BLOCK_SIZE = 1024
+
+
+def calculate_weight_float8(target_weights: torch.Tensor, original_weights: torch.Tensor) -> torch.Tensor:
+    result = _fused_add_round_launch(target_weights, original_weights, seed=0).to(target_weights.dtype)
+    target_weights.copy_(result, non_blocking=True)
+    return target_weights
+
+
+def _fused_add_round_launch(target_weight: torch.Tensor, original_weight: torch.Tensor, seed: int) -> torch.Tensor:
+    # Lazy import triton - only available on CUDA platforms
+    import triton  # noqa: PLC0415
+
+    from ltx_core.loader.kernels import fused_add_round_kernel  # noqa: PLC0415
+
+    if original_weight.dtype == torch.float8_e4m3fn:
+        exponent_bits, mantissa_bits, exponent_bias = 4, 3, 7
+    elif original_weight.dtype == torch.float8_e5m2:
+        exponent_bits, mantissa_bits, exponent_bias = 5, 2, 15  # noqa: F841
+    else:
+        raise ValueError("Unsupported dtype")
+
+    if target_weight.dtype != torch.bfloat16:
+        raise ValueError("target_weight dtype must be bfloat16")
+
+    # Calculate grid and block sizes
+    n_elements = original_weight.numel()
+    grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
+
+    # Launch kernel
+    fused_add_round_kernel[grid](
+        original_weight,
+        target_weight,
+        seed,
+        n_elements,
+        exponent_bias,
+        mantissa_bits,
+        BLOCK_SIZE,
+    )
+    return target_weight
+
+
+def _naive_weight_or_bias_downcast(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
+    """
+    Downcast the weight or bias to the float8_e4m3fn dtype.
+    """
+    return [KeyValueOperationResult(key, value.to(dtype=torch.float8_e4m3fn))]
+
+
+def _upcast_and_round(
+    weight: torch.Tensor, dtype: torch.dtype, with_stochastic_rounding: bool = False, seed: int = 0
+) -> torch.Tensor:
+    """
+    Upcast the weight to the given dtype and optionally apply stochastic rounding.
+    Input weight needs to have float8_e4m3fn or float8_e5m2 dtype.
+    """
+    if not with_stochastic_rounding:
+        return weight.to(dtype)
+    return _fused_add_round_launch(torch.zeros_like(weight, dtype=dtype), weight, seed)
+
+
+def _replace_fwd_with_upcast(layer: torch.nn.Linear, with_stochastic_rounding: bool = False, seed: int = 0) -> None:
+    """
+    Replace linear.forward and rms_norm.forward with a version that:
+      - upcasts weight and bias to input's dtype
+      - returns F.linear or F.rms_norm calculated in that dtype
+    """
+
+    layer.original_forward = layer.forward
+
+    def new_linear_forward(*args, **_kwargs) -> torch.Tensor:
+        # assume first arg is the input tensor
+        x = args[0]
+        w_up = _upcast_and_round(layer.weight, x.dtype, with_stochastic_rounding, seed)
+        b_up = None
+
+        if layer.bias is not None:
+            b_up = _upcast_and_round(layer.bias, x.dtype, with_stochastic_rounding, seed)
+
+        return torch.nn.functional.linear(x, w_up, b_up)
+
+    layer.forward = new_linear_forward
+
+
+def _amend_forward_with_upcast(
+    model: torch.nn.Module, with_stochastic_rounding: bool = False, seed: int = 0
+) -> torch.nn.Module:
+    """
+    Replace the forward method of the model's Linear and RMSNorm layers to forward
+    with upcast and optional stochastic rounding.
+    """
+    for m in model.modules():
+        if isinstance(m, (torch.nn.Linear)):
+            _replace_fwd_with_upcast(m, with_stochastic_rounding, seed)
+    return model
+
+
+TRANSFORMER_LINEAR_DOWNCAST_MAP = (
+    SDOps("TRANSFORMER_LINEAR_DOWNCAST_MAP")
+    .with_kv_operation(
+        key_prefix="transformer_blocks.", key_suffix=".to_q.weight", operation=_naive_weight_or_bias_downcast
+    )
+    .with_kv_operation(
+        key_prefix="transformer_blocks.", key_suffix=".to_q.bias", operation=_naive_weight_or_bias_downcast
+    )
+    .with_kv_operation(
+        key_prefix="transformer_blocks.", key_suffix=".to_k.weight", operation=_naive_weight_or_bias_downcast
+    )
+    .with_kv_operation(
+        key_prefix="transformer_blocks.", key_suffix=".to_k.bias", operation=_naive_weight_or_bias_downcast
+    )
+    .with_kv_operation(
+        key_prefix="transformer_blocks.", key_suffix=".to_v.weight", operation=_naive_weight_or_bias_downcast
+    )
+    .with_kv_operation(
+        key_prefix="transformer_blocks.", key_suffix=".to_v.bias", operation=_naive_weight_or_bias_downcast
+    )
+    .with_kv_operation(
+        key_prefix="transformer_blocks.", key_suffix=".to_out.0.weight", operation=_naive_weight_or_bias_downcast
+    )
+    .with_kv_operation(
+        key_prefix="transformer_blocks.", key_suffix=".to_out.0.bias", operation=_naive_weight_or_bias_downcast
+    )
+    .with_kv_operation(
+        key_prefix="transformer_blocks.", key_suffix=".ff.net.0.proj.weight", operation=_naive_weight_or_bias_downcast
+    )
+    .with_kv_operation(
+        key_prefix="transformer_blocks.", key_suffix=".ff.net.0.proj.bias", operation=_naive_weight_or_bias_downcast
+    )
+    .with_kv_operation(
+        key_prefix="transformer_blocks.", key_suffix=".ff.net.2.weight", operation=_naive_weight_or_bias_downcast
+    )
+    .with_kv_operation(
+        key_prefix="transformer_blocks.", key_suffix=".ff.net.2.bias", operation=_naive_weight_or_bias_downcast
+    )
+)
+
+UPCAST_DURING_INFERENCE = ModuleOps(
+    name="upcast_fp8_during_linear_forward",
+    matcher=lambda model: isinstance(model, LTXModel),
+    mutator=lambda model: _amend_forward_with_upcast(model, False),
+)
+
+
+class UpcastWithStochasticRounding(ModuleOps):
+    """
+    ModuleOps for upcasting the model's float8_e4m3fn weights and biases to the bfloat16 dtype
+    and applying stochastic rounding during linear forward.
+    """
+
+    def __new__(cls, seed: int = 0):
+        return super().__new__(
+            cls,
+            name="upcast_fp8_during_linear_forward_with_stochastic_rounding",
+            matcher=lambda model: isinstance(model, LTXModel),
+            mutator=lambda model: _amend_forward_with_upcast(model, True, seed),
+        )

packages/ltx-core/src/ltx_core/quantization/fp8_scaled_mm.py

@@ -0,0 +1,207 @@
+from typing import Callable
+
+import torch
+from torch import nn
+
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
+from ltx_core.model.transformer import LTXModel
+
+
+class FP8Linear(nn.Module):
+    """Linear layer with FP8 weight storage for scaled matrix multiplication."""
+
+    in_features: int
+    out_features: int
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        device: torch.device | str | None = None,
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+
+        fp8_shape = (in_features, out_features)
+        self.weight = nn.Parameter(torch.empty(fp8_shape, dtype=torch.float8_e4m3fn, device=device))
+        # Weight scale for FP8 dequantization (shape matches checkpoint format)
+        self.weight_scale = nn.Parameter(torch.empty((), dtype=torch.float32, device=device))
+        # Input scale for static quantization (pre-quantized checkpoints)
+        self.input_scale = nn.Parameter(torch.empty((), dtype=torch.float32, device=device))
+
+        if bias:
+            self.bias = nn.Parameter(torch.empty(out_features, device=device))
+        else:
+            self.register_parameter("bias", None)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        origin_shape = x.shape
+
+        # Static quantization: use pre-computed scale
+        qinput, cur_input_scale = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x, self.input_scale)
+
+        # Flatten to 2D for matmul
+        if qinput.dim() == 3:
+            qinput = qinput.reshape(-1, qinput.shape[-1])
+
+        # FP8 scaled matmul
+        output = torch.ops.trtllm.cublas_scaled_mm(
+            qinput,
+            self.weight,
+            scale_a=cur_input_scale,
+            scale_b=self.weight_scale,
+            bias=None,
+            out_dtype=x.dtype,
+        )
+
+        # Add bias
+        if self.bias is not None:
+            bias = self.bias
+            if bias.dtype != output.dtype:
+                bias = bias.to(output.dtype)
+            output = output + bias
+
+        # Restore original shape
+        if output.dim() != len(origin_shape):
+            output_shape = list(origin_shape)
+            output_shape[-1] = output.shape[-1]
+            output = output.reshape(output_shape)
+
+        return output
+
+
+def quantize_weight_to_fp8_per_tensor(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Quantize a weight tensor to FP8 (float8_e4m3fn) using per-tensor scaling.
+    Args:
+        weight: The weight tensor to quantize (any dtype, will be cast to float32)
+    Returns:
+        Tuple of (quantized_weight, weight_scale):
+        - quantized_weight: FP8 tensor, transposed for cublas_scaled_mm
+        - weight_scale: Per-tensor scale factor (reciprocal of quantization scale)
+    """
+    weight_fp32 = weight.to(torch.float32)
+
+    fp8_min = torch.finfo(torch.float8_e4m3fn).min
+    fp8_max = torch.finfo(torch.float8_e4m3fn).max
+
+    max_abs = torch.amax(torch.abs(weight_fp32))
+    scale = fp8_max / max_abs
+
+    @torch.compiler.disable
+    def _quantize(
+        weight_fp32: torch.Tensor, scale: torch.Tensor, fp8_min: torch.Tensor, fp8_max: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        quantized_weight = torch.clamp(weight_fp32 * scale, min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
+        quantized_weight = quantized_weight.t()
+        weight_scale = scale.reciprocal()
+        return quantized_weight, weight_scale
+
+    quantized_weight, weight_scale = _quantize(weight_fp32, scale, fp8_min, fp8_max)
+    return quantized_weight, weight_scale
+
+
+def _should_skip_layer(layer_name: str, excluded_layer_substrings: tuple[str, ...]) -> bool:
+    return any(substring in layer_name for substring in excluded_layer_substrings)
+
+
+EXCLUDED_LAYER_SUBSTRINGS = (
+    "patchify_proj",
+    "adaln_single",
+    "av_ca_video_scale_shift_adaln_single",
+    "av_ca_a2v_gate_adaln_single",
+    "caption_projection",
+    "proj_out",
+    "audio_patchify_proj",
+    "audio_adaln_single",
+    "av_ca_audio_scale_shift_adaln_single",
+    "av_ca_v2a_gate_adaln_single",
+    "audio_caption_projection",
+    "audio_proj_out",
+    "transformer_blocks.0.",
+    *[f"transformer_blocks.{i}." for i in range(43, 48)],
+)
+
+
+def _linear_to_fp8linear(layer: nn.Linear) -> FP8Linear:
+    """
+    Create an FP8Linear layer from an nn.Linear layer.
+    Args:
+        layer: The nn.Linear layer to convert (typically on meta device)
+    Returns:
+        A new FP8Linear with the same configuration
+    """
+    return FP8Linear(
+        in_features=layer.in_features,
+        out_features=layer.out_features,
+        bias=layer.bias is not None,
+        device=layer.weight.device,
+    )
+
+
+def _apply_fp8_prepare_to_model(model: nn.Module, excluded_layer_substrings: tuple[str, ...]) -> nn.Module:
+    """Replace nn.Linear layers with FP8Linear in the module tree."""
+    replacements: list[tuple[nn.Module, str, nn.Linear]] = []
+
+    for name, module in model.named_modules():
+        if not isinstance(module, nn.Linear) or isinstance(module, FP8Linear):
+            continue
+
+        if _should_skip_layer(name, excluded_layer_substrings):
+            continue
+
+        if "." in name:
+            parent_name, attr_name = name.rsplit(".", 1)
+            parent = model.get_submodule(parent_name)
+        else:
+            parent = model
+            attr_name = name
+
+        replacements.append((parent, attr_name, module))
+
+    for parent, attr_name, linear in replacements:
+        setattr(parent, attr_name, _linear_to_fp8linear(linear))
+
+    return model
+
+
+def _create_transpose_kv_operation(
+    excluded_layer_substrings: tuple[str, ...],
+) -> Callable[[str, torch.Tensor], list[KeyValueOperationResult]]:
+    def transpose_if_matches(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
+        # Only process .weight keys
+        if not key.endswith(".weight"):
+            return [KeyValueOperationResult(key, value)]
+
+        # Only transpose 2D FP8 tensors (Linear weights)
+        if value.dim() != 2 or value.dtype != torch.float8_e4m3fn:
+            return [KeyValueOperationResult(key, value)]
+
+        # Check if the layer is excluded
+        layer_name = key.rsplit(".weight", 1)[0]
+        if _should_skip_layer(layer_name, excluded_layer_substrings):
+            return [KeyValueOperationResult(key, value)]
+
+        # Transpose to cuBLAS layout (in, out)
+        transposed_weight = value.t()
+
+        return [KeyValueOperationResult(key, transposed_weight)]
+
+    return transpose_if_matches
+
+
+FP8_TRANSPOSE_SD_OPS = SDOps("fp8_transpose_weights").with_kv_operation(
+    _create_transpose_kv_operation(EXCLUDED_LAYER_SUBSTRINGS),
+    key_prefix="transformer_blocks.",
+    key_suffix=".weight",
+)
+
+
+FP8_PREPARE_MODULE_OPS = ModuleOps(
+    name="fp8_prepare_for_loading",
+    matcher=lambda model: isinstance(model, LTXModel),
+    mutator=lambda model: _apply_fp8_prepare_to_model(model, EXCLUDED_LAYER_SUBSTRINGS),
+)

packages/ltx-core/src/ltx_core/quantization/policy.py

@@ -0,0 +1,39 @@
+from dataclasses import dataclass
+
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.sd_ops import SDOps
+from ltx_core.quantization.fp8_cast import TRANSFORMER_LINEAR_DOWNCAST_MAP, UPCAST_DURING_INFERENCE
+from ltx_core.quantization.fp8_scaled_mm import FP8_PREPARE_MODULE_OPS, FP8_TRANSPOSE_SD_OPS
+
+
+@dataclass(frozen=True)
+class QuantizationPolicy:
+    """Configuration for model quantization during loading.
+    Attributes:
+        sd_ops: State dict operations for weight transformation.
+        module_ops: Post-load module transformations.
+    """
+
+    sd_ops: SDOps | None = None
+    module_ops: tuple[ModuleOps, ...] = ()
+
+    @classmethod
+    def fp8_cast(cls) -> "QuantizationPolicy":
+        """Create policy using FP8 casting with upcasting during inference."""
+        return cls(
+            sd_ops=TRANSFORMER_LINEAR_DOWNCAST_MAP,
+            module_ops=(UPCAST_DURING_INFERENCE,),
+        )
+
+    @classmethod
+    def fp8_scaled_mm(cls) -> "QuantizationPolicy":
+        """Create policy using FP8 scaled matrix multiplication."""
+        try:
+            import tensorrt_llm  # noqa: F401, PLC0415
+        except ImportError as e:
+            raise ImportError("tensorrt_llm is not installed, skipping FP8 scaled MM quantization") from e
+
+        return cls(
+            sd_ops=FP8_TRANSPOSE_SD_OPS,
+            module_ops=(FP8_PREPARE_MODULE_OPS,),
+        )

packages/ltx-core/src/ltx_core/text_encoders/gemma/embeddings_connector.py

@@ -19,6 +19,7 @@ class _BasicTransformerBlock1D(torch.nn.Module):
         heads: int,
         dim_head: int,
         rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+        apply_gated_attention: bool = False,
     ):
         super().__init__()
 
@@ -27,6 +28,7 @@ class _BasicTransformerBlock1D(torch.nn.Module):
             heads=heads,
             dim_head=dim_head,
             rope_type=rope_type,
+            apply_gated_attention=apply_gated_attention,
         )
 
         self.ff = FeedForward(
@@ -99,6 +101,7 @@ class Embeddings1DConnector(torch.nn.Module):
         num_learnable_registers: int | None = 128,
         rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
         double_precision_rope: bool = False,
+        apply_gated_attention: bool = False,
     ):
         super().__init__()
         self.num_attention_heads = num_attention_heads
@@ -117,6 +120,7 @@ class Embeddings1DConnector(torch.nn.Module):
                     heads=num_attention_heads,
                     dim_head=attention_head_dim,
                     rope_type=rope_type,
+                    apply_gated_attention=apply_gated_attention,
                 )
                 for _ in range(num_layers)
             ]
@@ -206,5 +210,6 @@ class Embeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConnector]
             positional_embedding_max_pos=pe_max_pos,
             rope_type=rope_type,
             double_precision_rope=double_precision_rope,
+            apply_gated_attention=config.get("connector_apply_gated_attention", False),
         )
         return connector

packages/ltx-pipelines/src/ltx_pipelines/distilled.py

@@ -11,6 +11,7 @@ from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
 from ltx_core.model.upsampler import upsample_video
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.model.video_vae import decode_video as vae_decode_video
+from ltx_core.quantization import QuantizationPolicy
 from ltx_core.text_encoders.gemma import encode_text
 from ltx_core.types import LatentState, VideoPixelShape
 from ltx_pipelines.utils import ModelLedger
@@ -50,7 +51,7 @@ class DistilledPipeline:
         spatial_upsampler_path: str,
         loras: list[LoraPathStrengthAndSDOps],
         device: torch.device = device,
-        fp8transformer: bool = False,
+        quantization: QuantizationPolicy | None = None,
     ):
         self.device = device
         self.dtype = torch.bfloat16
@@ -62,7 +63,7 @@ class DistilledPipeline:
             spatial_upsampler_path=spatial_upsampler_path,
             gemma_root_path=gemma_root,
             loras=loras,
-            fp8transformer=fp8transformer,
+            quantization=quantization,
         )
 
         self.pipeline_components = PipelineComponents(
@@ -204,7 +205,7 @@ def main() -> None:
         spatial_upsampler_path=args.spatial_upsampler_path,
         gemma_root=args.gemma_root,
         loras=args.lora,
-        fp8transformer=args.enable_fp8,
+        quantization=args.quantization,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)

packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py

@@ -13,6 +13,7 @@ from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
 from ltx_core.model.upsampler import upsample_video
 from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
 from ltx_core.model.video_vae import decode_video as vae_decode_video
+from ltx_core.quantization import QuantizationPolicy
 from ltx_core.text_encoders.gemma import encode_text
 from ltx_core.types import LatentState, VideoPixelShape
 from ltx_pipelines.utils import ModelLedger
@@ -55,7 +56,7 @@ class ICLoraPipeline:
         gemma_root: str,
         loras: list[LoraPathStrengthAndSDOps],
         device: torch.device = device,
-        fp8transformer: bool = False,
+        quantization: QuantizationPolicy | None = None,
     ):
         self.dtype = torch.bfloat16
         self.stage_1_model_ledger = ModelLedger(
@@ -65,7 +66,7 @@ class ICLoraPipeline:
             spatial_upsampler_path=spatial_upsampler_path,
             gemma_root_path=gemma_root,
             loras=loras,
-            fp8transformer=fp8transformer,
+            quantization=quantization,
         )
         self.stage_2_model_ledger = ModelLedger(
             dtype=self.dtype,
@@ -74,7 +75,7 @@ class ICLoraPipeline:
             spatial_upsampler_path=spatial_upsampler_path,
             gemma_root_path=gemma_root,
             loras=[],
-            fp8transformer=fp8transformer,
+            quantization=quantization,
         )
         self.pipeline_components = PipelineComponents(
             dtype=self.dtype,
@@ -314,7 +315,7 @@ def main() -> None:
         spatial_upsampler_path=args.spatial_upsampler_path,
         gemma_root=args.gemma_root,
         loras=args.lora,
-        fp8transformer=args.enable_fp8,
+        quantization=args.quantization,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)

packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py

@@ -13,6 +13,7 @@ from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
 from ltx_core.model.upsampler import upsample_video
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.model.video_vae import decode_video as vae_decode_video
+from ltx_core.quantization import QuantizationPolicy
 from ltx_core.text_encoders.gemma import encode_text
 from ltx_core.types import LatentState, VideoPixelShape
 from ltx_pipelines.utils import ModelLedger
@@ -54,7 +55,7 @@ class KeyframeInterpolationPipeline:
         gemma_root: str,
         loras: list[LoraPathStrengthAndSDOps],
         device: torch.device = device,
-        fp8transformer: bool = False,
+        quantization: QuantizationPolicy | None = None,
     ):
         self.device = device
         self.dtype = torch.bfloat16
@@ -65,7 +66,7 @@ class KeyframeInterpolationPipeline:
             spatial_upsampler_path=spatial_upsampler_path,
             gemma_root_path=gemma_root,
             loras=loras,
-            fp8transformer=fp8transformer,
+            quantization=quantization,
         )
         self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
             loras=distilled_lora,
@@ -248,7 +249,7 @@ def main() -> None:
         spatial_upsampler_path=args.spatial_upsampler_path,
         gemma_root=args.gemma_root,
         loras=args.lora,
-        fp8transformer=args.enable_fp8,
+        quantization=args.quantization,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)

packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py

@@ -11,6 +11,7 @@ from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
 from ltx_core.model.video_vae import decode_video as vae_decode_video
+from ltx_core.quantization import QuantizationPolicy
 from ltx_core.text_encoders.gemma import encode_text
 from ltx_core.types import LatentState, VideoPixelShape
 from ltx_pipelines.utils import ModelLedger
@@ -46,7 +47,7 @@ class TI2VidOneStagePipeline:
         gemma_root: str,
         loras: list[LoraPathStrengthAndSDOps],
         device: torch.device = device,
-        fp8transformer: bool = False,
+        quantization: QuantizationPolicy | None = None,
     ):
         self.dtype = torch.bfloat16
         self.device = device
@@ -56,7 +57,7 @@ class TI2VidOneStagePipeline:
             checkpoint_path=checkpoint_path,
             gemma_root_path=gemma_root,
             loras=loras,
-            fp8transformer=fp8transformer,
+            quantization=quantization,
         )
         self.pipeline_components = PipelineComponents(
             dtype=self.dtype,
@@ -169,7 +170,7 @@ def main() -> None:
         checkpoint_path=args.checkpoint_path,
         gemma_root=args.gemma_root,
         loras=args.lora,
-        fp8transformer=args.enable_fp8,
+        quantization=args.quantization,
     )
     video, audio = pipeline(
         prompt=args.prompt,

packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py

@@ -13,6 +13,7 @@ from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
 from ltx_core.model.upsampler import upsample_video
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.model.video_vae import decode_video as vae_decode_video
+from ltx_core.quantization import QuantizationPolicy
 from ltx_core.text_encoders.gemma import encode_text
 from ltx_core.types import LatentState, VideoPixelShape
 from ltx_pipelines.utils import ModelLedger
@@ -54,7 +55,7 @@ class TI2VidTwoStagesPipeline:
         gemma_root: str,
         loras: list[LoraPathStrengthAndSDOps],
         device: str = device,
-        fp8transformer: bool = False,
+        quantization: QuantizationPolicy | None = None,
     ):
         self.device = device
         self.dtype = torch.bfloat16
@@ -65,7 +66,7 @@ class TI2VidTwoStagesPipeline:
             gemma_root_path=gemma_root,
             spatial_upsampler_path=spatial_upsampler_path,
             loras=loras,
-            fp8transformer=fp8transformer,
+            quantization=quantization,
         )
 
         self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
@@ -251,7 +252,7 @@ def main() -> None:
         spatial_upsampler_path=args.spatial_upsampler_path,
         gemma_root=args.gemma_root,
         loras=args.lora,
-        fp8transformer=args.enable_fp8,
+        quantization=args.quantization,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)

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

@@ -2,6 +2,7 @@ import argparse
 from pathlib import Path
 
 from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
+from ltx_core.quantization import QuantizationPolicy
 from ltx_pipelines.utils.constants import (
     DEFAULT_1_STAGE_HEIGHT,
     DEFAULT_1_STAGE_WIDTH,
@@ -78,6 +79,40 @@ def resolve_path(path: str) -> str:
     return str(Path(path).expanduser().resolve().as_posix())
 
 
+QUANTIZATION_POLICIES = ("fp8-cast", "fp8-scaled-mm")
+
+
+class QuantizationAction(argparse.Action):
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,  # noqa: ARG002
+        namespace: argparse.Namespace,
+        values: list[str],
+        option_string: str | None = None,
+    ) -> None:
+        if len(values) > 2:
+            msg = (
+                f"{option_string} accepts at most 2 arguments (POLICY and optional AMAX_PATH), got {len(values)} values"
+            )
+            raise argparse.ArgumentError(self, msg)
+
+        policy_name = values[0]
+        if policy_name not in QUANTIZATION_POLICIES:
+            msg = f"Unknown quantization policy '{policy_name}'. Choose from: {', '.join(QUANTIZATION_POLICIES)}"
+            raise argparse.ArgumentError(self, msg)
+
+        if policy_name == "fp8-cast":
+            if len(values) > 1:
+                msg = f"{option_string} fp8-cast does not accept additional arguments"
+                raise argparse.ArgumentError(self, msg)
+            policy = QuantizationPolicy.fp8_cast()
+        elif policy_name == "fp8-scaled-mm":
+            amax_path = resolve_path(values[1]) if len(values) > 1 else None
+            policy = QuantizationPolicy.fp8_scaled_mm(amax_path)
+
+        setattr(namespace, self.dest, policy)
+
+
 def basic_arg_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -174,13 +209,22 @@ def basic_arg_parser() -> argparse.ArgumentParser:
             "Example: --lora path/to/lora1.safetensors 0.8 --lora path/to/lora2.safetensors"
         ),
     )
-    parser.add_argument(
-        "--enable-fp8",
-        action="store_true",
-        help="Enable FP8 mode to reduce memory footprint by keeping model in lower precision. "
-        "Note that calculations are still performed in bfloat16 precision.",
-    )
+
     parser.add_argument("--enhance-prompt", action="store_true")
+    parser.add_argument(
+        "--quantization",
+        dest="quantization",
+        action=QuantizationAction,
+        nargs="+",
+        metavar=("POLICY", "AMAX_PATH"),
+        default=None,
+        help=(
+            f"Quantization policy: {', '.join(QUANTIZATION_POLICIES)}. "
+            "fp8-cast uses FP8 casting with upcasting during inference. "
+            "fp8-scaled-mm uses FP8 scaled matrix multiplication (optionally provide amax calibration file path). "
+            "Example: --quantization fp8-cast or --quantization fp8-scaled-mm /path/to/amax.json"
+        ),
+    )
     return parser
 
 

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

@@ -1,3 +1,4 @@
+import logging
 import math
 from collections.abc import Generator, Iterator
 from fractions import Fraction
@@ -13,6 +14,8 @@ from tqdm import tqdm
 
 from ltx_pipelines.utils.constants import DEFAULT_IMAGE_CRF
 
+logger = logging.getLogger(__name__)
+
 
 def resize_aspect_ratio_preserving(image: torch.Tensor, long_side: int) -> torch.Tensor:
     """
@@ -227,6 +230,7 @@ def encode_video(
         _write_audio(container, audio_stream, audio, audio_sample_rate)
 
     container.close()
+    logger.info(f"Video saved to {output_path}")
 
 
 def decode_audio_from_file(path: str, device: torch.device) -> torch.Tensor | None:

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

@@ -2,6 +2,7 @@ from dataclasses import replace
 
 import torch
 
+from ltx_core.loader import SDOps
 from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import DummyRegistry, Registry
 from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
@@ -15,8 +16,6 @@ from ltx_core.model.audio_vae import (
 )
 from ltx_core.model.transformer import (
     LTXV_MODEL_COMFY_RENAMING_MAP,
-    LTXV_MODEL_COMFY_RENAMING_WITH_TRANSFORMER_LINEAR_DOWNCAST_MAP,
-    UPCAST_DURING_INFERENCE,
     LTXModelConfigurator,
     X0Model,
 )
@@ -29,6 +28,7 @@ from ltx_core.model.video_vae import (
     VideoEncoder,
     VideoEncoderConfigurator,
 )
+from ltx_core.quantization import QuantizationPolicy
 from ltx_core.text_encoders.gemma import (
     AV_GEMMA_TEXT_ENCODER_KEY_OPS,
     AVGemmaTextEncoderModel,
@@ -78,8 +78,9 @@ class ModelLedger:
     registry:
         Optional :class:`Registry` instance for weight caching across builders.
         Defaults to :class:`DummyRegistry` which performs no cross-builder caching.
-    fp8transformer:
-        If ``True``, builds the transformer with FP8 quantization and upcasting during inference.
+    quantization:
+        Optional :class:`QuantizationPolicy` controlling how transformer weights
+        are stored and how matmul is executed. Defaults to None, which means no quantization.
     ### Creating Variants
     Use :meth:`with_loras` to create a new ``ModelLedger`` instance that includes
     additional LoRA configurations while sharing the same registry for weight caching.
@@ -94,7 +95,7 @@ class ModelLedger:
         spatial_upsampler_path: str | None = None,
         loras: LoraPathStrengthAndSDOps | None = None,
         registry: Registry | None = None,
-        fp8transformer: bool = False,
+        quantization: QuantizationPolicy | None = None,
     ):
         self.dtype = dtype
         self.device = device
@@ -103,7 +104,7 @@ class ModelLedger:
         self.spatial_upsampler_path = spatial_upsampler_path
         self.loras = loras or ()
         self.registry = registry or DummyRegistry()
-        self.fp8transformer = fp8transformer
+        self.quantization = quantization
         self.build_model_builders()
 
     def build_model_builders(self) -> None:
@@ -179,7 +180,7 @@ class ModelLedger:
             spatial_upsampler_path=self.spatial_upsampler_path,
             loras=(*self.loras, *loras),
             registry=self.registry,
-            fp8transformer=self.fp8transformer,
+            quantization=self.quantization,
         )
 
     def transformer(self) -> X0Model:
@@ -187,19 +188,26 @@ class ModelLedger:
             raise ValueError(
                 "Transformer not initialized. Please provide a checkpoint path to the ModelLedger constructor."
             )
-        if self.fp8transformer:
-            fp8_builder = replace(
-                self.transformer_builder,
-                module_ops=(UPCAST_DURING_INFERENCE,),
-                model_sd_ops=LTXV_MODEL_COMFY_RENAMING_WITH_TRANSFORMER_LINEAR_DOWNCAST_MAP,
-            )
-            return X0Model(fp8_builder.build(device=self._target_device())).to(self.device).eval()
-        else:
+
+        if self.quantization is None:
             return (
                 X0Model(self.transformer_builder.build(device=self._target_device(), dtype=self.dtype))
                 .to(self.device)
                 .eval()
             )
+        else:
+            sd_ops = self.transformer_builder.model_sd_ops
+            if self.quantization.sd_ops is not None:
+                sd_ops = SDOps(
+                    name=f"sd_ops_chain_{sd_ops.name}+{self.quantization.sd_ops.name}",
+                    mapping=(*sd_ops.mapping, *self.quantization.sd_ops.mapping),
+                )
+            builder = replace(
+                self.transformer_builder,
+                module_ops=(*self.transformer_builder.module_ops, *self.quantization.module_ops),
+                model_sd_ops=sd_ops,
+            )
+            return X0Model(builder.build(device=self._target_device())).to(self.device).eval()
 
     def video_decoder(self) -> VideoDecoder:
         if not hasattr(self, "vae_decoder_builder"):

packages/ltx-trainer/pyproject.toml

@@ -26,12 +26,12 @@ dependencies = [
     "scenedetect>=0.6.5.2",
     "sentencepiece>=0.2.0",
     "torch>=2.6.0",
-    "torchaudio>=2.9.0",
+    "torchaudio>=2.7.0",
     "torchcodec>=0.8.1",
     "torchvision>=0.21.0",
     "typer>=0.15.1",
     "wandb>=0.19.11",
-    "setuptools>=80.9.0",
+    "setuptools>=79.0.0",
 ]
 
 [dependency-groups]