diffsynth/models/cog_dit.py

import torch
from einops import rearrange, repeat
from .sd3_dit import TimestepEmbeddings
from .attention import Attention
from .utils import load_state_dict_from_folder
from .tiler import TileWorker2Dto3D
import numpy as np



class CogPatchify(torch.nn.Module):
    def __init__(self, dim_in, dim_out, patch_size) -> None:
        super().__init__()
        self.proj = torch.nn.Conv3d(dim_in, dim_out, kernel_size=(1, patch_size, patch_size), stride=(1, patch_size, patch_size))

    def forward(self, hidden_states):
        hidden_states = self.proj(hidden_states)
        hidden_states = rearrange(hidden_states, "B C T H W -> B (T H W) C")
        return hidden_states
    


class CogAdaLayerNorm(torch.nn.Module):
    def __init__(self, dim, dim_cond, single=False):
        super().__init__()
        self.single = single
        self.linear = torch.nn.Linear(dim_cond, dim * (2 if single else 6))
        self.norm = torch.nn.LayerNorm(dim, elementwise_affine=True, eps=1e-5)


    def forward(self, hidden_states, prompt_emb, emb):
        emb = self.linear(torch.nn.functional.silu(emb))
        if self.single:
            shift, scale = emb.unsqueeze(1).chunk(2, dim=2)
            hidden_states = self.norm(hidden_states) * (1 + scale) + shift
            return hidden_states
        else:
            shift_a, scale_a, gate_a, shift_b, scale_b, gate_b = emb.unsqueeze(1).chunk(6, dim=2)
            hidden_states = self.norm(hidden_states) * (1 + scale_a) + shift_a
            prompt_emb = self.norm(prompt_emb) * (1 + scale_b) + shift_b
            return hidden_states, prompt_emb, gate_a, gate_b



class CogDiTBlock(torch.nn.Module):
    def __init__(self, dim, dim_cond, num_heads):
        super().__init__()
        self.norm1 = CogAdaLayerNorm(dim, dim_cond)
        self.attn1 = Attention(q_dim=dim, num_heads=48, head_dim=dim//num_heads, bias_q=True, bias_kv=True, bias_out=True)
        self.norm_q = torch.nn.LayerNorm((dim//num_heads,), eps=1e-06, elementwise_affine=True)
        self.norm_k = torch.nn.LayerNorm((dim//num_heads,), eps=1e-06, elementwise_affine=True)

        self.norm2 = CogAdaLayerNorm(dim, dim_cond)
        self.ff = torch.nn.Sequential(
            torch.nn.Linear(dim, dim*4),
            torch.nn.GELU(approximate="tanh"),
            torch.nn.Linear(dim*4, dim)
        )
    

    def apply_rotary_emb(self, x, freqs_cis):
        cos, sin = freqs_cis  # [S, D]
        cos = cos[None, None]
        sin = sin[None, None]
        cos, sin = cos.to(x.device), sin.to(x.device)
        x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
        x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
        out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
        return out
    

    def process_qkv(self, q, k, v, image_rotary_emb, text_seq_length):
        q = self.norm_q(q)
        k = self.norm_k(k)
        q[:, :, text_seq_length:] = self.apply_rotary_emb(q[:, :, text_seq_length:], image_rotary_emb)
        k[:, :, text_seq_length:] = self.apply_rotary_emb(k[:, :, text_seq_length:], image_rotary_emb)
        return q, k, v
        

    def forward(self, hidden_states, prompt_emb, time_emb, image_rotary_emb):
        # Attention
        norm_hidden_states, norm_encoder_hidden_states, gate_a, gate_b = self.norm1(
            hidden_states, prompt_emb, time_emb
        )
        attention_io = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1)
        attention_io = self.attn1(
            attention_io,
            qkv_preprocessor=lambda q, k, v: self.process_qkv(q, k, v, image_rotary_emb, prompt_emb.shape[1])
        )

        hidden_states = hidden_states + gate_a * attention_io[:, prompt_emb.shape[1]:]
        prompt_emb = prompt_emb + gate_b * attention_io[:, :prompt_emb.shape[1]]

        # Feed forward
        norm_hidden_states, norm_encoder_hidden_states, gate_a, gate_b = self.norm2(
            hidden_states, prompt_emb, time_emb
        )
        ff_io = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1)
        ff_io = self.ff(ff_io)

        hidden_states = hidden_states + gate_a * ff_io[:, prompt_emb.shape[1]:]
        prompt_emb = prompt_emb + gate_b * ff_io[:, :prompt_emb.shape[1]]

        return hidden_states, prompt_emb



class CogDiT(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.patchify = CogPatchify(16, 3072, 2)
        self.time_embedder = TimestepEmbeddings(3072, 512)
        self.context_embedder = torch.nn.Linear(4096, 3072)
        self.blocks = torch.nn.ModuleList([CogDiTBlock(3072, 512, 48) for _ in range(42)])
        self.norm_final = torch.nn.LayerNorm((3072,), eps=1e-05, elementwise_affine=True)
        self.norm_out = CogAdaLayerNorm(3072, 512, single=True)
        self.proj_out = torch.nn.Linear(3072, 64, bias=True)


    def get_resize_crop_region_for_grid(self, src, tgt_width, tgt_height):
        tw = tgt_width
        th = tgt_height
        h, w = src
        r = h / w
        if r > (th / tw):
            resize_height = th
            resize_width = int(round(th / h * w))
        else:
            resize_width = tw
            resize_height = int(round(tw / w * h))

        crop_top = int(round((th - resize_height) / 2.0))
        crop_left = int(round((tw - resize_width) / 2.0))

        return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width)
    

    def get_3d_rotary_pos_embed(
        self, embed_dim, crops_coords, grid_size, temporal_size, theta: int = 10000, use_real: bool = True
    ):
        start, stop = crops_coords
        grid_h = np.linspace(start[0], stop[0], grid_size[0], endpoint=False, dtype=np.float32)
        grid_w = np.linspace(start[1], stop[1], grid_size[1], endpoint=False, dtype=np.float32)
        grid_t = np.linspace(0, temporal_size, temporal_size, endpoint=False, dtype=np.float32)

        # Compute dimensions for each axis
        dim_t = embed_dim // 4
        dim_h = embed_dim // 8 * 3
        dim_w = embed_dim // 8 * 3

        # Temporal frequencies
        freqs_t = 1.0 / (theta ** (torch.arange(0, dim_t, 2).float() / dim_t))
        grid_t = torch.from_numpy(grid_t).float()
        freqs_t = torch.einsum("n , f -> n f", grid_t, freqs_t)
        freqs_t = freqs_t.repeat_interleave(2, dim=-1)

        # Spatial frequencies for height and width
        freqs_h = 1.0 / (theta ** (torch.arange(0, dim_h, 2).float() / dim_h))
        freqs_w = 1.0 / (theta ** (torch.arange(0, dim_w, 2).float() / dim_w))
        grid_h = torch.from_numpy(grid_h).float()
        grid_w = torch.from_numpy(grid_w).float()
        freqs_h = torch.einsum("n , f -> n f", grid_h, freqs_h)
        freqs_w = torch.einsum("n , f -> n f", grid_w, freqs_w)
        freqs_h = freqs_h.repeat_interleave(2, dim=-1)
        freqs_w = freqs_w.repeat_interleave(2, dim=-1)

        # Broadcast and concatenate tensors along specified dimension
        def broadcast(tensors, dim=-1):
            num_tensors = len(tensors)
            shape_lens = {len(t.shape) for t in tensors}
            assert len(shape_lens) == 1, "tensors must all have the same number of dimensions"
            shape_len = list(shape_lens)[0]
            dim = (dim + shape_len) if dim < 0 else dim
            dims = list(zip(*(list(t.shape) for t in tensors)))
            expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
            assert all(
                [*(len(set(t[1])) <= 2 for t in expandable_dims)]
            ), "invalid dimensions for broadcastable concatenation"
            max_dims = [(t[0], max(t[1])) for t in expandable_dims]
            expanded_dims = [(t[0], (t[1],) * num_tensors) for t in max_dims]
            expanded_dims.insert(dim, (dim, dims[dim]))
            expandable_shapes = list(zip(*(t[1] for t in expanded_dims)))
            tensors = [t[0].expand(*t[1]) for t in zip(tensors, expandable_shapes)]
            return torch.cat(tensors, dim=dim)

        freqs = broadcast((freqs_t[:, None, None, :], freqs_h[None, :, None, :], freqs_w[None, None, :, :]), dim=-1)

        t, h, w, d = freqs.shape
        freqs = freqs.view(t * h * w, d)

        # Generate sine and cosine components
        sin = freqs.sin()
        cos = freqs.cos()

        if use_real:
            return cos, sin
        else:
            freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
            return freqs_cis
    

    def prepare_rotary_positional_embeddings(
        self,
        height: int,
        width: int,
        num_frames: int,
        device: torch.device,
    ):
        grid_height = height // 2
        grid_width = width // 2
        base_size_width = 720 // (8 * 2)
        base_size_height = 480 // (8 * 2)

        grid_crops_coords = self.get_resize_crop_region_for_grid(
            (grid_height, grid_width), base_size_width, base_size_height
        )
        freqs_cos, freqs_sin = self.get_3d_rotary_pos_embed(
            embed_dim=64,
            crops_coords=grid_crops_coords,
            grid_size=(grid_height, grid_width),
            temporal_size=num_frames,
            use_real=True,
        )

        freqs_cos = freqs_cos.to(device=device)
        freqs_sin = freqs_sin.to(device=device)
        return freqs_cos, freqs_sin


    def unpatchify(self, hidden_states, height, width):
        hidden_states = rearrange(hidden_states, "B (T H W) (C P Q) -> B C T (H P) (W Q)", P=2, Q=2, H=height//2, W=width//2)
        return hidden_states
    

    def build_mask(self, T, H, W, dtype, device, is_bound):
        t = repeat(torch.arange(T), "T -> T H W", T=T, H=H, W=W)
        h = repeat(torch.arange(H), "H -> T H W", T=T, H=H, W=W)
        w = repeat(torch.arange(W), "W -> T H W", T=T, H=H, W=W)
        border_width = (H + W) // 4
        pad = torch.ones_like(h) * border_width
        mask = torch.stack([
            pad if is_bound[0] else t + 1,
            pad if is_bound[1] else T - t,
            pad if is_bound[2] else h + 1,
            pad if is_bound[3] else H - h,
            pad if is_bound[4] else w + 1,
            pad if is_bound[5] else W - w
        ]).min(dim=0).values
        mask = mask.clip(1, border_width)
        mask = (mask / border_width).to(dtype=dtype, device=device)
        mask = rearrange(mask, "T H W -> 1 1 T H W")
        return mask
    

    def tiled_forward(self, hidden_states, timestep, prompt_emb, tile_size=(60, 90), tile_stride=(30, 45)):
        B, C, T, H, W = hidden_states.shape
        value = torch.zeros((B, C, T, H, W), dtype=hidden_states.dtype, device=hidden_states.device)
        weight = torch.zeros((B, C, T, H, W), dtype=hidden_states.dtype, device=hidden_states.device)

        # Split tasks
        tasks = []
        for h in range(0, H, tile_stride):
            for w in range(0, W, tile_stride):
                if (h-tile_stride >= 0 and h-tile_stride+tile_size >= H) or (w-tile_stride >= 0 and w-tile_stride+tile_size >= W):
                    continue
                h_, w_ = h + tile_size, w + tile_size
                if h_ > H: h, h_ = max(H - tile_size, 0), H
                if w_ > W: w, w_ = max(W - tile_size, 0), W
                tasks.append((h, h_, w, w_))

        # Run
        for hl, hr, wl, wr in tasks:
            mask = self.build_mask(
                value.shape[2], (hr-hl), (wr-wl),
                hidden_states.dtype, hidden_states.device,
                is_bound=(True, True, hl==0, hr>=H, wl==0, wr>=W)
            )
            model_output = self.forward(hidden_states[:, :, :, hl:hr, wl:wr], timestep, prompt_emb)
            value[:, :, :, hl:hr, wl:wr] += model_output * mask
            weight[:, :, :, hl:hr, wl:wr] += mask
        value = value / weight

        return value


    def forward(self, hidden_states, timestep, prompt_emb, image_rotary_emb=None, tiled=False, tile_size=90, tile_stride=30, use_gradient_checkpointing=False):
        if tiled:
            return TileWorker2Dto3D().tiled_forward(
                forward_fn=lambda x: self.forward(x, timestep, prompt_emb),
                model_input=hidden_states,
                tile_size=tile_size, tile_stride=tile_stride,
                tile_device=hidden_states.device, tile_dtype=hidden_states.dtype,
                computation_device=self.context_embedder.weight.device, computation_dtype=self.context_embedder.weight.dtype
            )
        num_frames, height, width = hidden_states.shape[-3:]
        if image_rotary_emb is None:
            image_rotary_emb = self.prepare_rotary_positional_embeddings(height, width, num_frames, device=self.context_embedder.weight.device)
        hidden_states = self.patchify(hidden_states)
        time_emb = self.time_embedder(timestep, dtype=hidden_states.dtype)
        prompt_emb = self.context_embedder(prompt_emb)

        def create_custom_forward(module):
            def custom_forward(*inputs):
                return module(*inputs)
            return custom_forward
        
        for block in self.blocks:
            if self.training and use_gradient_checkpointing:
                hidden_states, prompt_emb = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(block),
                    hidden_states, prompt_emb, time_emb, image_rotary_emb,
                    use_reentrant=False,
                )
            else:
                hidden_states, prompt_emb = block(hidden_states, prompt_emb, time_emb, image_rotary_emb)

        hidden_states = torch.cat([prompt_emb, hidden_states], dim=1)
        hidden_states = self.norm_final(hidden_states)
        hidden_states = hidden_states[:, prompt_emb.shape[1]:]
        hidden_states = self.norm_out(hidden_states, prompt_emb, time_emb)
        hidden_states = self.proj_out(hidden_states)
        hidden_states = self.unpatchify(hidden_states, height, width)

        return hidden_states
    

    @staticmethod
    def state_dict_converter():
        return CogDiTStateDictConverter()
    

    @staticmethod
    def from_pretrained(file_path, torch_dtype=torch.bfloat16):
        model = CogDiT().to(torch_dtype)
        state_dict = load_state_dict_from_folder(file_path, torch_dtype=torch_dtype)
        state_dict = CogDiT.state_dict_converter().from_diffusers(state_dict)
        model.load_state_dict(state_dict)
        return model



class CogDiTStateDictConverter:
    def __init__(self):
        pass


    def from_diffusers(self, state_dict):
        rename_dict = {
            "patch_embed.proj.weight": "patchify.proj.weight",
            "patch_embed.proj.bias": "patchify.proj.bias",
            "patch_embed.text_proj.weight": "context_embedder.weight",
            "patch_embed.text_proj.bias": "context_embedder.bias",
            "time_embedding.linear_1.weight": "time_embedder.timestep_embedder.0.weight",
            "time_embedding.linear_1.bias": "time_embedder.timestep_embedder.0.bias",
            "time_embedding.linear_2.weight": "time_embedder.timestep_embedder.2.weight",
            "time_embedding.linear_2.bias": "time_embedder.timestep_embedder.2.bias",

            "norm_final.weight": "norm_final.weight",
            "norm_final.bias": "norm_final.bias",
            "norm_out.linear.weight": "norm_out.linear.weight",
            "norm_out.linear.bias": "norm_out.linear.bias",
            "norm_out.norm.weight": "norm_out.norm.weight",
            "norm_out.norm.bias": "norm_out.norm.bias",
            "proj_out.weight": "proj_out.weight",
            "proj_out.bias": "proj_out.bias",
        }
        suffix_dict = {
            "norm1.linear.weight": "norm1.linear.weight",
            "norm1.linear.bias": "norm1.linear.bias",
            "norm1.norm.weight": "norm1.norm.weight",
            "norm1.norm.bias": "norm1.norm.bias",
            "attn1.norm_q.weight": "norm_q.weight",
            "attn1.norm_q.bias": "norm_q.bias",
            "attn1.norm_k.weight": "norm_k.weight",
            "attn1.norm_k.bias": "norm_k.bias",
            "attn1.to_q.weight": "attn1.to_q.weight",
            "attn1.to_q.bias": "attn1.to_q.bias",
            "attn1.to_k.weight": "attn1.to_k.weight",
            "attn1.to_k.bias": "attn1.to_k.bias",
            "attn1.to_v.weight": "attn1.to_v.weight",
            "attn1.to_v.bias": "attn1.to_v.bias",
            "attn1.to_out.0.weight": "attn1.to_out.weight",
            "attn1.to_out.0.bias": "attn1.to_out.bias",
            "norm2.linear.weight": "norm2.linear.weight",
            "norm2.linear.bias": "norm2.linear.bias",
            "norm2.norm.weight": "norm2.norm.weight",
            "norm2.norm.bias": "norm2.norm.bias",
            "ff.net.0.proj.weight": "ff.0.weight",
            "ff.net.0.proj.bias": "ff.0.bias",
            "ff.net.2.weight": "ff.2.weight",
            "ff.net.2.bias": "ff.2.bias",
        }
        state_dict_ = {}
        for name, param in state_dict.items():
            if name in rename_dict:
                if name == "patch_embed.proj.weight":
                    param = param.unsqueeze(2)
                state_dict_[rename_dict[name]] = param
            else:
                names = name.split(".")
                if names[0] == "transformer_blocks":
                    suffix = ".".join(names[2:])
                    state_dict_[f"blocks.{names[1]}." + suffix_dict[suffix]] = param
        return state_dict_
    

    def from_civitai(self, state_dict):
        return self.from_diffusers(state_dict)