diff --git a/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py b/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py
index 9c7db02ff..91af3263f 100644
--- a/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py
+++ b/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py
@@ -1265,14 +1265,19 @@ class OmniGen2Pipeline(CFGParallelMixin, nn.Module, SupportsComponentDiscovery):
         if "ref_image_hidden_states" in set(inspect.signature(self.transformer.forward).parameters.keys()):
             optional_kwargs["ref_image_hidden_states"] = ref_image_hidden_states
 
-        model_pred = self.transformer(
-            latents,
-            timestep,
-            prompt_embeds,
-            freqs_cis,
-            prompt_attention_mask,
-            **optional_kwargs,
-        )
+        with torch.autocast(
+            device_type=self.device.type,
+            enabled=self.device.type != "cpu",
+            dtype=self.od_config.dtype,
+        ):
+            model_pred = self.transformer(
+                latents,
+                timestep,
+                prompt_embeds,
+                freqs_cis,
+                prompt_attention_mask,
+                **optional_kwargs,
+            )
         return model_pred
 
     def predict_noise(self, **kwargs):
diff --git a/vllm_omni/diffusion/offloader/module_collector.py b/vllm_omni/diffusion/offloader/module_collector.py
index 94cbb1fe3..320f9bf38 100644
--- a/vllm_omni/diffusion/offloader/module_collector.py
+++ b/vllm_omni/diffusion/offloader/module_collector.py
@@ -48,6 +48,7 @@ class ModuleDiscovery:
         "text_encoder_2",
         "text_encoder_3",
         "image_encoder",
+        "mllm",
     ]
     _FALLBACK_VAE_ATTRS = [
         "vae",