diffsynth-engine 0.1.0__py3-none-any.whl

diffsynth_engine/models/sd3/sd3_dit.py

@@ -0,0 +1,302 @@
+import json
+import torch
+import torch.nn as nn
+from typing import Dict
+from einops import rearrange
+
+from diffsynth_engine.models.basic.timestep import TimestepEmbeddings
+from diffsynth_engine.models.basic.transformer_helper import AdaLayerNorm
+from diffsynth_engine.models.base import PreTrainedModel, StateDictConverter
+from diffsynth_engine.models.utils import no_init_weights
+from diffsynth_engine.utils.constants import SD3_DIT_CONFIG_FILE
+from diffsynth_engine.utils import logging
+
+logger = logging.get_logger(__name__)
+
+with open(SD3_DIT_CONFIG_FILE, "r") as f:
+    config = json.load(f)
+
+
+class SD3DiTStateDictConverter(StateDictConverter):
+    def __init__(self):
+        pass
+
+    def _from_diffusers(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        rename_dict = config["diffusers"]["rename_dict"]
+        state_dict_ = {}
+        for name, param in state_dict.items():
+            if name in rename_dict:
+                if name == "pos_embed.pos_embed":
+                    param = param.reshape((1, 192, 192, 1536))
+                state_dict_[rename_dict[name]] = param
+            elif name.endswith(".weight") or name.endswith(".bias"):
+                suffix = ".weight" if name.endswith(".weight") else ".bias"
+                prefix = name[: -len(suffix)]
+                if prefix in rename_dict:
+                    state_dict_[rename_dict[prefix] + suffix] = param
+                elif prefix.startswith("transformer_blocks."):
+                    names = prefix.split(".")
+                    names[0] = "blocks"
+                    middle = ".".join(names[2:])
+                    if middle in rename_dict:
+                        name_ = ".".join(names[:2] + [rename_dict[middle]] + [suffix[1:]])
+                        state_dict_[name_] = param
+        return state_dict_
+
+    def _from_civitai(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        rename_dict = config["civitai"]["rename_dict"]
+        state_dict_ = {}
+        for name in state_dict:
+            if name in rename_dict:
+                param = state_dict[name]
+                if name.startswith("model.diffusion_model.joint_blocks.23.context_block.adaLN_modulation.1."):
+                    param = torch.concat([param[1536:], param[:1536]], axis=0)
+                elif name.startswith("model.diffusion_model.final_layer.adaLN_modulation.1."):
+                    param = torch.concat([param[1536:], param[:1536]], axis=0)
+                elif name == "model.diffusion_model.pos_embed":
+                    param = param.reshape((1, 192, 192, 1536))
+                if isinstance(rename_dict[name], str):
+                    state_dict_[rename_dict[name]] = param
+                else:
+                    name_ = rename_dict[name][0].replace(".a_to_q.", ".a_to_qkv.").replace(".b_to_q.", ".b_to_qkv.")
+                    state_dict_[name_] = param
+        return state_dict_
+
+    def convert(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        if "model.diffusion_model.context_embedder.weight" in state_dict:
+            state_dict = self._from_civitai(state_dict)
+            logger.info("use civitai format state dict")
+        elif "time_text_embed.timestep_embedder.linear_1" in state_dict:
+            state_dict = self._from_diffusers(state_dict)
+            logger.info("use diffusers format state dict")
+        else:
+            logger.info("use diffsynth format state dict")
+        return state_dict
+
+
+class PatchEmbed(nn.Module):
+    def __init__(
+        self,
+        patch_size=2,
+        in_channels=16,
+        embed_dim=1536,
+        pos_embed_max_size=192,
+        device: str = "cuda:0",
+        dtype: torch.dtype = torch.float16,
+    ):
+        super().__init__()
+        self.pos_embed_max_size = pos_embed_max_size
+        self.patch_size = patch_size
+
+        self.proj = nn.Conv2d(
+            in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, device=device, dtype=dtype
+        )
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, self.pos_embed_max_size, self.pos_embed_max_size, 1536, device=device, dtype=dtype)
+        )
+
+    def cropped_pos_embed(self, height, width):
+        height = height // self.patch_size
+        width = width // self.patch_size
+        top = (self.pos_embed_max_size - height) // 2
+        left = (self.pos_embed_max_size - width) // 2
+        spatial_pos_embed = self.pos_embed[:, top : top + height, left : left + width, :].flatten(1, 2)
+        return spatial_pos_embed
+
+    def forward(self, latent):
+        height, width = latent.shape[-2:]
+        latent = self.proj(latent)
+        latent = latent.flatten(2).transpose(1, 2)
+        pos_embed = self.cropped_pos_embed(height, width)
+        return latent + pos_embed
+
+
+class JointAttention(nn.Module):
+    def __init__(
+        self,
+        dim_a,
+        dim_b,
+        num_heads,
+        head_dim,
+        only_out_a=False,
+        device: str = "cuda:0",
+        dtype: torch.dtype = torch.float16,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.only_out_a = only_out_a
+
+        self.a_to_qkv = nn.Linear(dim_a, dim_a * 3, device=device, dtype=dtype)
+        self.b_to_qkv = nn.Linear(dim_b, dim_b * 3, device=device, dtype=dtype)
+
+        self.a_to_out = nn.Linear(dim_a, dim_a, device=device, dtype=dtype)
+        if not only_out_a:
+            self.b_to_out = nn.Linear(dim_b, dim_b, device=device, dtype=dtype)
+
+    def forward(self, hidden_states_a, hidden_states_b):
+        batch_size = hidden_states_a.shape[0]
+
+        qkv = torch.concat([self.a_to_qkv(hidden_states_a), self.b_to_qkv(hidden_states_b)], dim=1)
+        qkv = qkv.view(batch_size, -1, 3 * self.num_heads, self.head_dim).transpose(1, 2)
+        q, k, v = qkv.chunk(3, dim=1)
+
+        hidden_states = nn.functional.scaled_dot_product_attention(q, k, v)
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim)
+        hidden_states = hidden_states.to(q.dtype)
+        hidden_states_a, hidden_states_b = (
+            hidden_states[:, : hidden_states_a.shape[1]],
+            hidden_states[:, hidden_states_a.shape[1] :],
+        )
+        hidden_states_a = self.a_to_out(hidden_states_a)
+        if self.only_out_a:
+            return hidden_states_a
+        else:
+            hidden_states_b = self.b_to_out(hidden_states_b)
+            return hidden_states_a, hidden_states_b
+
+
+class JointTransformerBlock(nn.Module):
+    def __init__(self, dim, num_attention_heads, device: str = "cuda:0", dtype: torch.dtype = torch.float16):
+        super().__init__()
+        self.norm1_a = AdaLayerNorm(dim, device=device, dtype=dtype)
+        self.norm1_b = AdaLayerNorm(dim, device=device, dtype=dtype)
+
+        self.attn = JointAttention(
+            dim, dim, num_attention_heads, dim // num_attention_heads, device=device, dtype=dtype
+        )
+
+        self.norm2_a = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6, device=device, dtype=dtype)
+        self.ff_a = nn.Sequential(
+            nn.Linear(dim, dim * 4, device=device, dtype=dtype),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(dim * 4, dim, device=device, dtype=dtype),
+        )
+
+        self.norm2_b = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6, device=device, dtype=dtype)
+        self.ff_b = nn.Sequential(
+            nn.Linear(dim, dim * 4, device=device, dtype=dtype),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(dim * 4, dim, device=device, dtype=dtype),
+        )
+
+    def forward(self, hidden_states_a, hidden_states_b, temb):
+        norm_hidden_states_a, gate_msa_a, shift_mlp_a, scale_mlp_a, gate_mlp_a = self.norm1_a(hidden_states_a, emb=temb)
+        norm_hidden_states_b, gate_msa_b, shift_mlp_b, scale_mlp_b, gate_mlp_b = self.norm1_b(hidden_states_b, emb=temb)
+
+        # Attention
+        attn_output_a, attn_output_b = self.attn(norm_hidden_states_a, norm_hidden_states_b)
+
+        # Part A
+        hidden_states_a = hidden_states_a + gate_msa_a * attn_output_a
+        norm_hidden_states_a = self.norm2_a(hidden_states_a) * (1 + scale_mlp_a) + shift_mlp_a
+        hidden_states_a = hidden_states_a + gate_mlp_a * self.ff_a(norm_hidden_states_a)
+
+        # Part B
+        hidden_states_b = hidden_states_b + gate_msa_b * attn_output_b
+        norm_hidden_states_b = self.norm2_b(hidden_states_b) * (1 + scale_mlp_b) + shift_mlp_b
+        hidden_states_b = hidden_states_b + gate_mlp_b * self.ff_b(norm_hidden_states_b)
+
+        return hidden_states_a, hidden_states_b
+
+
+class JointTransformerFinalBlock(nn.Module):
+    def __init__(self, dim, num_attention_heads, device: str = "cuda:0", dtype: torch.dtype = torch.float16):
+        super().__init__()
+        self.norm1_a = AdaLayerNorm(dim, device=device, dtype=dtype)
+        self.norm1_b = AdaLayerNorm(dim, single=True, device=device, dtype=dtype)
+
+        self.attn = JointAttention(
+            dim, dim, num_attention_heads, dim // num_attention_heads, only_out_a=True, device=device, dtype=dtype
+        )
+
+        self.norm2_a = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6, device=device, dtype=dtype)
+        self.ff_a = nn.Sequential(
+            nn.Linear(dim, dim * 4, device=device, dtype=dtype),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(dim * 4, dim, device=device, dtype=dtype),
+        )
+
+    def forward(self, hidden_states_a, hidden_states_b, temb):
+        norm_hidden_states_a, gate_msa_a, shift_mlp_a, scale_mlp_a, gate_mlp_a = self.norm1_a(hidden_states_a, emb=temb)
+        norm_hidden_states_b = self.norm1_b(hidden_states_b, emb=temb)
+
+        # Attention
+        attn_output_a = self.attn(norm_hidden_states_a, norm_hidden_states_b)
+
+        # Part A
+        hidden_states_a = hidden_states_a + gate_msa_a * attn_output_a
+        norm_hidden_states_a = self.norm2_a(hidden_states_a) * (1 + scale_mlp_a) + shift_mlp_a
+        hidden_states_a = hidden_states_a + gate_mlp_a * self.ff_a(norm_hidden_states_a)
+
+        return hidden_states_a, hidden_states_b
+
+
+class SD3DiT(PreTrainedModel):
+    converter = SD3DiTStateDictConverter()
+
+    def __init__(self, device: str = "cuda:0", dtype: torch.dtype = torch.float16):
+        super().__init__()
+        self.pos_embedder = PatchEmbed(
+            patch_size=2, in_channels=16, embed_dim=1536, pos_embed_max_size=192, device=device, dtype=dtype
+        )
+        self.time_embedder = TimestepEmbeddings(256, 1536, device=device, dtype=dtype)
+        self.pooled_text_embedder = nn.Sequential(
+            nn.Linear(2048, 1536, device=device, dtype=dtype),
+            nn.SiLU(),
+            nn.Linear(1536, 1536, device=device, dtype=dtype),
+        )
+        self.context_embedder = nn.Linear(4096, 1536, device=device, dtype=dtype)
+        self.blocks = nn.ModuleList(
+            [JointTransformerBlock(1536, 24, device=device, dtype=dtype) for _ in range(23)]
+            + [JointTransformerFinalBlock(1536, 24, device=device, dtype=dtype)]
+        )
+        self.norm_out = AdaLayerNorm(1536, single=True, device=device, dtype=dtype)
+        self.proj_out = nn.Linear(1536, 64, device=device, dtype=dtype)
+
+    def forward(
+        self,
+        hidden_states,
+        timestep,
+        prompt_emb,
+        pooled_prompt_emb,
+        use_gradient_checkpointing=False,
+    ):
+        conditioning = self.time_embedder(timestep, hidden_states.dtype) + self.pooled_text_embedder(pooled_prompt_emb)
+        prompt_emb = self.context_embedder(prompt_emb)
+
+        height, width = hidden_states.shape[-2:]
+        hidden_states = self.pos_embedder(hidden_states)
+
+        def create_custom_forward(module):
+            def custom_forward(*inputs):
+                return module(*inputs)
+
+            return custom_forward
+
+        for block in self.blocks:
+            if self.training and use_gradient_checkpointing:
+                hidden_states, prompt_emb = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    prompt_emb,
+                    conditioning,
+                    use_reentrant=False,
+                )
+            else:
+                hidden_states, prompt_emb = block(hidden_states, prompt_emb, conditioning)
+
+        hidden_states = self.norm_out(hidden_states, conditioning)
+        hidden_states = self.proj_out(hidden_states)
+        hidden_states = rearrange(
+            hidden_states, "B (H W) (P Q C) -> B C (H P) (W Q)", P=2, Q=2, H=height // 2, W=width // 2
+        )
+        return hidden_states
+
+    @classmethod
+    def from_state_dict(cls, state_dict: Dict[str, torch.Tensor], device: str, dtype: torch.dtype):
+        with no_init_weights():
+            model = torch.nn.utils.skip_init(cls, device=device, dtype=dtype)
+        model.load_state_dict(state_dict, assign=True)
+        model.to(device=device, dtype=dtype, non_blocking=True)
+        return model

diffsynth_engine/models/sd3/sd3_text_encoder.py

@@ -0,0 +1,163 @@
+import json
+import torch
+from typing import Dict
+
+from diffsynth_engine.models.components.t5 import T5EncoderModel
+from diffsynth_engine.models.base import StateDictConverter
+from diffsynth_engine.models.sd import SDTextEncoder
+from diffsynth_engine.models.sdxl import SDXLTextEncoder2
+from diffsynth_engine.models.utils import no_init_weights
+from diffsynth_engine.utils.constants import SD3_TEXT_ENCODER_CONFIG_FILE
+from diffsynth_engine.utils import logging
+
+logger = logging.get_logger(__name__)
+
+with open(SD3_TEXT_ENCODER_CONFIG_FILE, "r") as f:
+    config = json.load(f)
+
+
+class SD3TextEncoder1StateDictConverter(StateDictConverter):
+    def __init__(self):
+        pass
+
+    def _from_diffusers(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        rename_dict = config["diffusers"]["te1_rename_dict"]
+        attn_rename_dict = config["diffusers"]["te1_attn_rename_dict"]
+        state_dict_ = {}
+        for name in state_dict:
+            if name in rename_dict:
+                param = state_dict[name]
+                if name == "text_model.embeddings.position_embedding.weight":
+                    param = param.reshape((1, param.shape[0], param.shape[1]))
+                state_dict_[rename_dict[name]] = param
+            elif name.startswith("text_model.encoder.layers."):
+                param = state_dict[name]
+                names = name.split(".")
+                layer_id, layer_type, tail = names[3], ".".join(names[4:-1]), names[-1]
+                name_ = ".".join(["encoders", layer_id, attn_rename_dict[layer_type], tail])
+                state_dict_[name_] = param
+        return state_dict_
+
+    def _from_civitai(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        rename_dict = config["civitai"]["te1_rename_dict"]
+        state_dict_ = {}
+        for name in state_dict:
+            if name in rename_dict:
+                param = state_dict[name]
+                if name == "text_encoders.clip_l.transformer.text_model.embeddings.position_embedding.weight":
+                    param = param.reshape((1, param.shape[0], param.shape[1]))
+                state_dict_[rename_dict[name]] = param
+        return state_dict_
+
+    def convert(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        if "text_encoders.clip_l.transformer.text_model.embeddings.position_embedding.weight" in state_dict:
+            state_dict = self._from_civitai(state_dict)
+            logger.info("use civitai format state dict")
+        elif "text_model.embeddings.token_embedding.weight" in state_dict:
+            state_dict = self._from_diffusers(state_dict)
+            logger.info("use diffusers format state dict")
+        else:
+            logger.info("use diffsynth format state dict")
+        return state_dict
+
+
+class SD3TextEncoder1(SDTextEncoder):
+    converter = SD3TextEncoder1StateDictConverter()
+
+    def __init__(self, vocab_size=49408, device: str = "cuda:0", dtype: torch.dtype = torch.float16):
+        super().__init__(vocab_size=vocab_size, device=device, dtype=dtype)
+
+    def forward(self, input_ids, clip_skip=2):
+        embeds = self.token_embedding(input_ids) + self.position_embeds
+        attn_mask = self.attn_mask.to(device=embeds.device, dtype=embeds.dtype)
+        for encoder_id, encoder in enumerate(self.encoders):
+            embeds = encoder(embeds, attn_mask=attn_mask)
+            if encoder_id + clip_skip == len(self.encoders):
+                hidden_states = embeds
+        embeds = self.final_layer_norm(embeds)
+        pooled_embeds = embeds[torch.arange(embeds.shape[0]), input_ids.to(dtype=torch.int).argmax(dim=-1)]
+        return hidden_states, pooled_embeds
+
+    @classmethod
+    def from_state_dict(
+        cls,
+        state_dict: Dict[str, torch.Tensor],
+        device: str,
+        dtype: torch.dtype,
+        vocab_size: int = 49408,
+    ):
+        with no_init_weights():
+            model = torch.nn.utils.skip_init(cls, device=device, dtype=dtype, vocab_size=vocab_size)
+        model.load_state_dict(state_dict)
+        return model
+
+
+class SD3TextEncoder2StateDictConverter(StateDictConverter):
+    def _from_diffusers(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        rename_dict = config["diffusers"]["te2_rename_dict"]
+        attn_rename_dict = config["diffusers"]["te2_attn_rename_dict"]
+        state_dict_ = {}
+        for name in state_dict:
+            if name in rename_dict:
+                param = state_dict[name]
+                if name == "text_model.embeddings.position_embedding.weight":
+                    param = param.reshape((1, param.shape[0], param.shape[1]))
+                state_dict_[rename_dict[name]] = param
+            elif name.startswith("text_model.encoder.layers."):
+                param = state_dict[name]
+                names = name.split(".")
+                layer_id, layer_type, tail = names[3], ".".join(names[4:-1]), names[-1]
+                name_ = ".".join(["encoders", layer_id, attn_rename_dict[layer_type], tail])
+                state_dict_[name_] = param
+        return state_dict_
+
+    def _from_civitai(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        rename_dict = config["civitai"]["te2_rename_dict"]
+        state_dict_ = {}
+        for name in state_dict:
+            if name in rename_dict:
+                param = state_dict[name]
+                if name == "text_encoders.clip_g.transformer.text_model.embeddings.position_embedding.weight":
+                    param = param.reshape((1, param.shape[0], param.shape[1]))
+                state_dict_[rename_dict[name]] = param
+        return state_dict_
+
+    def convert(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        if "text_encoders.clip_g.transformer.text_model.embeddings.position_embedding.weight" in state_dict:
+            state_dict = self._from_civitai(state_dict)
+            logger.info("use civitai format state dict")
+        elif "text_model.final_layer_norm.weight" in state_dict:
+            state_dict = self._from_diffusers(state_dict)
+            logger.info("use diffusers format state dict")
+        else:
+            logger.info("use diffsynth format state dict")
+        return state_dict
+
+
+class SD3TextEncoder2(SDXLTextEncoder2):
+    converter = SD3TextEncoder2StateDictConverter()
+
+    def __init__(self):
+        super().__init__()
+
+
+class SD3TextEncoder3(T5EncoderModel):
+    def __init__(self, device: str = "cuda:0", dtype: torch.dtype = torch.float16):
+        super().__init__(
+            embed_dim=4096,
+            vocab_size=32128,
+            num_encoder_layers=24,
+            d_ff=10240,
+            num_heads=64,
+            relative_attention_num_buckets=32,
+            relative_attention_max_distance=128,
+            dropout_rate=0.0,
+            eps=1e-6,
+            device=device,
+            dtype=dtype,
+        )
+
+    def forward(self, input_ids):
+        outputs = super().forward(input_ids=input_ids)
+        prompt_emb = outputs.last_hidden_state
+        return prompt_emb

diffsynth_engine/models/sd3/sd3_vae.py

@@ -0,0 +1,43 @@
+import torch
+from typing import Dict
+
+from diffsynth_engine.models.components.vae import VAEDecoder, VAEEncoder
+from diffsynth_engine.models.utils import no_init_weights
+
+
+class SD3VAEEncoder(VAEEncoder):
+    def __init__(self, device: str = "cuda:0", dtype: torch.dtype = torch.float32):
+        super().__init__(
+            latent_channels=16,
+            scaling_factor=1.5305,
+            shift_factor=0.0609,
+            use_quant_conv=False,
+            device=device,
+            dtype=dtype,
+        )
+
+    @classmethod
+    def from_state_dict(cls, state_dict: Dict[str, torch.Tensor], device: str, dtype: torch.dtype):
+        with no_init_weights():
+            model = torch.nn.utils.skip_init(cls, device=device, dtype=dtype)
+        model.load_state_dict(state_dict)
+        return model
+
+
+class SD3VAEDecoder(VAEDecoder):
+    def __init__(self, device: str = "cuda:0", dtype: torch.dtype = torch.float32):
+        super().__init__(
+            latent_channels=16,
+            scaling_factor=1.5305,
+            shift_factor=0.0609,
+            use_post_quant_conv=False,
+            device=device,
+            dtype=dtype,
+        )
+
+    @classmethod
+    def from_state_dict(cls, state_dict: Dict[str, torch.Tensor], device: str, dtype: torch.dtype):
+        with no_init_weights():
+            model = torch.nn.utils.skip_init(cls, device=device, dtype=dtype)
+        model.load_state_dict(state_dict)
+        return model

diffsynth_engine/models/sdxl/__init__.py

@@ -0,0 +1,13 @@
+from .sdxl_text_encoder import SDXLTextEncoder, SDXLTextEncoder2, config as sdxl_text_encoder_config
+from .sdxl_unet import SDXLUNet, config as sdxl_unet_config
+from .sdxl_vae import SDXLVAEDecoder, SDXLVAEEncoder
+
+__all__ = [
+    "SDXLTextEncoder",
+    "SDXLTextEncoder2",
+    "SDXLUNet",
+    "SDXLVAEDecoder",
+    "SDXLVAEEncoder",
+    "sdxl_text_encoder_config",
+    "sdxl_unet_config",
+]