Basic initial support for cosmos predict2 text to image 2B and 14B mo…

…dels. (comfyanonymous#8517)

comfy/ldm/cosmos/blocks.py

@@ -26,16 +26,6 @@
 from comfy.ldm.modules.attention import optimized_attention
 
 
-def apply_rotary_pos_emb(
-    t: torch.Tensor,
-    freqs: torch.Tensor,
-) -> torch.Tensor:
-    t_ = t.reshape(*t.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2).float()
-    t_out = freqs[..., 0] * t_[..., 0] + freqs[..., 1] * t_[..., 1]
-    t_out = t_out.movedim(-1, -2).reshape(*t.shape).type_as(t)
-    return t_out
-
-
 def get_normalization(name: str, channels: int, weight_args={}, operations=None):
     if name == "I":
         return nn.Identity()

comfy/ldm/cosmos/position_embedding.py

@@ -66,6 +66,7 @@ def __init__(
         h_extrapolation_ratio: float = 1.0,
         w_extrapolation_ratio: float = 1.0,
         t_extrapolation_ratio: float = 1.0,
+        enable_fps_modulation: bool = True,
         device=None,
         **kwargs,  # used for compatibility with other positional embeddings; unused in this class
     ):
@@ -75,6 +76,7 @@ def __init__(
         self.base_fps = base_fps
         self.max_h = len_h
         self.max_w = len_w
+        self.enable_fps_modulation = enable_fps_modulation
 
         dim = head_dim
         dim_h = dim // 6 * 2
@@ -143,7 +145,7 @@ def generate_embeddings(
         half_emb_w = torch.outer(self.seq[:W].to(device=device), w_spatial_freqs)
 
         # apply sequence scaling in temporal dimension
-        if fps is None:  # image case
+        if fps is None or self.enable_fps_modulation is False:  # image case
             half_emb_t = torch.outer(self.seq[:T].to(device=device), temporal_freqs)
         else:
             half_emb_t = torch.outer(self.seq[:T].to(device=device) / fps * self.base_fps, temporal_freqs)