update

.gitignore

@@ -1,3 +1,2 @@
 tmp.png
-diffusers/
 src

app.py

@@ -480,6 +480,7 @@ with gr.Blocks(
                 """
                 > ❗️ <strong>Note</strong>: 
                 > The Gradio apps use an accelerated version, which may result in a slight reduction in image generation quality.
+                > - This demo is the open-source version, utilizing [Qwen-Image](https://github.com/QwenLM/Qwen-Image) as the pre-trained model, while the more powerful closed-source version employs Kolors 2.1 as the pre-trained model and will soon be launched on the [KlingAI](https://app.klingai.com/global/?gad_source=1&gad_campaignid=22803840655&gbraid=0AAAAA_AcKMnNNjEHRRI1l9_5z1qK881dO).
                 """
                 )
 
@@ -489,8 +490,6 @@ with gr.Blocks(
             > - Adjust the <strong>Number of Prompts</strong> slider to add or remove input rows.
             > - Type your own prompts directly in the text boxes .
             > - You can click any template below to quickly load preset style code and prompts.
-            > - This model is the open-source version, utilizing [Qwen-Image](https://github.com/QwenLM/Qwen-Image) as the pre-trained model, while the more powerful closed-source version employs Kolors 2.1 as the pre-trained model and will soon be launched on the [KlingAI](https://app.klingai.com/global/?gad_source=1&gad_campaignid=22803840655&gbraid=0AAAAA_AcKMnNNjEHRRI1l9_5z1qK881dO).
-
             """
             )
 

lakonlab/models/architecture/__init__.py

@@ -1,5 +1,6 @@
 import sys
 sys.path.insert(0, '/home/user/app')
+print('=====insert=====')
 from .ddpm import *
 from .diffusers import *
 from .gmflow import *

lakonlab/models/architecture/diffusers/__init__.py

@@ -0,0 +1,15 @@
+from .pretrained import (
+    PretrainedVAE, PretrainedVAEDecoder, PretrainedVAEEncoder, PretrainedVAEQwenImage,
+    PretrainedFluxTextEncoder, PretrainedQwenImageTextEncoder, PretrainedStableDiffusion3TextEncoder)
+from .unet import UNet2DConditionModel
+from .flux import FluxTransformer2DModel
+from .dit import DiTTransformer2DModelMod
+from .sd3 import SD3Transformer2DModel
+from .qwen import QwenImageTransformer2DModel
+
+__all__ = [
+    'PretrainedVAE', 'PretrainedVAEDecoder', 'PretrainedVAEEncoder', 'PretrainedFluxTextEncoder',
+    'PretrainedQwenImageTextEncoder', 'UNet2DConditionModel', 'FluxTransformer2DModel',
+    'DiTTransformer2DModelMod', 'SD3Transformer2DModel',
+    'QwenImageTransformer2DModel', 'PretrainedVAEQwenImage', 'PretrainedStableDiffusion3TextEncoder',
+]

lakonlab/models/architecture/diffusers/dit.py

@@ -0,0 +1,428 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from typing import Any, Dict, Optional
+from diffusers.models import DiTTransformer2DModel, ModelMixin
+from diffusers.models.attention import BasicTransformerBlock, _chunked_feed_forward, Attention, FeedForward
+from diffusers.models.embeddings import (
+    PatchEmbed, Timesteps, CombinedTimestepLabelEmbeddings, TimestepEmbedding, LabelEmbedding)
+from diffusers.models.normalization import AdaLayerNormZero
+from diffusers.configuration_utils import register_to_config
+from mmcv.runner import load_checkpoint, _load_checkpoint, load_state_dict
+from mmcv.cnn import constant_init, xavier_init
+from mmgen.models.builder import MODULES
+from mmgen.utils import get_root_logger
+from ..utils import flex_freeze
+
+
+class LabelEmbeddingMod(LabelEmbedding):
+    def __init__(self, num_classes, hidden_size, dropout_prob=0.0, use_cfg_embedding=True):
+        super(LabelEmbedding, self).__init__()
+        if dropout_prob > 0:
+            assert use_cfg_embedding
+        self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)
+        self.num_classes = num_classes
+        self.dropout_prob = dropout_prob
+
+
+class CombinedTimestepLabelEmbeddingsMod(CombinedTimestepLabelEmbeddings):
+    """
+    Modified CombinedTimestepLabelEmbeddings for reproducing the original DiT (downscale_freq_shift=0).
+    """
+    def __init__(
+            self, num_classes, embedding_dim, class_dropout_prob=0.1, downscale_freq_shift=0, use_cfg_embedding=True):
+        super(CombinedTimestepLabelEmbeddings, self).__init__()
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=downscale_freq_shift)
+        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+        self.class_embedder = LabelEmbeddingMod(num_classes, embedding_dim, class_dropout_prob, use_cfg_embedding)
+
+
+class BasicTransformerBlockMod(BasicTransformerBlock):
+    """
+    Modified BasicTransformerBlock for reproducing the original DiT with shared time and class
+    embeddings across all layers.
+    """
+    def __init__(
+            self,
+            dim: int,
+            num_attention_heads: int,
+            attention_head_dim: int,
+            dropout=0.0,
+            cross_attention_dim: Optional[int] = None,
+            activation_fn: str = 'geglu',
+            num_embeds_ada_norm: Optional[int] = None,
+            attention_bias: bool = False,
+            only_cross_attention: bool = False,
+            double_self_attention: bool = False,
+            upcast_attention: bool = False,
+            norm_elementwise_affine: bool = True,
+            norm_type: str = 'layer_norm',
+            norm_eps: float = 1e-5,
+            final_dropout: bool = False,
+            attention_type: str = 'default',
+            ada_norm_continous_conditioning_embedding_dim: Optional[int] = None,
+            ada_norm_bias: Optional[int] = None,
+            ff_inner_dim: Optional[int] = None,
+            ff_bias: bool = True,
+            attention_out_bias: bool = True):
+        super(BasicTransformerBlock, self).__init__()
+        self.only_cross_attention = only_cross_attention
+        self.norm_type = norm_type
+        self.num_embeds_ada_norm = num_embeds_ada_norm
+
+        assert self.norm_type == 'ada_norm_zero'
+        self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)
+        self.attn1 = Attention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            dropout=dropout,
+            bias=attention_bias,
+            cross_attention_dim=cross_attention_dim if only_cross_attention else None,
+            upcast_attention=upcast_attention,
+            out_bias=attention_out_bias,
+        )
+
+        self.norm2 = None
+        self.attn2 = None
+
+        self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
+        self.ff = FeedForward(
+            dim,
+            dropout=dropout,
+            activation_fn=activation_fn,
+            final_dropout=final_dropout,
+            inner_dim=ff_inner_dim,
+            bias=ff_bias,
+        )
+
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            attention_mask: Optional[torch.Tensor] = None,
+            encoder_hidden_states: Optional[torch.Tensor] = None,
+            encoder_attention_mask: Optional[torch.Tensor] = None,
+            timestep: Optional[torch.LongTensor] = None,
+            cross_attention_kwargs: Dict[str, Any] = None,
+            class_labels: Optional[torch.LongTensor] = None,
+            emb: Optional[torch.Tensor] = None,
+            added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None) -> torch.Tensor:
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+            hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype, emb=emb)
+
+        if cross_attention_kwargs is None:
+            cross_attention_kwargs = dict()
+        attn_output = self.attn1(
+            norm_hidden_states,
+            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
+            attention_mask=attention_mask,
+            **cross_attention_kwargs)
+        attn_output = gate_msa.unsqueeze(1) * attn_output
+
+        hidden_states = attn_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+
+        norm_hidden_states = self.norm3(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+        if self._chunk_size is not None:
+            # "feed_forward_chunk_size" can be used to save memory
+            ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)
+        else:
+            ff_output = self.ff(norm_hidden_states)
+
+        ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+        hidden_states = ff_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+
+        return hidden_states
+
+
+class _DiTTransformer2DModelMod(DiTTransformer2DModel):
+
+    @register_to_config
+    def __init__(
+            self,
+            class_dropout_prob=0.0,
+            num_attention_heads: int = 16,
+            attention_head_dim: int = 72,
+            in_channels: int = 4,
+            out_channels: Optional[int] = None,
+            num_layers: int = 28,
+            dropout: float = 0.0,
+            norm_num_groups: int = 32,
+            attention_bias: bool = True,
+            sample_size: int = 32,
+            patch_size: int = 2,
+            activation_fn: str = 'gelu-approximate',
+            num_embeds_ada_norm: Optional[int] = 1000,
+            upcast_attention: bool = False,
+            norm_type: str = 'ada_norm_zero',
+            norm_elementwise_affine: bool = False,
+            norm_eps: float = 1e-5):
+
+        super(DiTTransformer2DModel, self).__init__()
+
+        # Validate inputs.
+        if norm_type != "ada_norm_zero":
+            raise NotImplementedError(
+                f"Forward pass is not implemented when `patch_size` is not None and `norm_type` is '{norm_type}'."
+            )
+        elif norm_type == "ada_norm_zero" and num_embeds_ada_norm is None:
+            raise ValueError(
+                f"When using a `patch_size` and this `norm_type` ({norm_type}), `num_embeds_ada_norm` cannot be None."
+            )
+
+        # Set some common variables used across the board.
+        self.attention_head_dim = attention_head_dim
+        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
+        self.out_channels = in_channels if out_channels is None else out_channels
+        self.gradient_checkpointing = False
+
+        # 2. Initialize the position embedding and transformer blocks.
+        self.height = self.config.sample_size
+        self.width = self.config.sample_size
+
+        self.patch_size = self.config.patch_size
+        self.pos_embed = PatchEmbed(
+            height=self.config.sample_size,
+            width=self.config.sample_size,
+            patch_size=self.config.patch_size,
+            in_channels=self.config.in_channels,
+            embed_dim=self.inner_dim)
+        self.emb = CombinedTimestepLabelEmbeddingsMod(
+            num_embeds_ada_norm, self.inner_dim, class_dropout_prob=0.0)
+
+        self.transformer_blocks = nn.ModuleList([
+            BasicTransformerBlockMod(
+                self.inner_dim,
+                self.config.num_attention_heads,
+                self.config.attention_head_dim,
+                dropout=self.config.dropout,
+                activation_fn=self.config.activation_fn,
+                num_embeds_ada_norm=None,
+                attention_bias=self.config.attention_bias,
+                upcast_attention=self.config.upcast_attention,
+                norm_type=norm_type,
+                norm_elementwise_affine=self.config.norm_elementwise_affine,
+                norm_eps=self.config.norm_eps)
+            for _ in range(self.config.num_layers)])
+
+        # 3. Output blocks.
+        self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim)
+        self.proj_out_2 = nn.Linear(
+            self.inner_dim, self.config.patch_size * self.config.patch_size * self.out_channels)
+
+    # https://github.com/facebookresearch/DiT/blob/main/models.py
+    def init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                xavier_init(m, distribution='uniform')
+            elif isinstance(m, nn.Embedding):
+                torch.nn.init.normal_(m.weight, mean=0.0, std=0.02)
+
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d)
+        w = self.pos_embed.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.pos_embed.proj.bias, 0)
+
+        # Zero-out adaLN modulation layers in DiT blocks
+        for m in self.modules():
+            if isinstance(m, AdaLayerNormZero):
+                constant_init(m.linear, val=0)
+
+        # Zero-out output layers
+        constant_init(self.proj_out_1, val=0)
+        constant_init(self.proj_out_2, val=0)
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            timestep: Optional[torch.LongTensor] = None,
+            class_labels: Optional[torch.LongTensor] = None,
+            cross_attention_kwargs: Dict[str, Any] = None):
+        # 1. Input
+        bs, _, h, w = hidden_states.size()
+        height, width = h // self.patch_size, w // self.patch_size
+        hidden_states = self.pos_embed(hidden_states)
+
+        cond_emb = self.emb(
+            timestep, class_labels, hidden_dtype=hidden_states.dtype)
+        dropout_enabled = self.config.class_dropout_prob > 0 and self.training
+        if dropout_enabled:
+            uncond_emb = self.emb(timestep, torch.full_like(
+                class_labels, self.config.num_embeds_ada_norm), hidden_dtype=hidden_states.dtype)
+
+        # 2. Blocks
+        for block in self.transformer_blocks:
+            if dropout_enabled:
+                dropout_mask = torch.rand((bs, 1), device=hidden_states.device) < self.config.class_dropout_prob
+                emb = torch.where(dropout_mask, uncond_emb, cond_emb)
+            else:
+                emb = cond_emb
+
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    None,
+                    None,
+                    None,
+                    timestep,
+                    cross_attention_kwargs,
+                    class_labels,
+                    emb,
+                    use_reentrant=False)
+
+            else:
+                hidden_states = block(
+                    hidden_states,
+                    attention_mask=None,
+                    encoder_hidden_states=None,
+                    encoder_attention_mask=None,
+                    timestep=timestep,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    class_labels=class_labels,
+                    emb=emb)
+
+        # 3. Output
+        if dropout_enabled:
+            dropout_mask = torch.rand((bs, 1), device=hidden_states.device) < self.config.class_dropout_prob
+            emb = torch.where(dropout_mask, uncond_emb, cond_emb)
+        else:
+            emb = cond_emb
+        shift, scale = self.proj_out_1(F.silu(emb)).chunk(2, dim=1)
+        hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
+        output = self.proj_out_2(hidden_states).reshape(
+                bs, height, width, self.patch_size, self.patch_size, self.out_channels
+            ).permute(0, 5, 1, 3, 2, 4).reshape(
+                bs, self.out_channels, height * self.patch_size, width * self.patch_size)
+
+        return output
+
+
+@MODULES.register_module()
+class DiTTransformer2DModelMod(_DiTTransformer2DModelMod):
+
+    def __init__(
+            self,
+            *args,
+            freeze=False,
+            freeze_exclude=[],
+            pretrained=None,
+            torch_dtype='float32',
+            autocast_dtype=None,
+            freeze_exclude_fp32=True,
+            freeze_exclude_autocast_dtype='float32',
+            checkpointing=True,
+            **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.init_weights(pretrained)
+
+        if autocast_dtype is not None:
+            assert torch_dtype == 'float32'
+        self.autocast_dtype = autocast_dtype
+
+        if torch_dtype is not None:
+            self.to(getattr(torch, torch_dtype))
+
+        self.freeze = freeze
+        if self.freeze:
+            flex_freeze(
+                self,
+                exclude_keys=freeze_exclude,
+                exclude_fp32=freeze_exclude_fp32,
+                exclude_autocast_dtype=freeze_exclude_autocast_dtype)
+
+        if checkpointing:
+            self.enable_gradient_checkpointing()
+
+    def init_weights(self, pretrained=None):
+        super().init_weights()
+        if pretrained is not None:
+            logger = get_root_logger()
+            # load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger)
+            checkpoint = _load_checkpoint(pretrained, map_location='cpu', logger=logger)
+            if 'state_dict' in checkpoint:
+                state_dict = checkpoint['state_dict']
+            else:
+                state_dict = checkpoint
+            # load from GMDiT V1 model with 1 Gaussian
+            p2 = self.config.patch_size * self.config.patch_size
+            ori_out_channels = p2 * self.out_channels
+            if 'proj_out_2.weight' in state_dict:
+                # if this is GMDiT V1 model with 1 Gaussian
+                if state_dict['proj_out_2.weight'].size(0) == p2 * (self.out_channels + 1):
+                    state_dict['proj_out_2.weight'] = state_dict['proj_out_2.weight'].reshape(
+                        p2, self.out_channels + 1, -1
+                    )[:, :-1].reshape(ori_out_channels, -1)
+                # if this is original DiT with variance prediction
+                if state_dict['proj_out_2.weight'].size(0) == 2 * ori_out_channels:
+                    state_dict['proj_out_2.weight'] = state_dict['proj_out_2.weight'].reshape(
+                        p2, 2 * self.out_channels, -1
+                    )[:, :self.out_channels].reshape(ori_out_channels, -1)
+            if 'proj_out_2.bias' in state_dict:
+                # if this is GMDiT V1 model with 1 Gaussian
+                if state_dict['proj_out_2.bias'].size(0) == p2 * (self.out_channels + 1):
+                    state_dict['proj_out_2.bias'] = state_dict['proj_out_2.bias'].reshape(
+                        p2, self.out_channels + 1
+                    )[:, :-1].reshape(ori_out_channels)
+                # if this is original DiT with variance prediction
+                if state_dict['proj_out_2.bias'].size(0) == 2 * ori_out_channels:
+                    state_dict['proj_out_2.bias'] = state_dict['proj_out_2.bias'].reshape(
+                        p2, 2 * self.out_channels
+                    )[:, :self.out_channels].reshape(ori_out_channels)
+            if 'emb.class_embedder.embedding_table.weight' not in state_dict \
+                    and 'transformer_blocks.0.norm1.emb.class_embedder.embedding_table.weight' in state_dict:
+                # convert original diffusers DiT model to our modified DiT model with shared embeddings
+                keys_to_remove = []
+                state_update = {}
+                for k, v in state_dict.items():
+                    if k.startswith('transformer_blocks.0.norm1.emb.'):
+                        new_k = k.replace('transformer_blocks.0.norm1.', '')
+                        state_update[new_k] = v
+                    if k.startswith('transformer_blocks.') and '.norm1.emb.' in k:
+                        keys_to_remove.append(k)
+                state_dict.update(state_update)
+                for k in keys_to_remove:
+                    del state_dict[k]
+            load_state_dict(self, state_dict, logger=logger)
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            timestep: Optional[torch.LongTensor] = None,
+            class_labels: Optional[torch.LongTensor] = None,
+            **kwargs):
+        if self.autocast_dtype is not None:
+            dtype = getattr(torch, self.autocast_dtype)
+        else:
+            dtype = hidden_states.dtype
+        with torch.autocast(
+                device_type='cuda',
+                enabled=self.autocast_dtype is not None,
+                dtype=dtype if self.autocast_dtype is not None else None):
+            return super().forward(
+                hidden_states.to(dtype),
+                timestep=timestep,
+                class_labels=class_labels,
+                **kwargs)

lakonlab/models/architecture/diffusers/flux.py

@@ -0,0 +1,156 @@
+import torch
+
+from typing import Optional
+from accelerate import init_empty_weights
+from diffusers.models import FluxTransformer2DModel as _FluxTransformer2DModel
+from peft import LoraConfig
+from mmgen.models.builder import MODULES
+from mmgen.utils import get_root_logger
+from ..utils import flex_freeze
+from lakonlab.runner.checkpoint import load_checkpoint, _load_checkpoint
+
+
+@MODULES.register_module()
+class FluxTransformer2DModel(_FluxTransformer2DModel):
+
+    def __init__(
+            self,
+            *args,
+            patch_size=2,
+            freeze=False,
+            freeze_exclude=[],
+            pretrained=None,
+            pretrained_lora=None,
+            pretrained_lora_scale=1.0,
+            torch_dtype='float32',
+            freeze_exclude_fp32=True,
+            freeze_exclude_autocast_dtype='float32',
+            checkpointing=True,
+            use_lora=False,
+            lora_target_modules=None,
+            lora_rank=16,
+            **kwargs):
+        with init_empty_weights():
+            super().__init__(patch_size=1, *args, **kwargs)
+        self.patch_size = patch_size
+
+        self.init_weights(pretrained, pretrained_lora, pretrained_lora_scale)
+
+        self.use_lora = use_lora
+        self.lora_target_modules = lora_target_modules
+        self.lora_rank = lora_rank
+        if self.use_lora:
+            transformer_lora_config = LoraConfig(
+                r=lora_rank,
+                lora_alpha=lora_rank,
+                init_lora_weights='gaussian',
+                target_modules=lora_target_modules,
+            )
+            self.add_adapter(transformer_lora_config)
+
+        if torch_dtype is not None:
+            self.to(getattr(torch, torch_dtype))
+
+        self.freeze = freeze
+        if self.freeze:
+            flex_freeze(
+                self,
+                exclude_keys=freeze_exclude,
+                exclude_fp32=freeze_exclude_fp32,
+                exclude_autocast_dtype=freeze_exclude_autocast_dtype)
+
+        if checkpointing:
+            self.enable_gradient_checkpointing()
+
+    def init_weights(self, pretrained=None, pretrained_lora=None, pretrained_lora_scale=1.0):
+        if pretrained is not None:
+            logger = get_root_logger()
+            load_checkpoint(
+                self, pretrained, map_location='cpu', strict=False, logger=logger, assign=True)
+            if pretrained_lora is not None:
+                if not isinstance(pretrained_lora, (list, tuple)):
+                    assert isinstance(pretrained_lora, str)
+                    pretrained_lora = [pretrained_lora]
+                if not isinstance(pretrained_lora_scale, (list, tuple)):
+                    assert isinstance(pretrained_lora_scale, (int, float))
+                    pretrained_lora_scale = [pretrained_lora_scale]
+                for pretrained_lora_single, pretrained_lora_scale_single in zip(pretrained_lora, pretrained_lora_scale):
+                    lora_state_dict = _load_checkpoint(
+                        pretrained_lora_single, map_location='cpu', logger=logger)
+                    self.load_lora_adapter(lora_state_dict)
+                    self.fuse_lora(lora_scale=pretrained_lora_scale_single)
+                    self.unload_lora()
+
+    @staticmethod
+    def _prepare_latent_image_ids(height, width, device, dtype):
+        """
+        Copied from Diffusers
+        """
+        latent_image_ids = torch.zeros(height, width, 3)
+        latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None]
+        latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :]
+
+        latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+
+        latent_image_ids = latent_image_ids.reshape(
+            latent_image_id_height * latent_image_id_width, latent_image_id_channels)
+
+        return latent_image_ids.to(device=device, dtype=dtype)
+
+    def patchify(self, latents):
+        if self.patch_size > 1:
+            bs, c, h, w = latents.size()
+            latents = latents.reshape(
+                bs, c, h // self.patch_size, self.patch_size, w // self.patch_size, self.patch_size
+            ).permute(
+                0, 1, 3, 5, 2, 4
+            ).reshape(
+                bs, c * self.patch_size * self.patch_size, h // self.patch_size, w // self.patch_size)
+        return latents
+
+    def unpatchify(self, latents):
+        if self.patch_size > 1:
+            bs, c, h, w = latents.size()
+            latents = latents.reshape(
+                bs, c // (self.patch_size * self.patch_size), self.patch_size, self.patch_size, h, w
+            ).permute(
+                0, 1, 4, 2, 5, 3
+            ).reshape(
+                bs, c // (self.patch_size * self.patch_size), h * self.patch_size, w * self.patch_size)
+        return latents
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            timestep: torch.Tensor,
+            encoder_hidden_states: torch.Tensor = None,
+            pooled_projections: torch.Tensor = None,
+            mask: Optional[torch.Tensor] = None,
+            masked_image_latents: Optional[torch.Tensor] = None,
+            **kwargs):
+        hidden_states = self.patchify(hidden_states)
+        bs, c, h, w = hidden_states.size()
+        dtype = hidden_states.dtype
+        device = hidden_states.device
+        hidden_states = hidden_states.reshape(bs, c, h * w).permute(0, 2, 1)
+        img_ids = self._prepare_latent_image_ids(
+            h, w, device, dtype)
+        txt_ids = img_ids.new_zeros((encoder_hidden_states.shape[-2], 3))
+
+        #  Flux fill
+        if mask is not None and masked_image_latents is not None:
+            hidden_states = torch.cat(
+                (hidden_states, masked_image_latents.to(dtype=dtype), mask.to(dtype=dtype)), dim=-1)
+
+        output = super().forward(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states.to(dtype),
+            pooled_projections=pooled_projections.to(dtype),
+            timestep=timestep,
+            img_ids=img_ids,
+            txt_ids=txt_ids,
+            return_dict=False,
+            **kwargs)[0]
+
+        output = output.permute(0, 2, 1).reshape(bs, self.out_channels, h, w)
+        return self.unpatchify(output)

lakonlab/models/architecture/diffusers/pretrained.py

@@ -0,0 +1,281 @@
+import logging
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.models import AutoencoderKL, AutoencoderKLQwenImage
+from diffusers.pipelines import FluxPipeline, QwenImagePipeline, StableDiffusion3Pipeline
+from mmgen.models.builder import MODULES
+
+# Suppress truncation warnings from transformers and diffusers
+for name in (
+        'transformers.tokenization_utils_base',
+        'transformers.tokenization_utils',
+        'transformers.tokenization_utils_fast'):
+    logging.getLogger(name).setLevel(logging.ERROR)
+
+for name, logger in logging.root.manager.loggerDict.items():
+    if isinstance(logger, logging.Logger) and (name.startswith('diffusers.pipelines.')):
+        logger.setLevel(logging.ERROR)
+
+
+@MODULES.register_module()
+class PretrainedVAE(nn.Module):
+    def __init__(self,
+                 from_pretrained=None,
+                 del_encoder=False,
+                 del_decoder=False,
+                 use_slicing=False,
+                 freeze=True,
+                 eval_mode=True,
+                 torch_dtype='float32',
+                 **kwargs):
+        super().__init__()
+        if torch_dtype is not None:
+            kwargs.update(torch_dtype=getattr(torch, torch_dtype))
+        self.vae = AutoencoderKL.from_pretrained(
+            from_pretrained, **kwargs)
+        if del_encoder:
+            del self.vae.encoder
+        if del_decoder:
+            del self.vae.decoder
+        if use_slicing:
+            self.vae.enable_slicing()
+        self.freeze = freeze
+        self.eval_mode = eval_mode
+        if self.freeze:
+            self.requires_grad_(False)
+        if self.eval_mode:
+            self.eval()
+        self.vae.set_use_memory_efficient_attention_xformers(
+            not hasattr(torch.nn.functional, 'scaled_dot_product_attention'))
+
+    def train(self, mode=True):
+        mode = mode and (not self.eval_mode)
+        return super().train(mode)
+
+    def forward(self, *args, **kwargs):
+        return self.vae(*args, return_dict=False, **kwargs)[0]
+
+    def encode(self, img):
+        scaling_factor = self.vae.config.scaling_factor
+        shift_factor = self.vae.config.shift_factor
+        if scaling_factor is None:
+            scaling_factor = 1.0
+        if shift_factor is None:
+            shift_factor = 0.0
+        return (self.vae.encode(img).latent_dist.sample() - shift_factor) * scaling_factor
+
+    def decode(self, code):
+        scaling_factor = self.vae.config.scaling_factor
+        shift_factor = self.vae.config.shift_factor
+        if scaling_factor is None:
+            scaling_factor = 1.0
+        if shift_factor is None:
+            shift_factor = 0.0
+        return self.vae.decode(code / scaling_factor + shift_factor, return_dict=False)[0]
+
+
+@MODULES.register_module()
+class PretrainedVAEDecoder(PretrainedVAE):
+    def __init__(self, **kwargs):
+        super().__init__(
+            del_encoder=True,
+            del_decoder=False,
+            **kwargs)
+
+    def forward(self, code):
+        return super().decode(code)
+
+
+@MODULES.register_module()
+class PretrainedVAEEncoder(PretrainedVAE):
+    def __init__(self, **kwargs):
+        super().__init__(
+            del_encoder=False,
+            del_decoder=True,
+            **kwargs)
+
+    def forward(self, img):
+        return super().encode(img)
+
+
+@MODULES.register_module()
+class PretrainedVAEQwenImage(nn.Module):
+    def __init__(self,
+                 from_pretrained=None,
+                 use_slicing=False,
+                 freeze=True,
+                 eval_mode=True,
+                 torch_dtype='float32',
+                 **kwargs):
+        super().__init__()
+        if torch_dtype is not None:
+            kwargs.update(torch_dtype=getattr(torch, torch_dtype))
+        self.vae = AutoencoderKLQwenImage.from_pretrained(
+            from_pretrained, **kwargs)
+        if use_slicing:
+            self.vae.enable_slicing()
+        self.freeze = freeze
+        self.eval_mode = eval_mode
+        if self.freeze:
+            self.requires_grad_(False)
+        if self.eval_mode:
+            self.eval()
+
+    def train(self, mode=True):
+        mode = mode and (not self.eval_mode)
+        return super().train(mode)
+
+    def forward(self, *args, **kwargs):
+        return self.vae(*args, return_dict=False, **kwargs)[0]
+
+    def encode(self, img):
+        device = img.device
+        dtype = img.dtype
+        latents_mean = torch.tensor(self.vae.config.latents_mean, device=device, dtype=dtype).view(
+            1, self.vae.config.z_dim, 1, 1, 1)
+        latents_std = torch.tensor(self.vae.config.latents_std, device=device, dtype=dtype).view(
+            1, self.vae.config.z_dim, 1, 1, 1)
+        return ((self.vae.encode(img.unsqueeze(-3)).latent_dist.sample() - latents_mean) / latents_std).squeeze(-3)
+
+    def decode(self, code):
+        device = code.device
+        dtype = code.dtype
+        latents_mean = torch.tensor(self.vae.config.latents_mean, device=device, dtype=dtype).view(
+            1, self.vae.config.z_dim, 1, 1, 1)
+        latents_std = torch.tensor(self.vae.config.latents_std, device=device, dtype=dtype).view(
+            1, self.vae.config.z_dim, 1, 1, 1)
+        return self.vae.decode(code.unsqueeze(-3) * latents_std + latents_mean, return_dict=False)[0].squeeze(-3)
+
+
+@MODULES.register_module()
+class PretrainedFluxTextEncoder(nn.Module):
+    def __init__(self,
+                 from_pretrained='black-forest-labs/FLUX.1-dev',
+                 freeze=True,
+                 eval_mode=True,
+                 torch_dtype='bfloat16',
+                 max_sequence_length=512,
+                 **kwargs):
+        super().__init__()
+        self.max_sequence_length = max_sequence_length
+        self.pipeline = FluxPipeline.from_pretrained(
+            from_pretrained,
+            scheduler=None,
+            vae=None,
+            transformer=None,
+            image_encoder=None,
+            feature_extractor=None,
+            torch_dtype=getattr(torch, torch_dtype),
+            **kwargs)
+        self.text_encoder = self.pipeline.text_encoder
+        self.text_encoder_2 = self.pipeline.text_encoder_2
+        self.freeze = freeze
+        self.eval_mode = eval_mode
+        if self.freeze:
+            self.requires_grad_(False)
+        if self.eval_mode:
+            self.eval()
+
+    def train(self, mode=True):
+        mode = mode and (not self.eval_mode)
+        return super().train(mode)
+
+    def forward(self, prompt, prompt_2=None):
+        prompt_embeds, pooled_prompt_embeds, text_ids = self.pipeline.encode_prompt(
+            prompt, prompt_2=prompt_2, max_sequence_length=self.max_sequence_length)
+        return dict(
+            encoder_hidden_states=prompt_embeds,
+            pooled_projections=pooled_prompt_embeds)
+
+
+@MODULES.register_module()
+class PretrainedQwenImageTextEncoder(nn.Module):
+    def __init__(self,
+                 from_pretrained='Qwen/Qwen-Image',
+                 freeze=True,
+                 eval_mode=True,
+                 torch_dtype='bfloat16',
+                 max_sequence_length=512,
+                 pad_seq_len=None,
+                 **kwargs):
+        super().__init__()
+        self.max_sequence_length = max_sequence_length
+        if pad_seq_len is not None:
+            assert pad_seq_len >= max_sequence_length
+        self.pad_seq_len = pad_seq_len
+        self.pipeline = QwenImagePipeline.from_pretrained(
+            from_pretrained,
+            scheduler=None,
+            vae=None,
+            transformer=None,
+            torch_dtype=getattr(torch, torch_dtype),
+            **kwargs)
+        self.text_encoder = self.pipeline.text_encoder
+        self.freeze = freeze
+        self.eval_mode = eval_mode
+        if self.freeze:
+            self.requires_grad_(False)
+        if self.eval_mode:
+            self.eval()
+
+    def train(self, mode=True):
+        mode = mode and (not self.eval_mode)
+        return super().train(mode)
+
+    def forward(self, prompt):
+        prompt_embeds, prompt_embeds_mask = self.pipeline.encode_prompt(
+            prompt, max_sequence_length=self.max_sequence_length)
+        if self.pad_seq_len is not None:
+            pad_len = self.pad_seq_len - prompt_embeds.size(1)
+            prompt_embeds = F.pad(
+                prompt_embeds, (0, 0, 0, pad_len), value=0.0)
+            prompt_embeds_mask = F.pad(
+                prompt_embeds_mask, (0, pad_len), value=0.0)
+        return dict(
+            encoder_hidden_states=prompt_embeds,
+            encoder_hidden_states_mask=prompt_embeds_mask)
+
+
+@MODULES.register_module()
+class PretrainedStableDiffusion3TextEncoder(nn.Module):
+    def __init__(self,
+                 from_pretrained='stabilityai/stable-diffusion-3.5-large',
+                 freeze=True,
+                 eval_mode=True,
+                 torch_dtype='float32',
+                 max_sequence_length=256,
+                 **kwargs):
+        super().__init__()
+        self.max_sequence_length = max_sequence_length
+        self.pipeline = StableDiffusion3Pipeline.from_pretrained(
+            from_pretrained,
+            scheduler=None,
+            vae=None,
+            transformer=None,
+            image_encoder=None,
+            feature_extractor=None,
+            torch_dtype=getattr(torch, torch_dtype),
+            **kwargs)
+        self.text_encoder = self.pipeline.text_encoder
+        self.text_encoder_2 = self.pipeline.text_encoder_2
+        self.text_encoder_3 = self.pipeline.text_encoder_3
+        self.freeze = freeze
+        self.eval_mode = eval_mode
+        if self.freeze:
+            self.requires_grad_(False)
+        if self.eval_mode:
+            self.eval()
+
+    def train(self, mode=True):
+        mode = mode and (not self.eval_mode)
+        return super().train(mode)
+
+    def forward(self, prompt, prompt_2=None, prompt_3=None):
+        prompt_embeds, _, pooled_prompt_embeds, _ = self.pipeline.encode_prompt(
+            prompt, prompt_2=prompt_2, prompt_3=prompt_3, do_classifier_free_guidance=False,
+            max_sequence_length=self.max_sequence_length)
+        return dict(
+            encoder_hidden_states=prompt_embeds,
+            pooled_projections=pooled_prompt_embeds)

lakonlab/models/architecture/diffusers/qwen.py

@@ -0,0 +1,139 @@
+import torch
+
+from accelerate import init_empty_weights
+from diffusers.models import QwenImageTransformer2DModel as _QwenImageTransformer2DModel
+from diffusers.loaders.lora_conversion_utils import _convert_non_diffusers_qwen_lora_to_diffusers
+from peft import LoraConfig
+from mmgen.models.builder import MODULES
+from mmgen.utils import get_root_logger
+from ..utils import flex_freeze
+from lakonlab.runner.checkpoint import load_checkpoint, _load_checkpoint
+
+
+@MODULES.register_module()
+class QwenImageTransformer2DModel(_QwenImageTransformer2DModel):
+
+    def __init__(
+            self,
+            *args,
+            patch_size=2,
+            freeze=False,
+            freeze_exclude=[],
+            pretrained=None,
+            pretrained_lora=None,
+            pretrained_lora_scale=1.0,
+            torch_dtype='float32',
+            freeze_exclude_fp32=True,
+            freeze_exclude_autocast_dtype='float32',
+            checkpointing=True,
+            use_lora=False,
+            lora_target_modules=None,
+            lora_rank=16,
+            **kwargs):
+        with init_empty_weights():
+            super().__init__(*args, patch_size=1, **kwargs)
+        self.patch_size = patch_size
+
+        self.init_weights(pretrained, pretrained_lora, pretrained_lora_scale)
+
+        self.use_lora = use_lora
+        self.lora_target_modules = lora_target_modules
+        self.lora_rank = lora_rank
+        if self.use_lora:
+            transformer_lora_config = LoraConfig(
+                r=lora_rank,
+                lora_alpha=lora_rank,
+                init_lora_weights='gaussian',
+                target_modules=lora_target_modules,
+            )
+            self.add_adapter(transformer_lora_config)
+
+        if torch_dtype is not None:
+            self.to(getattr(torch, torch_dtype))
+
+        self.freeze = freeze
+        if self.freeze:
+            flex_freeze(
+                self,
+                exclude_keys=freeze_exclude,
+                exclude_fp32=freeze_exclude_fp32,
+                exclude_autocast_dtype=freeze_exclude_autocast_dtype)
+
+        if checkpointing:
+            self.enable_gradient_checkpointing()
+
+    def init_weights(self, pretrained=None, pretrained_lora=None, pretrained_lora_scale=1.0):
+        if pretrained is not None:
+            logger = get_root_logger()
+            load_checkpoint(
+                self, pretrained, map_location='cpu', strict=False, logger=logger, assign=True)
+            if pretrained_lora is not None:
+                if not isinstance(pretrained_lora, (list, tuple)):
+                    assert isinstance(pretrained_lora, str)
+                    pretrained_lora = [pretrained_lora]
+                if not isinstance(pretrained_lora_scale, (list, tuple)):
+                    assert isinstance(pretrained_lora_scale, (int, float))
+                    pretrained_lora_scale = [pretrained_lora_scale]
+                for pretrained_lora_single, pretrained_lora_scale_single in zip(pretrained_lora, pretrained_lora_scale):
+                    lora_state_dict = _load_checkpoint(
+                        pretrained_lora_single, map_location='cpu', logger=logger)
+                    lora_state_dict = _convert_non_diffusers_qwen_lora_to_diffusers(lora_state_dict)
+                    self.load_lora_adapter(lora_state_dict)
+                    self.fuse_lora(lora_scale=pretrained_lora_scale_single)
+                    self.unload_lora()
+
+    def patchify(self, latents):
+        if self.patch_size > 1:
+            bs, c, h, w = latents.size()
+            latents = latents.reshape(
+                bs, c, h // self.patch_size, self.patch_size, w // self.patch_size, self.patch_size
+            ).permute(
+                0, 1, 3, 5, 2, 4
+            ).reshape(
+                bs, c * self.patch_size * self.patch_size, h // self.patch_size, w // self.patch_size)
+        return latents
+
+    def unpatchify(self, latents):
+        if self.patch_size > 1:
+            bs, c, h, w = latents.size()
+            latents = latents.reshape(
+                bs, c // (self.patch_size * self.patch_size), self.patch_size, self.patch_size, h, w
+            ).permute(
+                0, 1, 4, 2, 5, 3
+            ).reshape(
+                bs, c // (self.patch_size * self.patch_size), h * self.patch_size, w * self.patch_size)
+        return latents
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            timestep: torch.Tensor,
+            encoder_hidden_states: torch.Tensor = None,
+            encoder_hidden_states_mask: torch.Tensor = None,
+            **kwargs):
+        hidden_states = self.patchify(hidden_states)
+        bs, c, h, w = hidden_states.size()
+        dtype = hidden_states.dtype
+        hidden_states = hidden_states.reshape(bs, c, h * w).permute(0, 2, 1)
+        img_shapes = [[(1, h, w)]]
+        if encoder_hidden_states_mask is not None:
+            txt_seq_lens = encoder_hidden_states_mask.sum(dim=1)
+            max_txt_seq_len = txt_seq_lens.max()
+            encoder_hidden_states = encoder_hidden_states[:, :max_txt_seq_len]
+            encoder_hidden_states_mask = encoder_hidden_states_mask[:, :max_txt_seq_len]
+            txt_seq_lens = txt_seq_lens.tolist()
+        else:
+            txt_seq_lens = None
+
+        output = super().forward(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states.to(dtype),
+            encoder_hidden_states_mask=encoder_hidden_states_mask,
+            timestep=timestep,
+            img_shapes=img_shapes,
+            txt_seq_lens=txt_seq_lens,
+            return_dict=False,
+            **kwargs)[0]
+
+        output = output.permute(0, 2, 1).reshape(bs, self.out_channels, h, w)
+        return self.unpatchify(output)

lakonlab/models/architecture/diffusers/sd3.py

@@ -0,0 +1,80 @@
+import torch
+
+from accelerate import init_empty_weights
+from diffusers.models import SD3Transformer2DModel as _SD3Transformer2DModel
+from peft import LoraConfig
+from mmgen.models.builder import MODULES
+from mmgen.utils import get_root_logger
+from ..utils import flex_freeze
+from lakonlab.runner.checkpoint import load_checkpoint
+
+
+@MODULES.register_module()
+class SD3Transformer2DModel(_SD3Transformer2DModel):
+
+    def __init__(
+            self,
+            *args,
+            freeze=False,
+            freeze_exclude=[],
+            pretrained=None,
+            torch_dtype='float32',
+            freeze_exclude_fp32=True,
+            freeze_exclude_autocast_dtype='float32',
+            checkpointing=True,
+            use_lora=False,
+            lora_target_modules=None,
+            lora_rank=16,
+            **kwargs):
+        with init_empty_weights():
+            super().__init__(*args, **kwargs)
+        self.init_weights(pretrained)
+
+        self.use_lora = use_lora
+        self.lora_target_modules = lora_target_modules
+        self.lora_rank = lora_rank
+        if self.use_lora:
+            transformer_lora_config = LoraConfig(
+                r=lora_rank,
+                lora_alpha=lora_rank,
+                init_lora_weights='gaussian',
+                target_modules=lora_target_modules,
+            )
+            self.add_adapter(transformer_lora_config)
+
+        if torch_dtype is not None:
+            self.to(getattr(torch, torch_dtype))
+
+        self.freeze = freeze
+        if self.freeze:
+            flex_freeze(
+                self,
+                exclude_keys=freeze_exclude,
+                exclude_fp32=freeze_exclude_fp32,
+                exclude_autocast_dtype=freeze_exclude_autocast_dtype)
+
+        if checkpointing:
+            self.enable_gradient_checkpointing()
+
+    def init_weights(self, pretrained=None):
+        if pretrained is not None:
+            logger = get_root_logger()
+            load_checkpoint(
+                self, pretrained, map_location='cpu', strict=False, logger=logger, assign=True)
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            timestep: torch.Tensor,
+            encoder_hidden_states: torch.Tensor = None,
+            pooled_projections: torch.Tensor = None,
+            **kwargs):
+        dtype = hidden_states.dtype
+
+        return super().forward(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states.to(dtype),
+            pooled_projections=pooled_projections.to(dtype),
+            timestep=timestep,
+            return_dict=False,
+            **kwargs)[0]

lakonlab/models/architecture/diffusers/unet.py

@@ -0,0 +1,192 @@
+import torch
+import torch.nn.functional as F
+
+from typing import Dict, Any, Optional, Union, Tuple
+from collections import OrderedDict
+from diffusers.models import UNet2DConditionModel as _UNet2DConditionModel
+from mmcv.runner import _load_checkpoint, load_state_dict
+from mmgen.models.builder import MODULES
+from mmgen.utils import get_root_logger
+from ..utils import flex_freeze
+
+
+def ceildiv(a, b):
+    return -(a // -b)
+
+
+def unet_enc(
+        unet,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        added_cond_kwargs=None):
+    # 0. center input if necessary
+    if unet.config.center_input_sample:
+        sample = 2 * sample - 1.0
+
+    # 1. time
+    t_emb = unet.get_time_embed(sample=sample, timestep=timestep)
+    emb = unet.time_embedding(t_emb)
+    aug_emb = unet.get_aug_embed(
+        emb=emb, encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs)
+    emb = emb + aug_emb if aug_emb is not None else emb
+
+    if unet.time_embed_act is not None:
+        emb = unet.time_embed_act(emb)
+
+    encoder_hidden_states = unet.process_encoder_hidden_states(
+        encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs)
+
+    # 2. pre-process
+    sample = unet.conv_in(sample)
+
+    # 3. down
+    down_block_res_samples = (sample,)
+    for downsample_block in unet.down_blocks:
+        if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+            sample, res_samples = downsample_block(
+                hidden_states=sample,
+                temb=emb,
+                encoder_hidden_states=encoder_hidden_states,
+                cross_attention_kwargs=cross_attention_kwargs,
+            )
+        else:
+            sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+
+        down_block_res_samples += res_samples
+
+    return emb, down_block_res_samples, sample
+
+
+def unet_dec(
+        unet,
+        emb,
+        down_block_res_samples,
+        sample,
+        encoder_hidden_states: torch.Tensor,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        mid_block_additional_residual: Optional[torch.Tensor] = None):
+    is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
+
+    if is_controlnet:
+        new_down_block_res_samples = ()
+
+        for down_block_res_sample, down_block_additional_residual in zip(
+                down_block_res_samples, down_block_additional_residuals):
+            down_block_res_sample = down_block_res_sample + down_block_additional_residual
+            new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
+
+        down_block_res_samples = new_down_block_res_samples
+
+    # 4. mid
+    if unet.mid_block is not None:
+        if hasattr(unet.mid_block, "has_cross_attention") and unet.mid_block.has_cross_attention:
+            sample = unet.mid_block(
+                sample,
+                emb,
+                encoder_hidden_states=encoder_hidden_states,
+                cross_attention_kwargs=cross_attention_kwargs,
+            )
+        else:
+            sample = unet.mid_block(sample, emb)
+
+    if is_controlnet:
+        sample = sample + mid_block_additional_residual
+
+    # 5. up
+    for i, upsample_block in enumerate(unet.up_blocks):
+        res_samples = down_block_res_samples[-len(upsample_block.resnets):]
+        down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+        if hasattr(upsample_block, 'has_cross_attention') and upsample_block.has_cross_attention:
+            sample = upsample_block(
+                hidden_states=sample,
+                temb=emb,
+                res_hidden_states_tuple=res_samples,
+                encoder_hidden_states=encoder_hidden_states,
+                cross_attention_kwargs=cross_attention_kwargs,
+            )
+        else:
+            sample = upsample_block(
+                hidden_states=sample,
+                temb=emb,
+                res_hidden_states_tuple=res_samples,
+            )
+
+    # 6. post-process
+    if unet.conv_norm_out:
+        sample = unet.conv_norm_out(sample)
+        sample = unet.conv_act(sample)
+    sample = unet.conv_out(sample)
+
+    return sample
+
+
+@MODULES.register_module()
+class UNet2DConditionModel(_UNet2DConditionModel):
+    def __init__(self,
+                 *args,
+                 freeze=True,
+                 freeze_exclude=[],
+                 pretrained=None,
+                 torch_dtype='float32',
+                 freeze_exclude_fp32=True,
+                 freeze_exclude_autocast_dtype='float32',
+                 **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.init_weights(pretrained)
+        if torch_dtype is not None:
+            self.to(getattr(torch, torch_dtype))
+
+        self.set_use_memory_efficient_attention_xformers(
+            not hasattr(torch.nn.functional, 'scaled_dot_product_attention'))
+
+        self.freeze = freeze
+        if self.freeze:
+            flex_freeze(
+                self,
+                exclude_keys=freeze_exclude,
+                exclude_fp32=freeze_exclude_fp32,
+                exclude_autocast_dtype=freeze_exclude_autocast_dtype)
+
+    def init_weights(self, pretrained):
+        if pretrained is not None:
+            logger = get_root_logger()
+            # load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger)
+            checkpoint = _load_checkpoint(pretrained, map_location='cpu', logger=logger)
+            if 'state_dict' in checkpoint:
+                state_dict = checkpoint['state_dict']
+            else:
+                state_dict = checkpoint
+            metadata = getattr(state_dict, '_metadata', OrderedDict())
+            state_dict._metadata = metadata
+            assert self.conv_in.weight.shape[1] == self.conv_out.weight.shape[0]
+            if state_dict['conv_in.weight'].size() != self.conv_in.weight.size():
+                assert state_dict['conv_in.weight'].shape[1] == state_dict['conv_out.weight'].shape[0]
+                src_chn = state_dict['conv_in.weight'].shape[1]
+                tgt_chn = self.conv_in.weight.shape[1]
+                assert src_chn < tgt_chn
+                convert_mat_out = torch.tile(torch.eye(src_chn), (ceildiv(tgt_chn, src_chn), 1))
+                convert_mat_out = convert_mat_out[:tgt_chn]
+                convert_mat_in = F.normalize(convert_mat_out.pinverse(), dim=-1)
+                state_dict['conv_out.weight'] = torch.einsum(
+                    'ts,scxy->tcxy', convert_mat_out, state_dict['conv_out.weight'])
+                state_dict['conv_out.bias'] = torch.einsum(
+                    'ts,s->t', convert_mat_out, state_dict['conv_out.bias'])
+                state_dict['conv_in.weight'] = torch.einsum(
+                    'st,csxy->ctxy', convert_mat_in, state_dict['conv_in.weight'])
+            load_state_dict(self, state_dict, logger=logger)
+
+    def forward(self, sample, timestep, encoder_hidden_states, **kwargs):
+        dtype = sample.dtype
+        return super().forward(
+            sample, timestep, encoder_hidden_states, return_dict=False, **kwargs)[0].to(dtype)
+
+    def forward_enc(self, sample, timestep, encoder_hidden_states, **kwargs):
+        return unet_enc(self, sample, timestep, encoder_hidden_states, **kwargs)
+
+    def forward_dec(self, emb, down_block_res_samples, sample, encoder_hidden_states, **kwargs):
+        return unet_dec(self, emb, down_block_res_samples, sample, encoder_hidden_states, **kwargs)

models/lakonlab/models/architecture/diffusers/__init__.py

@@ -0,0 +1,15 @@
+from .pretrained import (
+    PretrainedVAE, PretrainedVAEDecoder, PretrainedVAEEncoder, PretrainedVAEQwenImage,
+    PretrainedFluxTextEncoder, PretrainedQwenImageTextEncoder, PretrainedStableDiffusion3TextEncoder)
+from .unet import UNet2DConditionModel
+from .flux import FluxTransformer2DModel
+from .dit import DiTTransformer2DModelMod
+from .sd3 import SD3Transformer2DModel
+from .qwen import QwenImageTransformer2DModel
+
+__all__ = [
+    'PretrainedVAE', 'PretrainedVAEDecoder', 'PretrainedVAEEncoder', 'PretrainedFluxTextEncoder',
+    'PretrainedQwenImageTextEncoder', 'UNet2DConditionModel', 'FluxTransformer2DModel',
+    'DiTTransformer2DModelMod', 'SD3Transformer2DModel',
+    'QwenImageTransformer2DModel', 'PretrainedVAEQwenImage', 'PretrainedStableDiffusion3TextEncoder',
+]

models/lakonlab/models/architecture/diffusers/dit.py

@@ -0,0 +1,428 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from typing import Any, Dict, Optional
+from diffusers.models import DiTTransformer2DModel, ModelMixin
+from diffusers.models.attention import BasicTransformerBlock, _chunked_feed_forward, Attention, FeedForward
+from diffusers.models.embeddings import (
+    PatchEmbed, Timesteps, CombinedTimestepLabelEmbeddings, TimestepEmbedding, LabelEmbedding)
+from diffusers.models.normalization import AdaLayerNormZero
+from diffusers.configuration_utils import register_to_config
+from mmcv.runner import load_checkpoint, _load_checkpoint, load_state_dict
+from mmcv.cnn import constant_init, xavier_init
+from mmgen.models.builder import MODULES
+from mmgen.utils import get_root_logger
+from ..utils import flex_freeze
+
+
+class LabelEmbeddingMod(LabelEmbedding):
+    def __init__(self, num_classes, hidden_size, dropout_prob=0.0, use_cfg_embedding=True):
+        super(LabelEmbedding, self).__init__()
+        if dropout_prob > 0:
+            assert use_cfg_embedding
+        self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)
+        self.num_classes = num_classes
+        self.dropout_prob = dropout_prob
+
+
+class CombinedTimestepLabelEmbeddingsMod(CombinedTimestepLabelEmbeddings):
+    """
+    Modified CombinedTimestepLabelEmbeddings for reproducing the original DiT (downscale_freq_shift=0).
+    """
+    def __init__(
+            self, num_classes, embedding_dim, class_dropout_prob=0.1, downscale_freq_shift=0, use_cfg_embedding=True):
+        super(CombinedTimestepLabelEmbeddings, self).__init__()
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=downscale_freq_shift)
+        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+        self.class_embedder = LabelEmbeddingMod(num_classes, embedding_dim, class_dropout_prob, use_cfg_embedding)
+
+
+class BasicTransformerBlockMod(BasicTransformerBlock):
+    """
+    Modified BasicTransformerBlock for reproducing the original DiT with shared time and class
+    embeddings across all layers.
+    """
+    def __init__(
+            self,
+            dim: int,
+            num_attention_heads: int,
+            attention_head_dim: int,
+            dropout=0.0,
+            cross_attention_dim: Optional[int] = None,
+            activation_fn: str = 'geglu',
+            num_embeds_ada_norm: Optional[int] = None,
+            attention_bias: bool = False,
+            only_cross_attention: bool = False,
+            double_self_attention: bool = False,
+            upcast_attention: bool = False,
+            norm_elementwise_affine: bool = True,
+            norm_type: str = 'layer_norm',
+            norm_eps: float = 1e-5,
+            final_dropout: bool = False,
+            attention_type: str = 'default',
+            ada_norm_continous_conditioning_embedding_dim: Optional[int] = None,
+            ada_norm_bias: Optional[int] = None,
+            ff_inner_dim: Optional[int] = None,
+            ff_bias: bool = True,
+            attention_out_bias: bool = True):
+        super(BasicTransformerBlock, self).__init__()
+        self.only_cross_attention = only_cross_attention
+        self.norm_type = norm_type
+        self.num_embeds_ada_norm = num_embeds_ada_norm
+
+        assert self.norm_type == 'ada_norm_zero'
+        self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)
+        self.attn1 = Attention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            dropout=dropout,
+            bias=attention_bias,
+            cross_attention_dim=cross_attention_dim if only_cross_attention else None,
+            upcast_attention=upcast_attention,
+            out_bias=attention_out_bias,
+        )
+
+        self.norm2 = None
+        self.attn2 = None
+
+        self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
+        self.ff = FeedForward(
+            dim,
+            dropout=dropout,
+            activation_fn=activation_fn,
+            final_dropout=final_dropout,
+            inner_dim=ff_inner_dim,
+            bias=ff_bias,
+        )
+
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            attention_mask: Optional[torch.Tensor] = None,
+            encoder_hidden_states: Optional[torch.Tensor] = None,
+            encoder_attention_mask: Optional[torch.Tensor] = None,
+            timestep: Optional[torch.LongTensor] = None,
+            cross_attention_kwargs: Dict[str, Any] = None,
+            class_labels: Optional[torch.LongTensor] = None,
+            emb: Optional[torch.Tensor] = None,
+            added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None) -> torch.Tensor:
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+            hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype, emb=emb)
+
+        if cross_attention_kwargs is None:
+            cross_attention_kwargs = dict()
+        attn_output = self.attn1(
+            norm_hidden_states,
+            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
+            attention_mask=attention_mask,
+            **cross_attention_kwargs)
+        attn_output = gate_msa.unsqueeze(1) * attn_output
+
+        hidden_states = attn_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+
+        norm_hidden_states = self.norm3(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+        if self._chunk_size is not None:
+            # "feed_forward_chunk_size" can be used to save memory
+            ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)
+        else:
+            ff_output = self.ff(norm_hidden_states)
+
+        ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+        hidden_states = ff_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+
+        return hidden_states
+
+
+class _DiTTransformer2DModelMod(DiTTransformer2DModel):
+
+    @register_to_config
+    def __init__(
+            self,
+            class_dropout_prob=0.0,
+            num_attention_heads: int = 16,
+            attention_head_dim: int = 72,
+            in_channels: int = 4,
+            out_channels: Optional[int] = None,
+            num_layers: int = 28,
+            dropout: float = 0.0,
+            norm_num_groups: int = 32,
+            attention_bias: bool = True,
+            sample_size: int = 32,
+            patch_size: int = 2,
+            activation_fn: str = 'gelu-approximate',
+            num_embeds_ada_norm: Optional[int] = 1000,
+            upcast_attention: bool = False,
+            norm_type: str = 'ada_norm_zero',
+            norm_elementwise_affine: bool = False,
+            norm_eps: float = 1e-5):
+
+        super(DiTTransformer2DModel, self).__init__()
+
+        # Validate inputs.
+        if norm_type != "ada_norm_zero":
+            raise NotImplementedError(
+                f"Forward pass is not implemented when `patch_size` is not None and `norm_type` is '{norm_type}'."
+            )
+        elif norm_type == "ada_norm_zero" and num_embeds_ada_norm is None:
+            raise ValueError(
+                f"When using a `patch_size` and this `norm_type` ({norm_type}), `num_embeds_ada_norm` cannot be None."
+            )
+
+        # Set some common variables used across the board.
+        self.attention_head_dim = attention_head_dim
+        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
+        self.out_channels = in_channels if out_channels is None else out_channels
+        self.gradient_checkpointing = False
+
+        # 2. Initialize the position embedding and transformer blocks.
+        self.height = self.config.sample_size
+        self.width = self.config.sample_size
+
+        self.patch_size = self.config.patch_size
+        self.pos_embed = PatchEmbed(
+            height=self.config.sample_size,
+            width=self.config.sample_size,
+            patch_size=self.config.patch_size,
+            in_channels=self.config.in_channels,
+            embed_dim=self.inner_dim)
+        self.emb = CombinedTimestepLabelEmbeddingsMod(
+            num_embeds_ada_norm, self.inner_dim, class_dropout_prob=0.0)
+
+        self.transformer_blocks = nn.ModuleList([
+            BasicTransformerBlockMod(
+                self.inner_dim,
+                self.config.num_attention_heads,
+                self.config.attention_head_dim,
+                dropout=self.config.dropout,
+                activation_fn=self.config.activation_fn,
+                num_embeds_ada_norm=None,
+                attention_bias=self.config.attention_bias,
+                upcast_attention=self.config.upcast_attention,
+                norm_type=norm_type,
+                norm_elementwise_affine=self.config.norm_elementwise_affine,
+                norm_eps=self.config.norm_eps)
+            for _ in range(self.config.num_layers)])
+
+        # 3. Output blocks.
+        self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim)
+        self.proj_out_2 = nn.Linear(
+            self.inner_dim, self.config.patch_size * self.config.patch_size * self.out_channels)
+
+    # https://github.com/facebookresearch/DiT/blob/main/models.py
+    def init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                xavier_init(m, distribution='uniform')
+            elif isinstance(m, nn.Embedding):
+                torch.nn.init.normal_(m.weight, mean=0.0, std=0.02)
+
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d)
+        w = self.pos_embed.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.pos_embed.proj.bias, 0)
+
+        # Zero-out adaLN modulation layers in DiT blocks
+        for m in self.modules():
+            if isinstance(m, AdaLayerNormZero):
+                constant_init(m.linear, val=0)
+
+        # Zero-out output layers
+        constant_init(self.proj_out_1, val=0)
+        constant_init(self.proj_out_2, val=0)
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            timestep: Optional[torch.LongTensor] = None,
+            class_labels: Optional[torch.LongTensor] = None,
+            cross_attention_kwargs: Dict[str, Any] = None):
+        # 1. Input
+        bs, _, h, w = hidden_states.size()
+        height, width = h // self.patch_size, w // self.patch_size
+        hidden_states = self.pos_embed(hidden_states)
+
+        cond_emb = self.emb(
+            timestep, class_labels, hidden_dtype=hidden_states.dtype)
+        dropout_enabled = self.config.class_dropout_prob > 0 and self.training
+        if dropout_enabled:
+            uncond_emb = self.emb(timestep, torch.full_like(
+                class_labels, self.config.num_embeds_ada_norm), hidden_dtype=hidden_states.dtype)
+
+        # 2. Blocks
+        for block in self.transformer_blocks:
+            if dropout_enabled:
+                dropout_mask = torch.rand((bs, 1), device=hidden_states.device) < self.config.class_dropout_prob
+                emb = torch.where(dropout_mask, uncond_emb, cond_emb)
+            else:
+                emb = cond_emb
+
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    None,
+                    None,
+                    None,
+                    timestep,
+                    cross_attention_kwargs,
+                    class_labels,
+                    emb,
+                    use_reentrant=False)
+
+            else:
+                hidden_states = block(
+                    hidden_states,
+                    attention_mask=None,
+                    encoder_hidden_states=None,
+                    encoder_attention_mask=None,
+                    timestep=timestep,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    class_labels=class_labels,
+                    emb=emb)
+
+        # 3. Output
+        if dropout_enabled:
+            dropout_mask = torch.rand((bs, 1), device=hidden_states.device) < self.config.class_dropout_prob
+            emb = torch.where(dropout_mask, uncond_emb, cond_emb)
+        else:
+            emb = cond_emb
+        shift, scale = self.proj_out_1(F.silu(emb)).chunk(2, dim=1)
+        hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
+        output = self.proj_out_2(hidden_states).reshape(
+                bs, height, width, self.patch_size, self.patch_size, self.out_channels
+            ).permute(0, 5, 1, 3, 2, 4).reshape(
+                bs, self.out_channels, height * self.patch_size, width * self.patch_size)
+
+        return output
+
+
+@MODULES.register_module()
+class DiTTransformer2DModelMod(_DiTTransformer2DModelMod):
+
+    def __init__(
+            self,
+            *args,
+            freeze=False,
+            freeze_exclude=[],
+            pretrained=None,
+            torch_dtype='float32',
+            autocast_dtype=None,
+            freeze_exclude_fp32=True,
+            freeze_exclude_autocast_dtype='float32',
+            checkpointing=True,
+            **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.init_weights(pretrained)
+
+        if autocast_dtype is not None:
+            assert torch_dtype == 'float32'
+        self.autocast_dtype = autocast_dtype
+
+        if torch_dtype is not None:
+            self.to(getattr(torch, torch_dtype))
+
+        self.freeze = freeze
+        if self.freeze:
+            flex_freeze(
+                self,
+                exclude_keys=freeze_exclude,
+                exclude_fp32=freeze_exclude_fp32,
+                exclude_autocast_dtype=freeze_exclude_autocast_dtype)
+
+        if checkpointing:
+            self.enable_gradient_checkpointing()
+
+    def init_weights(self, pretrained=None):
+        super().init_weights()
+        if pretrained is not None:
+            logger = get_root_logger()
+            # load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger)
+            checkpoint = _load_checkpoint(pretrained, map_location='cpu', logger=logger)
+            if 'state_dict' in checkpoint:
+                state_dict = checkpoint['state_dict']
+            else:
+                state_dict = checkpoint
+            # load from GMDiT V1 model with 1 Gaussian
+            p2 = self.config.patch_size * self.config.patch_size
+            ori_out_channels = p2 * self.out_channels
+            if 'proj_out_2.weight' in state_dict:
+                # if this is GMDiT V1 model with 1 Gaussian
+                if state_dict['proj_out_2.weight'].size(0) == p2 * (self.out_channels + 1):
+                    state_dict['proj_out_2.weight'] = state_dict['proj_out_2.weight'].reshape(
+                        p2, self.out_channels + 1, -1
+                    )[:, :-1].reshape(ori_out_channels, -1)
+                # if this is original DiT with variance prediction
+                if state_dict['proj_out_2.weight'].size(0) == 2 * ori_out_channels:
+                    state_dict['proj_out_2.weight'] = state_dict['proj_out_2.weight'].reshape(
+                        p2, 2 * self.out_channels, -1
+                    )[:, :self.out_channels].reshape(ori_out_channels, -1)
+            if 'proj_out_2.bias' in state_dict:
+                # if this is GMDiT V1 model with 1 Gaussian
+                if state_dict['proj_out_2.bias'].size(0) == p2 * (self.out_channels + 1):
+                    state_dict['proj_out_2.bias'] = state_dict['proj_out_2.bias'].reshape(
+                        p2, self.out_channels + 1
+                    )[:, :-1].reshape(ori_out_channels)
+                # if this is original DiT with variance prediction
+                if state_dict['proj_out_2.bias'].size(0) == 2 * ori_out_channels:
+                    state_dict['proj_out_2.bias'] = state_dict['proj_out_2.bias'].reshape(
+                        p2, 2 * self.out_channels
+                    )[:, :self.out_channels].reshape(ori_out_channels)
+            if 'emb.class_embedder.embedding_table.weight' not in state_dict \
+                    and 'transformer_blocks.0.norm1.emb.class_embedder.embedding_table.weight' in state_dict:
+                # convert original diffusers DiT model to our modified DiT model with shared embeddings
+                keys_to_remove = []
+                state_update = {}
+                for k, v in state_dict.items():
+                    if k.startswith('transformer_blocks.0.norm1.emb.'):
+                        new_k = k.replace('transformer_blocks.0.norm1.', '')
+                        state_update[new_k] = v
+                    if k.startswith('transformer_blocks.') and '.norm1.emb.' in k:
+                        keys_to_remove.append(k)
+                state_dict.update(state_update)
+                for k in keys_to_remove:
+                    del state_dict[k]
+            load_state_dict(self, state_dict, logger=logger)
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            timestep: Optional[torch.LongTensor] = None,
+            class_labels: Optional[torch.LongTensor] = None,
+            **kwargs):
+        if self.autocast_dtype is not None:
+            dtype = getattr(torch, self.autocast_dtype)
+        else:
+            dtype = hidden_states.dtype
+        with torch.autocast(
+                device_type='cuda',
+                enabled=self.autocast_dtype is not None,
+                dtype=dtype if self.autocast_dtype is not None else None):
+            return super().forward(
+                hidden_states.to(dtype),
+                timestep=timestep,
+                class_labels=class_labels,
+                **kwargs)

models/lakonlab/models/architecture/diffusers/flux.py

@@ -0,0 +1,156 @@
+import torch
+
+from typing import Optional
+from accelerate import init_empty_weights
+from diffusers.models import FluxTransformer2DModel as _FluxTransformer2DModel
+from peft import LoraConfig
+from mmgen.models.builder import MODULES
+from mmgen.utils import get_root_logger
+from ..utils import flex_freeze
+from lakonlab.runner.checkpoint import load_checkpoint, _load_checkpoint
+
+
+@MODULES.register_module()
+class FluxTransformer2DModel(_FluxTransformer2DModel):
+
+    def __init__(
+            self,
+            *args,
+            patch_size=2,
+            freeze=False,
+            freeze_exclude=[],
+            pretrained=None,
+            pretrained_lora=None,
+            pretrained_lora_scale=1.0,
+            torch_dtype='float32',
+            freeze_exclude_fp32=True,
+            freeze_exclude_autocast_dtype='float32',
+            checkpointing=True,
+            use_lora=False,
+            lora_target_modules=None,
+            lora_rank=16,
+            **kwargs):
+        with init_empty_weights():
+            super().__init__(patch_size=1, *args, **kwargs)
+        self.patch_size = patch_size
+
+        self.init_weights(pretrained, pretrained_lora, pretrained_lora_scale)
+
+        self.use_lora = use_lora
+        self.lora_target_modules = lora_target_modules
+        self.lora_rank = lora_rank
+        if self.use_lora:
+            transformer_lora_config = LoraConfig(
+                r=lora_rank,
+                lora_alpha=lora_rank,
+                init_lora_weights='gaussian',
+                target_modules=lora_target_modules,
+            )
+            self.add_adapter(transformer_lora_config)
+
+        if torch_dtype is not None:
+            self.to(getattr(torch, torch_dtype))
+
+        self.freeze = freeze
+        if self.freeze:
+            flex_freeze(
+                self,
+                exclude_keys=freeze_exclude,
+                exclude_fp32=freeze_exclude_fp32,
+                exclude_autocast_dtype=freeze_exclude_autocast_dtype)
+
+        if checkpointing:
+            self.enable_gradient_checkpointing()
+
+    def init_weights(self, pretrained=None, pretrained_lora=None, pretrained_lora_scale=1.0):
+        if pretrained is not None:
+            logger = get_root_logger()
+            load_checkpoint(
+                self, pretrained, map_location='cpu', strict=False, logger=logger, assign=True)
+            if pretrained_lora is not None:
+                if not isinstance(pretrained_lora, (list, tuple)):
+                    assert isinstance(pretrained_lora, str)
+                    pretrained_lora = [pretrained_lora]
+                if not isinstance(pretrained_lora_scale, (list, tuple)):
+                    assert isinstance(pretrained_lora_scale, (int, float))
+                    pretrained_lora_scale = [pretrained_lora_scale]
+                for pretrained_lora_single, pretrained_lora_scale_single in zip(pretrained_lora, pretrained_lora_scale):
+                    lora_state_dict = _load_checkpoint(
+                        pretrained_lora_single, map_location='cpu', logger=logger)
+                    self.load_lora_adapter(lora_state_dict)
+                    self.fuse_lora(lora_scale=pretrained_lora_scale_single)
+                    self.unload_lora()
+
+    @staticmethod
+    def _prepare_latent_image_ids(height, width, device, dtype):
+        """
+        Copied from Diffusers
+        """
+        latent_image_ids = torch.zeros(height, width, 3)
+        latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None]
+        latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :]
+
+        latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+
+        latent_image_ids = latent_image_ids.reshape(
+            latent_image_id_height * latent_image_id_width, latent_image_id_channels)
+
+        return latent_image_ids.to(device=device, dtype=dtype)
+
+    def patchify(self, latents):
+        if self.patch_size > 1:
+            bs, c, h, w = latents.size()
+            latents = latents.reshape(
+                bs, c, h // self.patch_size, self.patch_size, w // self.patch_size, self.patch_size
+            ).permute(
+                0, 1, 3, 5, 2, 4
+            ).reshape(
+                bs, c * self.patch_size * self.patch_size, h // self.patch_size, w // self.patch_size)
+        return latents
+
+    def unpatchify(self, latents):
+        if self.patch_size > 1:
+            bs, c, h, w = latents.size()
+            latents = latents.reshape(
+                bs, c // (self.patch_size * self.patch_size), self.patch_size, self.patch_size, h, w
+            ).permute(
+                0, 1, 4, 2, 5, 3
+            ).reshape(
+                bs, c // (self.patch_size * self.patch_size), h * self.patch_size, w * self.patch_size)
+        return latents
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            timestep: torch.Tensor,
+            encoder_hidden_states: torch.Tensor = None,
+            pooled_projections: torch.Tensor = None,
+            mask: Optional[torch.Tensor] = None,
+            masked_image_latents: Optional[torch.Tensor] = None,
+            **kwargs):
+        hidden_states = self.patchify(hidden_states)
+        bs, c, h, w = hidden_states.size()
+        dtype = hidden_states.dtype
+        device = hidden_states.device
+        hidden_states = hidden_states.reshape(bs, c, h * w).permute(0, 2, 1)
+        img_ids = self._prepare_latent_image_ids(
+            h, w, device, dtype)
+        txt_ids = img_ids.new_zeros((encoder_hidden_states.shape[-2], 3))
+
+        #  Flux fill
+        if mask is not None and masked_image_latents is not None:
+            hidden_states = torch.cat(
+                (hidden_states, masked_image_latents.to(dtype=dtype), mask.to(dtype=dtype)), dim=-1)
+
+        output = super().forward(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states.to(dtype),
+            pooled_projections=pooled_projections.to(dtype),
+            timestep=timestep,
+            img_ids=img_ids,
+            txt_ids=txt_ids,
+            return_dict=False,
+            **kwargs)[0]
+
+        output = output.permute(0, 2, 1).reshape(bs, self.out_channels, h, w)
+        return self.unpatchify(output)

models/lakonlab/models/architecture/diffusers/pretrained.py

@@ -0,0 +1,281 @@
+import logging
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.models import AutoencoderKL, AutoencoderKLQwenImage
+from diffusers.pipelines import FluxPipeline, QwenImagePipeline, StableDiffusion3Pipeline
+from mmgen.models.builder import MODULES
+
+# Suppress truncation warnings from transformers and diffusers
+for name in (
+        'transformers.tokenization_utils_base',
+        'transformers.tokenization_utils',
+        'transformers.tokenization_utils_fast'):
+    logging.getLogger(name).setLevel(logging.ERROR)
+
+for name, logger in logging.root.manager.loggerDict.items():
+    if isinstance(logger, logging.Logger) and (name.startswith('diffusers.pipelines.')):
+        logger.setLevel(logging.ERROR)
+
+
+@MODULES.register_module()
+class PretrainedVAE(nn.Module):
+    def __init__(self,
+                 from_pretrained=None,
+                 del_encoder=False,
+                 del_decoder=False,
+                 use_slicing=False,
+                 freeze=True,
+                 eval_mode=True,
+                 torch_dtype='float32',
+                 **kwargs):
+        super().__init__()
+        if torch_dtype is not None:
+            kwargs.update(torch_dtype=getattr(torch, torch_dtype))
+        self.vae = AutoencoderKL.from_pretrained(
+            from_pretrained, **kwargs)
+        if del_encoder:
+            del self.vae.encoder
+        if del_decoder:
+            del self.vae.decoder
+        if use_slicing:
+            self.vae.enable_slicing()
+        self.freeze = freeze
+        self.eval_mode = eval_mode
+        if self.freeze:
+            self.requires_grad_(False)
+        if self.eval_mode:
+            self.eval()
+        self.vae.set_use_memory_efficient_attention_xformers(
+            not hasattr(torch.nn.functional, 'scaled_dot_product_attention'))
+
+    def train(self, mode=True):
+        mode = mode and (not self.eval_mode)
+        return super().train(mode)
+
+    def forward(self, *args, **kwargs):
+        return self.vae(*args, return_dict=False, **kwargs)[0]
+
+    def encode(self, img):
+        scaling_factor = self.vae.config.scaling_factor
+        shift_factor = self.vae.config.shift_factor
+        if scaling_factor is None:
+            scaling_factor = 1.0
+        if shift_factor is None:
+            shift_factor = 0.0
+        return (self.vae.encode(img).latent_dist.sample() - shift_factor) * scaling_factor
+
+    def decode(self, code):
+        scaling_factor = self.vae.config.scaling_factor
+        shift_factor = self.vae.config.shift_factor
+        if scaling_factor is None:
+            scaling_factor = 1.0
+        if shift_factor is None:
+            shift_factor = 0.0
+        return self.vae.decode(code / scaling_factor + shift_factor, return_dict=False)[0]
+
+
+@MODULES.register_module()
+class PretrainedVAEDecoder(PretrainedVAE):
+    def __init__(self, **kwargs):
+        super().__init__(
+            del_encoder=True,
+            del_decoder=False,
+            **kwargs)
+
+    def forward(self, code):
+        return super().decode(code)
+
+
+@MODULES.register_module()
+class PretrainedVAEEncoder(PretrainedVAE):
+    def __init__(self, **kwargs):
+        super().__init__(
+            del_encoder=False,
+            del_decoder=True,
+            **kwargs)
+
+    def forward(self, img):
+        return super().encode(img)
+
+
+@MODULES.register_module()
+class PretrainedVAEQwenImage(nn.Module):
+    def __init__(self,
+                 from_pretrained=None,
+                 use_slicing=False,
+                 freeze=True,
+                 eval_mode=True,
+                 torch_dtype='float32',
+                 **kwargs):
+        super().__init__()
+        if torch_dtype is not None:
+            kwargs.update(torch_dtype=getattr(torch, torch_dtype))
+        self.vae = AutoencoderKLQwenImage.from_pretrained(
+            from_pretrained, **kwargs)
+        if use_slicing:
+            self.vae.enable_slicing()
+        self.freeze = freeze
+        self.eval_mode = eval_mode
+        if self.freeze:
+            self.requires_grad_(False)
+        if self.eval_mode:
+            self.eval()
+
+    def train(self, mode=True):
+        mode = mode and (not self.eval_mode)
+        return super().train(mode)
+
+    def forward(self, *args, **kwargs):
+        return self.vae(*args, return_dict=False, **kwargs)[0]
+
+    def encode(self, img):
+        device = img.device
+        dtype = img.dtype
+        latents_mean = torch.tensor(self.vae.config.latents_mean, device=device, dtype=dtype).view(
+            1, self.vae.config.z_dim, 1, 1, 1)
+        latents_std = torch.tensor(self.vae.config.latents_std, device=device, dtype=dtype).view(
+            1, self.vae.config.z_dim, 1, 1, 1)
+        return ((self.vae.encode(img.unsqueeze(-3)).latent_dist.sample() - latents_mean) / latents_std).squeeze(-3)
+
+    def decode(self, code):
+        device = code.device
+        dtype = code.dtype
+        latents_mean = torch.tensor(self.vae.config.latents_mean, device=device, dtype=dtype).view(
+            1, self.vae.config.z_dim, 1, 1, 1)
+        latents_std = torch.tensor(self.vae.config.latents_std, device=device, dtype=dtype).view(
+            1, self.vae.config.z_dim, 1, 1, 1)
+        return self.vae.decode(code.unsqueeze(-3) * latents_std + latents_mean, return_dict=False)[0].squeeze(-3)
+
+
+@MODULES.register_module()
+class PretrainedFluxTextEncoder(nn.Module):
+    def __init__(self,
+                 from_pretrained='black-forest-labs/FLUX.1-dev',
+                 freeze=True,
+                 eval_mode=True,
+                 torch_dtype='bfloat16',
+                 max_sequence_length=512,
+                 **kwargs):
+        super().__init__()
+        self.max_sequence_length = max_sequence_length
+        self.pipeline = FluxPipeline.from_pretrained(
+            from_pretrained,
+            scheduler=None,
+            vae=None,
+            transformer=None,
+            image_encoder=None,
+            feature_extractor=None,
+            torch_dtype=getattr(torch, torch_dtype),
+            **kwargs)
+        self.text_encoder = self.pipeline.text_encoder
+        self.text_encoder_2 = self.pipeline.text_encoder_2
+        self.freeze = freeze
+        self.eval_mode = eval_mode
+        if self.freeze:
+            self.requires_grad_(False)
+        if self.eval_mode:
+            self.eval()
+
+    def train(self, mode=True):
+        mode = mode and (not self.eval_mode)
+        return super().train(mode)
+
+    def forward(self, prompt, prompt_2=None):
+        prompt_embeds, pooled_prompt_embeds, text_ids = self.pipeline.encode_prompt(
+            prompt, prompt_2=prompt_2, max_sequence_length=self.max_sequence_length)
+        return dict(
+            encoder_hidden_states=prompt_embeds,
+            pooled_projections=pooled_prompt_embeds)
+
+
+@MODULES.register_module()
+class PretrainedQwenImageTextEncoder(nn.Module):
+    def __init__(self,
+                 from_pretrained='Qwen/Qwen-Image',
+                 freeze=True,
+                 eval_mode=True,
+                 torch_dtype='bfloat16',
+                 max_sequence_length=512,
+                 pad_seq_len=None,
+                 **kwargs):
+        super().__init__()
+        self.max_sequence_length = max_sequence_length
+        if pad_seq_len is not None:
+            assert pad_seq_len >= max_sequence_length
+        self.pad_seq_len = pad_seq_len
+        self.pipeline = QwenImagePipeline.from_pretrained(
+            from_pretrained,
+            scheduler=None,
+            vae=None,
+            transformer=None,
+            torch_dtype=getattr(torch, torch_dtype),
+            **kwargs)
+        self.text_encoder = self.pipeline.text_encoder
+        self.freeze = freeze
+        self.eval_mode = eval_mode
+        if self.freeze:
+            self.requires_grad_(False)
+        if self.eval_mode:
+            self.eval()
+
+    def train(self, mode=True):
+        mode = mode and (not self.eval_mode)
+        return super().train(mode)
+
+    def forward(self, prompt):
+        prompt_embeds, prompt_embeds_mask = self.pipeline.encode_prompt(
+            prompt, max_sequence_length=self.max_sequence_length)
+        if self.pad_seq_len is not None:
+            pad_len = self.pad_seq_len - prompt_embeds.size(1)
+            prompt_embeds = F.pad(
+                prompt_embeds, (0, 0, 0, pad_len), value=0.0)
+            prompt_embeds_mask = F.pad(
+                prompt_embeds_mask, (0, pad_len), value=0.0)
+        return dict(
+            encoder_hidden_states=prompt_embeds,
+            encoder_hidden_states_mask=prompt_embeds_mask)
+
+
+@MODULES.register_module()
+class PretrainedStableDiffusion3TextEncoder(nn.Module):
+    def __init__(self,
+                 from_pretrained='stabilityai/stable-diffusion-3.5-large',
+                 freeze=True,
+                 eval_mode=True,
+                 torch_dtype='float32',
+                 max_sequence_length=256,
+                 **kwargs):
+        super().__init__()
+        self.max_sequence_length = max_sequence_length
+        self.pipeline = StableDiffusion3Pipeline.from_pretrained(
+            from_pretrained,
+            scheduler=None,
+            vae=None,
+            transformer=None,
+            image_encoder=None,
+            feature_extractor=None,
+            torch_dtype=getattr(torch, torch_dtype),
+            **kwargs)
+        self.text_encoder = self.pipeline.text_encoder
+        self.text_encoder_2 = self.pipeline.text_encoder_2
+        self.text_encoder_3 = self.pipeline.text_encoder_3
+        self.freeze = freeze
+        self.eval_mode = eval_mode
+        if self.freeze:
+            self.requires_grad_(False)
+        if self.eval_mode:
+            self.eval()
+
+    def train(self, mode=True):
+        mode = mode and (not self.eval_mode)
+        return super().train(mode)
+
+    def forward(self, prompt, prompt_2=None, prompt_3=None):
+        prompt_embeds, _, pooled_prompt_embeds, _ = self.pipeline.encode_prompt(
+            prompt, prompt_2=prompt_2, prompt_3=prompt_3, do_classifier_free_guidance=False,
+            max_sequence_length=self.max_sequence_length)
+        return dict(
+            encoder_hidden_states=prompt_embeds,
+            pooled_projections=pooled_prompt_embeds)

models/lakonlab/models/architecture/diffusers/qwen.py

@@ -0,0 +1,139 @@
+import torch
+
+from accelerate import init_empty_weights
+from diffusers.models import QwenImageTransformer2DModel as _QwenImageTransformer2DModel
+from diffusers.loaders.lora_conversion_utils import _convert_non_diffusers_qwen_lora_to_diffusers
+from peft import LoraConfig
+from mmgen.models.builder import MODULES
+from mmgen.utils import get_root_logger
+from ..utils import flex_freeze
+from lakonlab.runner.checkpoint import load_checkpoint, _load_checkpoint
+
+
+@MODULES.register_module()
+class QwenImageTransformer2DModel(_QwenImageTransformer2DModel):
+
+    def __init__(
+            self,
+            *args,
+            patch_size=2,
+            freeze=False,
+            freeze_exclude=[],
+            pretrained=None,
+            pretrained_lora=None,
+            pretrained_lora_scale=1.0,
+            torch_dtype='float32',
+            freeze_exclude_fp32=True,
+            freeze_exclude_autocast_dtype='float32',
+            checkpointing=True,
+            use_lora=False,
+            lora_target_modules=None,
+            lora_rank=16,
+            **kwargs):
+        with init_empty_weights():
+            super().__init__(*args, patch_size=1, **kwargs)
+        self.patch_size = patch_size
+
+        self.init_weights(pretrained, pretrained_lora, pretrained_lora_scale)
+
+        self.use_lora = use_lora
+        self.lora_target_modules = lora_target_modules
+        self.lora_rank = lora_rank
+        if self.use_lora:
+            transformer_lora_config = LoraConfig(
+                r=lora_rank,
+                lora_alpha=lora_rank,
+                init_lora_weights='gaussian',
+                target_modules=lora_target_modules,
+            )
+            self.add_adapter(transformer_lora_config)
+
+        if torch_dtype is not None:
+            self.to(getattr(torch, torch_dtype))
+
+        self.freeze = freeze
+        if self.freeze:
+            flex_freeze(
+                self,
+                exclude_keys=freeze_exclude,
+                exclude_fp32=freeze_exclude_fp32,
+                exclude_autocast_dtype=freeze_exclude_autocast_dtype)
+
+        if checkpointing:
+            self.enable_gradient_checkpointing()
+
+    def init_weights(self, pretrained=None, pretrained_lora=None, pretrained_lora_scale=1.0):
+        if pretrained is not None:
+            logger = get_root_logger()
+            load_checkpoint(
+                self, pretrained, map_location='cpu', strict=False, logger=logger, assign=True)
+            if pretrained_lora is not None:
+                if not isinstance(pretrained_lora, (list, tuple)):
+                    assert isinstance(pretrained_lora, str)
+                    pretrained_lora = [pretrained_lora]
+                if not isinstance(pretrained_lora_scale, (list, tuple)):
+                    assert isinstance(pretrained_lora_scale, (int, float))
+                    pretrained_lora_scale = [pretrained_lora_scale]
+                for pretrained_lora_single, pretrained_lora_scale_single in zip(pretrained_lora, pretrained_lora_scale):
+                    lora_state_dict = _load_checkpoint(
+                        pretrained_lora_single, map_location='cpu', logger=logger)
+                    lora_state_dict = _convert_non_diffusers_qwen_lora_to_diffusers(lora_state_dict)
+                    self.load_lora_adapter(lora_state_dict)
+                    self.fuse_lora(lora_scale=pretrained_lora_scale_single)
+                    self.unload_lora()
+
+    def patchify(self, latents):
+        if self.patch_size > 1:
+            bs, c, h, w = latents.size()
+            latents = latents.reshape(
+                bs, c, h // self.patch_size, self.patch_size, w // self.patch_size, self.patch_size
+            ).permute(
+                0, 1, 3, 5, 2, 4
+            ).reshape(
+                bs, c * self.patch_size * self.patch_size, h // self.patch_size, w // self.patch_size)
+        return latents
+
+    def unpatchify(self, latents):
+        if self.patch_size > 1:
+            bs, c, h, w = latents.size()
+            latents = latents.reshape(
+                bs, c // (self.patch_size * self.patch_size), self.patch_size, self.patch_size, h, w
+            ).permute(
+                0, 1, 4, 2, 5, 3
+            ).reshape(
+                bs, c // (self.patch_size * self.patch_size), h * self.patch_size, w * self.patch_size)
+        return latents
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            timestep: torch.Tensor,
+            encoder_hidden_states: torch.Tensor = None,
+            encoder_hidden_states_mask: torch.Tensor = None,
+            **kwargs):
+        hidden_states = self.patchify(hidden_states)
+        bs, c, h, w = hidden_states.size()
+        dtype = hidden_states.dtype
+        hidden_states = hidden_states.reshape(bs, c, h * w).permute(0, 2, 1)
+        img_shapes = [[(1, h, w)]]
+        if encoder_hidden_states_mask is not None:
+            txt_seq_lens = encoder_hidden_states_mask.sum(dim=1)
+            max_txt_seq_len = txt_seq_lens.max()
+            encoder_hidden_states = encoder_hidden_states[:, :max_txt_seq_len]
+            encoder_hidden_states_mask = encoder_hidden_states_mask[:, :max_txt_seq_len]
+            txt_seq_lens = txt_seq_lens.tolist()
+        else:
+            txt_seq_lens = None
+
+        output = super().forward(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states.to(dtype),
+            encoder_hidden_states_mask=encoder_hidden_states_mask,
+            timestep=timestep,
+            img_shapes=img_shapes,
+            txt_seq_lens=txt_seq_lens,
+            return_dict=False,
+            **kwargs)[0]
+
+        output = output.permute(0, 2, 1).reshape(bs, self.out_channels, h, w)
+        return self.unpatchify(output)

models/lakonlab/models/architecture/diffusers/sd3.py

@@ -0,0 +1,80 @@
+import torch
+
+from accelerate import init_empty_weights
+from diffusers.models import SD3Transformer2DModel as _SD3Transformer2DModel
+from peft import LoraConfig
+from mmgen.models.builder import MODULES
+from mmgen.utils import get_root_logger
+from ..utils import flex_freeze
+from lakonlab.runner.checkpoint import load_checkpoint
+
+
+@MODULES.register_module()
+class SD3Transformer2DModel(_SD3Transformer2DModel):
+
+    def __init__(
+            self,
+            *args,
+            freeze=False,
+            freeze_exclude=[],
+            pretrained=None,
+            torch_dtype='float32',
+            freeze_exclude_fp32=True,
+            freeze_exclude_autocast_dtype='float32',
+            checkpointing=True,
+            use_lora=False,
+            lora_target_modules=None,
+            lora_rank=16,
+            **kwargs):
+        with init_empty_weights():
+            super().__init__(*args, **kwargs)
+        self.init_weights(pretrained)
+
+        self.use_lora = use_lora
+        self.lora_target_modules = lora_target_modules
+        self.lora_rank = lora_rank
+        if self.use_lora:
+            transformer_lora_config = LoraConfig(
+                r=lora_rank,
+                lora_alpha=lora_rank,
+                init_lora_weights='gaussian',
+                target_modules=lora_target_modules,
+            )
+            self.add_adapter(transformer_lora_config)
+
+        if torch_dtype is not None:
+            self.to(getattr(torch, torch_dtype))
+
+        self.freeze = freeze
+        if self.freeze:
+            flex_freeze(
+                self,
+                exclude_keys=freeze_exclude,
+                exclude_fp32=freeze_exclude_fp32,
+                exclude_autocast_dtype=freeze_exclude_autocast_dtype)
+
+        if checkpointing:
+            self.enable_gradient_checkpointing()
+
+    def init_weights(self, pretrained=None):
+        if pretrained is not None:
+            logger = get_root_logger()
+            load_checkpoint(
+                self, pretrained, map_location='cpu', strict=False, logger=logger, assign=True)
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            timestep: torch.Tensor,
+            encoder_hidden_states: torch.Tensor = None,
+            pooled_projections: torch.Tensor = None,
+            **kwargs):
+        dtype = hidden_states.dtype
+
+        return super().forward(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states.to(dtype),
+            pooled_projections=pooled_projections.to(dtype),
+            timestep=timestep,
+            return_dict=False,
+            **kwargs)[0]

models/lakonlab/models/architecture/diffusers/unet.py

@@ -0,0 +1,192 @@
+import torch
+import torch.nn.functional as F
+
+from typing import Dict, Any, Optional, Union, Tuple
+from collections import OrderedDict
+from diffusers.models import UNet2DConditionModel as _UNet2DConditionModel
+from mmcv.runner import _load_checkpoint, load_state_dict
+from mmgen.models.builder import MODULES
+from mmgen.utils import get_root_logger
+from ..utils import flex_freeze
+
+
+def ceildiv(a, b):
+    return -(a // -b)
+
+
+def unet_enc(
+        unet,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        added_cond_kwargs=None):
+    # 0. center input if necessary
+    if unet.config.center_input_sample:
+        sample = 2 * sample - 1.0
+
+    # 1. time
+    t_emb = unet.get_time_embed(sample=sample, timestep=timestep)
+    emb = unet.time_embedding(t_emb)
+    aug_emb = unet.get_aug_embed(
+        emb=emb, encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs)
+    emb = emb + aug_emb if aug_emb is not None else emb
+
+    if unet.time_embed_act is not None:
+        emb = unet.time_embed_act(emb)
+
+    encoder_hidden_states = unet.process_encoder_hidden_states(
+        encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs)
+
+    # 2. pre-process
+    sample = unet.conv_in(sample)
+
+    # 3. down
+    down_block_res_samples = (sample,)
+    for downsample_block in unet.down_blocks:
+        if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+            sample, res_samples = downsample_block(
+                hidden_states=sample,
+                temb=emb,
+                encoder_hidden_states=encoder_hidden_states,
+                cross_attention_kwargs=cross_attention_kwargs,
+            )
+        else:
+            sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+
+        down_block_res_samples += res_samples
+
+    return emb, down_block_res_samples, sample
+
+
+def unet_dec(
+        unet,
+        emb,
+        down_block_res_samples,
+        sample,
+        encoder_hidden_states: torch.Tensor,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        mid_block_additional_residual: Optional[torch.Tensor] = None):
+    is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
+
+    if is_controlnet:
+        new_down_block_res_samples = ()
+
+        for down_block_res_sample, down_block_additional_residual in zip(
+                down_block_res_samples, down_block_additional_residuals):
+            down_block_res_sample = down_block_res_sample + down_block_additional_residual
+            new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
+
+        down_block_res_samples = new_down_block_res_samples
+
+    # 4. mid
+    if unet.mid_block is not None:
+        if hasattr(unet.mid_block, "has_cross_attention") and unet.mid_block.has_cross_attention:
+            sample = unet.mid_block(
+                sample,
+                emb,
+                encoder_hidden_states=encoder_hidden_states,
+                cross_attention_kwargs=cross_attention_kwargs,
+            )
+        else:
+            sample = unet.mid_block(sample, emb)
+
+    if is_controlnet:
+        sample = sample + mid_block_additional_residual
+
+    # 5. up
+    for i, upsample_block in enumerate(unet.up_blocks):
+        res_samples = down_block_res_samples[-len(upsample_block.resnets):]
+        down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+        if hasattr(upsample_block, 'has_cross_attention') and upsample_block.has_cross_attention:
+            sample = upsample_block(
+                hidden_states=sample,
+                temb=emb,
+                res_hidden_states_tuple=res_samples,
+                encoder_hidden_states=encoder_hidden_states,
+                cross_attention_kwargs=cross_attention_kwargs,
+            )
+        else:
+            sample = upsample_block(
+                hidden_states=sample,
+                temb=emb,
+                res_hidden_states_tuple=res_samples,
+            )
+
+    # 6. post-process
+    if unet.conv_norm_out:
+        sample = unet.conv_norm_out(sample)
+        sample = unet.conv_act(sample)
+    sample = unet.conv_out(sample)
+
+    return sample
+
+
+@MODULES.register_module()
+class UNet2DConditionModel(_UNet2DConditionModel):
+    def __init__(self,
+                 *args,
+                 freeze=True,
+                 freeze_exclude=[],
+                 pretrained=None,
+                 torch_dtype='float32',
+                 freeze_exclude_fp32=True,
+                 freeze_exclude_autocast_dtype='float32',
+                 **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.init_weights(pretrained)
+        if torch_dtype is not None:
+            self.to(getattr(torch, torch_dtype))
+
+        self.set_use_memory_efficient_attention_xformers(
+            not hasattr(torch.nn.functional, 'scaled_dot_product_attention'))
+
+        self.freeze = freeze
+        if self.freeze:
+            flex_freeze(
+                self,
+                exclude_keys=freeze_exclude,
+                exclude_fp32=freeze_exclude_fp32,
+                exclude_autocast_dtype=freeze_exclude_autocast_dtype)
+
+    def init_weights(self, pretrained):
+        if pretrained is not None:
+            logger = get_root_logger()
+            # load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger)
+            checkpoint = _load_checkpoint(pretrained, map_location='cpu', logger=logger)
+            if 'state_dict' in checkpoint:
+                state_dict = checkpoint['state_dict']
+            else:
+                state_dict = checkpoint
+            metadata = getattr(state_dict, '_metadata', OrderedDict())
+            state_dict._metadata = metadata
+            assert self.conv_in.weight.shape[1] == self.conv_out.weight.shape[0]
+            if state_dict['conv_in.weight'].size() != self.conv_in.weight.size():
+                assert state_dict['conv_in.weight'].shape[1] == state_dict['conv_out.weight'].shape[0]
+                src_chn = state_dict['conv_in.weight'].shape[1]
+                tgt_chn = self.conv_in.weight.shape[1]
+                assert src_chn < tgt_chn
+                convert_mat_out = torch.tile(torch.eye(src_chn), (ceildiv(tgt_chn, src_chn), 1))
+                convert_mat_out = convert_mat_out[:tgt_chn]
+                convert_mat_in = F.normalize(convert_mat_out.pinverse(), dim=-1)
+                state_dict['conv_out.weight'] = torch.einsum(
+                    'ts,scxy->tcxy', convert_mat_out, state_dict['conv_out.weight'])
+                state_dict['conv_out.bias'] = torch.einsum(
+                    'ts,s->t', convert_mat_out, state_dict['conv_out.bias'])
+                state_dict['conv_in.weight'] = torch.einsum(
+                    'st,csxy->ctxy', convert_mat_in, state_dict['conv_in.weight'])
+            load_state_dict(self, state_dict, logger=logger)
+
+    def forward(self, sample, timestep, encoder_hidden_states, **kwargs):
+        dtype = sample.dtype
+        return super().forward(
+            sample, timestep, encoder_hidden_states, return_dict=False, **kwargs)[0].to(dtype)
+
+    def forward_enc(self, sample, timestep, encoder_hidden_states, **kwargs):
+        return unet_enc(self, sample, timestep, encoder_hidden_states, **kwargs)
+
+    def forward_dec(self, emb, down_block_res_samples, sample, encoder_hidden_states, **kwargs):
+        return unet_dec(self, emb, down_block_res_samples, sample, encoder_hidden_states, **kwargs)

piFlow/lakonlab/models/architecture/diffusers/__init__.py

@@ -0,0 +1,15 @@
+from .pretrained import (
+    PretrainedVAE, PretrainedVAEDecoder, PretrainedVAEEncoder, PretrainedVAEQwenImage,
+    PretrainedFluxTextEncoder, PretrainedQwenImageTextEncoder, PretrainedStableDiffusion3TextEncoder)
+from .unet import UNet2DConditionModel
+from .flux import FluxTransformer2DModel
+from .dit import DiTTransformer2DModelMod
+from .sd3 import SD3Transformer2DModel
+from .qwen import QwenImageTransformer2DModel
+
+__all__ = [
+    'PretrainedVAE', 'PretrainedVAEDecoder', 'PretrainedVAEEncoder', 'PretrainedFluxTextEncoder',
+    'PretrainedQwenImageTextEncoder', 'UNet2DConditionModel', 'FluxTransformer2DModel',
+    'DiTTransformer2DModelMod', 'SD3Transformer2DModel',
+    'QwenImageTransformer2DModel', 'PretrainedVAEQwenImage', 'PretrainedStableDiffusion3TextEncoder',
+]

piFlow/lakonlab/models/architecture/diffusers/dit.py

@@ -0,0 +1,428 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from typing import Any, Dict, Optional
+from diffusers.models import DiTTransformer2DModel, ModelMixin
+from diffusers.models.attention import BasicTransformerBlock, _chunked_feed_forward, Attention, FeedForward
+from diffusers.models.embeddings import (
+    PatchEmbed, Timesteps, CombinedTimestepLabelEmbeddings, TimestepEmbedding, LabelEmbedding)
+from diffusers.models.normalization import AdaLayerNormZero
+from diffusers.configuration_utils import register_to_config
+from mmcv.runner import load_checkpoint, _load_checkpoint, load_state_dict
+from mmcv.cnn import constant_init, xavier_init
+from mmgen.models.builder import MODULES
+from mmgen.utils import get_root_logger
+from ..utils import flex_freeze
+
+
+class LabelEmbeddingMod(LabelEmbedding):
+    def __init__(self, num_classes, hidden_size, dropout_prob=0.0, use_cfg_embedding=True):
+        super(LabelEmbedding, self).__init__()
+        if dropout_prob > 0:
+            assert use_cfg_embedding
+        self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)
+        self.num_classes = num_classes
+        self.dropout_prob = dropout_prob
+
+
+class CombinedTimestepLabelEmbeddingsMod(CombinedTimestepLabelEmbeddings):
+    """
+    Modified CombinedTimestepLabelEmbeddings for reproducing the original DiT (downscale_freq_shift=0).
+    """
+    def __init__(
+            self, num_classes, embedding_dim, class_dropout_prob=0.1, downscale_freq_shift=0, use_cfg_embedding=True):
+        super(CombinedTimestepLabelEmbeddings, self).__init__()
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=downscale_freq_shift)
+        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+        self.class_embedder = LabelEmbeddingMod(num_classes, embedding_dim, class_dropout_prob, use_cfg_embedding)
+
+
+class BasicTransformerBlockMod(BasicTransformerBlock):
+    """
+    Modified BasicTransformerBlock for reproducing the original DiT with shared time and class
+    embeddings across all layers.
+    """
+    def __init__(
+            self,
+            dim: int,
+            num_attention_heads: int,
+            attention_head_dim: int,
+            dropout=0.0,
+            cross_attention_dim: Optional[int] = None,
+            activation_fn: str = 'geglu',
+            num_embeds_ada_norm: Optional[int] = None,
+            attention_bias: bool = False,
+            only_cross_attention: bool = False,
+            double_self_attention: bool = False,
+            upcast_attention: bool = False,
+            norm_elementwise_affine: bool = True,
+            norm_type: str = 'layer_norm',
+            norm_eps: float = 1e-5,
+            final_dropout: bool = False,
+            attention_type: str = 'default',
+            ada_norm_continous_conditioning_embedding_dim: Optional[int] = None,
+            ada_norm_bias: Optional[int] = None,
+            ff_inner_dim: Optional[int] = None,
+            ff_bias: bool = True,
+            attention_out_bias: bool = True):
+        super(BasicTransformerBlock, self).__init__()
+        self.only_cross_attention = only_cross_attention
+        self.norm_type = norm_type
+        self.num_embeds_ada_norm = num_embeds_ada_norm
+
+        assert self.norm_type == 'ada_norm_zero'
+        self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)
+        self.attn1 = Attention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            dropout=dropout,
+            bias=attention_bias,
+            cross_attention_dim=cross_attention_dim if only_cross_attention else None,
+            upcast_attention=upcast_attention,
+            out_bias=attention_out_bias,
+        )
+
+        self.norm2 = None
+        self.attn2 = None
+
+        self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
+        self.ff = FeedForward(
+            dim,
+            dropout=dropout,
+            activation_fn=activation_fn,
+            final_dropout=final_dropout,
+            inner_dim=ff_inner_dim,
+            bias=ff_bias,
+        )
+
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            attention_mask: Optional[torch.Tensor] = None,
+            encoder_hidden_states: Optional[torch.Tensor] = None,
+            encoder_attention_mask: Optional[torch.Tensor] = None,
+            timestep: Optional[torch.LongTensor] = None,
+            cross_attention_kwargs: Dict[str, Any] = None,
+            class_labels: Optional[torch.LongTensor] = None,
+            emb: Optional[torch.Tensor] = None,
+            added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None) -> torch.Tensor:
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+            hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype, emb=emb)
+
+        if cross_attention_kwargs is None:
+            cross_attention_kwargs = dict()
+        attn_output = self.attn1(
+            norm_hidden_states,
+            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
+            attention_mask=attention_mask,
+            **cross_attention_kwargs)
+        attn_output = gate_msa.unsqueeze(1) * attn_output
+
+        hidden_states = attn_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+
+        norm_hidden_states = self.norm3(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+        if self._chunk_size is not None:
+            # "feed_forward_chunk_size" can be used to save memory
+            ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)
+        else:
+            ff_output = self.ff(norm_hidden_states)
+
+        ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+        hidden_states = ff_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+
+        return hidden_states
+
+
+class _DiTTransformer2DModelMod(DiTTransformer2DModel):
+
+    @register_to_config
+    def __init__(
+            self,
+            class_dropout_prob=0.0,
+            num_attention_heads: int = 16,
+            attention_head_dim: int = 72,
+            in_channels: int = 4,
+            out_channels: Optional[int] = None,
+            num_layers: int = 28,
+            dropout: float = 0.0,
+            norm_num_groups: int = 32,
+            attention_bias: bool = True,
+            sample_size: int = 32,
+            patch_size: int = 2,
+            activation_fn: str = 'gelu-approximate',
+            num_embeds_ada_norm: Optional[int] = 1000,
+            upcast_attention: bool = False,
+            norm_type: str = 'ada_norm_zero',
+            norm_elementwise_affine: bool = False,
+            norm_eps: float = 1e-5):
+
+        super(DiTTransformer2DModel, self).__init__()
+
+        # Validate inputs.
+        if norm_type != "ada_norm_zero":
+            raise NotImplementedError(
+                f"Forward pass is not implemented when `patch_size` is not None and `norm_type` is '{norm_type}'."
+            )
+        elif norm_type == "ada_norm_zero" and num_embeds_ada_norm is None:
+            raise ValueError(
+                f"When using a `patch_size` and this `norm_type` ({norm_type}), `num_embeds_ada_norm` cannot be None."
+            )
+
+        # Set some common variables used across the board.
+        self.attention_head_dim = attention_head_dim
+        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
+        self.out_channels = in_channels if out_channels is None else out_channels
+        self.gradient_checkpointing = False
+
+        # 2. Initialize the position embedding and transformer blocks.
+        self.height = self.config.sample_size
+        self.width = self.config.sample_size
+
+        self.patch_size = self.config.patch_size
+        self.pos_embed = PatchEmbed(
+            height=self.config.sample_size,
+            width=self.config.sample_size,
+            patch_size=self.config.patch_size,
+            in_channels=self.config.in_channels,
+            embed_dim=self.inner_dim)
+        self.emb = CombinedTimestepLabelEmbeddingsMod(
+            num_embeds_ada_norm, self.inner_dim, class_dropout_prob=0.0)
+
+        self.transformer_blocks = nn.ModuleList([
+            BasicTransformerBlockMod(
+                self.inner_dim,
+                self.config.num_attention_heads,
+                self.config.attention_head_dim,
+                dropout=self.config.dropout,
+                activation_fn=self.config.activation_fn,
+                num_embeds_ada_norm=None,
+                attention_bias=self.config.attention_bias,
+                upcast_attention=self.config.upcast_attention,
+                norm_type=norm_type,
+                norm_elementwise_affine=self.config.norm_elementwise_affine,
+                norm_eps=self.config.norm_eps)
+            for _ in range(self.config.num_layers)])
+
+        # 3. Output blocks.
+        self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim)
+        self.proj_out_2 = nn.Linear(
+            self.inner_dim, self.config.patch_size * self.config.patch_size * self.out_channels)
+
+    # https://github.com/facebookresearch/DiT/blob/main/models.py
+    def init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                xavier_init(m, distribution='uniform')
+            elif isinstance(m, nn.Embedding):
+                torch.nn.init.normal_(m.weight, mean=0.0, std=0.02)
+
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d)
+        w = self.pos_embed.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.pos_embed.proj.bias, 0)
+
+        # Zero-out adaLN modulation layers in DiT blocks
+        for m in self.modules():
+            if isinstance(m, AdaLayerNormZero):
+                constant_init(m.linear, val=0)
+
+        # Zero-out output layers
+        constant_init(self.proj_out_1, val=0)
+        constant_init(self.proj_out_2, val=0)
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            timestep: Optional[torch.LongTensor] = None,
+            class_labels: Optional[torch.LongTensor] = None,
+            cross_attention_kwargs: Dict[str, Any] = None):
+        # 1. Input
+        bs, _, h, w = hidden_states.size()
+        height, width = h // self.patch_size, w // self.patch_size
+        hidden_states = self.pos_embed(hidden_states)
+
+        cond_emb = self.emb(
+            timestep, class_labels, hidden_dtype=hidden_states.dtype)
+        dropout_enabled = self.config.class_dropout_prob > 0 and self.training
+        if dropout_enabled:
+            uncond_emb = self.emb(timestep, torch.full_like(
+                class_labels, self.config.num_embeds_ada_norm), hidden_dtype=hidden_states.dtype)
+
+        # 2. Blocks
+        for block in self.transformer_blocks:
+            if dropout_enabled:
+                dropout_mask = torch.rand((bs, 1), device=hidden_states.device) < self.config.class_dropout_prob
+                emb = torch.where(dropout_mask, uncond_emb, cond_emb)
+            else:
+                emb = cond_emb
+
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    None,
+                    None,
+                    None,
+                    timestep,
+                    cross_attention_kwargs,
+                    class_labels,
+                    emb,
+                    use_reentrant=False)
+
+            else:
+                hidden_states = block(
+                    hidden_states,
+                    attention_mask=None,
+                    encoder_hidden_states=None,
+                    encoder_attention_mask=None,
+                    timestep=timestep,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    class_labels=class_labels,
+                    emb=emb)
+
+        # 3. Output
+        if dropout_enabled:
+            dropout_mask = torch.rand((bs, 1), device=hidden_states.device) < self.config.class_dropout_prob
+            emb = torch.where(dropout_mask, uncond_emb, cond_emb)
+        else:
+            emb = cond_emb
+        shift, scale = self.proj_out_1(F.silu(emb)).chunk(2, dim=1)
+        hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
+        output = self.proj_out_2(hidden_states).reshape(
+                bs, height, width, self.patch_size, self.patch_size, self.out_channels
+            ).permute(0, 5, 1, 3, 2, 4).reshape(
+                bs, self.out_channels, height * self.patch_size, width * self.patch_size)
+
+        return output
+
+
+@MODULES.register_module()
+class DiTTransformer2DModelMod(_DiTTransformer2DModelMod):
+
+    def __init__(
+            self,
+            *args,
+            freeze=False,
+            freeze_exclude=[],
+            pretrained=None,
+            torch_dtype='float32',
+            autocast_dtype=None,
+            freeze_exclude_fp32=True,
+            freeze_exclude_autocast_dtype='float32',
+            checkpointing=True,
+            **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.init_weights(pretrained)
+
+        if autocast_dtype is not None:
+            assert torch_dtype == 'float32'
+        self.autocast_dtype = autocast_dtype
+
+        if torch_dtype is not None:
+            self.to(getattr(torch, torch_dtype))
+
+        self.freeze = freeze
+        if self.freeze:
+            flex_freeze(
+                self,
+                exclude_keys=freeze_exclude,
+                exclude_fp32=freeze_exclude_fp32,
+                exclude_autocast_dtype=freeze_exclude_autocast_dtype)
+
+        if checkpointing:
+            self.enable_gradient_checkpointing()
+
+    def init_weights(self, pretrained=None):
+        super().init_weights()
+        if pretrained is not None:
+            logger = get_root_logger()
+            # load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger)
+            checkpoint = _load_checkpoint(pretrained, map_location='cpu', logger=logger)
+            if 'state_dict' in checkpoint:
+                state_dict = checkpoint['state_dict']
+            else:
+                state_dict = checkpoint
+            # load from GMDiT V1 model with 1 Gaussian
+            p2 = self.config.patch_size * self.config.patch_size
+            ori_out_channels = p2 * self.out_channels
+            if 'proj_out_2.weight' in state_dict:
+                # if this is GMDiT V1 model with 1 Gaussian
+                if state_dict['proj_out_2.weight'].size(0) == p2 * (self.out_channels + 1):
+                    state_dict['proj_out_2.weight'] = state_dict['proj_out_2.weight'].reshape(
+                        p2, self.out_channels + 1, -1
+                    )[:, :-1].reshape(ori_out_channels, -1)
+                # if this is original DiT with variance prediction
+                if state_dict['proj_out_2.weight'].size(0) == 2 * ori_out_channels:
+                    state_dict['proj_out_2.weight'] = state_dict['proj_out_2.weight'].reshape(
+                        p2, 2 * self.out_channels, -1
+                    )[:, :self.out_channels].reshape(ori_out_channels, -1)
+            if 'proj_out_2.bias' in state_dict:
+                # if this is GMDiT V1 model with 1 Gaussian
+                if state_dict['proj_out_2.bias'].size(0) == p2 * (self.out_channels + 1):
+                    state_dict['proj_out_2.bias'] = state_dict['proj_out_2.bias'].reshape(
+                        p2, self.out_channels + 1
+                    )[:, :-1].reshape(ori_out_channels)
+                # if this is original DiT with variance prediction
+                if state_dict['proj_out_2.bias'].size(0) == 2 * ori_out_channels:
+                    state_dict['proj_out_2.bias'] = state_dict['proj_out_2.bias'].reshape(
+                        p2, 2 * self.out_channels
+                    )[:, :self.out_channels].reshape(ori_out_channels)
+            if 'emb.class_embedder.embedding_table.weight' not in state_dict \
+                    and 'transformer_blocks.0.norm1.emb.class_embedder.embedding_table.weight' in state_dict:
+                # convert original diffusers DiT model to our modified DiT model with shared embeddings
+                keys_to_remove = []
+                state_update = {}
+                for k, v in state_dict.items():
+                    if k.startswith('transformer_blocks.0.norm1.emb.'):
+                        new_k = k.replace('transformer_blocks.0.norm1.', '')
+                        state_update[new_k] = v
+                    if k.startswith('transformer_blocks.') and '.norm1.emb.' in k:
+                        keys_to_remove.append(k)
+                state_dict.update(state_update)
+                for k in keys_to_remove:
+                    del state_dict[k]
+            load_state_dict(self, state_dict, logger=logger)
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            timestep: Optional[torch.LongTensor] = None,
+            class_labels: Optional[torch.LongTensor] = None,
+            **kwargs):
+        if self.autocast_dtype is not None:
+            dtype = getattr(torch, self.autocast_dtype)
+        else:
+            dtype = hidden_states.dtype
+        with torch.autocast(
+                device_type='cuda',
+                enabled=self.autocast_dtype is not None,
+                dtype=dtype if self.autocast_dtype is not None else None):
+            return super().forward(
+                hidden_states.to(dtype),
+                timestep=timestep,
+                class_labels=class_labels,
+                **kwargs)

piFlow/lakonlab/models/architecture/diffusers/flux.py

@@ -0,0 +1,156 @@
+import torch
+
+from typing import Optional
+from accelerate import init_empty_weights
+from diffusers.models import FluxTransformer2DModel as _FluxTransformer2DModel
+from peft import LoraConfig
+from mmgen.models.builder import MODULES
+from mmgen.utils import get_root_logger
+from ..utils import flex_freeze
+from lakonlab.runner.checkpoint import load_checkpoint, _load_checkpoint
+
+
+@MODULES.register_module()
+class FluxTransformer2DModel(_FluxTransformer2DModel):
+
+    def __init__(
+            self,
+            *args,
+            patch_size=2,
+            freeze=False,
+            freeze_exclude=[],
+            pretrained=None,
+            pretrained_lora=None,
+            pretrained_lora_scale=1.0,
+            torch_dtype='float32',
+            freeze_exclude_fp32=True,
+            freeze_exclude_autocast_dtype='float32',
+            checkpointing=True,
+            use_lora=False,
+            lora_target_modules=None,
+            lora_rank=16,
+            **kwargs):
+        with init_empty_weights():
+            super().__init__(patch_size=1, *args, **kwargs)
+        self.patch_size = patch_size
+
+        self.init_weights(pretrained, pretrained_lora, pretrained_lora_scale)
+
+        self.use_lora = use_lora
+        self.lora_target_modules = lora_target_modules
+        self.lora_rank = lora_rank
+        if self.use_lora:
+            transformer_lora_config = LoraConfig(
+                r=lora_rank,
+                lora_alpha=lora_rank,
+                init_lora_weights='gaussian',
+                target_modules=lora_target_modules,
+            )
+            self.add_adapter(transformer_lora_config)
+
+        if torch_dtype is not None:
+            self.to(getattr(torch, torch_dtype))
+
+        self.freeze = freeze
+        if self.freeze:
+            flex_freeze(
+                self,
+                exclude_keys=freeze_exclude,
+                exclude_fp32=freeze_exclude_fp32,
+                exclude_autocast_dtype=freeze_exclude_autocast_dtype)
+
+        if checkpointing:
+            self.enable_gradient_checkpointing()
+
+    def init_weights(self, pretrained=None, pretrained_lora=None, pretrained_lora_scale=1.0):
+        if pretrained is not None:
+            logger = get_root_logger()
+            load_checkpoint(
+                self, pretrained, map_location='cpu', strict=False, logger=logger, assign=True)
+            if pretrained_lora is not None:
+                if not isinstance(pretrained_lora, (list, tuple)):
+                    assert isinstance(pretrained_lora, str)
+                    pretrained_lora = [pretrained_lora]
+                if not isinstance(pretrained_lora_scale, (list, tuple)):
+                    assert isinstance(pretrained_lora_scale, (int, float))
+                    pretrained_lora_scale = [pretrained_lora_scale]
+                for pretrained_lora_single, pretrained_lora_scale_single in zip(pretrained_lora, pretrained_lora_scale):
+                    lora_state_dict = _load_checkpoint(
+                        pretrained_lora_single, map_location='cpu', logger=logger)
+                    self.load_lora_adapter(lora_state_dict)
+                    self.fuse_lora(lora_scale=pretrained_lora_scale_single)
+                    self.unload_lora()
+
+    @staticmethod
+    def _prepare_latent_image_ids(height, width, device, dtype):
+        """
+        Copied from Diffusers
+        """
+        latent_image_ids = torch.zeros(height, width, 3)
+        latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None]
+        latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :]
+
+        latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+
+        latent_image_ids = latent_image_ids.reshape(
+            latent_image_id_height * latent_image_id_width, latent_image_id_channels)
+
+        return latent_image_ids.to(device=device, dtype=dtype)
+
+    def patchify(self, latents):
+        if self.patch_size > 1:
+            bs, c, h, w = latents.size()
+            latents = latents.reshape(
+                bs, c, h // self.patch_size, self.patch_size, w // self.patch_size, self.patch_size
+            ).permute(
+                0, 1, 3, 5, 2, 4
+            ).reshape(
+                bs, c * self.patch_size * self.patch_size, h // self.patch_size, w // self.patch_size)
+        return latents
+
+    def unpatchify(self, latents):
+        if self.patch_size > 1:
+            bs, c, h, w = latents.size()
+            latents = latents.reshape(
+                bs, c // (self.patch_size * self.patch_size), self.patch_size, self.patch_size, h, w
+            ).permute(
+                0, 1, 4, 2, 5, 3
+            ).reshape(
+                bs, c // (self.patch_size * self.patch_size), h * self.patch_size, w * self.patch_size)
+        return latents
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            timestep: torch.Tensor,
+            encoder_hidden_states: torch.Tensor = None,
+            pooled_projections: torch.Tensor = None,
+            mask: Optional[torch.Tensor] = None,
+            masked_image_latents: Optional[torch.Tensor] = None,
+            **kwargs):
+        hidden_states = self.patchify(hidden_states)
+        bs, c, h, w = hidden_states.size()
+        dtype = hidden_states.dtype
+        device = hidden_states.device
+        hidden_states = hidden_states.reshape(bs, c, h * w).permute(0, 2, 1)
+        img_ids = self._prepare_latent_image_ids(
+            h, w, device, dtype)
+        txt_ids = img_ids.new_zeros((encoder_hidden_states.shape[-2], 3))
+
+        #  Flux fill
+        if mask is not None and masked_image_latents is not None:
+            hidden_states = torch.cat(
+                (hidden_states, masked_image_latents.to(dtype=dtype), mask.to(dtype=dtype)), dim=-1)
+
+        output = super().forward(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states.to(dtype),
+            pooled_projections=pooled_projections.to(dtype),
+            timestep=timestep,
+            img_ids=img_ids,
+            txt_ids=txt_ids,
+            return_dict=False,
+            **kwargs)[0]
+
+        output = output.permute(0, 2, 1).reshape(bs, self.out_channels, h, w)
+        return self.unpatchify(output)

piFlow/lakonlab/models/architecture/diffusers/pretrained.py

@@ -0,0 +1,281 @@
+import logging
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.models import AutoencoderKL, AutoencoderKLQwenImage
+from diffusers.pipelines import FluxPipeline, QwenImagePipeline, StableDiffusion3Pipeline
+from mmgen.models.builder import MODULES
+
+# Suppress truncation warnings from transformers and diffusers
+for name in (
+        'transformers.tokenization_utils_base',
+        'transformers.tokenization_utils',
+        'transformers.tokenization_utils_fast'):
+    logging.getLogger(name).setLevel(logging.ERROR)
+
+for name, logger in logging.root.manager.loggerDict.items():
+    if isinstance(logger, logging.Logger) and (name.startswith('diffusers.pipelines.')):
+        logger.setLevel(logging.ERROR)
+
+
+@MODULES.register_module()
+class PretrainedVAE(nn.Module):
+    def __init__(self,
+                 from_pretrained=None,
+                 del_encoder=False,
+                 del_decoder=False,
+                 use_slicing=False,
+                 freeze=True,
+                 eval_mode=True,
+                 torch_dtype='float32',
+                 **kwargs):
+        super().__init__()
+        if torch_dtype is not None:
+            kwargs.update(torch_dtype=getattr(torch, torch_dtype))
+        self.vae = AutoencoderKL.from_pretrained(
+            from_pretrained, **kwargs)
+        if del_encoder:
+            del self.vae.encoder
+        if del_decoder:
+            del self.vae.decoder
+        if use_slicing:
+            self.vae.enable_slicing()
+        self.freeze = freeze
+        self.eval_mode = eval_mode
+        if self.freeze:
+            self.requires_grad_(False)
+        if self.eval_mode:
+            self.eval()
+        self.vae.set_use_memory_efficient_attention_xformers(
+            not hasattr(torch.nn.functional, 'scaled_dot_product_attention'))
+
+    def train(self, mode=True):
+        mode = mode and (not self.eval_mode)
+        return super().train(mode)
+
+    def forward(self, *args, **kwargs):
+        return self.vae(*args, return_dict=False, **kwargs)[0]
+
+    def encode(self, img):
+        scaling_factor = self.vae.config.scaling_factor
+        shift_factor = self.vae.config.shift_factor
+        if scaling_factor is None:
+            scaling_factor = 1.0
+        if shift_factor is None:
+            shift_factor = 0.0
+        return (self.vae.encode(img).latent_dist.sample() - shift_factor) * scaling_factor
+
+    def decode(self, code):
+        scaling_factor = self.vae.config.scaling_factor
+        shift_factor = self.vae.config.shift_factor
+        if scaling_factor is None:
+            scaling_factor = 1.0
+        if shift_factor is None:
+            shift_factor = 0.0
+        return self.vae.decode(code / scaling_factor + shift_factor, return_dict=False)[0]
+
+
+@MODULES.register_module()
+class PretrainedVAEDecoder(PretrainedVAE):
+    def __init__(self, **kwargs):
+        super().__init__(
+            del_encoder=True,
+            del_decoder=False,
+            **kwargs)
+
+    def forward(self, code):
+        return super().decode(code)
+
+
+@MODULES.register_module()
+class PretrainedVAEEncoder(PretrainedVAE):
+    def __init__(self, **kwargs):
+        super().__init__(
+            del_encoder=False,
+            del_decoder=True,
+            **kwargs)
+
+    def forward(self, img):
+        return super().encode(img)
+
+
+@MODULES.register_module()
+class PretrainedVAEQwenImage(nn.Module):
+    def __init__(self,
+                 from_pretrained=None,
+                 use_slicing=False,
+                 freeze=True,
+                 eval_mode=True,
+                 torch_dtype='float32',
+                 **kwargs):
+        super().__init__()
+        if torch_dtype is not None:
+            kwargs.update(torch_dtype=getattr(torch, torch_dtype))
+        self.vae = AutoencoderKLQwenImage.from_pretrained(
+            from_pretrained, **kwargs)
+        if use_slicing:
+            self.vae.enable_slicing()
+        self.freeze = freeze
+        self.eval_mode = eval_mode
+        if self.freeze:
+            self.requires_grad_(False)
+        if self.eval_mode:
+            self.eval()
+
+    def train(self, mode=True):
+        mode = mode and (not self.eval_mode)
+        return super().train(mode)
+
+    def forward(self, *args, **kwargs):
+        return self.vae(*args, return_dict=False, **kwargs)[0]
+
+    def encode(self, img):
+        device = img.device
+        dtype = img.dtype
+        latents_mean = torch.tensor(self.vae.config.latents_mean, device=device, dtype=dtype).view(
+            1, self.vae.config.z_dim, 1, 1, 1)
+        latents_std = torch.tensor(self.vae.config.latents_std, device=device, dtype=dtype).view(
+            1, self.vae.config.z_dim, 1, 1, 1)
+        return ((self.vae.encode(img.unsqueeze(-3)).latent_dist.sample() - latents_mean) / latents_std).squeeze(-3)
+
+    def decode(self, code):
+        device = code.device
+        dtype = code.dtype
+        latents_mean = torch.tensor(self.vae.config.latents_mean, device=device, dtype=dtype).view(
+            1, self.vae.config.z_dim, 1, 1, 1)
+        latents_std = torch.tensor(self.vae.config.latents_std, device=device, dtype=dtype).view(
+            1, self.vae.config.z_dim, 1, 1, 1)
+        return self.vae.decode(code.unsqueeze(-3) * latents_std + latents_mean, return_dict=False)[0].squeeze(-3)
+
+
+@MODULES.register_module()
+class PretrainedFluxTextEncoder(nn.Module):
+    def __init__(self,
+                 from_pretrained='black-forest-labs/FLUX.1-dev',
+                 freeze=True,
+                 eval_mode=True,
+                 torch_dtype='bfloat16',
+                 max_sequence_length=512,
+                 **kwargs):
+        super().__init__()
+        self.max_sequence_length = max_sequence_length
+        self.pipeline = FluxPipeline.from_pretrained(
+            from_pretrained,
+            scheduler=None,
+            vae=None,
+            transformer=None,
+            image_encoder=None,
+            feature_extractor=None,
+            torch_dtype=getattr(torch, torch_dtype),
+            **kwargs)
+        self.text_encoder = self.pipeline.text_encoder
+        self.text_encoder_2 = self.pipeline.text_encoder_2
+        self.freeze = freeze
+        self.eval_mode = eval_mode
+        if self.freeze:
+            self.requires_grad_(False)
+        if self.eval_mode:
+            self.eval()
+
+    def train(self, mode=True):
+        mode = mode and (not self.eval_mode)
+        return super().train(mode)
+
+    def forward(self, prompt, prompt_2=None):
+        prompt_embeds, pooled_prompt_embeds, text_ids = self.pipeline.encode_prompt(
+            prompt, prompt_2=prompt_2, max_sequence_length=self.max_sequence_length)
+        return dict(
+            encoder_hidden_states=prompt_embeds,
+            pooled_projections=pooled_prompt_embeds)
+
+
+@MODULES.register_module()
+class PretrainedQwenImageTextEncoder(nn.Module):
+    def __init__(self,
+                 from_pretrained='Qwen/Qwen-Image',
+                 freeze=True,
+                 eval_mode=True,
+                 torch_dtype='bfloat16',
+                 max_sequence_length=512,
+                 pad_seq_len=None,
+                 **kwargs):
+        super().__init__()
+        self.max_sequence_length = max_sequence_length
+        if pad_seq_len is not None:
+            assert pad_seq_len >= max_sequence_length
+        self.pad_seq_len = pad_seq_len
+        self.pipeline = QwenImagePipeline.from_pretrained(
+            from_pretrained,
+            scheduler=None,
+            vae=None,
+            transformer=None,
+            torch_dtype=getattr(torch, torch_dtype),
+            **kwargs)
+        self.text_encoder = self.pipeline.text_encoder
+        self.freeze = freeze
+        self.eval_mode = eval_mode
+        if self.freeze:
+            self.requires_grad_(False)
+        if self.eval_mode:
+            self.eval()
+
+    def train(self, mode=True):
+        mode = mode and (not self.eval_mode)
+        return super().train(mode)
+
+    def forward(self, prompt):
+        prompt_embeds, prompt_embeds_mask = self.pipeline.encode_prompt(
+            prompt, max_sequence_length=self.max_sequence_length)
+        if self.pad_seq_len is not None:
+            pad_len = self.pad_seq_len - prompt_embeds.size(1)
+            prompt_embeds = F.pad(
+                prompt_embeds, (0, 0, 0, pad_len), value=0.0)
+            prompt_embeds_mask = F.pad(
+                prompt_embeds_mask, (0, pad_len), value=0.0)
+        return dict(
+            encoder_hidden_states=prompt_embeds,
+            encoder_hidden_states_mask=prompt_embeds_mask)
+
+
+@MODULES.register_module()
+class PretrainedStableDiffusion3TextEncoder(nn.Module):
+    def __init__(self,
+                 from_pretrained='stabilityai/stable-diffusion-3.5-large',
+                 freeze=True,
+                 eval_mode=True,
+                 torch_dtype='float32',
+                 max_sequence_length=256,
+                 **kwargs):
+        super().__init__()
+        self.max_sequence_length = max_sequence_length
+        self.pipeline = StableDiffusion3Pipeline.from_pretrained(
+            from_pretrained,
+            scheduler=None,
+            vae=None,
+            transformer=None,
+            image_encoder=None,
+            feature_extractor=None,
+            torch_dtype=getattr(torch, torch_dtype),
+            **kwargs)
+        self.text_encoder = self.pipeline.text_encoder
+        self.text_encoder_2 = self.pipeline.text_encoder_2
+        self.text_encoder_3 = self.pipeline.text_encoder_3
+        self.freeze = freeze
+        self.eval_mode = eval_mode
+        if self.freeze:
+            self.requires_grad_(False)
+        if self.eval_mode:
+            self.eval()
+
+    def train(self, mode=True):
+        mode = mode and (not self.eval_mode)
+        return super().train(mode)
+
+    def forward(self, prompt, prompt_2=None, prompt_3=None):
+        prompt_embeds, _, pooled_prompt_embeds, _ = self.pipeline.encode_prompt(
+            prompt, prompt_2=prompt_2, prompt_3=prompt_3, do_classifier_free_guidance=False,
+            max_sequence_length=self.max_sequence_length)
+        return dict(
+            encoder_hidden_states=prompt_embeds,
+            pooled_projections=pooled_prompt_embeds)

piFlow/lakonlab/models/architecture/diffusers/qwen.py

@@ -0,0 +1,139 @@
+import torch
+
+from accelerate import init_empty_weights
+from diffusers.models import QwenImageTransformer2DModel as _QwenImageTransformer2DModel
+from diffusers.loaders.lora_conversion_utils import _convert_non_diffusers_qwen_lora_to_diffusers
+from peft import LoraConfig
+from mmgen.models.builder import MODULES
+from mmgen.utils import get_root_logger
+from ..utils import flex_freeze
+from lakonlab.runner.checkpoint import load_checkpoint, _load_checkpoint
+
+
+@MODULES.register_module()
+class QwenImageTransformer2DModel(_QwenImageTransformer2DModel):
+
+    def __init__(
+            self,
+            *args,
+            patch_size=2,
+            freeze=False,
+            freeze_exclude=[],
+            pretrained=None,
+            pretrained_lora=None,
+            pretrained_lora_scale=1.0,
+            torch_dtype='float32',
+            freeze_exclude_fp32=True,
+            freeze_exclude_autocast_dtype='float32',
+            checkpointing=True,
+            use_lora=False,
+            lora_target_modules=None,
+            lora_rank=16,
+            **kwargs):
+        with init_empty_weights():
+            super().__init__(*args, patch_size=1, **kwargs)
+        self.patch_size = patch_size
+
+        self.init_weights(pretrained, pretrained_lora, pretrained_lora_scale)
+
+        self.use_lora = use_lora
+        self.lora_target_modules = lora_target_modules
+        self.lora_rank = lora_rank
+        if self.use_lora:
+            transformer_lora_config = LoraConfig(
+                r=lora_rank,
+                lora_alpha=lora_rank,
+                init_lora_weights='gaussian',
+                target_modules=lora_target_modules,
+            )
+            self.add_adapter(transformer_lora_config)
+
+        if torch_dtype is not None:
+            self.to(getattr(torch, torch_dtype))
+
+        self.freeze = freeze
+        if self.freeze:
+            flex_freeze(
+                self,
+                exclude_keys=freeze_exclude,
+                exclude_fp32=freeze_exclude_fp32,
+                exclude_autocast_dtype=freeze_exclude_autocast_dtype)
+
+        if checkpointing:
+            self.enable_gradient_checkpointing()
+
+    def init_weights(self, pretrained=None, pretrained_lora=None, pretrained_lora_scale=1.0):
+        if pretrained is not None:
+            logger = get_root_logger()
+            load_checkpoint(
+                self, pretrained, map_location='cpu', strict=False, logger=logger, assign=True)
+            if pretrained_lora is not None:
+                if not isinstance(pretrained_lora, (list, tuple)):
+                    assert isinstance(pretrained_lora, str)
+                    pretrained_lora = [pretrained_lora]
+                if not isinstance(pretrained_lora_scale, (list, tuple)):
+                    assert isinstance(pretrained_lora_scale, (int, float))
+                    pretrained_lora_scale = [pretrained_lora_scale]
+                for pretrained_lora_single, pretrained_lora_scale_single in zip(pretrained_lora, pretrained_lora_scale):
+                    lora_state_dict = _load_checkpoint(
+                        pretrained_lora_single, map_location='cpu', logger=logger)
+                    lora_state_dict = _convert_non_diffusers_qwen_lora_to_diffusers(lora_state_dict)
+                    self.load_lora_adapter(lora_state_dict)
+                    self.fuse_lora(lora_scale=pretrained_lora_scale_single)
+                    self.unload_lora()
+
+    def patchify(self, latents):
+        if self.patch_size > 1:
+            bs, c, h, w = latents.size()
+            latents = latents.reshape(
+                bs, c, h // self.patch_size, self.patch_size, w // self.patch_size, self.patch_size
+            ).permute(
+                0, 1, 3, 5, 2, 4
+            ).reshape(
+                bs, c * self.patch_size * self.patch_size, h // self.patch_size, w // self.patch_size)
+        return latents
+
+    def unpatchify(self, latents):
+        if self.patch_size > 1:
+            bs, c, h, w = latents.size()
+            latents = latents.reshape(
+                bs, c // (self.patch_size * self.patch_size), self.patch_size, self.patch_size, h, w
+            ).permute(
+                0, 1, 4, 2, 5, 3
+            ).reshape(
+                bs, c // (self.patch_size * self.patch_size), h * self.patch_size, w * self.patch_size)
+        return latents
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            timestep: torch.Tensor,
+            encoder_hidden_states: torch.Tensor = None,
+            encoder_hidden_states_mask: torch.Tensor = None,
+            **kwargs):
+        hidden_states = self.patchify(hidden_states)
+        bs, c, h, w = hidden_states.size()
+        dtype = hidden_states.dtype
+        hidden_states = hidden_states.reshape(bs, c, h * w).permute(0, 2, 1)
+        img_shapes = [[(1, h, w)]]
+        if encoder_hidden_states_mask is not None:
+            txt_seq_lens = encoder_hidden_states_mask.sum(dim=1)
+            max_txt_seq_len = txt_seq_lens.max()
+            encoder_hidden_states = encoder_hidden_states[:, :max_txt_seq_len]
+            encoder_hidden_states_mask = encoder_hidden_states_mask[:, :max_txt_seq_len]
+            txt_seq_lens = txt_seq_lens.tolist()
+        else:
+            txt_seq_lens = None
+
+        output = super().forward(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states.to(dtype),
+            encoder_hidden_states_mask=encoder_hidden_states_mask,
+            timestep=timestep,
+            img_shapes=img_shapes,
+            txt_seq_lens=txt_seq_lens,
+            return_dict=False,
+            **kwargs)[0]
+
+        output = output.permute(0, 2, 1).reshape(bs, self.out_channels, h, w)
+        return self.unpatchify(output)

piFlow/lakonlab/models/architecture/diffusers/sd3.py

@@ -0,0 +1,80 @@
+import torch
+
+from accelerate import init_empty_weights
+from diffusers.models import SD3Transformer2DModel as _SD3Transformer2DModel
+from peft import LoraConfig
+from mmgen.models.builder import MODULES
+from mmgen.utils import get_root_logger
+from ..utils import flex_freeze
+from lakonlab.runner.checkpoint import load_checkpoint
+
+
+@MODULES.register_module()
+class SD3Transformer2DModel(_SD3Transformer2DModel):
+
+    def __init__(
+            self,
+            *args,
+            freeze=False,
+            freeze_exclude=[],
+            pretrained=None,
+            torch_dtype='float32',
+            freeze_exclude_fp32=True,
+            freeze_exclude_autocast_dtype='float32',
+            checkpointing=True,
+            use_lora=False,
+            lora_target_modules=None,
+            lora_rank=16,
+            **kwargs):
+        with init_empty_weights():
+            super().__init__(*args, **kwargs)
+        self.init_weights(pretrained)
+
+        self.use_lora = use_lora
+        self.lora_target_modules = lora_target_modules
+        self.lora_rank = lora_rank
+        if self.use_lora:
+            transformer_lora_config = LoraConfig(
+                r=lora_rank,
+                lora_alpha=lora_rank,
+                init_lora_weights='gaussian',
+                target_modules=lora_target_modules,
+            )
+            self.add_adapter(transformer_lora_config)
+
+        if torch_dtype is not None:
+            self.to(getattr(torch, torch_dtype))
+
+        self.freeze = freeze
+        if self.freeze:
+            flex_freeze(
+                self,
+                exclude_keys=freeze_exclude,
+                exclude_fp32=freeze_exclude_fp32,
+                exclude_autocast_dtype=freeze_exclude_autocast_dtype)
+
+        if checkpointing:
+            self.enable_gradient_checkpointing()
+
+    def init_weights(self, pretrained=None):
+        if pretrained is not None:
+            logger = get_root_logger()
+            load_checkpoint(
+                self, pretrained, map_location='cpu', strict=False, logger=logger, assign=True)
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            timestep: torch.Tensor,
+            encoder_hidden_states: torch.Tensor = None,
+            pooled_projections: torch.Tensor = None,
+            **kwargs):
+        dtype = hidden_states.dtype
+
+        return super().forward(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states.to(dtype),
+            pooled_projections=pooled_projections.to(dtype),
+            timestep=timestep,
+            return_dict=False,
+            **kwargs)[0]

piFlow/lakonlab/models/architecture/diffusers/unet.py

@@ -0,0 +1,192 @@
+import torch
+import torch.nn.functional as F
+
+from typing import Dict, Any, Optional, Union, Tuple
+from collections import OrderedDict
+from diffusers.models import UNet2DConditionModel as _UNet2DConditionModel
+from mmcv.runner import _load_checkpoint, load_state_dict
+from mmgen.models.builder import MODULES
+from mmgen.utils import get_root_logger
+from ..utils import flex_freeze
+
+
+def ceildiv(a, b):
+    return -(a // -b)
+
+
+def unet_enc(
+        unet,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        added_cond_kwargs=None):
+    # 0. center input if necessary
+    if unet.config.center_input_sample:
+        sample = 2 * sample - 1.0
+
+    # 1. time
+    t_emb = unet.get_time_embed(sample=sample, timestep=timestep)
+    emb = unet.time_embedding(t_emb)
+    aug_emb = unet.get_aug_embed(
+        emb=emb, encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs)
+    emb = emb + aug_emb if aug_emb is not None else emb
+
+    if unet.time_embed_act is not None:
+        emb = unet.time_embed_act(emb)
+
+    encoder_hidden_states = unet.process_encoder_hidden_states(
+        encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs)
+
+    # 2. pre-process
+    sample = unet.conv_in(sample)
+
+    # 3. down
+    down_block_res_samples = (sample,)
+    for downsample_block in unet.down_blocks:
+        if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+            sample, res_samples = downsample_block(
+                hidden_states=sample,
+                temb=emb,
+                encoder_hidden_states=encoder_hidden_states,
+                cross_attention_kwargs=cross_attention_kwargs,
+            )
+        else:
+            sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+
+        down_block_res_samples += res_samples
+
+    return emb, down_block_res_samples, sample
+
+
+def unet_dec(
+        unet,
+        emb,
+        down_block_res_samples,
+        sample,
+        encoder_hidden_states: torch.Tensor,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        mid_block_additional_residual: Optional[torch.Tensor] = None):
+    is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
+
+    if is_controlnet:
+        new_down_block_res_samples = ()
+
+        for down_block_res_sample, down_block_additional_residual in zip(
+                down_block_res_samples, down_block_additional_residuals):
+            down_block_res_sample = down_block_res_sample + down_block_additional_residual
+            new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
+
+        down_block_res_samples = new_down_block_res_samples
+
+    # 4. mid
+    if unet.mid_block is not None:
+        if hasattr(unet.mid_block, "has_cross_attention") and unet.mid_block.has_cross_attention:
+            sample = unet.mid_block(
+                sample,
+                emb,
+                encoder_hidden_states=encoder_hidden_states,
+                cross_attention_kwargs=cross_attention_kwargs,
+            )
+        else:
+            sample = unet.mid_block(sample, emb)
+
+    if is_controlnet:
+        sample = sample + mid_block_additional_residual
+
+    # 5. up
+    for i, upsample_block in enumerate(unet.up_blocks):
+        res_samples = down_block_res_samples[-len(upsample_block.resnets):]
+        down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+        if hasattr(upsample_block, 'has_cross_attention') and upsample_block.has_cross_attention:
+            sample = upsample_block(
+                hidden_states=sample,
+                temb=emb,
+                res_hidden_states_tuple=res_samples,
+                encoder_hidden_states=encoder_hidden_states,
+                cross_attention_kwargs=cross_attention_kwargs,
+            )
+        else:
+            sample = upsample_block(
+                hidden_states=sample,
+                temb=emb,
+                res_hidden_states_tuple=res_samples,
+            )
+
+    # 6. post-process
+    if unet.conv_norm_out:
+        sample = unet.conv_norm_out(sample)
+        sample = unet.conv_act(sample)
+    sample = unet.conv_out(sample)
+
+    return sample
+
+
+@MODULES.register_module()
+class UNet2DConditionModel(_UNet2DConditionModel):
+    def __init__(self,
+                 *args,
+                 freeze=True,
+                 freeze_exclude=[],
+                 pretrained=None,
+                 torch_dtype='float32',
+                 freeze_exclude_fp32=True,
+                 freeze_exclude_autocast_dtype='float32',
+                 **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.init_weights(pretrained)
+        if torch_dtype is not None:
+            self.to(getattr(torch, torch_dtype))
+
+        self.set_use_memory_efficient_attention_xformers(
+            not hasattr(torch.nn.functional, 'scaled_dot_product_attention'))
+
+        self.freeze = freeze
+        if self.freeze:
+            flex_freeze(
+                self,
+                exclude_keys=freeze_exclude,
+                exclude_fp32=freeze_exclude_fp32,
+                exclude_autocast_dtype=freeze_exclude_autocast_dtype)
+
+    def init_weights(self, pretrained):
+        if pretrained is not None:
+            logger = get_root_logger()
+            # load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger)
+            checkpoint = _load_checkpoint(pretrained, map_location='cpu', logger=logger)
+            if 'state_dict' in checkpoint:
+                state_dict = checkpoint['state_dict']
+            else:
+                state_dict = checkpoint
+            metadata = getattr(state_dict, '_metadata', OrderedDict())
+            state_dict._metadata = metadata
+            assert self.conv_in.weight.shape[1] == self.conv_out.weight.shape[0]
+            if state_dict['conv_in.weight'].size() != self.conv_in.weight.size():
+                assert state_dict['conv_in.weight'].shape[1] == state_dict['conv_out.weight'].shape[0]
+                src_chn = state_dict['conv_in.weight'].shape[1]
+                tgt_chn = self.conv_in.weight.shape[1]
+                assert src_chn < tgt_chn
+                convert_mat_out = torch.tile(torch.eye(src_chn), (ceildiv(tgt_chn, src_chn), 1))
+                convert_mat_out = convert_mat_out[:tgt_chn]
+                convert_mat_in = F.normalize(convert_mat_out.pinverse(), dim=-1)
+                state_dict['conv_out.weight'] = torch.einsum(
+                    'ts,scxy->tcxy', convert_mat_out, state_dict['conv_out.weight'])
+                state_dict['conv_out.bias'] = torch.einsum(
+                    'ts,s->t', convert_mat_out, state_dict['conv_out.bias'])
+                state_dict['conv_in.weight'] = torch.einsum(
+                    'st,csxy->ctxy', convert_mat_in, state_dict['conv_in.weight'])
+            load_state_dict(self, state_dict, logger=logger)
+
+    def forward(self, sample, timestep, encoder_hidden_states, **kwargs):
+        dtype = sample.dtype
+        return super().forward(
+            sample, timestep, encoder_hidden_states, return_dict=False, **kwargs)[0].to(dtype)
+
+    def forward_enc(self, sample, timestep, encoder_hidden_states, **kwargs):
+        return unet_enc(self, sample, timestep, encoder_hidden_states, **kwargs)
+
+    def forward_dec(self, emb, down_block_res_samples, sample, encoder_hidden_states, **kwargs):
+        return unet_dec(self, emb, down_block_res_samples, sample, encoder_hidden_states, **kwargs)