diffsynth-engine 0.0.0__py3-none-any.whl

diffsynth_engine/models/utils.py

@@ -0,0 +1,54 @@
+import torch
+import torch.nn as nn
+from contextlib import contextmanager
+
+
+# mofified from transformers.modeling_utils
+TORCH_INIT_FUNCTIONS = {
+    "uniform_": nn.init.uniform_,
+    "normal_": nn.init.normal_,
+    "trunc_normal_": nn.init.trunc_normal_,
+    "constant_": nn.init.constant_,
+    "xavier_uniform_": nn.init.xavier_uniform_,
+    "xavier_normal_": nn.init.xavier_normal_,
+    "kaiming_uniform_": nn.init.kaiming_uniform_,
+    "kaiming_normal_": nn.init.kaiming_normal_,
+    "uniform": nn.init.uniform,
+    "normal": nn.init.normal,
+    "xavier_uniform": nn.init.xavier_uniform,
+    "xavier_normal": nn.init.xavier_normal,
+    "kaiming_uniform": nn.init.kaiming_uniform,
+    "kaiming_normal": nn.init.kaiming_normal,
+}
+
+_init_weights = True
+
+
+@contextmanager
+def no_init_weights():
+    """
+    Context manager to globally disable weight initialization to speed up loading large models.
+    """
+    global _init_weights
+    old_init_weights = _init_weights
+
+    def _skip_init(*args, **kwargs):
+        pass
+
+    _init_weights = False
+    # Save the original initialization functions
+    for name, init_func in TORCH_INIT_FUNCTIONS.items():
+        setattr(torch.nn.init, name, _skip_init)
+    try:
+        yield
+    finally:
+        _init_weights = old_init_weights
+        # Restore the original initialization functions
+        for name, init_func in TORCH_INIT_FUNCTIONS.items():
+            setattr(torch.nn.init, name, init_func)
+
+
+def zero_module(module: nn.Module):
+    for p in module.parameters():
+        nn.init.zeros_(p)
+    return module

diffsynth_engine/models/wan/wan_dit.py

@@ -0,0 +1,497 @@
+import math
+import json
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+from typing import Tuple, Optional
+from einops import rearrange
+
+from diffsynth_engine.models.base import StateDictConverter, PreTrainedModel
+from diffsynth_engine.models.basic.attention import attention, long_context_attention
+from diffsynth_engine.models.utils import no_init_weights
+from diffsynth_engine.utils.constants import (
+    WAN_DIT_1_3B_T2V_CONFIG_FILE,
+    WAN_DIT_14B_I2V_CONFIG_FILE,
+    WAN_DIT_14B_T2V_CONFIG_FILE,
+)
+from diffsynth_engine.utils.gguf import gguf_inference
+from diffsynth_engine.utils.parallel import (
+    get_sp_group,
+    get_sp_world_size,
+    get_sp_rank,
+)
+
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor):
+    return x * (1 + scale) + shift
+
+
+def sinusoidal_embedding_1d(dim, position):
+    sinusoid = torch.outer(
+        position.type(torch.float64),
+        torch.pow(10000, -torch.arange(dim // 2, dtype=torch.float64, device=position.device).div(dim // 2)),
+    )
+    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+    return x.to(position.dtype)
+
+
+def precompute_freqs_cis_3d(dim: int, end: int = 1024, theta: float = 10000.0):
+    # 3d rope precompute
+    f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta)
+    h_freqs_cis = precompute_freqs_cis(dim // 3, end, theta)
+    w_freqs_cis = precompute_freqs_cis(dim // 3, end, theta)
+    return f_freqs_cis, h_freqs_cis, w_freqs_cis
+
+
+def precompute_freqs_cis(dim: int, end: int = 1024, theta: float = 10000.0):
+    # 1d rope precompute
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].double() / dim))
+    freqs = torch.outer(torch.arange(end, device=freqs.device), freqs)
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+
+def rope_apply(x, freqs):
+    x_out = torch.view_as_complex(x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2))
+    x_out = torch.view_as_real(x_out * freqs)
+    return x_out.to(x.dtype).flatten(3)
+
+
+class RMSNorm(nn.Module):
+    def __init__(
+        self,
+        dim,
+        eps=1e-5,
+        device: str = "cuda:0",
+        dtype: torch.dtype = torch.bfloat16,
+    ):
+        super().__init__()
+        self.eps = eps
+        self.dim = dim
+        self.weight = nn.Parameter(torch.ones(dim, device=device, dtype=dtype))
+
+    def norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        return self.norm(x.float()).to(x.dtype) * self.weight
+
+
+class SelfAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        eps: float = 1e-6,
+        attn_impl: Optional[str] = None,
+        device: str = "cuda:0",
+        dtype: torch.dtype = torch.bfloat16,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.head_dim = dim // num_heads
+        self.q = nn.Linear(dim, dim, device=device, dtype=dtype)
+        self.k = nn.Linear(dim, dim, device=device, dtype=dtype)
+        self.v = nn.Linear(dim, dim, device=device, dtype=dtype)
+        self.o = nn.Linear(dim, dim, device=device, dtype=dtype)
+        self.norm_q = RMSNorm(dim, eps=eps, device=device, dtype=dtype)
+        self.norm_k = RMSNorm(dim, eps=eps, device=device, dtype=dtype)
+        self.attn_impl = attn_impl
+
+    def forward(self, x, freqs):
+        q, k, v = self.norm_q(self.q(x)), self.norm_k(self.k(x)), self.v(x)
+        num_heads = q.shape[2] // self.head_dim
+        q = rearrange(q, "b s (n d) -> b s n d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b s n d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b s n d", n=num_heads)
+        if getattr(self, "use_usp", False):
+            x = long_context_attention(
+                q=rope_apply(q, freqs),
+                k=rope_apply(k, freqs),
+                v=v,
+                attn_impl=self.attn_impl,
+            )
+        else:
+            x = attention(
+                q=rope_apply(q, freqs),
+                k=rope_apply(k, freqs),
+                v=v,
+                attn_impl=self.attn_impl,
+            )
+        x = x.flatten(2)
+        return self.o(x)
+
+
+class CrossAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        eps: float = 1e-6,
+        has_image_input: bool = False,
+        attn_impl: Optional[str] = None,
+        device: str = "cuda:0",
+        dtype: torch.dtype = torch.bfloat16,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.head_dim = dim // num_heads
+
+        self.q = nn.Linear(dim, dim, device=device, dtype=dtype)
+        self.k = nn.Linear(dim, dim, device=device, dtype=dtype)
+        self.v = nn.Linear(dim, dim, device=device, dtype=dtype)
+        self.o = nn.Linear(dim, dim, device=device, dtype=dtype)
+        self.norm_q = RMSNorm(dim, eps=eps, device=device, dtype=dtype)
+        self.norm_k = RMSNorm(dim, eps=eps, device=device, dtype=dtype)
+        self.has_image_input = has_image_input
+        if has_image_input:
+            self.k_img = nn.Linear(dim, dim, device=device, dtype=dtype)
+            self.v_img = nn.Linear(dim, dim, device=device, dtype=dtype)
+            self.norm_k_img = RMSNorm(dim, eps=eps, device=device, dtype=dtype)
+        self.attn_impl = attn_impl
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        if self.has_image_input:
+            img = y[:, :257]
+            ctx = y[:, 257:]
+        else:
+            ctx = y
+        q, k, v = self.norm_q(self.q(x)), self.norm_k(self.k(ctx)), self.v(ctx)
+        num_heads = q.shape[2] // self.head_dim
+        q = rearrange(q, "b s (n d) -> b s n d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b s n d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b s n d", n=num_heads)
+
+        x = attention(q, k, v, attn_impl=self.attn_impl).flatten(2)
+        if self.has_image_input:
+            k_img, v_img = self.norm_k_img(self.k_img(img)), self.v_img(img)
+            k_img = rearrange(k_img, "b s (n d) -> b s n d", n=num_heads)
+            y = attention(q, k_img, v_img, attn_impl=self.attn_impl).flatten(2)
+            x = x + y
+        return self.o(x)
+
+
+class DiTBlock(nn.Module):
+    def __init__(
+        self,
+        has_image_input: bool,
+        dim: int,
+        num_heads: int,
+        ffn_dim: int,
+        eps: float = 1e-6,
+        attn_impl: Optional[str] = None,
+        device: str = "cuda:0",
+        dtype: torch.dtype = torch.bfloat16,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.ffn_dim = ffn_dim
+        self.self_attn = SelfAttention(dim, num_heads, eps, attn_impl=attn_impl, device=device, dtype=dtype)
+        self.cross_attn = CrossAttention(
+            dim, num_heads, eps, has_image_input=has_image_input, attn_impl=attn_impl, device=device, dtype=dtype
+        )
+        self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False, device=device, dtype=dtype)
+        self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False, device=device, dtype=dtype)
+        self.norm3 = nn.LayerNorm(dim, eps=eps, device=device, dtype=dtype)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, ffn_dim, device=device, dtype=dtype),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(ffn_dim, dim, device=device, dtype=dtype),
+        )
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim, device=device, dtype=dtype) / dim**0.5)
+
+    def forward(self, x, context, t_mod, freqs):
+        # msa: multi-head self-attention  mlp: multi-layer perceptron
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.modulation + t_mod).chunk(6, dim=1)
+        input_x = modulate(self.norm1(x), shift_msa, scale_msa)
+        x = x + gate_msa * self.self_attn(input_x, freqs)
+        x = x + self.cross_attn(self.norm3(x), context)
+        input_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
+        x = x + gate_mlp * self.ffn(input_x)
+        return x
+
+
+class MLP(torch.nn.Module):
+    def __init__(
+        self,
+        in_dim,
+        out_dim,
+        device: str = "cuda:0",
+        dtype: torch.dtype = torch.bfloat16,
+    ):
+        super().__init__()
+        self.proj = torch.nn.Sequential(
+            nn.LayerNorm(in_dim, device=device, dtype=dtype),
+            nn.Linear(in_dim, in_dim, device=device, dtype=dtype),
+            nn.GELU(),
+            nn.Linear(in_dim, out_dim, device=device, dtype=dtype),
+            nn.LayerNorm(out_dim, device=device, dtype=dtype),
+        )
+
+    def forward(self, x):
+        return self.proj(x)
+
+
+class Head(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        out_dim: int,
+        patch_size: Tuple[int, int, int],
+        eps: float,
+        device: str = "cuda:0",
+        dtype: torch.dtype = torch.bfloat16,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.patch_size = patch_size
+        self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False, device=device, dtype=dtype)
+        self.head = nn.Linear(dim, out_dim * math.prod(patch_size), device=device, dtype=dtype)
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim, device=device, dtype=dtype) / dim**0.5)
+
+    def forward(self, x, t_mod):
+        shift, scale = (self.modulation + t_mod).chunk(2, dim=1)
+        x = self.head(self.norm(x) * (1 + scale) + shift)
+        return x
+
+
+class WanDiTStateDictConverter(StateDictConverter):
+    def convert(self, state_dict):
+        return state_dict
+
+
+class WanDiT(PreTrainedModel):
+    converter = WanDiTStateDictConverter()
+
+    def __init__(
+        self,
+        dim: int,
+        in_dim: int,
+        ffn_dim: int,
+        out_dim: int,
+        text_dim: int,
+        freq_dim: int,
+        eps: float,
+        patch_size: Tuple[int, int, int],
+        num_heads: int,
+        num_layers: int,
+        has_image_input: bool,
+        attn_impl: Optional[str] = None,
+        use_usp: bool = False,
+        device: str = "cpu",
+        dtype: torch.dtype = torch.bfloat16,
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.freq_dim = freq_dim
+        self.has_image_input = has_image_input
+        self.patch_size = patch_size
+
+        self.patch_embedding = nn.Conv3d(
+            in_dim, dim, kernel_size=patch_size, stride=patch_size, device=device, dtype=dtype
+        )
+        self.text_embedding = nn.Sequential(
+            nn.Linear(text_dim, dim, device=device, dtype=dtype),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(dim, dim, device=device, dtype=dtype),
+        )
+        self.time_embedding = nn.Sequential(
+            nn.Linear(freq_dim, dim, device=device, dtype=dtype),
+            nn.SiLU(),
+            nn.Linear(dim, dim, device=device, dtype=dtype),
+        )
+
+        self.time_projection = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim, dim * 6, device=device, dtype=dtype),
+        )
+        self.blocks = nn.ModuleList(
+            [
+                DiTBlock(has_image_input, dim, num_heads, ffn_dim, eps, attn_impl, device=device, dtype=dtype)
+                for _ in range(num_layers)
+            ]
+        )
+        self.head = Head(dim, out_dim, patch_size, eps, device=device, dtype=dtype)
+
+        head_dim = dim // num_heads
+        self.freqs = precompute_freqs_cis_3d(head_dim)
+
+        if has_image_input:
+            self.img_emb = MLP(1280, dim, device=device, dtype=dtype)  # clip_feature_dim = 1280
+
+        if use_usp:
+            setattr(self, "use_usp", True)
+            for block in self.blocks:
+                setattr(block.self_attn, "use_usp", True)
+
+    def patchify(self, x: torch.Tensor):
+        x = self.patch_embedding(x)  # b c f h w -> b 4c f h/2 w/2
+        grid_size = x.shape[2:]
+        x = rearrange(x, "b c f h w -> b (f h w) c").contiguous()
+        return x, grid_size  # x, grid_size: (f, h, w)
+
+    def unpatchify(self, x: torch.Tensor, grid_size: torch.Tensor):
+        return rearrange(
+            x,
+            "b (f h w) (x y z c) -> b c (f x) (h y) (w z)",
+            f=grid_size[0],
+            h=grid_size[1],
+            w=grid_size[2],
+            x=self.patch_size[0],
+            y=self.patch_size[1],
+            z=self.patch_size[2],
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor,
+        timestep: torch.Tensor,
+        clip_feature: Optional[torch.Tensor] = None,  # clip_vision_encoder(img)
+        y: Optional[torch.Tensor] = None,  # vae_encoder(img)
+    ):
+        with gguf_inference():
+            t = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, timestep))
+            t_mod = self.time_projection(t).unflatten(1, (6, self.dim))
+            context = self.text_embedding(context)
+            if self.has_image_input:
+                x = torch.cat([x, y], dim=1)  # (b, c_x + c_y, f, h, w)
+                clip_embdding = self.img_emb(clip_feature)
+                context = torch.cat([clip_embdding, context], dim=1)  # (b, s1 + s2, d)
+            x, (f, h, w) = self.patchify(x)
+            freqs = (
+                torch.cat(
+                    [
+                        self.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+                        self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+                        self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1),
+                    ],
+                    dim=-1,
+                )
+                .reshape(f * h * w, 1, -1)
+                .to(x.device)
+            )
+            if getattr(self, "use_usp", False):
+                s, p = x.size(1), get_sp_world_size()  # (sequence_length, parallelism)
+                split_size = [s // p + 1 if i < s % p else s // p for i in range(p)]
+                x = torch.split(x, split_size, dim=1)[get_sp_rank()]
+                freqs = torch.split(freqs, split_size, dim=0)[get_sp_rank()]
+
+            for block in self.blocks:
+                x = block(x, context, t_mod, freqs)
+            x = self.head(x, t)
+
+            if getattr(self, "use_usp", False):
+                b, d = x.size(0), x.size(2)  # (batch_size, out_dim)
+                xs = [torch.zeros((b, s, d), dtype=x.dtype, device=x.device) for s in split_size]
+                dist.all_gather(xs, x, group=get_sp_group())
+                x = torch.concat(xs, dim=1)
+            x = self.unpatchify(x, (f, h, w))
+            return x
+
+    @classmethod
+    def from_state_dict(
+        cls,
+        state_dict,
+        device,
+        dtype,
+        model_type="1.3b-t2v",
+        attn_impl: Optional[str] = None,
+        use_usp=False,
+        assign=True,
+    ):
+        if model_type == "1.3b-t2v":
+            config = json.load(open(WAN_DIT_1_3B_T2V_CONFIG_FILE, "r"))
+        elif model_type == "14b-t2v":
+            config = json.load(open(WAN_DIT_14B_T2V_CONFIG_FILE, "r"))
+        elif model_type == "14b-i2v":
+            config = json.load(open(WAN_DIT_14B_I2V_CONFIG_FILE, "r"))
+        else:
+            raise ValueError(f"Unsupported model type: {model_type}")
+        with no_init_weights():
+            model = torch.nn.utils.skip_init(
+                cls, **config, device=device, dtype=dtype, attn_impl=attn_impl, use_usp=use_usp
+            )
+            model = model.requires_grad_(False)
+        model.load_state_dict(state_dict, assign=assign)
+        model.to(device=device, dtype=dtype)
+        return model
+
+    def get_tp_plan(self):
+        from torch.distributed.tensor.parallel import (
+            ColwiseParallel,
+            RowwiseParallel,
+            PrepareModuleInput,
+            PrepareModuleOutput,
+        )
+        from torch.distributed.tensor import Replicate, Shard
+
+        tp_plan = {
+            "text_embedding.0": ColwiseParallel(),
+            "text_embedding.2": RowwiseParallel(),
+            "time_embedding.0": ColwiseParallel(),
+            "time_embedding.2": RowwiseParallel(),
+            "time_projection.1": ColwiseParallel(output_layouts=Replicate()),
+            "blocks.0": PrepareModuleInput(
+                input_layouts=(Replicate(), None, None, None),
+                desired_input_layouts=(Shard(1), None, None, None),  # sequence parallel
+                use_local_output=True,
+            ),
+            "head": PrepareModuleOutput(
+                output_layouts=Shard(1),
+                desired_output_layouts=Replicate(),
+                use_local_output=True,
+            ),
+        }
+        for idx in range(len(self.blocks)):
+            tp_plan.update(
+                {
+                    f"blocks.{idx}.self_attn": PrepareModuleInput(
+                        input_layouts=(Shard(1), None),
+                        desired_input_layouts=(Replicate(), None),
+                    ),
+                    f"blocks.{idx}.self_attn.q": ColwiseParallel(output_layouts=Shard(1)),
+                    f"blocks.{idx}.self_attn.k": ColwiseParallel(output_layouts=Shard(1)),
+                    f"blocks.{idx}.self_attn.v": ColwiseParallel(),
+                    f"blocks.{idx}.self_attn.o": RowwiseParallel(output_layouts=Shard(1)),
+                    f"blocks.{idx}.self_attn.norm_q": PrepareModuleOutput(
+                        output_layouts=Shard(1),
+                        desired_output_layouts=Shard(-1),
+                    ),
+                    f"blocks.{idx}.self_attn.norm_k": PrepareModuleOutput(
+                        output_layouts=Shard(1),
+                        desired_output_layouts=Shard(-1),
+                    ),
+                    f"blocks.{idx}.cross_attn": PrepareModuleInput(
+                        input_layouts=(Shard(1), None),
+                        desired_input_layouts=(Replicate(), None),
+                    ),
+                    f"blocks.{idx}.cross_attn.q": ColwiseParallel(output_layouts=Shard(1)),
+                    f"blocks.{idx}.cross_attn.k": ColwiseParallel(output_layouts=Shard(1)),
+                    f"blocks.{idx}.cross_attn.v": ColwiseParallel(),
+                    f"blocks.{idx}.cross_attn.o": RowwiseParallel(output_layouts=Shard(1)),
+                    f"blocks.{idx}.cross_attn.norm_q": PrepareModuleOutput(
+                        output_layouts=Shard(1),
+                        desired_output_layouts=Shard(-1),
+                    ),
+                    f"blocks.{idx}.cross_attn.norm_k": PrepareModuleOutput(
+                        output_layouts=Shard(1),
+                        desired_output_layouts=Shard(-1),
+                    ),
+                    f"blocks.{idx}.cross_attn.k_img": ColwiseParallel(output_layouts=Shard(1)),
+                    f"blocks.{idx}.cross_attn.v_img": ColwiseParallel(),
+                    f"blocks.{idx}.cross_attn.norm_k_img": PrepareModuleOutput(
+                        output_layouts=Shard(1),
+                        desired_output_layouts=Shard(-1),
+                    ),
+                    f"blocks.{idx}.ffn": PrepareModuleInput(
+                        input_layouts=(Shard(1),),
+                        desired_input_layouts=(Replicate(),),
+                    ),
+                    f"blocks.{idx}.ffn.0": ColwiseParallel(),
+                    f"blocks.{idx}.ffn.2": RowwiseParallel(output_layouts=Shard(1)),
+                }
+            )
+        return tp_plan