diffsynth-engine 0.0.0__py3-none-any.whl

diffsynth_engine/models/__init__.py

@@ -0,0 +1,7 @@
+from .base import PreTrainedModel, StateDictConverter
+
+
+__all__ = [
+    "PreTrainedModel",
+    "StateDictConverter",
+]

diffsynth_engine/models/base.py

@@ -0,0 +1,64 @@
+import os
+import torch
+import torch.nn as nn
+from typing import Dict, List, Union
+from safetensors.torch import load_file
+
+from diffsynth_engine.models.basic.lora import LoRALinear, LoRAConv2d
+from diffsynth_engine.models.utils import no_init_weights
+
+
+class StateDictConverter:
+    def convert(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        return state_dict
+
+
+class PreTrainedModel(nn.Module):
+    converter = StateDictConverter()
+
+    def load_state_dict(self, state_dict: Dict[str, torch.Tensor], strict: bool = True, assign: bool = False):
+        state_dict = self.converter.convert(state_dict)
+        super().load_state_dict(state_dict, strict=strict, assign=assign)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_path: Union[str, os.PathLike], device: str, dtype: torch.dtype, **kwargs):
+        state_dict = load_file(pretrained_model_path, device=device)
+        return cls.from_state_dict(state_dict, device=device, dtype=dtype, **kwargs)
+
+    @classmethod
+    def from_state_dict(cls, state_dict: Dict[str, torch.Tensor], device: str, dtype: torch.dtype, **kwargs):
+        with no_init_weights():
+            model = torch.nn.utils.skip_init(cls, device=device, dtype=dtype, **kwargs)
+        model.load_state_dict(state_dict)
+        model.to(device=device, dtype=dtype, non_blocking=True)
+        return model
+
+    def load_loras(self, lora_args: List[Dict[str, any]], fused: bool = True):
+        for args in lora_args:
+            key = args["name"]
+            module = self.get_submodule(key)
+            if not isinstance(module, (LoRALinear, LoRAConv2d)):
+                raise ValueError(f"Unsupported lora key: {key}")
+            if fused:
+                module.add_frozen_lora(**args)
+            else:
+                module.add_lora(**args)
+
+    def unload_loras(self):
+        for module in self.modules():
+            if isinstance(module, (LoRALinear, LoRAConv2d)):
+                module.clear()
+
+
+def split_suffix(name: str):
+    suffix_list = [
+        ".lora_up.weight",
+        ".lora_down.weight",
+        ".weight",
+        ".bias",
+        ".alpha",
+    ]
+    for suffix in suffix_list:
+        if name.endswith(suffix):
+            return name.replace(suffix, ""), suffix
+    return name, ""

diffsynth_engine/models/basic/attention.py

@@ -0,0 +1,217 @@
+import torch
+import torch.nn as nn
+from einops import rearrange
+from typing import Optional
+from yunchang import LongContextAttention
+from yunchang.kernels import AttnType
+
+from diffsynth_engine.utils import logging
+from diffsynth_engine.utils.flag import (
+    FLASH_ATTN_3_AVAILABLE,
+    FLASH_ATTN_2_AVAILABLE,
+    XFORMERS_AVAILABLE,
+    SDPA_AVAILABLE,
+    SAGE_ATTN_AVAILABLE,
+    SPARGE_ATTN_AVAILABLE,
+)
+
+logger = logging.get_logger(__name__)
+
+
+if FLASH_ATTN_3_AVAILABLE:
+    from flash_attn_interface import flash_attn_func as flash_attn3
+if FLASH_ATTN_2_AVAILABLE:
+    from flash_attn import flash_attn_func as flash_attn2
+if XFORMERS_AVAILABLE:
+    from xformers.ops import memory_efficient_attention as xformers_attn
+if SDPA_AVAILABLE:
+
+    def sdpa_attn(q, k, v, attn_mask=None, scale=None):
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, scale=scale)
+        return out.transpose(1, 2)
+
+
+if SAGE_ATTN_AVAILABLE:
+    from sageattention import sageattn
+
+    def sage_attn(q, k, v, attn_mask=None, scale=None):
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        out = sageattn(q, k, v, attn_mask=attn_mask, sm_scale=scale)
+        return out.transpose(1, 2)
+
+
+if SPARGE_ATTN_AVAILABLE:
+    from spas_sage_attn import spas_sage2_attn_meansim_cuda
+
+    def sparge_attn(self, q, k, v, attn_mask=None, scale=None):
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        out = spas_sage2_attn_meansim_cuda(q, k, v, attn_mask=attn_mask, scale=scale)
+        return out.transpose(1, 2)
+
+
+def eager_attn(q, k, v, attn_mask=None, scale=None):
+    q = q.transpose(1, 2)
+    k = k.transpose(1, 2)
+    v = v.transpose(1, 2)
+    scale = 1 / q.shape[-1] ** 0.5 if scale is None else scale
+    q = q * scale
+    attn = torch.matmul(q, k.transpose(-2, -1))
+    if attn_mask is not None:
+        attn = attn + attn_mask
+    attn = attn.softmax(-1)
+    out = attn @ v
+    return out.transpose(1, 2)
+
+
+def attention(
+    q,
+    k,
+    v,
+    attn_impl: Optional[str] = None,
+    attn_mask: Optional[torch.Tensor] = None,
+    scale: Optional[float] = None,
+):
+    """
+    q: [B, Lq, Nq, C1]
+    k: [B, Lk, Nk, C1]
+    v: [B, Lk, Nk, C2]
+    """
+    assert attn_impl in [
+        None,
+        "auto",
+        "eager",
+        "flash_attn_2",
+        "flash_attn_3",
+        "xformers",
+        "sdpa",
+        "sage_attn",
+        "sparge_attn",
+    ]
+    if attn_impl is None or attn_impl == "auto":
+        if FLASH_ATTN_3_AVAILABLE:
+            return flash_attn3(q, k, v, softmax_scale=scale)
+        elif FLASH_ATTN_2_AVAILABLE:
+            return flash_attn2(q, k, v, softmax_scale=scale)
+        elif XFORMERS_AVAILABLE:
+            return xformers_attn(q, k, v, attn_bias=attn_mask, scale=scale)
+        elif SDPA_AVAILABLE:
+            return sdpa_attn(q, k, v, attn_mask=attn_mask, scale=scale)
+        else:
+            return eager_attn(q, k, v, attn_mask=attn_mask, scale=scale)
+    else:
+        if attn_impl == "eager":
+            return eager_attn(q, k, v, attn_mask=attn_mask, scale=scale)
+        elif attn_impl == "flash_attn_3":
+            return flash_attn3(q, k, v, softmax_scale=scale)
+        elif attn_impl == "flash_attn_2":
+            return flash_attn2(q, k, v, softmax_scale=scale)
+        elif attn_impl == "xformers":
+            return xformers_attn(q, k, v, attn_bias=attn_mask, scale=scale)
+        elif attn_impl == "sdpa":
+            return sdpa_attn(q, k, v, attn_mask=attn_mask, scale=scale)
+        elif attn_impl == "sage_attn":
+            return sage_attn(q, k, v, attn_mask=attn_mask, scale=scale)
+        elif attn_impl == "sparge_attn":
+            return sparge_attn(q, k, v, attn_mask=attn_mask, scale=scale)
+        else:
+            raise ValueError(f"Invalid attention implementation: {attn_impl}")
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        q_dim,
+        num_heads,
+        head_dim,
+        kv_dim=None,
+        bias_q=False,
+        bias_kv=False,
+        bias_out=False,
+        scale=None,
+        attn_impl: Optional[str] = None,
+        device: str = "cuda:0",
+        dtype: torch.dtype = torch.float16,
+    ):
+        super().__init__()
+        dim_inner = head_dim * num_heads
+        kv_dim = kv_dim if kv_dim is not None else q_dim
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+
+        self.to_q = nn.Linear(q_dim, dim_inner, bias=bias_q, device=device, dtype=dtype)
+        self.to_k = nn.Linear(kv_dim, dim_inner, bias=bias_kv, device=device, dtype=dtype)
+        self.to_v = nn.Linear(kv_dim, dim_inner, bias=bias_kv, device=device, dtype=dtype)
+        self.to_out = nn.Linear(dim_inner, q_dim, bias=bias_out, device=device, dtype=dtype)
+        self.attn_impl = attn_impl
+        self.scale = scale
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        y: Optional[torch.Tensor] = None,
+        attn_mask: Optional[torch.Tensor] = None,
+    ):
+        if y is None:
+            y = x
+        q = rearrange(self.to_q(x), "b s (n d) -> b s n d", n=self.num_heads)
+        k = rearrange(self.to_k(y), "b s (n d) -> b s n d", n=self.num_heads)
+        v = rearrange(self.to_v(y), "b s (n d) -> b s n d", n=self.num_heads)
+        out = attention(q, k, v, attn_mask=attn_mask, attn_impl=self.attn_impl, scale=self.scale)
+        out = rearrange(out, "b s n d -> b s (n d)", n=self.num_heads)
+        return self.to_out(out)
+
+
+def long_context_attention(
+    q,
+    k,
+    v,
+    attn_impl: Optional[str] = None,
+    attn_mask: Optional[torch.Tensor] = None,
+    scale: Optional[float] = None,
+):
+    """
+    q: [B, Lq, Nq, C1]
+    k: [B, Lk, Nk, C1]
+    v: [B, Lk, Nk, C2]
+    """
+    assert attn_impl in [
+        None,
+        "auto",
+        "eager",
+        "flash_attn_2",
+        "flash_attn_3",
+        "xformers",
+        "sdpa",
+        "sage_attn",
+        "sparge_attn",
+    ]
+    if attn_impl is None or attn_impl == "auto":
+        if FLASH_ATTN_3_AVAILABLE:
+            attn_func = LongContextAttention(attn_type=AttnType.FA3)
+        elif FLASH_ATTN_2_AVAILABLE:
+            attn_func = LongContextAttention(attn_type=AttnType.FA)
+        elif SDPA_AVAILABLE:
+            attn_func = LongContextAttention(attn_type=AttnType.TORCH)
+        else:
+            raise ValueError("No available long context attention implementation")
+    else:
+        if attn_impl == "flash_attn_3":
+            attn_func = LongContextAttention(attn_type=AttnType.FA3)
+        elif attn_impl == "flash_attn_2":
+            attn_func = LongContextAttention(attn_type=AttnType.FA)
+        elif attn_impl == "sdpa":
+            attn_func = LongContextAttention(attn_type=AttnType.TORCH)
+        elif attn_impl == "sage_attn":
+            attn_func = LongContextAttention(attn_type=AttnType.SAGE_FP8)
+        elif attn_impl == "sparge_attn":
+            attn_func = LongContextAttention(attn_type=AttnType.SPARSE_SAGE)
+        else:
+            raise ValueError(f"Invalid long context attention implementation: {attn_impl}")
+    return attn_func(q, k, v, softmax_scale=scale)

diffsynth_engine/models/basic/lora.py

@@ -0,0 +1,293 @@
+import torch
+import torch.nn as nn
+from torch.nn.common_types import _size_2_t
+from typing import Union
+from collections import OrderedDict
+from contextlib import contextmanager
+
+
+class LoRA(nn.Module):
+    def __init__(
+        self,
+        scale: float,
+        rank: int,
+        alpha: int,
+        up: Union[nn.Linear, nn.Conv2d, torch.Tensor],
+        down: Union[nn.Linear, nn.Conv2d, torch.Tensor],
+        device: str,
+        dtype: torch.dtype,
+    ):
+        super().__init__()
+        self.device = device
+        self.dtype = dtype
+        self.scale = scale
+        self.rank = rank
+        self.alpha = alpha.item() if isinstance(alpha, torch.Tensor) else alpha
+        self.up = up.to(device=device, dtype=dtype)
+        self.down = down.to(device=device, dtype=dtype)
+
+    def forward(self, x):
+        if isinstance(self.up, torch.Tensor) and isinstance(self.down, torch.Tensor):
+            return self.scale * (self.alpha / self.rank) * (x @ self.down.T @ self.up.T)
+        return self.scale * (self.alpha / self.rank) * (self.up(self.down(x)))
+
+    def apply_to(self, w: Union[nn.Linear, nn.Conv2d, nn.Parameter, torch.Tensor]):
+        if isinstance(self.up, torch.Tensor) and isinstance(self.down, torch.Tensor):
+            delta_w = self.scale * (self.alpha / self.rank) * (self.up @ self.down)
+        else:
+            delta_w = self.scale * (self.alpha / self.rank) * (self.up.weight @ self.down.weight)
+        if isinstance(w, (nn.Linear, nn.Conv2d)):
+            delta_w = delta_w.to(device=w.weight.data.device, dtype=w.weight.data.dtype)
+            w.weight.data.add_(delta_w)
+        elif isinstance(w, nn.Parameter):
+            delta_w = delta_w.to(device=w.data.device, dtype=w.data.dtype)
+            w.data.add_(delta_w)
+        elif isinstance(w, torch.Tensor):
+            delta_w = delta_w.to(device=w.device, dtype=w.dtype)
+            w.add_(delta_w)
+
+
+class LoRALinear(nn.Linear):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        device=None,
+        dtype=None,
+    ) -> None:
+        super().__init__(in_features, out_features, bias, device, dtype)
+        # LoRA
+        self._lora_dict = OrderedDict()
+        # Frozen LoRA
+        self._frozen_lora_list = []
+        self.register_buffer("_original_weight", None)
+
+    @staticmethod
+    def from_linear(linear: nn.Linear):
+        lora_linear = torch.nn.utils.skip_init(
+            LoRALinear,
+            linear.in_features,
+            linear.out_features,
+            linear.bias is not None,
+            device=linear.weight.device,
+            dtype=linear.weight.dtype,
+        )
+        lora_linear.weight = linear.weight
+        lora_linear.bias = linear.bias
+        return lora_linear
+
+    def add_lora(
+        self,
+        name: str,
+        scale: float,
+        rank: int,
+        alpha: int,
+        up: torch.Tensor,
+        down: torch.Tensor,
+        device: str,
+        dtype: torch.dtype,
+        **kwargs,
+    ):
+        up_linear = torch.nn.utils.skip_init(
+            nn.Linear, up.shape[1], up.shape[0], bias=False, device=device, dtype=dtype
+        )
+        down_linear = torch.nn.utils.skip_init(
+            nn.Linear, down.shape[0], down.shape[1], bias=False, device=device, dtype=dtype
+        )
+        up_linear.weight.data = up
+        down_linear.weight.data = down
+        lora = LoRA(scale, rank, alpha, up_linear, down_linear, device, dtype)
+        self._lora_dict[name] = lora
+
+    def modify_scale(self, name: str, scale: float):
+        if name not in self._lora_dict:
+            raise ValueError(f"LoRA name {name} not found in LoRALinear {self.__class__.__name__}")
+        self._lora_dict[name].scale = scale
+
+    def add_frozen_lora(
+        self,
+        name: str,
+        scale: float,
+        rank: int,
+        alpha: int,
+        up: torch.Tensor,
+        down: torch.Tensor,
+        device: str,
+        dtype: torch.dtype,
+        save_original_weight: bool = True,
+    ):
+        if save_original_weight and self._original_weight is None:
+            self._original_weight = self.weight.clone()
+        lora = LoRA(scale, rank, alpha, up, down, device, dtype)
+        lora.apply_to(self)
+        self._frozen_lora_list.append(lora)
+
+    def clear(self):
+        if self._original_weight is None and len(self._frozen_lora_list) > 0:
+            raise RuntimeError(
+                "Current LoRALinear has patched by frozen LoRA, but original weight is not saved, so you cannot clear LoRA."
+            )
+        self._lora_dict.clear()
+        self._frozen_lora_list = []
+        if self._original_weight is not None:
+            self.weight.data = self._original_weight
+            self._original_weight = None
+
+    def forward(self, x):
+        w_x = super().forward(x)
+        for name, lora in self._lora_dict.items():
+            w_x += lora(x)
+        return w_x
+
+
+class LoRAConv2d(nn.Conv2d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: _size_2_t,
+        stride: _size_2_t = 1,
+        padding: Union[str, _size_2_t] = 0,
+        dilation: _size_2_t = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: str = "zeros",  # TODO: refine this type
+        device=None,
+        dtype=None,
+    ) -> None:
+        super().__init__(
+            in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias, padding_mode, device, dtype
+        )
+        # LoRA
+        self._lora_dict = OrderedDict()
+        # Frozen LoRA
+        self._frozen_lora_list = []
+        self._original_weight = None
+
+    @staticmethod
+    def from_conv2d(conv2d: nn.Conv2d):
+        lora_conv2d = torch.nn.utils.skip_init(
+            LoRAConv2d,
+            conv2d.in_channels,
+            conv2d.out_channels,
+            conv2d.kernel_size,
+            conv2d.stride,
+            conv2d.padding,
+            conv2d.dilation,
+            conv2d.groups,
+            conv2d.bias is not None,
+            conv2d.padding_mode,
+            device=conv2d.weight.device,
+            dtype=conv2d.weight.dtype,
+        )
+        lora_conv2d.weight = conv2d.weight
+        lora_conv2d.bias = conv2d.bias
+        return lora_conv2d
+
+    def _construct_lora(
+        self,
+        name: str,
+        scale: float,
+        rank: int,
+        alpha: int,
+        up: torch.Tensor,
+        down: torch.Tensor,
+        device: str,
+        dtype: torch.dtype,
+    ):
+        down_conv = torch.nn.utils.skip_init(
+            nn.Conv2d,
+            self.in_channels,
+            rank,
+            kernel_size=self.kernel_size,
+            stride=self.stride,
+            padding=self.padding,
+            bias=False,
+            device=device,
+            dtype=dtype,
+        )
+        down_conv.weight.data = down
+        # according to the official kohya_ss trainer kernel_size are always fixed for the up layer
+        # see: https://github.com/bmaltais/kohya_ss/blob/2accb1305979ba62f5077a23aabac23b4c37e935/networks/lora_diffusers.py#L129
+        # refer from diffusers
+        up_conv = torch.nn.utils.skip_init(
+            nn.Conv2d,
+            rank,
+            self.out_channels,
+            kernel_size=(1, 1),
+            stride=(1, 1),
+            bias=False,
+            device=device,
+            dtype=dtype,
+        )
+        up_conv.weight.data = up
+
+        lora = LoRA(scale, rank, alpha, up_conv, down_conv, device, dtype)
+        return lora
+
+    def add_lora(
+        self,
+        name: str,
+        scale: float,
+        rank: int,
+        alpha: int,
+        up: torch.Tensor,
+        down: torch.Tensor,
+        device: str,
+        dtype: torch.dtype,
+        **kwargs,
+    ):
+        self._lora_dict[name] = self._construct_lora(name, scale, rank, alpha, up, down, device, dtype)
+
+    def modify_scale(self, name: str, scale: float):
+        if name not in self._lora_dict:
+            raise ValueError(f"LoRA name {name} not found in LoRAConv2d {self.__class__.__name__}")
+        self._lora_dict[name].scale = scale
+
+    def add_frozen_lora(
+        self,
+        name: str,
+        scale: float,
+        rank: int,
+        alpha: int,
+        up: torch.Tensor,
+        down: torch.Tensor,
+        device: str,
+        dtype: torch.dtype,
+        save_original_weight: bool = True,
+    ):
+        if save_original_weight and self._original_weight is None:
+            self._original_weight = self.weight.clone()
+        lora = self._construct_lora(name, scale, rank, alpha, up, down, device, dtype)
+        lora.apply_to(self)
+        self._frozen_lora_list.append(lora)
+
+    def clear(self):
+        if self._original_weight is None and len(self._frozen_lora_list) > 0:
+            raise RuntimeError(
+                "Current LoRALinear has patched by frozen LoRA, but original weight is not saved, so you cannot clear LoRA."
+            )
+        self._lora_dict.clear()
+        self._frozen_lora_list = []
+        if self._original_weight is not None:
+            self.weight.copy_(self._original_weight)
+            self._original_weight = None
+
+    def forward(self, x):
+        w_x = super().forward(x)
+        for name, lora in self._lora_dict.items():
+            w_x += lora(x)
+        return w_x
+
+
+@contextmanager
+def LoRAContext():
+    origin_linear = torch.nn.Linear
+    origin_conv2d = torch.nn.Conv2d
+
+    torch.nn.Linear = LoRALinear
+    torch.nn.Conv2d = LoRAConv2d
+    yield
+    torch.nn.Linear = origin_linear
+    torch.nn.Conv2d = origin_conv2d

diffsynth_engine/models/basic/relative_position_emb.py

@@ -0,0 +1,56 @@
+import torch
+import torch.nn as nn
+import math
+
+
+def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128):
+    relative_buckets = 0
+    if bidirectional:
+        num_buckets //= 2
+        relative_buckets += (relative_position > 0).to(torch.long) * num_buckets
+        relative_position = torch.abs(relative_position)
+    else:
+        relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))
+    # now relative_position is in the range [0, inf)
+
+    # half of the buckets are for exact increments in positions
+    max_exact = num_buckets // 2
+    is_small = relative_position < max_exact
+
+    # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance
+    relative_position_if_large = max_exact + (
+        torch.log(relative_position.float() / max_exact)
+        / math.log(max_distance / max_exact)
+        * (num_buckets - max_exact)
+    ).to(torch.long)
+    relative_position_if_large = torch.min(
+        relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1)
+    )
+
+    relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)
+    return relative_buckets
+
+
+class RelativePositionEmbedding(nn.Module):
+    def __init__(
+        self, num_buckets, max_distance, num_heads, device: str = "cuda:0", dtype: torch.dtype = torch.float16
+    ):
+        super().__init__()
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        self.relative_attention_bias = nn.Embedding(num_buckets, num_heads, device=device, dtype=dtype)
+
+    def forward(self, query_length, key_length):
+        device = self.relative_attention_bias.weight.device
+        context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
+        memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
+        relative_position = memory_position - context_position  # shape (query_length, key_length)
+        relative_position_bucket = _relative_position_bucket(
+            relative_position,  # shape (query_length, key_length)
+            bidirectional=True,
+            num_buckets=self.num_buckets,
+            max_distance=self.max_distance,
+        )
+        values = self.relative_attention_bias(relative_position_bucket)  # shape (query_length, key_length, num_heads)
+        values = values.permute([2, 0, 1]).unsqueeze(0)  # shape (1, num_heads, query_length, key_length)
+        return values