diff --git a/hy3dshape/hy3dshape/models/denoisers/hunyuandit_mlx.py b/hy3dshape/hy3dshape/models/denoisers/hunyuandit_mlx.py index 45658b6..73b0bae 100644 --- a/hy3dshape/hy3dshape/models/denoisers/hunyuandit_mlx.py +++ b/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: