"unexpected keyword argument" Error when passing "return_dict=True" to HiDreamImageTransformer2DModel

[Nerogar]

Describe the bug

The HiDreamImageTransformer2DModel returns Transformer2DModelOutput(sample=output, mask=hidden_states_masks) if return_dict is set to True, but mask is not a valid keyword argument, which results in an error

Reproduction

Just pass return_dict=True to HiDreamImageTransformer2DModel.forward()

Logs

output truncated because it's not relevant to the error

...
  File "H:\stable-diffusion\one-trainer\venv\src\diffusers\src\diffusers\models\transformers\transformer_hidream_image.py", line 922, in forward
    return Transformer2DModelOutput(sample=output, mask=hidden_states_masks)
TypeError: Transformer2DModelOutput.__init__() got an unexpected keyword argument 'mask'

System Info

Diffusers commit bbd0c16, everything else shouldn't be relevant

Who can help?

No response

[yiyixuxu]

indeed, thanks, we will fix