mlx-tune: T2 whole-DiT mx.compile (same HY3D_MLX_COMPILE flag; parity 7.3e-06, 1.01x micro — GEMM-bound, kept as free insurance)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

hy3dshape/hy3dshape/models/denoisers/hunyuandit_mlx.py

@@ -16,6 +16,8 @@ import math
 import re
 from typing import Optional
 
+import os
+
 import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
@@ -513,6 +515,17 @@ class HunYuanDiTPlain(nn.Module):
         self.final_layer = FinalLayer(hidden_size, self.out_channels)
 
     def __call__(self, x: mx.array, t: mx.array, contexts: dict, **kwargs) -> mx.array:
+        # HY3D_MLX_COMPILE=1: whole-DiT mx.compile — legal now that
+        # StaticMoELayer removed the data-dependent .item() skips (T3).
+        if os.environ.get("HY3D_MLX_COMPILE", "0") == "1":
+            fn = getattr(self, "_compiled_forward", None)
+            if fn is None:
+                fn = mx.compile(self._forward)
+                self._compiled_forward = fn
+            return fn(x, t, contexts, **kwargs)
+        return self._forward(x, t, contexts, **kwargs)
+
+    def _forward(self, x: mx.array, t: mx.array, contexts: dict, **kwargs) -> mx.array:
         """Forward pass.
 
         Args: