update output for Hidream transformer (#11366)

up

update output for Hidream transformer (#11366)
up
5a2e0f71 · YiYi Xu · GitHub · ef47726e · 5a2e0f71
Unverified Commit 5a2e0f71 authored Apr 18, 2025 by YiYi Xu Committed by GitHub Apr 18, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

src/diffusers/models/transformers/transformer_hidream_image.py ...iffusers/models/transformers/transformer_hidream_image.py +2 -2

No files found.
--- a/src/diffusers/models/transformers/transformer_hidream_image.py
+++ b/src/diffusers/models/transformers/transformer_hidream_image.py
@@ -918,5 +918,5 @@ class HiDreamImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
            unscale_lora_layers(self, lora_scale)

        if not return_dict:
-            return (output, hidden_states_masks)
-        return Transformer2DModelOutput(sample=output, mask=hidden_states_masks)
+            return (output,)
+        return Transformer2DModelOutput(sample=output)