Invalid API call in Cosmos VAE Encoder

[akhilg-nv]

Describe the bug

At this line, the reshape appears to be called incorrectly. I believe it should either be hidden_states.reshape(*) or torch.reshape(hidden_states, *)

https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py#L171

It appears like this part of the model is not used in the text2world pipeline so it doesn't raise an issue, but trying to run this line in local testing results in an error.

Reproduction

>>> hidden_states = torch.reshape(batch_size, num_channels, num_frames // p, p, height // p, p, width // p, p)
Traceback (most recent call last):                                                                                                     
  File "<stdin>", line 1, in <module>
TypeError: reshape() takes 2 positional arguments but 8 were given

Logs

System Info

Python 3.9

Who can help?

@DN6 @sayakpaul

[a-r-r-o-w]

This should indeed be hidden_states.reshape instead of torch.reshape. Would you like to open a PR with the fix?

[tanuj-rai]

Hi @a-r-r-o-w, can I fix this?

[akhilg-nv]

Regardless of who submits the fix, can I ask why this wasn't caught earlier @a-r-r-o-w? Ideally there should be some test/pipeline to validate the rest of the code is error-free as well.

[a-r-r-o-w]

@tanuj-rai Yes, please feel free to take it up :)

@akhilg-nv The branch of code is not used in any model (the VAEs we tested when converting to diffusers; Cosmos released many other VAE models, but it's not supported by us yet). Our tests are only minimal configs to initialize small models somewhat representative of the actual supported models. The branch of code should ideally be removed but since it's been in a release, we'll have to keep maintaining it, or deprecate and remove later.