import math

import torch
import torch.amp as amp
import torch.nn as nn
from util.model_util import hash_state_dict_keys
from einops import rearrange

try:
    import flash_attn_interface
    FLASH_ATTN_3_AVAILABLE = True
except ModuleNotFoundError:
    FLASH_ATTN_3_AVAILABLE = False

try:
    import flash_attn
    FLASH_ATTN_2_AVAILABLE = True
except ModuleNotFoundError:
    FLASH_ATTN_2_AVAILABLE = False

try:
    from sageattention import sageattn
    SAGE_ATTN_AVAILABLE = True
except ModuleNotFoundError:
    SAGE_ATTN_AVAILABLE = False

import warnings


__all__ = ['WanModel']


def flash_attention(
    q,
    k,
    v,
    q_lens=None,
    k_lens=None,
    dropout_p=0.,
    softmax_scale=None,
    q_scale=None,
    causal=False,
    window_size=(-1, -1),
    deterministic=False,
    dtype=torch.bfloat16,
    version=None,
):
    """
    q:              [B, Lq, Nq, C1].
    k:              [B, Lk, Nk, C1].
    v:              [B, Lk, Nk, C2]. Nq must be divisible by Nk.
    q_lens:         [B].
    k_lens:         [B].
    dropout_p:      float. Dropout probability.
    softmax_scale:  float. The scaling of QK^T before applying softmax.
    causal:         bool. Whether to apply causal attention mask.
    window_size:    (left right). If not (-1, -1), apply sliding window local attention.
    deterministic:  bool. If True, slightly slower and uses more memory.
    dtype:          torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
    """
    half_dtypes = (torch.float16, torch.bfloat16)
    assert dtype in half_dtypes
    assert q.device.type == 'cuda' and q.size(-1) <= 256

    # params
    b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype

    def half(x):
        return x if x.dtype in half_dtypes else x.to(dtype)

    # preprocess query
    if q_lens is None:
        q = half(q.flatten(0, 1))
        q_lens = torch.tensor(
            [lq] * b, dtype=torch.int32).to(
                device=q.device, non_blocking=True)
    else:
        q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))

    # preprocess key, value
    if k_lens is None:
        k = half(k.flatten(0, 1))
        v = half(v.flatten(0, 1))
        k_lens = torch.tensor(
            [lk] * b, dtype=torch.int32).to(
                device=k.device, non_blocking=True)
    else:
        k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
        v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))

    q = q.to(v.dtype)
    k = k.to(v.dtype)

    if q_scale is not None:
        q = q * q_scale

    if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
        warnings.warn(
            'Flash attention 3 is not available, use flash attention 2 instead.'
        )

    # apply attention
    if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
        # Note: dropout_p, window_size are not supported in FA3 now.
        x = flash_attn_interface.flash_attn_varlen_func(
            q=q,
            k=k,
            v=v,
            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
                0, dtype=torch.int32).to(q.device, non_blocking=True),
            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
                0, dtype=torch.int32).to(q.device, non_blocking=True),
            seqused_q=None,
            seqused_k=None,
            max_seqlen_q=lq,
            max_seqlen_k=lk,
            softmax_scale=softmax_scale,
            causal=causal,
            deterministic=deterministic)[0].unflatten(0, (b, lq))
    elif FLASH_ATTN_2_AVAILABLE:
        x = flash_attn.flash_attn_varlen_func(
            q=q,
            k=k,
            v=v,
            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
                0, dtype=torch.int32).to(q.device, non_blocking=True),
            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
                0, dtype=torch.int32).to(q.device, non_blocking=True),
            max_seqlen_q=lq,
            max_seqlen_k=lk,
            dropout_p=dropout_p,
            softmax_scale=softmax_scale,
            causal=causal,
            window_size=window_size,
            deterministic=deterministic).unflatten(0, (b, lq))
    elif SAGE_ATTN_AVAILABLE:
        q = q.unsqueeze(0).transpose(1, 2).to(dtype)
        k = k.unsqueeze(0).transpose(1, 2).to(dtype)
        v = v.unsqueeze(0).transpose(1, 2).to(dtype)
        x = sageattn(q, k, v, dropout_p=dropout_p, is_causal=causal)
        x = x.transpose(1, 2).contiguous()
    else:
        q = q.unsqueeze(0).transpose(1, 2).to(dtype)
        k = k.unsqueeze(0).transpose(1, 2).to(dtype)
        v = v.unsqueeze(0).transpose(1, 2).to(dtype)
        x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
        x = x.transpose(1, 2).contiguous()

    # output
    return x.type(out_dtype)


def create_sdpa_mask(q, k, q_lens, k_lens, causal=False):
    b, lq, lk = q.size(0), q.size(1), k.size(1)
    if q_lens is None:
        q_lens = torch.tensor([lq] * b, dtype=torch.int32)
    if k_lens is None:
        k_lens = torch.tensor([lk] * b, dtype=torch.int32)
    attn_mask = torch.zeros((b, lq, lk), dtype=torch.bool)
    for i in range(b):
        q_len, k_len = q_lens[i], k_lens[i]
        attn_mask[i, q_len:, :] = True
        attn_mask[i, :, k_len:] = True
        
        if causal:
            causal_mask = torch.triu(torch.ones((lq, lk), dtype=torch.bool), diagonal=1)
            attn_mask[i, :, :] = torch.logical_or(attn_mask[i, :, :], causal_mask)

    attn_mask = attn_mask.logical_not().to(q.device, non_blocking=True)
    return attn_mask


def attention(
    q,
    k,
    v,
    q_lens=None,
    k_lens=None,
    dropout_p=0.,
    softmax_scale=None,
    q_scale=None,
    causal=False,
    window_size=(-1, -1),
    deterministic=False,
    dtype=torch.bfloat16,
    fa_version=None,
):
    if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
        return flash_attention(
            q=q,
            k=k,
            v=v,
            q_lens=q_lens,
            k_lens=k_lens,
            dropout_p=dropout_p,
            softmax_scale=softmax_scale,
            q_scale=q_scale,
            causal=causal,
            window_size=window_size,
            deterministic=deterministic,
            dtype=dtype,
            version=fa_version,
        )
    else:
        if q_lens is not None or k_lens is not None:
            warnings.warn('Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.')
        attn_mask = None

        q = q.transpose(1, 2).to(dtype)
        k = k.transpose(1, 2).to(dtype)
        v = v.transpose(1, 2).to(dtype)

        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p)

        out = out.transpose(1, 2).contiguous()
        return out



def sinusoidal_embedding_1d(dim, position):
    # preprocess
    assert dim % 2 == 0
    half = dim // 2
    position = position.type(torch.float64)

    # calculation
    sinusoid = torch.outer(
        position, torch.pow(10000, -torch.arange(half).to(position).div(half)))
    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
    return x


@amp.autocast(enabled=False, device_type="cuda")
def rope_params(max_seq_len, dim, theta=10000):
    assert dim % 2 == 0
    freqs = torch.outer(
        torch.arange(max_seq_len),
        1.0 / torch.pow(theta,
                        torch.arange(0, dim, 2).to(torch.float64).div(dim)))
    freqs = torch.polar(torch.ones_like(freqs), freqs)
    return freqs


@amp.autocast(enabled=False, device_type="cuda")
def rope_apply(x, grid_sizes, freqs, sequence_cond_compressed_indices=None):
    batch, seq_len_actual, n, c = x.size(0), x.size(1), x.size(2), x.size(3) // 2
    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
    output = []
    assert len(grid_sizes) == batch, "grid_sizes must have the same length as the batch size ([b, 3=[f, h, w])"
    for i, (f, h, w) in enumerate(grid_sizes.tolist()):
        seq_len = f * h * w
        x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(
            seq_len, n, -1, 2))
        freqs_i = torch.cat([
            freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
            freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
            freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
        ], dim=-1).reshape(seq_len, 1, -1)
        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
        
        if seq_len_actual > seq_len:
            sequence_cond_seq_length = seq_len_actual - seq_len
            if sequence_cond_seq_length == seq_len:
                x_i_sequence_cond = torch.view_as_complex(x[i, seq_len:].to(torch.float64).reshape(seq_len_actual - seq_len, n, -1, 2))
                x_i_sequence_cond = torch.view_as_real(x_i_sequence_cond * freqs_i).flatten(2)
            else:
                sequence_cond_compressed_index = sequence_cond_compressed_indices[i]
                sequence_cond_t_length = len(sequence_cond_compressed_index)
                assert sequence_cond_t_length * h * w == sequence_cond_seq_length, "`sequence_cond_t_length * h * w` must be equal to `sequence_cond_seq_length`"
                x_i_sequence_cond = torch.view_as_complex(x[i, seq_len:].to(torch.float64).reshape(sequence_cond_seq_length, n, -1, 2))
                freqs_i_sequence_cond = torch.cat([
                    freqs[0][sequence_cond_compressed_index].view(sequence_cond_t_length, 1, 1, -1).expand(sequence_cond_t_length, h, w, -1),
                    freqs[1][:h].view(1, h, 1, -1).expand(sequence_cond_t_length, h, w, -1),
                    freqs[2][:w].view(1, 1, w, -1).expand(sequence_cond_t_length, h, w, -1)
                ], dim=-1).reshape(sequence_cond_seq_length, 1, -1)
                x_i_sequence_cond = torch.view_as_real(x_i_sequence_cond * freqs_i_sequence_cond).flatten(2)
            x_i = torch.cat([x_i, x_i_sequence_cond])
        
        output.append(x_i)
    return torch.stack(output).float()


class WanRMSNorm(nn.Module):

    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return self._norm(x.float()).type_as(x) * self.weight

    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)


class WanLayerNorm(nn.LayerNorm):

    def __init__(self, dim, eps=1e-6, elementwise_affine=False):
        super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)

    def forward(self, x):
        return super().forward(x.float()).type_as(x)


class WanSelfAttention(nn.Module):

    def __init__(self,
                 dim,
                 num_heads,
                 window_size=(-1, -1),
                 qk_norm=True,
                 eps=1e-6):
        assert dim % num_heads == 0
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.window_size = window_size
        self.qk_norm = qk_norm
        self.eps = eps

        self.q = nn.Linear(dim, dim)
        self.k = nn.Linear(dim, dim)
        self.v = nn.Linear(dim, dim)
        self.o = nn.Linear(dim, dim)
        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
        self.visualize_attention = False

    def forward(self, x, seq_lens, grid_sizes, freqs, sequence_cond_compressed_indices):
        """
        Args:
            x:              [B, L, C].
            seq_lens:       [B].
            grid_sizes:     [B, 3=[f, h, w]].
            freqs:          [L, 2].
            sequence_cond_compressed_indices: [B, T_sequence_condITION].
            
        `f` in `grid_sizes` can less than the actual seq_lens (L), 
        which indicates full in-context condition (when L=2*f) or 
        sparse in-context condition (when `f` < L < 2*f and `sequence_cond_compressed_indices` is not None) is used.
        """
        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim

        def qkv_fn(x):
            q = self.norm_q(self.q(x)).view(b, s, n, d)
            k = self.norm_k(self.k(x)).view(b, s, n, d)
            v = self.v(x).view(b, s, n, d)
            return q, k, v

        q, k, v = qkv_fn(x)
        
        q_rope = rope_apply(q, grid_sizes, freqs, sequence_cond_compressed_indices)
        k_rope = rope_apply(k, grid_sizes, freqs, sequence_cond_compressed_indices)
        
        if self.visualize_attention:
            with torch.no_grad():
                self._last_attn_maps = self._compute_attention_for_visualization(q_rope, k_rope) # CPU tesnor of [S, S]
                self._last_grid_sizes = grid_sizes
                self._last_seq_lens = seq_lens

        x = flash_attention(
            q=q_rope,
            k=k_rope,
            v=v,
            k_lens=None,
            window_size=self.window_size)

        # output
        x = x.flatten(2)
        x = self.o(x)
        return x
    
    def _compute_attention_for_visualization(self, q, k):
        """Compute attention maps for visualization purposes"""
        # b, _, n, d = q.shape
        print("Computing attention maps for visualization")
        # Reshape for attention computation
        q = q.permute(0, 2, 1, 3)  # [b, n, s, d]
        k = k.permute(0, 2, 1, 3)  # [b, n, s, d]
        # query: b, n, s, d
        print("q.shape=", q.shape)
        print("k.shape=", k.shape)
        attention_probs_list = []
        for i in range(0, q.shape[1], 20):
            print(f"Computing attention for head {i} to {i+20}")
            query_attention = q[-1][i : i + 20]
            key_attention = k[-1][i : i + 20]
            identity_matrix = torch.eye(
                query_attention.shape[-2],
                device=query_attention.device,
                dtype=query_attention.dtype,
            ) # shape=[s]
            attention_probs_temp = torch.nn.functional.scaled_dot_product_attention(
                query_attention,
                key_attention,
                identity_matrix,
                attn_mask=None,
                dropout_p=0.0,
                is_causal=False,
            )
            attention_probs_list.append(attention_probs_temp.detach().cpu())
            del (
                query_attention,
                key_attention,
                identity_matrix,
                attention_probs_temp,
            )
        attention_probs = torch.mean(torch.cat(attention_probs_list), dim=0).float().numpy()
        print("Attention maps computed. Shape=", attention_probs.shape)
        # Only keep attention maps, don't compute the output
        return attention_probs # [s, s]


class WanT2VCrossAttention(WanSelfAttention):

    def forward(self, x, context, context_lens):
        """
        x:              [B, L1, C].
        context:        [B, L2, C].
        context_lens:   [B].
        """
        b, n, d = x.size(0), self.num_heads, self.head_dim

        # compute query, key, value
        q = self.norm_q(self.q(x)).view(b, -1, n, d)
        k = self.norm_k(self.k(context)).view(b, -1, n, d)
        v = self.v(context).view(b, -1, n, d)

        # compute attention
        x = flash_attention(q, k, v, k_lens=context_lens)

        # output
        x = x.flatten(2)
        x = self.o(x)
        return x


class WanI2VCrossAttention(WanSelfAttention):

    def __init__(self,
                 dim,
                 num_heads,
                 window_size=(-1, -1),
                 qk_norm=True,
                 eps=1e-6):
        super().__init__(dim, num_heads, window_size, qk_norm, eps)

        self.k_img = nn.Linear(dim, dim)
        self.v_img = nn.Linear(dim, dim)
        # self.alpha = nn.Parameter(torch.zeros((1, )))
        self.norm_k_img = WanRMSNorm(
            dim, eps=eps) if qk_norm else nn.Identity()

    def forward(self, x, context, context_lens):
        """
        x:              [B, L1, C].
        context:        [B, L2, C].
        context_lens:   [B].
        """
        context_img = context[:, :257]
        context = context[:, 257:]
        b, n, d = x.size(0), self.num_heads, self.head_dim

        # compute query, key, value
        q = self.norm_q(self.q(x)).view(b, -1, n, d)
        k = self.norm_k(self.k(context)).view(b, -1, n, d)
        v = self.v(context).view(b, -1, n, d)
        k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d)
        v_img = self.v_img(context_img).view(b, -1, n, d)
        img_x = flash_attention(q, k_img, v_img, k_lens=None)
        # compute attention
        x = flash_attention(q, k, v, k_lens=context_lens)

        # output
        x = x.flatten(2)
        img_x = img_x.flatten(2)
        x = x + img_x
        x = self.o(x)
        return x


WANX_CROSSATTENTION_CLASSES = {
    't2v_cross_attn': WanT2VCrossAttention,
    'i2v_cross_attn': WanI2VCrossAttention,
}


class WanAttentionBlock(nn.Module):

    def __init__(self,
                 cross_attn_type,
                 dim,
                 ffn_dim,
                 num_heads,
                 window_size=(-1, -1),
                 qk_norm=True,
                 cross_attn_norm=False,
                 eps=1e-6,
                 use_local_lora=False,
                 use_dera=False,
                 dera_rank=None,
                 use_dera_spatial=True,
                 use_dera_temporal=True):
        super().__init__()
        self.dim = dim
        self.ffn_dim = ffn_dim
        self.num_heads = num_heads
        self.window_size = window_size
        self.qk_norm = qk_norm
        self.cross_attn_norm = cross_attn_norm
        self.eps = eps

        # layers
        self.norm1 = WanLayerNorm(dim, eps)
        self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, eps)
        self.norm3 = WanLayerNorm(
            dim, eps,
            elementwise_affine=True) if cross_attn_norm else nn.Identity()
        self.cross_attn = WANX_CROSSATTENTION_CLASSES[cross_attn_type](
            dim, num_heads, (-1, -1), qk_norm, eps)
        self.norm2 = WanLayerNorm(dim, eps)
        self.ffn = nn.Sequential(
            nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
            nn.Linear(ffn_dim, dim))

        # modulation
        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
        
        self.use_local_lora = use_local_lora
        if use_local_lora:
            from .local_lora import LocalLoRA
            self.local_lora = LocalLoRA(dim=dim, rank=64, kernel_size=(3, 3), stride=(1, 1))
            
        self.use_dera = use_dera
        if use_dera:
            from .dera import DeRA
            self.dera = DeRA(dim, rank=dera_rank, use_spatial=use_dera_spatial, use_temporal=use_dera_temporal)

    def forward(
        self,
        x,
        e,
        seq_lens,
        grid_sizes,
        freqs,
        context,
        context_lens,
        sequence_cond_compressed_indices,
        dera_freqs=None
    ):
        assert e.dtype == torch.float32
        with amp.autocast(dtype=torch.float32, device_type="cuda"):
            e = (self.modulation.to(dtype=e.dtype, device=e.device) + e).chunk(6, dim=1)
        assert e[0].dtype == torch.float32

        # self-attention
        x_self_attn_input = self.norm1(x).float() * (1 + e[1]) + e[0]
        y = self.self_attn(x_self_attn_input, seq_lens, grid_sizes, freqs, sequence_cond_compressed_indices)
        if self.use_local_lora:
            y = y + self.local_lora(x_self_attn_input, grid_sizes)
        
        if self.use_dera:
            y = y + self.dera(x_self_attn_input, seq_lens, grid_sizes, dera_freqs, sequence_cond_compressed_indices)
            
        with amp.autocast(dtype=torch.float32, device_type="cuda"):
            x = x + y * e[2]
            
        def cross_attn_ffn(x, context, context_lens, e):
            x = x + self.cross_attn(self.norm3(x), context, context_lens)
            y = self.ffn(self.norm2(x).float() * (1 + e[4]) + e[3])
            with amp.autocast(dtype=torch.float32, device_type="cuda"):
                x = x + y * e[5]
            return x

        x = cross_attn_ffn(x, context, context_lens, e)
        return x


class Head(nn.Module):

    def __init__(self, dim, out_dim, patch_size, eps=1e-6):
        super().__init__()
        self.dim = dim
        self.out_dim = out_dim
        self.patch_size = patch_size
        self.eps = eps

        # layers
        out_dim = math.prod(patch_size) * out_dim
        self.norm = WanLayerNorm(dim, eps)
        self.head = nn.Linear(dim, out_dim)

        # modulation
        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)

    def forward(self, x, e):
        assert e.dtype == torch.float32
        with amp.autocast(dtype=torch.float32, device_type="cuda"):
            e = (self.modulation.to(dtype=e.dtype, device=e.device) + e.unsqueeze(1)).chunk(2, dim=1)
            x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))
        return x


class MLPProj(torch.nn.Module):

    def __init__(self, in_dim, out_dim):
        super().__init__()

        self.proj = torch.nn.Sequential(
            torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim),
            torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim),
            torch.nn.LayerNorm(out_dim))

    def forward(self, image_embeds):
        clip_extra_context_tokens = self.proj(image_embeds)
        return clip_extra_context_tokens


class WanModel(nn.Module):

    def __init__(self,
                 model_type='t2v',
                 patch_size=(1, 2, 2),
                 text_len=512,
                 in_dim=16,
                 dim=2048,
                 ffn_dim=8192,
                 freq_dim=256,
                 text_dim=4096,
                 out_dim=16,
                 num_heads=16,
                 num_layers=32,
                 window_size=(-1, -1),
                 qk_norm=True,
                 cross_attn_norm=False,
                 eps=1e-6,
                 use_local_lora=False,
                 use_dera=False,
                 dera_rank=None,
                 use_dera_spatial=True,
                 use_dera_temporal=True,
                 use_sequence_cond=False,
                 sequence_cond_in_dim=None,
                 sequence_cond_mode=None,
                 use_channel_cond=False,
                 channel_cond_in_dim=None,
                 use_sequence_cond_position_aware_residual=False,
                 use_sequence_cond_loss=False
                 ):
        super().__init__()

        assert model_type in ['t2v', 'i2v']
        self.model_type = model_type

        self.patch_size = patch_size
        self.text_len = text_len
        self.in_dim = in_dim
        self.dim = dim
        self.ffn_dim = ffn_dim
        self.freq_dim = freq_dim
        self.text_dim = text_dim
        self.out_dim = out_dim
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.window_size = window_size
        self.qk_norm = qk_norm
        self.cross_attn_norm = cross_attn_norm
        self.eps = eps
        
        self.use_local_lora = use_local_lora
        self.use_dera = use_dera

        # embeddings
        self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size)
        self.text_embedding = nn.Sequential(
            nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),
            nn.Linear(dim, dim))

        self.time_embedding = nn.Sequential(
            nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
        self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))

        # blocks
        cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'
        self.blocks = nn.ModuleList([
            WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,
                               window_size, qk_norm, cross_attn_norm, eps, use_local_lora=use_local_lora, 
                               use_dera=use_dera, dera_rank=dera_rank, use_dera_spatial=use_dera_spatial, use_dera_temporal=use_dera_temporal)
            for _ in range(num_layers)
        ])

        # head
        self.head = Head(dim, out_dim, patch_size, eps)

        # buffers (don't use register_buffer otherwise dtype will be changed in to())
        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
        d = dim // num_heads
        self.freqs = torch.cat([
            rope_params(1024, d - 4 * (d // 6)),
            rope_params(1024, 2 * (d // 6)),
            rope_params(1024, 2 * (d // 6))
        ], dim=1)
        
        if self.use_dera:
            dera_d = dera_rank // 4  # (18)
            self.dera_freqs = torch.cat([
                rope_params(1024, dera_d - 4 * (dera_d // 6)),
                rope_params(1024, 2 * (dera_d // 6)),
                rope_params(1024, 2 * (dera_d // 6))
            ], dim=1)
        else:
            self.dera_freqs = None

        if model_type == 'i2v':
            self.img_emb = MLPProj(1280, dim)
            
        self.init_weights()
        
        self.use_sequence_cond = use_sequence_cond
        self.sequence_cond_in_dim = sequence_cond_in_dim
        self.sequence_cond_mode = sequence_cond_mode
        if use_sequence_cond:
            assert sequence_cond_in_dim is not None, "`sequence_cond_in_dim` must be provided when `use_sequence_cond` is True"
            self.sequence_cond_patch_embedding = nn.Conv3d(sequence_cond_in_dim, dim, kernel_size=patch_size, stride=patch_size)
            self.sequence_cond_identifier = nn.Parameter(torch.randn(1, 1, dim) / dim**0.5)
        
        self.use_channel_cond = use_channel_cond
        self.channel_cond_in_dim = channel_cond_in_dim
        if use_channel_cond:
            assert channel_cond_in_dim is not None, "`channel_cond_in_dim` must be provided when `use_channel_cond` is True"
        self.use_sequence_cond_position_aware_residual = use_sequence_cond_position_aware_residual
        if use_sequence_cond_position_aware_residual:
            self.sequence_cond_residual_proj = nn.Linear(dim, dim, bias=False)
            self.sequence_cond_residual_proj.weight.data.zero_()
        
        self.use_sequence_cond_loss = use_sequence_cond_loss
        if self.use_sequence_cond_loss:
            self.sequence_latent_to_cond_proj = nn.Linear(dim, dim, bias=False)
            self.sequence_latent_to_cond_proj.weight.data.zero_()
            self.head_sequence_cond_out = nn.Linear(dim, math.prod(patch_size) * out_dim)
    
    def copy_sequence_cond_patch_embedding_weights(self):
        size_patch_embedding = self.patch_embedding.weight.size(1)
        size_sequence_cond_patch_embedding = self.sequence_cond_patch_embedding.weight.size(1)
        self.sequence_cond_patch_embedding.weight.data = self.patch_embedding.weight.data[:, size_patch_embedding - size_sequence_cond_patch_embedding:, :, :, :].clone()
        if self.patch_embedding.bias is not None:
            self.sequence_cond_patch_embedding.bias.data = self.patch_embedding.bias.data.clone()
    
    def copy_patch_embedding_weights_for_channel_cond(self):
        original_patch_in_channels = self.patch_embedding.in_channels
        new_patch_embedding = nn.Conv3d(in_channels=original_patch_in_channels + self.channel_cond_in_dim, 
                                        out_channels=self.dim, kernel_size=self.patch_size, stride=self.patch_size)
        new_patch_embedding.weight.data[:, :original_patch_in_channels, :, :, :] = self.patch_embedding.weight.data.clone()
        if self.patch_embedding.bias is not None:
            new_patch_embedding.bias.data = self.patch_embedding.bias.data.clone()
        del self.patch_embedding
        self.patch_embedding = new_patch_embedding

    def forward(
        self,
        x,
        timestep,
        context,
        seq_len,
        clip_fea=None,
        y=None,
        use_gradient_checkpointing=False,
        sequence_cond=None,
        sequence_cond_compressed_indices=None,
        channel_cond=None,
        sequence_cond_residual_scale=1.0,
        **kwargs,
    ):
        """
        x:              A list of videos each with shape [C, T, H, W].
        t:              [B].
        context:        A list of text embeddings each with shape [L, C].
        sequence_cond: A list of conditional frames each with shape [C, T_sequence_cond, H, W].
        sequence_cond_compressed_indices: [B, T_sequence_cond] Indices for any additional conditional information, where T_sequence_cond < T. For sparse mode only.
        
        
            Note:
                sequence_cond will be injected into the model as an additional input sequence, i.e., sequence dimension.
                channel_cond will be injected into the model in the input' channel dimension.
                
            Examples:
                1) for extra cond case:
                    # given x: [B, C, T, H, W] ----> [B, L=T*H*W, C] --patch_embedding--> [B, L, D]
                    # sequence_cond: [B, C_sequence_cond, T_sequence_cond, H, W] ----> [B, L_sequence_cond=T_sequence_cond*H*W, C_sequence_cond] --sequence_cond_embedding--> [B, L_sequence_cond, D]
                    x = torch.concat([x, sequence_cond], dim=2) # Concat on sequence dimension after patch/extra cond embedding
                    # after concat, x: [B, L+L_sequence_cond, D]
                2) for channel cond case:
                    given x: [B, C, T, H, W]
                    channel_cond: [B, C_CHANNEL_COND, T, H, W]
                    x = torch.concat([x, channel_cond], dim=1) # Concat on channel dimension before patch/extra cond embedding    
                    # x: [B, C + C_CHANNEL_COND, T, H, W] --patch_embedding(requires param copy and tuning)--> [B, L=T*H*W, D]
        """
        if self.model_type == 'i2v':
            assert clip_fea is not None and y is not None
        # params
        device = x[0].device
        if self.freqs.device != device:
            self.freqs = self.freqs.to(device)
        if self.dera_freqs is not None and self.dera_freqs.device != device:
            self.dera_freqs = self.dera_freqs.to(device)
            
        if y is not None:
            x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
        
        if channel_cond is not None:
            assert self.use_channel_cond, "forward argument `channel_cond` is provided but model property `self.use_channel_cond` is False"
            x = [torch.cat([u, v], dim=0) for u, v in zip(x, channel_cond)]

        # embeddings
        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
        grid_sizes = torch.stack(
            [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
        x = [u.flatten(2).transpose(1, 2) for u in x]
        x = torch.cat(x, dim=0)
        
        if sequence_cond is not None:
            assert self.use_sequence_cond, "forward argument `sequence_cond` is provided but model property `self.use_sequence_cond` is False"
            sequence_cond = [self.sequence_cond_patch_embedding(u.unsqueeze(0)) for u in sequence_cond]
            sequence_cond = [u.flatten(2).transpose(1, 2) + self.sequence_cond_identifier for u in sequence_cond]
            sequence_cond = torch.concat(sequence_cond, dim=0)
            
            x = torch.concat([x, sequence_cond], dim=1)
        
        actual_seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
        
        # time embeddings
        with amp.autocast(dtype=torch.float32, device_type="cuda"):
            e = self.time_embedding(
                sinusoidal_embedding_1d(self.freq_dim, timestep).float())
            e0 = self.time_projection(e).unflatten(1, (6, self.dim))
            assert e.dtype == torch.float32 and e0.dtype == torch.float32

        # context
        context_lens = None
        context = self.text_embedding(
            torch.stack([
                torch.cat(
                    [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
                for u in context
            ]))

        if clip_fea is not None:
            context_clip = self.img_emb(clip_fea)  # bs x 257 x dim
            context = torch.concat([context_clip, context], dim=1)

        # arguments
        kwargs = dict(e=e0, seq_lens=actual_seq_lens, grid_sizes=grid_sizes,
            freqs=self.freqs, context=context, context_lens=context_lens,
            sequence_cond_compressed_indices=sequence_cond_compressed_indices, dera_freqs=self.dera_freqs)
        
        def create_custom_forward(module):
            def custom_forward(*inputs, **kwargs):
                return module(*inputs, **kwargs)
            return custom_forward

        for block_idx, block in enumerate(self.blocks):
            if self.training and use_gradient_checkpointing:
                x = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(block),
                    x, **kwargs,
                    use_reentrant=False,
                )
            else:
                x = block(x, **kwargs)
                
            if self.use_sequence_cond_loss and block_idx == len(self.blocks) - 3:
                # This this function, the context length will be extended from (N+C) to 2N, where C is the length of the sparse sequence cond.
                x_ori = x[:, :seq_len, :]
                x_ori_projected = self.sequence_latent_to_cond_proj(x_ori)
                x_seq_cond = x[:, seq_len:, :]
                seq_cond_length = len(sequence_cond_compressed_indices[0])
                x_ori_projected = rearrange(x_ori_projected, 'b (t h w) c -> b c t h w', t=grid_sizes[0, 0], h=grid_sizes[0, 1], w=grid_sizes[0, 2])
                x_seq_cond = rearrange(x_seq_cond, 'b (t h w) c -> b c t h w', t=seq_cond_length, h=grid_sizes[0, 1], w=grid_sizes[0, 2])
                x_ori_projected[:, :, sequence_cond_compressed_indices[0], :, :] += x_seq_cond
                x_ori_projected = rearrange(x_ori_projected, 'b c t h w -> b (t h w) c')
                x = torch.concat([x_ori, x_ori_projected], dim=1)
                # Let the later blocks generate sketches at the full seqeuence length
                
            if self.use_sequence_cond_position_aware_residual and block_idx < len(self.blocks) - 1:
                # Apply the sequence condition position-aware residual for all blocks except the last one
                x_ori = x[:, :seq_len, :]
                x_seq_cond = x[:, seq_len:, :]
                x_seq_cond_porjected = self.sequence_cond_residual_proj(x_seq_cond)
                assert x_ori.shape[0] == 1, "Only support batch size 1 for `sequence_cond_position_aware_residual`."
                seq_cond_length = len(sequence_cond_compressed_indices[0])
                x_ori = rearrange(x_ori, 'b (t h w) c -> b c t h w', t=grid_sizes[0, 0], h=grid_sizes[0, 1], w=grid_sizes[0, 2])
                x_seq_cond_porjected = rearrange(x_seq_cond_porjected, 'b (t h w) c -> b c t h w', t=seq_cond_length, h=grid_sizes[0, 1], w=grid_sizes[0, 2])
                
                x_ori[:, :, sequence_cond_compressed_indices[0], :, :] = x_ori[:, :, sequence_cond_compressed_indices[0], :, :] + x_seq_cond_porjected * sequence_cond_residual_scale
                x_ori = rearrange(x_ori, 'b c t h w -> b (t h w) c')
                x = torch.concat([x_ori, x_seq_cond], dim=1)
                
        if sequence_cond is not None:
            if self.use_sequence_cond_loss:
                sequence_cond_out = x[:, seq_len:, :]
                sequence_cond_out = self.unpatchify(sequence_cond_out, grid_sizes)  # sequence_cond_grid_sizes
                sequence_cond_out = torch.stack(sequence_cond_out).float()  # b, c, t, h, w
            else:
                sequence_cond_out = None
            x = x[:, :seq_len, :]
        # head
        x = self.head(x, e)

        # unpatchify
        x = self.unpatchify(x, grid_sizes)
        x = torch.stack(x).float()
        if sequence_cond is not None and self.use_sequence_cond_loss:
            return x, sequence_cond_out
        return x

    def unpatchify(self, x, grid_sizes):
        c = self.out_dim
        out = []
        for u, v in zip(x, grid_sizes.tolist()):
            u = u[:math.prod(v)].view(*v, *self.patch_size, c)
            u = torch.einsum('fhwpqrc->cfphqwr', u)
            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
            out.append(u)
        return out

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

        # init embeddings
        nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
        for m in self.text_embedding.modules():
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, std=.02)
        for m in self.time_embedding.modules():
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, std=.02)

        # init output layer
        nn.init.zeros_(self.head.head.weight)

    @staticmethod
    def state_dict_converter():
        return WanModelStateDictConverter()
    
    
class WanModelStateDictConverter:
    def __init__(self):
        pass

    def from_diffusers(self, state_dict):
        rename_dict = {"blocks.0.attn1.norm_k.weight": "blocks.0.self_attn.norm_k.weight",
            "blocks.0.attn1.norm_q.weight": "blocks.0.self_attn.norm_q.weight",
            "blocks.0.attn1.to_k.bias": "blocks.0.self_attn.k.bias",
            "blocks.0.attn1.to_k.weight": "blocks.0.self_attn.k.weight",
            "blocks.0.attn1.to_out.0.bias": "blocks.0.self_attn.o.bias",
            "blocks.0.attn1.to_out.0.weight": "blocks.0.self_attn.o.weight",
            "blocks.0.attn1.to_q.bias": "blocks.0.self_attn.q.bias",
            "blocks.0.attn1.to_q.weight": "blocks.0.self_attn.q.weight",
            "blocks.0.attn1.to_v.bias": "blocks.0.self_attn.v.bias",
            "blocks.0.attn1.to_v.weight": "blocks.0.self_attn.v.weight",
            "blocks.0.attn2.norm_k.weight": "blocks.0.cross_attn.norm_k.weight",
            "blocks.0.attn2.norm_q.weight": "blocks.0.cross_attn.norm_q.weight",
            "blocks.0.attn2.to_k.bias": "blocks.0.cross_attn.k.bias",
            "blocks.0.attn2.to_k.weight": "blocks.0.cross_attn.k.weight",
            "blocks.0.attn2.to_out.0.bias": "blocks.0.cross_attn.o.bias",
            "blocks.0.attn2.to_out.0.weight": "blocks.0.cross_attn.o.weight",
            "blocks.0.attn2.to_q.bias": "blocks.0.cross_attn.q.bias",
            "blocks.0.attn2.to_q.weight": "blocks.0.cross_attn.q.weight",
            "blocks.0.attn2.to_v.bias": "blocks.0.cross_attn.v.bias",
            "blocks.0.attn2.to_v.weight": "blocks.0.cross_attn.v.weight",
            "blocks.0.ffn.net.0.proj.bias": "blocks.0.ffn.0.bias",
            "blocks.0.ffn.net.0.proj.weight": "blocks.0.ffn.0.weight",
            "blocks.0.ffn.net.2.bias": "blocks.0.ffn.2.bias",
            "blocks.0.ffn.net.2.weight": "blocks.0.ffn.2.weight",
            "blocks.0.norm2.bias": "blocks.0.norm3.bias",
            "blocks.0.norm2.weight": "blocks.0.norm3.weight",
            "blocks.0.scale_shift_table": "blocks.0.modulation",
            "condition_embedder.text_embedder.linear_1.bias": "text_embedding.0.bias",
            "condition_embedder.text_embedder.linear_1.weight": "text_embedding.0.weight",
            "condition_embedder.text_embedder.linear_2.bias": "text_embedding.2.bias",
            "condition_embedder.text_embedder.linear_2.weight": "text_embedding.2.weight",
            "condition_embedder.time_embedder.linear_1.bias": "time_embedding.0.bias",
            "condition_embedder.time_embedder.linear_1.weight": "time_embedding.0.weight",
            "condition_embedder.time_embedder.linear_2.bias": "time_embedding.2.bias",
            "condition_embedder.time_embedder.linear_2.weight": "time_embedding.2.weight",
            "condition_embedder.time_proj.bias": "time_projection.1.bias",
            "condition_embedder.time_proj.weight": "time_projection.1.weight",
            "patch_embedding.bias": "patch_embedding.bias",
            "patch_embedding.weight": "patch_embedding.weight",
            "scale_shift_table": "head.modulation",
            "proj_out.bias": "head.head.bias",
            "proj_out.weight": "head.head.weight",
        }
        state_dict_ = {}
        for name, param in state_dict.items():
            if name in rename_dict:
                state_dict_[rename_dict[name]] = param
            else:
                name_ = ".".join(name.split(".")[:1] + ["0"] + name.split(".")[2:])
                if name_ in rename_dict:
                    name_ = rename_dict[name_]
                    name_ = ".".join(name_.split(".")[:1] + [name.split(".")[1]] + name_.split(".")[2:])
                    state_dict_[name_] = param
        if hash_state_dict_keys(state_dict) == "cb104773c6c2cb6df4f9529ad5c60d0b":
            config = {
                "model_type": "t2v",
                "patch_size": (1, 2, 2),
                "text_len": 512,
                "in_dim": 16,
                "dim": 5120,
                "ffn_dim": 13824,
                "freq_dim": 256,
                "text_dim": 4096,
                "out_dim": 16,
                "num_heads": 40,
                "num_layers": 40,
                "window_size": (-1, -1),
                "qk_norm": True,
                "cross_attn_norm": True,
                "eps": 1e-6,
            }
        else:
            config = {}
        return state_dict_, config
    
    def from_civitai(self, state_dict):
        if hash_state_dict_keys(state_dict) == "9269f8db9040a9d860eaca435be61814":
            config = {
                "model_type": "t2v",
                "patch_size": (1, 2, 2),
                "text_len": 512,
                "in_dim": 16,
                "dim": 1536,
                "ffn_dim": 8960,
                "freq_dim": 256,
                "text_dim": 4096,
                "out_dim": 16,
                "num_heads": 12,
                "num_layers": 30,
                "window_size": (-1, -1),
                "qk_norm": True,
                "cross_attn_norm": True,
                "eps": 1e-6,
            }
        elif hash_state_dict_keys(state_dict) == "aafcfd9672c3a2456dc46e1cb6e52c70":
            config = {
                "model_type": "t2v",
                "patch_size": (1, 2, 2),
                "text_len": 512,
                "in_dim": 16,
                "dim": 5120,
                "ffn_dim": 13824,
                "freq_dim": 256,
                "text_dim": 4096,
                "out_dim": 16,
                "num_heads": 40,
                "num_layers": 40,
                "window_size": (-1, -1),
                "qk_norm": True,
                "cross_attn_norm": True,
                "eps": 1e-6,
            }
        elif hash_state_dict_keys(state_dict) == "6bfcfb3b342cb286ce886889d519a77e":
            config = {
                "model_type": "i2v",
                "patch_size": (1, 2, 2),
                "text_len": 512,
                "in_dim": 36,
                "dim": 5120,
                "ffn_dim": 13824,
                "freq_dim": 256,
                "text_dim": 4096,
                "out_dim": 16,
                "num_heads": 40,
                "num_layers": 40,
                "window_size": (-1, -1),
                "qk_norm": True,
                "cross_attn_norm": True,
                "eps": 1e-6,
            }
        else:
            config = {}
        return state_dict, config