Clean up

Author: kevin314

fastvideo/configs/models/dits/cosmos.py

@@ -91,8 +91,6 @@ class CosmosArchConfig(DiTArchConfig):
     eps: float = 1e-6
     exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
 
-    # Attention backend selection - use torch to match diffusers behavior
-    attention_backend: str = "torch"  # Options: "distributed", "torch"
 
     def __post_init__(self):
         super().__post_init__()

fastvideo/models/dits/cosmos.py

@@ -168,52 +168,26 @@ def __init__(self,
                  num_heads: int,
                  qk_norm=True,
                  eps=1e-6,
-                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
-                 | None = None,
-                 prefix: str = "",
-                 attention_backend: str = "distributed") -> None:
+                 prefix: str = "") -> None:
         assert dim % num_heads == 0
         super().__init__()
         self.dim = dim
         self.num_heads = num_heads
         self.head_dim = dim // num_heads
         self.qk_norm = qk_norm
         self.eps = eps
-        self.attention_backend = attention_backend
 
         # layers - use standard PyTorch layers when using torch backend
-        if attention_backend == "torch":
-            self.to_q = nn.Linear(dim, dim, bias=False)
-            self.to_k = nn.Linear(dim, dim, bias=False)
-            self.to_v = nn.Linear(dim, dim, bias=False)
-            self.to_out = nn.Linear(dim, dim, bias=False)
-        else:
-            self.to_q = ReplicatedLinear(dim, dim, bias=False)
-            self.to_k = ReplicatedLinear(dim, dim, bias=False)
-            self.to_v = ReplicatedLinear(dim, dim, bias=False)
-            self.to_out = ReplicatedLinear(dim, dim, bias=False)
+        self.to_q = nn.Linear(dim, dim, bias=False)
+        self.to_k = nn.Linear(dim, dim, bias=False)
+        self.to_v = nn.Linear(dim, dim, bias=False)
+        self.to_out = nn.Linear(dim, dim, bias=False)
 
         self.norm_q = RMSNorm(self.head_dim,
                               eps=eps) if qk_norm else nn.Identity()
         self.norm_k = RMSNorm(self.head_dim,
                               eps=eps) if qk_norm else nn.Identity()
 
-        # Attention mechanism - select backend
-        if attention_backend == "torch":
-            self.use_torch_attention = True
-        elif attention_backend == "distributed":
-            self.attn = DistributedAttention(
-                num_heads=num_heads,
-                head_size=self.head_dim,
-                dropout_rate=0,
-                softmax_scale=None,
-                causal=False,
-                supported_attention_backends=supported_attention_backends,
-                prefix=prefix)
-            self.use_torch_attention = False
-        else:
-            raise ValueError(f"Unsupported attention backend: {attention_backend}")
-
     def forward(self,
                 hidden_states: torch.Tensor,
                 encoder_hidden_states: torch.Tensor | None = None,
@@ -224,14 +198,9 @@ def forward(self,
             encoder_hidden_states = hidden_states
 
         # Get QKV
-        if self.attention_backend == "torch":
-            query = self.to_q(hidden_states)
-            key = self.to_k(encoder_hidden_states)
-            value = self.to_v(encoder_hidden_states)
-        else:
-            query, _ = self.to_q(hidden_states)
-            key, _ = self.to_k(encoder_hidden_states)
-            value, _ = self.to_v(encoder_hidden_states)
+        query = self.to_q(hidden_states)
+        key = self.to_k(encoder_hidden_states)
+        value = self.to_v(encoder_hidden_states)
 
         # Reshape for multi-head attention
         query = query.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
@@ -256,21 +225,15 @@ def forward(self,
                                    use_real_unbind_dim=-2)
 
         # Attention computation
-        if self.use_torch_attention:
-            # Use standard PyTorch scaled dot product attention
-            attn_output = torch.nn.functional.scaled_dot_product_attention(
-                query, key, value, attn_mask=attention_mask, dropout_p=0.0
-            )
-        else:
-            attn_output, _ = self.attn(query, key, value)
-            
+        # Use standard PyTorch scaled dot product attention
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0
+        )
         attn_output = attn_output.transpose(1, 2).flatten(2, 3).type_as(query)
 
         # Output projection
-        if self.attention_backend == "torch":
-            attn_output = self.to_out(attn_output)
-        else:
-            attn_output, _ = self.to_out(attn_output)
+        attn_output = self.to_out(attn_output)
+
         return attn_output
 
 
@@ -282,10 +245,7 @@ def __init__(self,
                  num_heads: int,
                  qk_norm=True,
                  eps=1e-6,
-                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
-                 | None = None,
-                 prefix: str = "",
-                 attention_backend: str = "distributed") -> None:
+                 prefix: str = "") -> None:
         assert dim % num_heads == 0
         super().__init__()
         self.dim = dim
@@ -294,65 +254,33 @@ def __init__(self,
         self.head_dim = dim // num_heads
         self.qk_norm = qk_norm
         self.eps = eps
-        self.attention_backend = attention_backend
 
         # layers - use standard PyTorch layers when using torch backend
-        if attention_backend == "torch":
-            self.to_q = nn.Linear(dim, dim, bias=False)
-            self.to_k = nn.Linear(cross_attention_dim, dim, bias=False)
-            self.to_v = nn.Linear(cross_attention_dim, dim, bias=False)
-            self.to_out = nn.Linear(dim, dim, bias=False)
-        else:
-            self.to_q = ReplicatedLinear(dim, dim, bias=False)
-            self.to_k = ReplicatedLinear(cross_attention_dim, dim, bias=False)
-            self.to_v = ReplicatedLinear(cross_attention_dim, dim, bias=False)
-            self.to_out = ReplicatedLinear(dim, dim, bias=False)
+        self.to_q = nn.Linear(dim, dim, bias=False)
+        self.to_k = nn.Linear(cross_attention_dim, dim, bias=False)
+        self.to_v = nn.Linear(cross_attention_dim, dim, bias=False)
+        self.to_out = nn.Linear(dim, dim, bias=False)
 
         self.norm_q = RMSNorm(self.head_dim,
                               eps=eps) if qk_norm else nn.Identity()
         self.norm_k = RMSNorm(self.head_dim,
                               eps=eps) if qk_norm else nn.Identity()
 
-        # Attention mechanism - select backend
-        if attention_backend == "torch":
-            self.use_torch_attention = True
-        elif attention_backend == "distributed":
-            self.attn = LocalAttention(
-                num_heads=num_heads,
-                head_size=self.head_dim,
-                dropout_rate=0,
-                softmax_scale=None,
-                causal=False,
-                supported_attention_backends=supported_attention_backends)
-            self.use_torch_attention = False
-        else:
-            raise ValueError(f"Unsupported attention backend: {attention_backend}")
-
     def forward(self,
                 hidden_states: torch.Tensor,
                 encoder_hidden_states: torch.Tensor,
                 attention_mask: torch.Tensor | None = None) -> torch.Tensor:
 
         # Get QKV
-        if self.attention_backend == "torch":
-            query = self.to_q(hidden_states)
-            key = self.to_k(encoder_hidden_states)
-            value = self.to_v(encoder_hidden_states)
-        else:
-            query, _ = self.to_q(hidden_states)
-            key, _ = self.to_k(encoder_hidden_states)
-            value, _ = self.to_v(encoder_hidden_states)
+        query = self.to_q(hidden_states)
+        key = self.to_k(encoder_hidden_states)
+        value = self.to_v(encoder_hidden_states)
 
         # Reshape for multi-head attention
-        if self.use_torch_attention:
             # Standard PyTorch attention expects [batch, num_heads, seq_len, head_dim]
-            query = query.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
-            key = key.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
-            value = value.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
-        else:
-            query = query.unflatten(2, (self.num_heads, -1))
-            key = key.unflatten(2, (self.num_heads, -1))
-            value = value.unflatten(2, (self.num_heads, -1))
+        query = query.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
+        key = key.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
+        value = value.unflatten(2, (self.num_heads, -1)).transpose(1, 2)
 
         # Apply normalization
         if self.norm_q is not None:
@@ -361,20 +289,14 @@ def forward(self,
             key = self.norm_k.forward_native(key)
 
         # Attention computation
-        if self.use_torch_attention:
-            attn_output = torch.nn.functional.scaled_dot_product_attention(
-                query, key, value, attn_mask=attention_mask, dropout_p=0.0
-            )
-            attn_output = attn_output.transpose(1, 2).flatten(2, 3).type_as(query)
-        else:
-            attn_output = self.attn(query, key, value)
-            attn_output = attn_output.flatten(2, 3).type_as(query)
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0
+        )
+        attn_output = attn_output.transpose(1, 2).flatten(2, 3).type_as(query)
 
         # Output projection
-        if self.attention_backend == "torch":
-            attn_output = self.to_out(attn_output)
-        else:
-            attn_output, _ = self.to_out(attn_output)
+        attn_output = self.to_out(attn_output)
+
         return attn_output
 
 
@@ -389,10 +311,7 @@ def __init__(
         adaln_lora_dim: int = 256,
         qk_norm: str = "rms_norm",
         out_bias: bool = False,
-        supported_attention_backends: tuple[AttentionBackendEnum, ...]
-        | None = None,
         prefix: str = "",
-        attention_backend: str = "distributed",
     ) -> None:
         super().__init__()
 
@@ -404,9 +323,7 @@ def __init__(
             dim=hidden_size,
             num_heads=num_attention_heads,
             qk_norm=(qk_norm == "rms_norm"),
-            supported_attention_backends=supported_attention_backends,
-            prefix=f"{prefix}.attn1",
-            attention_backend=attention_backend)
+            prefix=f"{prefix}.attn1")
 
         self.norm2 = CosmosAdaLayerNormZero(in_features=hidden_size,
                                             hidden_features=adaln_lora_dim)
@@ -415,9 +332,7 @@ def __init__(
             cross_attention_dim=cross_attention_dim,
             num_heads=num_attention_heads,
             qk_norm=(qk_norm == "rms_norm"),
-            supported_attention_backends=supported_attention_backends,
-            prefix=f"{prefix}.attn2",
-            attention_backend=attention_backend)
+            prefix=f"{prefix}.attn2")
 
         self.norm3 = CosmosAdaLayerNormZero(in_features=hidden_size,
                                             hidden_features=adaln_lora_dim)
@@ -598,7 +513,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 class CosmosTransformer3DModel(BaseDiT):
     _fsdp_shard_conditions = CosmosVideoConfig()._fsdp_shard_conditions
     _compile_conditions = CosmosVideoConfig()._compile_conditions
-    _supported_attention_backends = CosmosVideoConfig()._supported_attention_backends
+    # _supported_attention_backends = CosmosVideoConfig()._supported_attention_backends
     param_names_mapping = CosmosVideoConfig().param_names_mapping
     lora_param_names_mapping = CosmosVideoConfig().lora_param_names_mapping
 
@@ -651,9 +566,9 @@ def __init__(self, config: CosmosVideoConfig, hf_config: dict[str, Any]) -> None
                 adaln_lora_dim=config.adaln_lora_dim,
                 qk_norm=config.qk_norm,
                 out_bias=False,
-                supported_attention_backends=self._supported_attention_backends,
+                # supported_attention_backends=self._supported_attention_backends,
                 prefix=f"{config.prefix}.transformer_blocks.{i}",
-                attention_backend=config.arch_config.attention_backend,
+                #attention_backend=config.arch_config.attention_backend,
             ) for i in range(config.num_layers)
         ])