[Bugfix] fix OmniGen2 offload and dtype mismatch (#2560)

Signed-off-by: Lancer <maruixiang6688@gmail.com>
Signed-off-by: Lancer <402430575@qq.com>

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):

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",