Add NVIDIA Cosmos #36476
Add NVIDIA Cosmos #36476
Conversation
|
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. When it is ready for review, please click the |
…6, then upcasted back ?
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
Ready for review. Note that most parts of the model are similar to existing Emu3 but we cannot copy with modular. Modular can't handle cases with |
ArthurZucker
left a comment
ArthurZucker
left a comment
There was a problem hiding this comment.
Late review but stil! Code should be simpler with our updates!˜
| channel_multiplier: List[int] = [2, 4, 4], | ||
| num_res_blocks: int = 2, | ||
| hidden_size: int = 512, | ||
| patch_size: int = 4, | ||
| levels: List[int] = [8, 8, 8, 5, 5, 5], |
There was a problem hiding this comment.
list defaults are never good!
| """ | ||
|
|
||
|
|
||
| class CosmosVQVAEVectorQuantizer(nn.Module): |
There was a problem hiding this comment.
as we have other VQVAEs now in transformers make sur to standardize here, checking if we can't re-use / re-factor old blocks! 🤗
| self.conv1 = ( | ||
| CosmosCausalConv3d(in_channels, in_channels, kernel_size=(3, 1, 1), stride=1, time_stride=1, padding=0) | ||
| if temporal_up | ||
| else nn.Identity() | ||
| ) |
There was a problem hiding this comment.
we should use if else!
| hidden_states = F.pad( | ||
| hidden_states, | ||
| pad=( | ||
| max(0, wavelet_len - 2), | ||
| wavelet_len - 1, | ||
| wavelet_len - 2, | ||
| wavelet_len - 1, | ||
| wavelet_len - 2, | ||
| wavelet_len - 1, | ||
| ), | ||
| mode=mode, | ||
| ).to(dtype) | ||
| hidden_states_low = F.conv3d( | ||
| hidden_states, wavelets_low.unsqueeze(3).unsqueeze(4), groups=seq_len, stride=(2, 1, 1) | ||
| ) | ||
| hidden_states_high = F.conv3d( | ||
| hidden_states, wavelets_high.unsqueeze(3).unsqueeze(4), groups=seq_len, stride=(2, 1, 1) | ||
| ) | ||
|
|
||
| # Handles spatial axes. | ||
| hidden_low_low = F.conv3d( | ||
| hidden_states_low, wavelets_low.unsqueeze(2).unsqueeze(4), groups=seq_len, stride=(1, 2, 1) | ||
| ) | ||
| hidden_low_hight = F.conv3d( | ||
| hidden_states_low, wavelets_high.unsqueeze(2).unsqueeze(4), groups=seq_len, stride=(1, 2, 1) | ||
| ) | ||
| hidden_high_low = F.conv3d( | ||
| hidden_states_high, wavelets_low.unsqueeze(2).unsqueeze(4), groups=seq_len, stride=(1, 2, 1) | ||
| ) | ||
| hidden_high_high = F.conv3d( | ||
| hidden_states_high, wavelets_high.unsqueeze(2).unsqueeze(4), groups=seq_len, stride=(1, 2, 1) | ||
| ) | ||
|
|
||
| hidden_lll = F.conv3d(hidden_low_low, wavelets_low.unsqueeze(2).unsqueeze(3), groups=seq_len, stride=(1, 1, 2)) | ||
| hidden_llh = F.conv3d( | ||
| hidden_low_low, wavelets_high.unsqueeze(2).unsqueeze(3), groups=seq_len, stride=(1, 1, 2) | ||
| ) | ||
| hidden_lhl = F.conv3d( | ||
| hidden_low_hight, wavelets_low.unsqueeze(2).unsqueeze(3), groups=seq_len, stride=(1, 1, 2) | ||
| ) | ||
| hidden_lhh = F.conv3d( | ||
| hidden_low_hight, wavelets_high.unsqueeze(2).unsqueeze(3), groups=seq_len, stride=(1, 1, 2) | ||
| ) | ||
| hidden_hll = F.conv3d( | ||
| hidden_high_low, wavelets_low.unsqueeze(2).unsqueeze(3), groups=seq_len, stride=(1, 1, 2) | ||
| ) | ||
| hidden_hlh = F.conv3d( | ||
| hidden_high_low, wavelets_high.unsqueeze(2).unsqueeze(3), groups=seq_len, stride=(1, 1, 2) | ||
| ) | ||
| hidden_hhl = F.conv3d( | ||
| hidden_high_high, wavelets_low.unsqueeze(2).unsqueeze(3), groups=seq_len, stride=(1, 1, 2) | ||
| ) | ||
| hidden_hhh = F.conv3d( | ||
| hidden_high_high, wavelets_high.unsqueeze(2).unsqueeze(3), groups=seq_len, stride=(1, 1, 2) | ||
| ) |
There was a problem hiding this comment.
wow that's ugly 🤣 we can probably make this simpler !
| xlll, xllh, xlhl, xlhh, xhll, xhlh, xhhl, xhhh = torch.chunk(hidden_states, 8, dim=1) | ||
|
|
||
| # Height height transposed convolutions. | ||
| xll = F.conv_transpose3d(xlll, hl.unsqueeze(2).unsqueeze(3), groups=groups, stride=(1, 1, 2)) |
There was a problem hiding this comment.
[..., None, None] is already more lightweight!
| position_ids = torch.stack([position_temporal, position_height, position_width], dim=1).to(device) | ||
| return position_ids | ||
|
|
||
| @add_start_docstrings_to_model_forward(COSMOS_TEXT_INPUTS_DOCSTRING) |
There was a problem hiding this comment.
use check_model_inputs
| output_attentions: Optional[bool] = None, | ||
| output_hidden_states: Optional[bool] = None, | ||
| return_dict: Optional[bool] = None, |
| output_hidden_states: Optional[bool] = None, | ||
| return_dict: Optional[bool] = None, | ||
| cache_position: Optional[torch.LongTensor] = None, | ||
| **flash_attn_kwargs: Unpack[FlashAttentionKwargs], |
There was a problem hiding this comment.
use TransformersKwargs
| hidden_states = self.absolute_position_emb(hidden_states, position_ids) | ||
|
|
||
| if self.gradient_checkpointing and self.training: | ||
| layer_outputs = self._gradient_checkpointing_func( |
There was a problem hiding this comment.
use gradient checkpointing layer!
| ) | ||
| return output if return_dict else output.to_tuple() | ||
|
|
||
| def _update_causal_mask( |
There was a problem hiding this comment.
I think not needed anymore
|
Is this PR still going on? I am willing to help, if updates on this are still needed. @ArthurZucker @zucchini-nlp |
|
@1himan the model got a bit low in priority and there has been almost no community request to add it. So I put it aside. Do you want to take over the PR? |
4 participants
What does this PR do?
Fixes #35565
Conversion and generation with video only model is tested on 4B checkpoint. Next steps: