diffusers/pr_14261/en/_app/immutable/nodes/58.DtyCmBRv.js

import"../chunks/DsnmJJEf.js";import{i as O,h as b,C as y,H as c,a as k,D as _,E as C,s as N}from"../chunks/BtE7mKSK.js";import{p as q,o as J,s as e,f as I,a as f,b as j,c as p,d as g,n as u,r as h}from"../chunks/jDjavuwI.js";const M='{"title":"Cosmos3OmniTransformer","local":"cosmos3omnitransformer","sections":[{"title":"Cosmos3OmniTransformer","local":"diffusers.Cosmos3OmniTransformer","sections":[],"depth":2}],"depth":1}';var z=g('<meta name="hf:doc:metadata"/>'),Z=g('<p></p> <!> <!> <p>A Mixture-of-Transformer (MoT) joint vision-language transformer introduced as part of NVIDIA’s Cosmos3 world foundation model family. The model runs two parallel computation pathways over a packed joint sequence:</p> <ul><li>a <strong>causal “understanding” pathway</strong> that self-attends over text tokens with causal masking, and</li> <li>a <strong>bi-directional “generation” pathway</strong> that cross-attends from generation tokens (vision + optional sound latents) over the full understanding-plus-generation key/value set.</li></ul> <p>The two pathways share the same hidden size and number of layers but maintain <strong>separate Q/K/V/O projections, MLPs, and RMSNorm parameters</strong>, which is what makes the architecture a Mixture-of-Transformer rather than a standard Mixture-of-Experts. Position information is supplied through a 3D multimodal RoPE (mRoPE) that interleaves temporal / height / width frequencies for video latents and reuses the temporal axis for text and audio.</p> <p>The model can be loaded as follows.</p> <!> <!> <div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8"><!> <div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8"><!> <p>Run a full denoising-step forward pass.</p></div></div> <!> <p></p>',1);function E(v,x){q(x,!1),J(()=>{new URLSearchParams(window.location.search).get("fw")}),O();var s=Z();b("8krje0",d=>{var l=z();N(l,"content",M),f(d,l)});var o=e(I(s),2);y(o,{containerStyle:"float: right; margin-left: 10px; display: inline-flex; position: relative; z-index: 10;"});var i=e(o,2);c(i,{title:"Cosmos3OmniTransformer",local:"cosmos3omnitransformer",headingTag:"h1"});var t=e(i,10);k(t,{code:"aW1wb3J0JTIwdG9yY2glMEFmcm9tJTIwZGlmZnVzZXJzJTIwaW1wb3J0JTIwQ29zbW9zM09tbmlUcmFuc2Zvcm1lciUwQSUwQXRyYW5zZm9ybWVyJTIwJTNEJTIwQ29zbW9zM09tbmlUcmFuc2Zvcm1lci5mcm9tX3ByZXRyYWluZWQoJTBBJTIwJTIwJTIwJTIwJTIybnZpZGlhJTJGQ29zbW9zMy1OYW5vJTIyJTJDJTIwc3ViZm9sZGVyJTNEJTIydHJhbnNmb3JtZXIlMjIlMkMlMjB0b3JjaF9kdHlwZSUzRHRvcmNoLmJmbG9hdDE2JTBBKQ==",highlighted:`<span class="hljs-keyword">import</span> torch
<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> Cosmos3OmniTransformer

transformer = Cosmos3OmniTransformer.from_pretrained(
    <span class="hljs-string">&quot;nvidia/Cosmos3-Nano&quot;</span>, subfolder=<span class="hljs-string">&quot;transformer&quot;</span>, torch_dtype=torch.bfloat16
)`,lang:"python",wrap:!1});var r=e(t,2);c(r,{title:"Cosmos3OmniTransformer",local:"diffusers.Cosmos3OmniTransformer",headingTag:"h2"});var n=e(r,2),a=p(n);_(a,{name:"class diffusers.Cosmos3OmniTransformer",anchor:"diffusers.Cosmos3OmniTransformer",source:"https://github.com/huggingface/diffusers/blob/vr_14261/src/diffusers/models/transformers/transformer_cosmos3.py#L373",parameters:[{name:"attention_bias",val:": bool = False"},{name:"attention_dropout",val:": float = 0.0"},{name:"dtype",val:": str = 'bfloat16'"},{name:"head_dim",val:": int = 128"},{name:"hidden_size",val:": int = 4096"},{name:"intermediate_size",val:": int = 12288"},{name:"base_fps",val:": int = 24"},{name:"enable_fps_modulation",val:": bool = True"},{name:"latent_channel",val:": int = 48"},{name:"unified_3d_mrope_reset_spatial_ids",val:": bool = True"},{name:"unified_3d_mrope_temporal_modality_margin",val:": int = 15000"},{name:"latent_patch_size",val:": int = 2"},{name:"num_attention_heads",val:": int = 32"},{name:"num_hidden_layers",val:": int = 36"},{name:"num_key_value_heads",val:": int = 8"},{name:"patch_latent_dim",val:": int = 192"},{name:"rms_norm_eps",val:": float = 1e-06"},{name:"rope_scaling",val:": dict | None = None"},{name:"rope_theta",val:": float = 5000000.0"},{name:"action_dim",val:": int | None = None"},{name:"action_gen",val:": bool = False"},{name:"num_embodiment_domains",val:": int = 32"},{name:"sound_dim",val:": int | None = None"},{name:"sound_gen",val:": bool = False"},{name:"sound_latent_fps",val:": float = 25.0"},{name:"timestep_scale",val:": float = 0.001"},{name:"vocab_size",val:": int = 151936"},{name:"hidden_act",val:": str = 'silu'"},{name:"qk_norm_for_text",val:": bool = True"},{name:"use_und_k_norm_for_gen",val:": bool = False"},{name:"rope_axes_dim",val:": tuple[int, int, int] | list[int] | None = None"}]});var m=e(a,2),T=p(m);_(T,{name:"forward",anchor:"diffusers.Cosmos3OmniTransformer.forward",source:"https://github.com/huggingface/diffusers/blob/vr_14261/src/diffusers/models/transformers/transformer_cosmos3.py#L654",parameters:[{name:"input_ids",val:": Tensor"},{name:"text_indexes",val:": Tensor"},{name:"position_ids",val:": Tensor"},{name:"und_len",val:": int"},{name:"sequence_length",val:": int"},{name:"vision_tokens",val:": list"},{name:"vision_token_shapes",val:": list"},{name:"vision_sequence_indexes",val:": Tensor"},{name:"vision_mse_loss_indexes",val:": Tensor"},{name:"vision_timesteps",val:": Tensor"},{name:"vision_noisy_frame_indexes",val:": list"},{name:"sound_tokens",val:": list[torch.Tensor] | None = None"},{name:"sound_token_shapes",val:": list[tuple[int, int, int]] | None = None"},{name:"sound_sequence_indexes",val:": typing.Optional[torch.Tensor] = None"},{name:"sound_mse_loss_indexes",val:": typing.Optional[torch.Tensor] = None"},{name:"sound_timesteps",val:": typing.Optional[torch.Tensor] = None"},{name:"sound_noisy_frame_indexes",val:": list[torch.Tensor] | None = None"},{name:"action_tokens",val:": list[torch.Tensor] | None = None"},{name:"action_token_shapes",val:": list[tuple[int, int, int]] | None = None"},{name:"action_sequence_indexes",val:": typing.Optional[torch.Tensor] = None"},{name:"action_mse_loss_indexes",val:": typing.Optional[torch.Tensor] = None"},{name:"action_timesteps",val:": typing.Optional[torch.Tensor] = None"},{name:"action_noisy_frame_indexes",val:": list[torch.Tensor] | None = None"},{name:"action_domain_ids",val:": list[torch.Tensor] | None = None"},{name:"return_dict",val:": bool = True"}],parametersDescription:[{anchor:"diffusers.Cosmos3OmniTransformer.forward.input_ids",description:"<strong>input_ids</strong> &#x2014; Text token IDs placed at <code>text_indexes</code> in the joint sequence.",name:"input_ids"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.text_indexes",description:"<strong>text_indexes</strong> &#x2014; Indices of text tokens in the joint sequence.",name:"text_indexes"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.position_ids",description:"<strong>position_ids</strong> &#x2014; <code>[3, sequence_length]</code> mRoPE position IDs for the full joint sequence.",name:"position_ids"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.und_len",description:"<strong>und_len</strong> &#x2014; Length of the causal text (understanding) prefix; generation tokens follow.",name:"und_len"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.sequence_length",description:"<strong>sequence_length</strong> &#x2014; Total length of the joint packed sequence.",name:"sequence_length"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.vision_tokens",description:"<strong>vision_tokens</strong> &#x2014; Per-item vision latent tensors before patchify.",name:"vision_tokens"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.vision_token_shapes",description:"<strong>vision_token_shapes</strong> &#x2014; Patch grid shapes <code>(T, H, W)</code> per vision item.",name:"vision_token_shapes"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.vision_sequence_indexes",description:"<strong>vision_sequence_indexes</strong> &#x2014; Indices of vision tokens in the joint sequence.",name:"vision_sequence_indexes"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.vision_mse_loss_indexes",description:"<strong>vision_mse_loss_indexes</strong> &#x2014; Indices used to read vision predictions after the backbone.",name:"vision_mse_loss_indexes"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.vision_timesteps",description:"<strong>vision_timesteps</strong> &#x2014; Per-patch diffusion timesteps for vision tokens.",name:"vision_timesteps"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.vision_noisy_frame_indexes",description:"<strong>vision_noisy_frame_indexes</strong> &#x2014; Noisy frame indices per vision item.",name:"vision_noisy_frame_indexes"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.sound_tokens",description:"<strong>sound_tokens</strong> &#x2014; Optional sound latent tensors before packing.",name:"sound_tokens"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.sound_token_shapes",description:"<strong>sound_token_shapes</strong> &#x2014; Optional patch grid shapes for sound items.",name:"sound_token_shapes"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.sound_sequence_indexes",description:"<strong>sound_sequence_indexes</strong> &#x2014; Optional indices of sound tokens in the joint sequence.",name:"sound_sequence_indexes"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.sound_mse_loss_indexes",description:"<strong>sound_mse_loss_indexes</strong> &#x2014; Optional indices used to read sound predictions.",name:"sound_mse_loss_indexes"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.sound_timesteps",description:"<strong>sound_timesteps</strong> &#x2014; Optional per-token diffusion timesteps for sound.",name:"sound_timesteps"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.sound_noisy_frame_indexes",description:"<strong>sound_noisy_frame_indexes</strong> &#x2014; Optional noisy frame indices per sound item.",name:"sound_noisy_frame_indexes"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.action_tokens",description:"<strong>action_tokens</strong> &#x2014; Optional action latent tensors before packing.",name:"action_tokens"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.action_token_shapes",description:"<strong>action_token_shapes</strong> &#x2014; Optional patch grid shapes <code>(T, H, W)</code> per action item.",name:"action_token_shapes"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.action_sequence_indexes",description:"<strong>action_sequence_indexes</strong> &#x2014; Optional indices of action tokens in the joint sequence.",name:"action_sequence_indexes"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.action_mse_loss_indexes",description:"<strong>action_mse_loss_indexes</strong> &#x2014; Optional indices used to read action predictions after the backbone.",name:"action_mse_loss_indexes"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.action_timesteps",description:"<strong>action_timesteps</strong> &#x2014; Optional per-token diffusion timesteps for action tokens.",name:"action_timesteps"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.action_noisy_frame_indexes",description:"<strong>action_noisy_frame_indexes</strong> &#x2014; Optional noisy frame indices per action item.",name:"action_noisy_frame_indexes"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.action_domain_ids",description:"<strong>action_domain_ids</strong> &#x2014; Optional per-item domain IDs selecting the action head weights.",name:"action_domain_ids"},{anchor:"diffusers.Cosmos3OmniTransformer.forward.return_dict",description:"<strong>return_dict</strong> &#x2014; Whether to return a [<em>Cosmos3OmniTransformerOutput</em>] instead of a tuple.",name:"return_dict"}],returnDescription:`<script context="module">export const metadata = 'undefined';<\/script>


<p>A [<em>Cosmos3OmniTransformerOutput</em>] or a tuple of per-modality prediction lists. Optional modalities return
<code>None</code> when their inputs are omitted.</p>
`}),u(2),h(m),h(n);var w=e(n,2);C(w,{source:"https://github.com/huggingface/diffusers/blob/main/docs/source/en/api/models/cosmos3_omni_transformer.md"}),u(2),f(v,s),j()}export{E as component};