From a5d53293c0d86a2e6078facc5226b4e20b3a492f Mon Sep 17 00:00:00 2001
From: Noah Peterson <NoahBPeterson1997@gmail.com>
Date: Fri, 1 May 2026 14:16:21 -0500
Subject: [PATCH] generate.py: add --dit-dtype {bfloat16,float16,float32} flag
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Lets users opt into fp16 DiT compute on Apple Silicon for ~25-32% sampling
wall-clock savings, matching the proven win from upstream-ish ports of
TRELLIS.2 to MLX. Apple's Metal SDPA + matmul kernels run noticeably
faster on fp16 than bf16; visual quality stays essentially identical
(sub-pixel mesh deviation), but single-seed numerical parity with
upstream is sacrificed.

Implementation:

  - Use the model's own `convert_to(dtype)` method when available rather
    than a plain `.to(dtype)`. `convert_to` casts the transformer blocks'
    parameters AND updates `self.dtype`, which the forward pass uses to
    drive `manual_cast(x, self.dtype)` on intermediates. A bare
    `.to(dtype)` would cast the parameters but leave the manual-cast
    targets at bf16 — silently undoing the speedup.
  - Falls back to `.to(dtype)` + setting `m.dtype = dtype` for any model
    that doesn't expose `convert_to`.
  - Iterates over all five flow-model keys (SS + shape@512 + shape@1024 +
    tex@512 + tex@1024) so the flag is consistent across pipeline_type.
  - VAE decoders are intentionally NOT recast — they already ship as fp16
    and accurate intermediates matter more there.

Default remains bfloat16 (matches upstream training/inference dtype).

Inspect:

    Pre:  {'torch.float32': 17.5M, 'torch.bfloat16': 1.27B}
    Post: {'torch.float32': 17.5M, 'torch.float16':  1.27B}

That is, only the transformer torso is recast; input/output layers stay
at fp32. The 17.5M fp32 params are normalization / time-embed / projection
weights where we want maximum precision regardless.
---
 generate.py | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/generate.py b/generate.py
index b3f6718..776a39c 100644
--- a/generate.py
+++ b/generate.py
@@ -63,6 +63,15 @@ def main():
         "--steps", type=int, default=None,
         help="Override sampler steps for all three flow phases (default: pipeline JSON, usually 12)",
     )
+    parser.add_argument(
+        "--dit-dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"],
+        help=(
+            "Compute dtype for the three flow DiTs (default: bfloat16, matches upstream "
+            "training/inference). 'float16' is ~25-32%% faster on Apple Silicon SDPA + "
+            "matmul kernels with sub-pixel mesh deviation; visual quality is essentially "
+            "identical. Single-seed numerical parity with upstream is sacrificed."
+        ),
+    )
     args = parser.parse_args()
 
     if not os.path.exists(args.image):
@@ -85,6 +94,43 @@ def main():
     pipeline.to(torch.device("mps"))
     print("Device: MPS")
 
+    # Optionally cast the three DiT flow models to a smaller compute dtype.
+    # Upstream weights ship as bf16; on Apple Silicon, fp16 SDPA + matmul are
+    # ~1.3x faster than bf16, so fp16 gives roughly 25-32% wall-clock savings
+    # on the sampling phase with sub-pixel mesh deviation. The shape and tex
+    # VAE decoders ship as fp16 already and are NOT recast — they're tiny
+    # relative to the DiTs and accurate intermediates matter more there.
+    #
+    # We use the model's own `convert_to(dtype)` method when available rather
+    # than a plain `.to(dtype)`, because `convert_to` also updates `self.dtype`
+    # which the forward pass uses to drive `manual_cast(x, self.dtype)` on
+    # intermediates. A bare `.to(dtype)` would cast the parameters but leave
+    # the manual-cast targets at bf16 — silently undoing the speedup.
+    target_dtype = {
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+        "float32": torch.float32,
+    }[args.dit_dtype]
+    if target_dtype is not torch.bfloat16:
+        for k in (
+            "sparse_structure_flow_model",
+            "shape_slat_flow_model_512",
+            "shape_slat_flow_model_1024",
+            "tex_slat_flow_model_512",
+            "tex_slat_flow_model_1024",
+        ):
+            m = pipeline.models.get(k)
+            if m is None:
+                continue
+            convert_to = getattr(m, "convert_to", None)
+            if callable(convert_to):
+                convert_to(target_dtype)
+            else:
+                m.to(target_dtype)
+                if hasattr(m, "dtype"):
+                    m.dtype = target_dtype
+        print(f"DiT dtype: {args.dit_dtype}")
+
     # Load image
     img = PILImage.open(args.image)
     print(f"Input: {args.image} ({img.size[0]}x{img.size[1]})")