From 6c2ba374ae4702436b7ad656e05e9b87025485cb Mon Sep 17 00:00:00 2001
From: CrossProduct <michael.chung@databites.ca>
Date: Thu, 13 Feb 2025 18:03:26 +0000
Subject: [PATCH] Todo fix the issue with the conditioning we do not re adjust
 the values into the model.

---
 .../LTXVideoDoubleCrossAttn3DModel.py         | 334 ++++++++++++++++++
 .../conditioning/condition_cross_attention.py |  43 +++
 ...conditioned_residual_adapter_bottleneck.py |  87 -----
 finetrainers/trainer.py                       |   2 +-
 4 files changed, 378 insertions(+), 88 deletions(-)
 create mode 100644 finetrainers/conditioning/LTXVideoDoubleCrossAttn3DModel.py
 create mode 100644 finetrainers/conditioning/condition_cross_attention.py
 delete mode 100644 finetrainers/conditioning/conditioned_residual_adapter_bottleneck.py

diff --git a/finetrainers/conditioning/LTXVideoDoubleCrossAttn3DModel.py b/finetrainers/conditioning/LTXVideoDoubleCrossAttn3DModel.py
new file mode 100644
index 0000000..c8a7575
--- /dev/null
+++ b/finetrainers/conditioning/LTXVideoDoubleCrossAttn3DModel.py
@@ -0,0 +1,334 @@
+
+import math
+from typing import Any, Dict, Optional, Tuple
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.utils import logging
+from diffusers import LTXVideoTransformer3DModel
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.utils import is_torch_version
+from diffusers.models.normalization import AdaLayerNormSingle, RMSNorm
+from diffusers.models import Attention
+from diffusers.models.transformers.transformer_ltx import LTXVideoAttentionProcessor2_0
+from diffusers.models.attention import FeedForward
+
+logger = logging.get_logger(__name__)
+
+@maybe_allow_in_graph
+class LTXVideoTransformerDoubleCrossAttn3DModel(LTXVideoTransformer3DModel):
+    r"""
+    A Transformer model for video-like data used in [LTX](https://huggingface.co/Lightricks/LTX-Video).
+
+    Args:
+        in_channels (`int`, defaults to `128`):
+            The number of channels in the input.
+        out_channels (`int`, defaults to `128`):
+            The number of channels in the output.
+        patch_size (`int`, defaults to `1`):
+            The size of the spatial patches to use in the patch embedding layer.
+        patch_size_t (`int`, defaults to `1`):
+            The size of the tmeporal patches to use in the patch embedding layer.
+        num_attention_heads (`int`, defaults to `32`):
+            The number of heads to use for multi-head attention.
+        attention_head_dim (`int`, defaults to `64`):
+            The number of channels in each head.
+        cross_attention_dim (`int`, defaults to `2048 `):
+            The number of channels for cross attention heads.
+        num_layers (`int`, defaults to `28`):
+            The number of layers of Transformer blocks to use.
+        activation_fn (`str`, defaults to `"gelu-approximate"`):
+            Activation function to use in feed-forward.
+        qk_norm (`str`, defaults to `"rms_norm_across_heads"`):
+            The normalization layer to use.
+    """
+
+    def __init__(self,
+        in_channels: int = 128,
+        out_channels: int = 128,
+        patch_size: int = 1,
+        patch_size_t: int = 1,
+        num_attention_heads: int = 32,
+        attention_head_dim: int = 64,
+        cross_attention_dim: int = 2048,
+        num_layers: int = 28,
+        activation_fn: str = "gelu-approximate",
+        qk_norm: str = "rms_norm_across_heads",
+        norm_elementwise_affine: bool = False,
+        norm_eps: float = 1e-6,
+        caption_channels: int = 4096,
+        attention_bias: bool = True,
+        attention_out_bias: bool = True):
+
+        super().__init__(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    patch_size=patch_size,
+                    patch_size_t=patch_size_t,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    cross_attention_dim=cross_attention_dim,
+                    num_layers=num_layers,
+                    activation_fn=activation_fn,
+                    qk_norm=qk_norm,
+                    norm_elementwise_affine=norm_elementwise_affine,
+                    norm_eps=norm_eps,
+                    caption_channels=caption_channels,
+                    attention_bias=attention_bias,
+                    attention_out_bias=attention_out_bias
+                )
+        
+        inner_dim = num_attention_heads * attention_head_dim
+        
+        self.transformer_blocks = nn.ModuleList(
+            [
+                LTXVideoTransformerDoubleCrossAttentionBlock(
+                    dim=inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    cross_attention_dim=cross_attention_dim,
+                    qk_norm=qk_norm,
+                    activation_fn=activation_fn,
+                    attention_bias=attention_bias,
+                    attention_out_bias=attention_out_bias,
+                    eps=norm_eps,
+                    elementwise_affine=norm_elementwise_affine,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+
+        self.gradient_checkpointing = False
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            encoder_hidden_states: torch.Tensor,
+            timestep: torch.LongTensor,
+            encoder_attention_mask: torch.Tensor,
+            num_frames: int,
+            height: int,
+            width: int,
+            rope_interpolation_scale: Optional[Tuple[float, float, float]] = None,
+            attention_kwargs: Optional[Dict[str, Any]] = None,
+            return_dict: bool = True,
+        ) -> torch.Tensor:
+          
+            if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
+                logger.warning(
+                    "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
+                )
+
+            image_rotary_emb = self.rope(hidden_states, num_frames, height, width, rope_interpolation_scale)
+
+            # convert encoder_attention_mask to a bias the same way we do for attention_mask
+            if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
+                encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
+                encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+
+            batch_size = hidden_states.size(0)
+            hidden_states = self.proj_in(hidden_states)
+
+            temb, embedded_timestep = self.time_embed(
+                timestep.flatten(),
+                batch_size=batch_size,
+                hidden_dtype=hidden_states.dtype,
+            )
+
+            temb = temb.view(batch_size, -1, temb.size(-1))
+            embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.size(-1))
+
+            encoder_hidden_states = self.caption_projection(encoder_hidden_states)
+            encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1))
+
+            for block in self.transformer_blocks:
+                if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                    def create_custom_forward(module, return_dict=None):
+                        def custom_forward(*inputs):
+                            if return_dict is not None:
+                                return module(*inputs, return_dict=return_dict)
+                            else:
+                                return module(*inputs)
+
+                        return custom_forward
+
+                    ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(block),
+                        hidden_states,
+                        encoder_hidden_states,
+                        temb,
+                        image_rotary_emb,
+                        encoder_attention_mask,
+                        **ckpt_kwargs,
+                    )
+                else:
+                    hidden_states = block(
+                        hidden_states=hidden_states,
+                        encoder_hidden_states=encoder_hidden_states,
+                        temb=temb,
+                        image_rotary_emb=image_rotary_emb,
+                        encoder_attention_mask=encoder_attention_mask,
+                    )
+
+            scale_shift_values = self.scale_shift_table[None, None] + embedded_timestep[:, :, None]
+            shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
+
+            hidden_states = self.norm_out(hidden_states)
+            hidden_states = hidden_states * (1 + scale) + shift
+            output = self.proj_out(hidden_states)
+
+
+            if not return_dict:
+                return (output,)
+            return Transformer2DModelOutput(sample=output)
+    
+
+
+
+def apply_rotary_emb(x, freqs):
+    cos, sin = freqs
+    x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1)  # [B, S, H, D // 2]
+    x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
+    out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+    return out
+
+
+
+@maybe_allow_in_graph
+class LTXVideoTransformerDoubleCrossAttentionBlock(nn.Module):
+    r"""
+    Transformer block used in [LTX](https://huggingface.co/Lightricks/LTX-Video).
+
+    Args:
+        dim (`int`):
+            The number of channels in the input and output.
+        num_attention_heads (`int`):
+            The number of heads to use for multi-head attention.
+        attention_head_dim (`int`):
+            The number of channels in each head.
+        qk_norm (`str`, defaults to `"rms_norm"`):
+            The normalization layer to use.
+        activation_fn (`str`, defaults to `"gelu-approximate"`):
+            Activation function to use in feed-forward.
+        eps (`float`, defaults to `1e-6`):
+            Epsilon value for normalization layers.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        cross_attention_dim: int,
+        qk_norm: str = "rms_norm_across_heads",
+        activation_fn: str = "gelu-approximate",
+        attention_bias: bool = True,
+        attention_out_bias: bool = True,
+        eps: float = 1e-6,
+        elementwise_affine: bool = False,
+    ):
+        super().__init__()
+
+        self.norm1 = RMSNorm(dim, eps=eps, elementwise_affine=elementwise_affine)
+        self.attn1 = Attention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            kv_heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            bias=attention_bias,
+            cross_attention_dim=None,
+            out_bias=attention_out_bias,
+            qk_norm=qk_norm,
+            processor=LTXVideoAttentionProcessor2_0(),
+        )
+
+        self.norm2 = RMSNorm(dim, eps=eps, elementwise_affine=elementwise_affine)
+
+        self.attn2 = Attention(
+            query_dim=dim,
+            cross_attention_dim=cross_attention_dim,
+            heads=num_attention_heads,
+            kv_heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            bias=attention_bias,
+            out_bias=attention_out_bias,
+            qk_norm=qk_norm,
+            processor=LTXVideoAttentionProcessor2_0(),
+        )
+
+        # Pose condition + img ref block
+        self.norm3 = RMSNorm(dim, eps=eps, elementwise_affine=elementwise_affine)
+        self.attn3 = Attention(query_dim=dim,
+            cross_attention_dim=cross_attention_dim,
+            heads=num_attention_heads,
+            kv_heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            bias=attention_bias,
+            out_bias=attention_out_bias,
+            qk_norm=qk_norm,
+            processor=LTXVideoAttentionProcessor2_0(),
+            )
+
+        self.ff = FeedForward(dim, activation_fn=activation_fn)
+
+        self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        pose_img_ref_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        batch_size = hidden_states.size(0)
+        norm_hidden_states = self.norm1(hidden_states)
+
+        num_ada_params = self.scale_shift_table.shape[0]
+        ada_values = self.scale_shift_table[None, None] + temb.reshape(batch_size, temb.size(1), num_ada_params, -1)
+
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(dim=2)
+
+        norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
+
+        attn_hidden_states = self.attn1(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=None,
+            image_rotary_emb=image_rotary_emb,
+        )
+        hidden_states = hidden_states + attn_hidden_states * gate_msa
+
+        attn_hidden_states_text_cond = self.attn2(
+            hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            image_rotary_emb=None,
+            attention_mask=encoder_attention_mask,
+        )
+
+        attn_hidden_states_pose_img_ref_cond = self.attn3(
+            hidden_states,
+            encoder_hidden_states=pose_img_ref_hidden_states,
+            image_rotary_emb=None,
+            attention_mask=None,
+        )
+
+        hidden_states_text_cond = hidden_states + attn_hidden_states_text_cond 
+        norm_hidden_states_text_cond = self.norm2(hidden_states_text_cond) * (1 + scale_mlp) + shift_mlp
+
+        hidden_states_pose_img_ref_cond = hidden_states + attn_hidden_states_pose_img_ref_cond
+        norm_hidden_states_imf_ref_cond = hidden_states + self.norm3(hidden_states_pose_img_ref_cond)
+
+        ff_output = self.ff(norm_hidden_states_text_cond + norm_hidden_states_imf_ref_cond)
+        hidden_states = hidden_states + ff_output * gate_mlp
+
+        return hidden_states
\ No newline at end of file
diff --git a/finetrainers/conditioning/condition_cross_attention.py b/finetrainers/conditioning/condition_cross_attention.py
new file mode 100644
index 0000000..71b1939
--- /dev/null
+++ b/finetrainers/conditioning/condition_cross_attention.py
@@ -0,0 +1,43 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class CrossAttentionBlock(nn.Module):
+    """
+    Cross-Attention Block:
+      - Query comes from 'x' (the main hidden states)
+      - Key/Value come from 'context' (the conditioning or encoder output)
+    """
+    def __init__(self, embed_dim, num_heads, dropout=0.0):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        
+        self.attn = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout, batch_first=True)
+        self.layernorm = nn.LayerNorm(embed_dim)
+        self.dropout_layer = nn.Dropout(dropout)
+
+    def forward(self, x, context):
+        """
+        Args:
+            x: [batch_size, seq_len, embed_dim] - queries
+            context: [batch_size, context_len, embed_dim] - keys & values (conditioning)
+        Returns:
+            out: The cross-attended hidden states (same shape as x)
+        """
+        # Apply LayerNorm before attention (pre-norm variant)
+        x_norm = self.layernorm(x)
+        
+        # In cross attention:
+        #   Q = x_norm
+        #   K,V = context
+        attn_out, _ = self.attn(
+            query=x_norm, 
+            key=context, 
+            value=context
+        )
+
+        # Residual connection
+        out = x + self.dropout_layer(attn_out)
+        return out
\ No newline at end of file
diff --git a/finetrainers/conditioning/conditioned_residual_adapter_bottleneck.py b/finetrainers/conditioning/conditioned_residual_adapter_bottleneck.py
deleted file mode 100644
index 26d0938..0000000
--- a/finetrainers/conditioning/conditioned_residual_adapter_bottleneck.py
+++ /dev/null
@@ -1,87 +0,0 @@
-import torch.nn as nn
-import torch
-
-# This class takes the conditioned multichannel video tensor in patchified 
-# Downsamples it using a bottle neck layer
-# To adapt the input into the diffusion transformer.
-# You want to use the patchified tensor non conditioned tensor to pass in as a residual before passing this off to the transformer.
-class ConditionedResidualAdapterBottleneck(nn.Module):
-    def __init__(
-        self,
-        input_dim: int,
-        output_dim: int,  # New parameter for output dimension
-        bottleneck_dim: int = None,
-        adapter_dropout: float = 0.1,
-        adapter_init_scale: float = 1e-3,
-    ):
-        """
-        Args:
-            input_dim: Size of input dimension
-            output_dim: Size of output dimension
-            adapter_dropout: Dropout probability
-            adapter_init_scale: Initial scale for adapter layer parameters
-            use_residual: Whether to use residual connection (only if input_dim == output_dim)
-        """
-        super().__init__()
-            
-        # Down projection
-        self.down_proj = nn.Linear(input_dim, bottleneck_dim)
-        
-        # Activation and dropout
-        self.activation = nn.GELU()
-        self.dropout = nn.Dropout(adapter_dropout)
-        
-        # Up projection (now projects to output_dim instead of input_dim)
-        self.up_proj = nn.Linear(bottleneck_dim, output_dim)
-        
-        # Initialize weights
-        self.down_proj.weight.data.normal_(mean=0.0, std=adapter_init_scale)
-        self.down_proj.bias.data.zero_()
-        self.up_proj.weight.data.normal_(mean=0.0, std=adapter_init_scale)
-        self.up_proj.bias.data.zero_()
-
-    def forward(self,residual_x:torch.tensor, conditioned_x: torch.Tensor) -> torch.Tensor:    
-        
-        # Down projection
-        hidden_states = self.down_proj(conditioned_x)
-        
-        # Activation and dropout
-        hidden_states = self.activation(hidden_states)
-        hidden_states = self.dropout(hidden_states)
-        
-        # Up projection to new dimension
-        hidden_states = self.up_proj(hidden_states)
-        
-        if residual_x is None:
-            return hidden_states
-        else:
-            hidden_states = hidden_states + residual_x
-            
-        return hidden_states
-
-# Example usage:
-def example_usage():
-    # Create a sample input tensor
-    batch_size, seq_length, input_dim = 1, 4086, 256
-   
-    output_dim = 128  # Reduced output dimension
-    x = torch.randn(batch_size, seq_length, input_dim)
-    
-    res_x = torch.randn(1,4086,128)
-    # Initialize adapter
-    adapter = ConditionedResidualAdapterBottleneck(
-        input_dim=input_dim,
-        output_dim=output_dim,
-        bottleneck_dim=64,
-        adapter_dropout=0.1,
-        adapter_init_scale=1e-3
-    )
-    adapter.requires_grad_(True)
-
-    # Forward pass
-    output = adapter(residual_x=res_x,conditioned_x=x)
-    print(f"Input shape: {x.shape}")  
-    print(f"Output shape: {output.shape}") 
-
-if __name__ == "__main__":
-    example_usage()
\ No newline at end of file
diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py
index eaffea8..c14a555 100644
--- a/finetrainers/trainer.py
+++ b/finetrainers/trainer.py
@@ -813,7 +813,7 @@ class Trainer:
                             noisy_latents = noisy_latents.to(target_video_latents.dtype)
 
                             # add pose information to both channels
-                            pose_noisy_latents = noisy_latents + pose_video_latents["latents"]
+                            pose_noisy_latents = noisy_latents
                             pose_img_ref_latents = img_refs_latents["latents"] + pose_video_latents["latents"]
                             
                             # expand channel information # B x 2C  latent will be projected to adapter to scale it back to 128d using adapter.