wehub-resource-sync
57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
import torch.nn as nn
|
|
|
|
from .model import WanAttentionBlock, WanLinearAttentionModel, WanModel
|
|
|
|
|
|
class SanaVideoMSBlock(WanAttentionBlock):
|
|
pass
|
|
|
|
|
|
class SanaWanModel(WanModel):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
cross_attn_type = "t2v_cross_attn" if self.model_type == "t2v" else "i2v_cross_attn"
|
|
self.blocks = nn.ModuleList(
|
|
[
|
|
SanaVideoMSBlock(
|
|
cross_attn_type,
|
|
self.dim,
|
|
self.ffn_dim,
|
|
self.num_heads,
|
|
self.window_size,
|
|
self.qk_norm,
|
|
self.cross_attn_norm,
|
|
self.eps,
|
|
)
|
|
for _ in range(self.num_layers)
|
|
]
|
|
)
|
|
|
|
|
|
class SanaWanLinearAttentionModel(WanLinearAttentionModel):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
cross_attn_type = "t2v_cross_attn" if self.model_type == "t2v" else "i2v_cross_attn"
|
|
self_attn_types = ["flash"] * self.num_layers
|
|
ffn_types = ["mlp"] * self.num_layers
|
|
|
|
self.blocks = nn.ModuleList(
|
|
[
|
|
SanaVideoMSBlock(
|
|
cross_attn_type,
|
|
self.dim,
|
|
self.ffn_dim,
|
|
self.num_heads,
|
|
self.window_size,
|
|
self.qk_norm,
|
|
self.cross_attn_norm,
|
|
self.eps,
|
|
self_attn_types[i],
|
|
self.rope_after,
|
|
self.power,
|
|
ffn_types[i],
|
|
)
|
|
for i in range(self.num_layers)
|
|
]
|
|
)
|