Automated PR - 2026-02-09

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