revert fp32 upcast and support bs > 1

Author: Ting-Yun Chang

src/diffusers/models/transformers/transformer_cosmos.py

@@ -74,8 +74,8 @@ def __init__(self, embedding_dim: int, condition_dim: int) -> None:
         self.t_embedder = CosmosTimestepEmbedding(embedding_dim, condition_dim)
         self.norm = RMSNorm(embedding_dim, eps=1e-6, elementwise_affine=True)
 
-    def forward(self, hidden_states: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
-        timesteps_proj = self.time_proj(timestep.float())
+    def forward(self, hidden_states: torch.Tensor, timestep: torch.LongTensor) -> torch.Tensor:
+        timesteps_proj = self.time_proj(timestep).type_as(hidden_states)
         temb = self.t_embedder(timesteps_proj)
         embedded_timestep = self.norm(timesteps_proj)
         return temb, embedded_timestep
@@ -102,7 +102,6 @@ def forward(
             embedded_timestep = embedded_timestep + temb[..., : 2 * self.embedding_dim]
 
         shift, scale = embedded_timestep.chunk(2, dim=-1)
-
         hidden_states = self.norm(hidden_states)
 
         if embedded_timestep.ndim == 2:
@@ -132,16 +131,14 @@ def forward(
         embedded_timestep: torch.Tensor,
         temb: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        original_dtype = hidden_states.dtype
-        embedded_timestep = self.activation(embedded_timestep.float())
+        embedded_timestep = self.activation(embedded_timestep)
         embedded_timestep = self.linear_1(embedded_timestep)
         embedded_timestep = self.linear_2(embedded_timestep)
+
         if temb is not None:
-            embedded_timestep = embedded_timestep + temb.float()
+            embedded_timestep = embedded_timestep + temb
+
         shift, scale, gate = embedded_timestep.chunk(3, dim=-1)
-        shift = shift.to(original_dtype)
-        scale = scale.to(original_dtype)
-        gate = gate.to(original_dtype)
         hidden_states = self.norm(hidden_states)
 
         if embedded_timestep.ndim == 2:
@@ -184,11 +181,8 @@ def __call__(
         if image_rotary_emb is not None:
             from ..embeddings import apply_rotary_emb
 
-            original_dtype = query.dtype
-            query = apply_rotary_emb(query.to(torch.float32), image_rotary_emb, use_real=True, use_real_unbind_dim=-2)
-            key = apply_rotary_emb(key.to(torch.float32), image_rotary_emb, use_real=True, use_real_unbind_dim=-2)
-            query = query.to(original_dtype)
-            key = key.to(original_dtype)
+            query = apply_rotary_emb(query, image_rotary_emb, use_real=True, use_real_unbind_dim=-2)
+            key = apply_rotary_emb(key, image_rotary_emb, use_real=True, use_real_unbind_dim=-2)
 
         # 4. Prepare for GQA
         if torch.onnx.is_in_onnx_export():
@@ -254,11 +248,8 @@ def __call__(
         if image_rotary_emb is not None:
             from ..embeddings import apply_rotary_emb
 
-            original_dtype = query.dtype
-            query = apply_rotary_emb(query.to(torch.float32), image_rotary_emb, use_real=True, use_real_unbind_dim=-2)
-            key = apply_rotary_emb(key.to(torch.float32), image_rotary_emb, use_real=True, use_real_unbind_dim=-2)
-            query = query.to(original_dtype)
-            key = key.to(original_dtype)
+            query = apply_rotary_emb(query, image_rotary_emb, use_real=True, use_real_unbind_dim=-2)
+            key = apply_rotary_emb(key, image_rotary_emb, use_real=True, use_real_unbind_dim=-2)
 
         if torch.onnx.is_in_onnx_export():
             query_idx = torch.tensor(query.size(3), device=query.device)
@@ -608,7 +599,7 @@ class CosmosTransformer3DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin,
     _supports_gradient_checkpointing = True
     _skip_layerwise_casting_patterns = ["patch_embed", "final_layer", "norm"]
     _no_split_modules = ["CosmosTransformerBlock"]
-    _keep_in_fp32_modules = ["learnable_pos_embed", "time_embed", "norm1", "norm2", "norm3", "norm_out", "proj_out"]
+    _keep_in_fp32_modules = ["learnable_pos_embed"]
 
     @register_to_config
     def __init__(
@@ -806,7 +797,7 @@ def forward(
                 )
 
         # 8. Output norm & projection & unpatchify
-        hidden_states = self.norm_out(hidden_states.float(), embedded_timestep, temb)
+        hidden_states = self.norm_out(hidden_states, embedded_timestep, temb)
         hidden_states = self.proj_out(hidden_states)
         hidden_states = hidden_states.unflatten(2, (p_h, p_w, p_t, -1))
         hidden_states = hidden_states.unflatten(1, (post_patch_num_frames, post_patch_height, post_patch_width))

src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py

@@ -781,7 +781,9 @@ def __call__(
                 # NOTE: assumes sigma(t) \in [0, 1]
                 sigma_t = self.scheduler.sigmas[i].expand(batch_size).to(device=device, dtype=torch.float32)
                 if conditional_frame_timestep >= 0:
-                    in_timestep = cond_indicator * conditional_frame_timestep + (1 - cond_indicator) * sigma_t
+                    in_timestep = cond_indicator * conditional_frame_timestep + (1 - cond_indicator) * sigma_t.view(
+                        batch_size, 1, 1, 1, 1
+                    )
                 else:
                     in_timestep = sigma_t
                 in_latents = cond_mask * cond_latent + (1 - cond_mask) * latents