AFM-XL-2-1NFE-guided/generator/modeling_afm.py

"""Self-contained AFM generator for Hugging Face diffusers checkpoints."""

from __future__ import annotations

import math
from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.modeling_utils import ModelMixin
from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
from torch import nn


def build_dit_config(
    *,
    depth: int = 28,
    hidden_size: int = 1152,
    patch_size: int = 2,
    num_heads: int = 16,
    learn_sigma: bool = False,
    class_dropout_prob: float = 0.0,
    input_size: int = 32,
    num_classes: int = 1000,
    in_channels: int = 4,
) -> Dict[str, Any]:
    attention_head_dim = hidden_size // num_heads
    num_embeds_ada_norm = num_classes + (1 if class_dropout_prob > 0 else 0)
    out_channels = in_channels * 2 if learn_sigma else in_channels
    return {
        "sample_size": input_size,
        "num_layers": depth,
        "num_attention_heads": num_heads,
        "attention_head_dim": attention_head_dim,
        "in_channels": in_channels,
        "out_channels": out_channels,
        "patch_size": patch_size,
        "attention_bias": True,
        "activation_fn": "gelu-approximate",
        "num_embeds_ada_norm": num_embeds_ada_norm,
        "norm_type": "ada_norm_zero",
        "norm_elementwise_affine": False,
        "dropout": 0.0,
        "norm_num_groups": 32,
        "norm_eps": 1e-5,
        "upcast_attention": False,
    }


_PREFIX_KEYS = ("transformer.", "t_tgt_embedder.")


def _wrap_dit_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    return {k if k.startswith(_PREFIX_KEYS) else f"transformer.{k}": v for k, v in state_dict.items()}


def forward_dit_output(
    transformer: DiTTransformer2DModel,
    hidden_states: torch.Tensor,
    timestep: Optional[torch.Tensor],
    class_labels: torch.Tensor,
    extra_cond: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    hidden_states = transformer.pos_embed(hidden_states)

    for block in transformer.transformer_blocks:
        if torch.is_grad_enabled() and transformer.gradient_checkpointing:
            hidden_states = transformer._gradient_checkpointing_func(
                block,
                hidden_states,
                None,
                None,
                None,
                timestep,
                None,
                class_labels,
            )
        else:
            hidden_states = block(
                hidden_states,
                attention_mask=None,
                encoder_hidden_states=None,
                encoder_attention_mask=None,
                timestep=timestep,
                cross_attention_kwargs=None,
                class_labels=class_labels,
            )

    conditioning = transformer.transformer_blocks[0].norm1.emb(
        timestep, class_labels, hidden_dtype=hidden_states.dtype
    )
    if extra_cond is not None:
        conditioning = conditioning + extra_cond
    shift, scale = transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)
    hidden_states = transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
    hidden_states = transformer.proj_out_2(hidden_states)

    height = width = int(hidden_states.shape[1] ** 0.5)
    hidden_states = hidden_states.reshape(
        shape=(-1, height, width, transformer.patch_size, transformer.patch_size, transformer.out_channels)
    )
    hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
    return hidden_states.reshape(
        shape=(-1, transformer.out_channels, height * transformer.patch_size, width * transformer.patch_size)
    )


def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor:
    if output.shape[1] == in_channels:
        return output
    output, _ = output.chunk(2, dim=1)
    return output


class TimestepEmbedder(nn.Module):
    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size, bias=True),
        )
        self.frequency_embedding_size = frequency_embedding_size

    @staticmethod
    def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
        half = dim // 2
        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
            device=t.device
        )
        args = t[:, None].float() * freqs[None]
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        if dim % 2:
            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
        return embedding

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        emb = self.timestep_embedding(t.float(), self.frequency_embedding_size)
        weight_dtype = self.mlp[0].weight.dtype
        return self.mlp(emb.to(dtype=weight_dtype))


class _AFMCheckpointMixin(ModelMixin):
    def _fix_state_dict_keys_on_load(self, state_dict):
        wrapped = _wrap_dit_state_dict(state_dict)
        state_dict.clear()
        state_dict.update(wrapped)
        super()._fix_state_dict_keys_on_load(state_dict)


class Generator(nn.Module):
    def __init__(
        self,
        *,
        use_t_src: bool = False,
        use_t_tgt: bool = False,
        depth: int = 28,
        hidden_size: int = 1152,
        patch_size: int = 2,
        num_heads: int = 16,
        learn_sigma: bool = False,
        class_dropout_prob: float = 0.0,
        input_size: int = 32,
        num_classes: int = 1000,
        in_channels: int = 4,
        **_,
    ):
        super().__init__()
        self.use_t_src = use_t_src
        self.use_t_tgt = use_t_tgt
        self.transformer = DiTTransformer2DModel(
            **build_dit_config(
                depth=depth,
                hidden_size=hidden_size,
                patch_size=patch_size,
                num_heads=num_heads,
                learn_sigma=learn_sigma,
                class_dropout_prob=class_dropout_prob,
                input_size=input_size,
                num_classes=num_classes,
                in_channels=in_channels,
            )
        )
        self.in_channels = in_channels
        if use_t_tgt:
            self.t_tgt_embedder = TimestepEmbedder(hidden_size)
            nn.init.normal_(self.t_tgt_embedder.mlp[0].weight, std=0.02)
            nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02)

    def load_state_dict(self, state_dict, strict: bool = True):
        if self.use_t_tgt and not any(key.startswith("t_tgt_embedder.") for key in state_dict):
            strict = False
        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)

    def forward(self, x, y, t_src=None, t_tgt=None):
        if self.use_t_src and t_src is not None:
            timestep = (t_src * 1000).long()
        else:
            timestep = torch.zeros(x.shape[0], device=x.device, dtype=torch.long)

        extra_cond = None
        if self.use_t_tgt and t_tgt is not None:
            extra_cond = self.t_tgt_embedder(t_tgt * 1000)

        output = forward_dit_output(self.transformer, x, timestep, y, extra_cond=extra_cond)
        return split_output_channels(output, self.in_channels)


class GeneratorDeep(nn.Module):
    def __init__(
        self,
        *,
        repeat: int,
        depth: int = 28,
        hidden_size: int = 1152,
        patch_size: int = 2,
        num_heads: int = 16,
        learn_sigma: bool = False,
        class_dropout_prob: float = 0.0,
        input_size: int = 32,
        num_classes: int = 1000,
        in_channels: int = 4,
        **_,
    ):
        super().__init__()
        self.repeat = repeat
        self.transformer = DiTTransformer2DModel(
            **build_dit_config(
                depth=depth,
                hidden_size=hidden_size,
                patch_size=patch_size,
                num_heads=num_heads,
                learn_sigma=learn_sigma,
                class_dropout_prob=class_dropout_prob,
                input_size=input_size,
                num_classes=num_classes,
                in_channels=in_channels,
            )
        )
        self.in_channels = in_channels

    def load_state_dict(self, state_dict, strict: bool = True):
        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)

    def forward(self, x, y, *args, **kwargs):
        hidden_states = self.transformer.pos_embed(x)
        for t in torch.arange(1, 0, -1 / self.repeat).tolist():
            timestep = torch.full([x.shape[0]], int(t * 1000), device=x.device, dtype=torch.long)
            for block in self.transformer.transformer_blocks:
                hidden_states = block(
                    hidden_states,
                    attention_mask=None,
                    encoder_hidden_states=None,
                    encoder_attention_mask=None,
                    timestep=timestep,
                    cross_attention_kwargs=None,
                    class_labels=y,
                )

        conditioning = self.transformer.transformer_blocks[0].norm1.emb(
            torch.zeros(x.shape[0], device=x.device, dtype=torch.long),
            y,
            hidden_dtype=hidden_states.dtype,
        )
        shift, scale = self.transformer.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)
        hidden_states = self.transformer.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
        hidden_states = self.transformer.proj_out_2(hidden_states)

        height = width = int(hidden_states.shape[1] ** 0.5)
        hidden_states = hidden_states.reshape(
            shape=(
                -1,
                height,
                width,
                self.transformer.patch_size,
                self.transformer.patch_size,
                self.transformer.out_channels,
            )
        )
        hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
        output = hidden_states.reshape(
            shape=(-1, self.transformer.out_channels, height * self.transformer.patch_size, width * self.transformer.patch_size)
        )
        return split_output_channels(output, self.in_channels)


class AFMGenerator2DModel(_AFMCheckpointMixin, Generator, ConfigMixin):
    config_name = "config.json"

    @register_to_config
    def __init__(
        self,
        model_type: str = "AFM-XL/2",
        architecture: str = "standard",
        repeat: int = 1,
        use_t_src: bool = False,
        use_t_tgt: bool = False,
        depth: int = 28,
        hidden_size: int = 1152,
        patch_size: int = 2,
        num_heads: int = 16,
        learn_sigma: bool = False,
        class_dropout_prob: float = 0.0,
        input_size: int = 32,
        num_classes: int = 1000,
        in_channels: int = 4,
        pred_type: str = "x",
        num_inference_steps: int = 1,
    ) -> None:
        if architecture == "deep":
            raise ValueError("Use AFMGeneratorDeep2DModel for deep AFM checkpoints.")
        Generator.__init__(
            self,
            use_t_src=use_t_src,
            use_t_tgt=use_t_tgt,
            depth=depth,
            hidden_size=hidden_size,
            patch_size=patch_size,
            num_heads=num_heads,
            learn_sigma=learn_sigma,
            class_dropout_prob=class_dropout_prob,
            input_size=input_size,
            num_classes=num_classes,
            in_channels=in_channels,
        )


class AFMGeneratorDeep2DModel(_AFMCheckpointMixin, GeneratorDeep, ConfigMixin):
    config_name = "config.json"

    @register_to_config
    def __init__(
        self,
        model_type: str = "AFM-XL/2",
        architecture: str = "deep",
        repeat: int = 2,
        use_t_src: bool = False,
        use_t_tgt: bool = False,
        depth: int = 28,
        hidden_size: int = 1152,
        patch_size: int = 2,
        num_heads: int = 16,
        learn_sigma: bool = False,
        class_dropout_prob: float = 0.0,
        input_size: int = 32,
        num_classes: int = 1000,
        in_channels: int = 4,
        pred_type: str = "x",
        num_inference_steps: int = 1,
    ) -> None:
        GeneratorDeep.__init__(
            self,
            repeat=repeat,
            depth=depth,
            hidden_size=hidden_size,
            patch_size=patch_size,
            num_heads=num_heads,
            learn_sigma=learn_sigma,
            class_dropout_prob=class_dropout_prob,
            input_size=input_size,
            num_classes=num_classes,
            in_channels=in_channels,
        )


__all__ = ["AFMGenerator2DModel", "AFMGeneratorDeep2DModel"]