---
title: cosmos
---

<div class="python-reference">
<a id="fastvideo.train.models.cosmos-classes"></a>
<a id="fastvideo.train.models.cosmos"></a>
<p>Cosmos model plugin package.</p>
<h2 id="fastvideo.train.models.cosmos.CosmosModel"><code>fastvideo.train.models.cosmos.CosmosModel</code></h2>
<pre class="doc-signature"><code>CosmosModel(*, init_from: str, training_config: TrainingConfig, trainable: bool = True, disable_custom_init_weights: bool = False, flow_shift: float = 1.0, enable_gradient_checkpointing_type: str | None = None, transformer_override_safetensor: str | None = None)</code></pre>
<p>Cosmos 2.5 per-role model.</p>
<p>Inherits most behaviour from WanModel (noise scheduler,
timestep sampling, attention metadata, backward).  Overrides
only the pieces that differ for Cosmos 2.5.</p>
<p>Cosmos 2.5 uses:
- Cosmos25Transformer3DModel (velocity prediction)
- EDM noise schedule: x_t = x_0 + sigma * eps
- No input/output preconditioning (raw latents)
- Timestep = raw sigma value
- Model output = velocity ≈ noise</p>

<h3 id="fastvideo.train.models.cosmos.CosmosModel.prepare_batch"><code>fastvideo.train.models.cosmos.CosmosModel.prepare_batch</code></h3>
<pre class="doc-signature"><code>prepare_batch(raw_batch: dict&#91;str, Any&#93;, *, generator: torch.Generator, latents_source: Literal&#91;&#x27;data&#x27;, &#x27;zeros&#x27;&#93; = &#x27;data&#x27;) -&gt; TrainingBatch</code></pre>
<p>Same flow as Wan, but uses Cosmos VAE
normalisation.</p>
<h3 id="fastvideo.train.models.cosmos.CosmosModel.predict_noise"><code>fastvideo.train.models.cosmos.CosmosModel.predict_noise</code></h3>
<pre class="doc-signature"><code>predict_noise(noisy_latents: torch.Tensor, timestep: torch.Tensor, batch: TrainingBatch, *, conditional: bool, cfg_uncond: dict&#91;str, Any&#93; | None = None, attn_kind: Literal&#91;&#x27;dense&#x27;, &#x27;vsa&#x27;&#93; = &#x27;dense&#x27;) -&gt; torch.Tensor</code></pre>

<h3 id="fastvideo.train.models.cosmos.CosmosModel.ensure_negative_conditioning"><code>fastvideo.train.models.cosmos.CosmosModel.ensure_negative_conditioning</code></h3>
<pre class="doc-signature"><code>ensure_negative_conditioning() -&gt; None</code></pre>
<p>Create negative (unconditional) prompt embeddings.</p>
<p>Cosmos 2.5 uses Reason1 (Qwen2.5-VL) which is expensive
to load.  This method only supports <code>training_cfg_rate=0</code>
(no classifier-free guidance dropout), in which case the
negative embedding is never used and a zero placeholder
sized to match the text embedding dimension is sufficient.
<code>training_cfg_rate&gt;0</code> would require real Reason1 negative
embeddings and is rejected here to avoid silently training
with zero-vector "unconditional" inputs.</p>
</div>

<h2>Submodules</h2>
<ul>
<li><a href="/FastVideo/api/fastvideo/train/models/cosmos/cosmos">cosmos</a></li>
</ul>
