AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-27 17:22:53 +03:00
Code Activity

Update autoencoder_kl_cosmos.py (#12045)

Browse Source 
* Update autoencoder_kl_cosmos.py

* Apply style fixes

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Aryan <aryan@huggingface.co>
This commit is contained in:
Tanuj Rai
2025-08-02 20:24:01 +05:30
committed by GitHub
parent 6febc08bfc
commit 359b605f4b
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py
View File

@@ -168,7 +168,9 @@ class CosmosPatchEmbed3d(nn.Module):
batch_size, num_channels, num_frames, height, width = hidden_states.shape
p = self.patch_size
hidden_states = torch.reshape(batch_size, num_channels, num_frames // p, p, height // p, p, width // p, p)
hidden_states = hidden_states.reshape(
batch_size, num_channels, num_frames // p, p, height // p, p, width // p, p
)
hidden_states = hidden_states.permute(0, 1, 3, 5, 7, 2, 4, 6).flatten(1, 4).contiguous()
return hidden_states