# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""Flux model flavors."""

from copy import deepcopy
from functools import partial

import torch.nn as nn

from torchtitan.config.transform import (
    ModelConfigConverter,
    validate_converter_compatibility,
)

from torchtitan.models.common.linear import Linear
from torchtitan.models.common.nn_modules import RMSNorm

from .model.autoencoder import AutoEncoder
from .model.hf_embedder import FluxEmbedder
from .model.layers import (
    DoubleStreamBlock,
    EmbedND,
    LastLayer,
    MLPEmbedder,
    Modulation,
    QKNorm,
    SelfAttention,
    SingleStreamBlock,
)
from .model.model import FluxModel

__all__ = [
    "FluxModel",
    "MODEL_FLAVORS",
    "build_model_config",
]

_ZERO_LINEAR = {"weight": nn.init.zeros_, "bias": nn.init.zeros_}
_XAVIER_LINEAR = {"weight": nn.init.xavier_uniform_, "bias": nn.init.zeros_}
_NORMAL_02 = {"weight": partial(nn.init.normal_, std=0.02), "bias": nn.init.zeros_}
_NORM_INIT = {"weight": nn.init.ones_}

_T5_ENCODER_VERSION = "google/t5-v1_1-xxl"
_CLIP_ENCODER_VERSION = "openai/clip-vit-large-patch14"


def _make_double_block_config(
    hidden_size: int,
    num_heads: int,
    mlp_ratio: float,
    qkv_bias: bool,
) -> DoubleStreamBlock.Config:
    hs = hidden_size
    head_dim = hs // num_heads
    mlp_hidden_dim = int(hs * mlp_ratio)
    return DoubleStreamBlock.Config(
        hidden_size=hs,
        num_heads=num_heads,
        mlp_ratio=mlp_ratio,
        qkv_bias=qkv_bias,
        img_mod=Modulation.Config(
            double=True,
            lin=Linear.Config(
                in_features=hs,
                out_features=6 * hs,
                bias=True,
                param_init=_ZERO_LINEAR,
            ),
        ),
        txt_mod=Modulation.Config(
            double=True,
            lin=Linear.Config(
                in_features=hs,
                out_features=6 * hs,
                bias=True,
                param_init=_ZERO_LINEAR,
            ),
        ),
        img_attn=SelfAttention.Config(
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qkv=Linear.Config(
                in_features=hs,
                out_features=hs * 3,
                bias=qkv_bias,
                param_init=_XAVIER_LINEAR,
            ),
            proj=Linear.Config(
                in_features=hs,
                out_features=hs,
                bias=True,
                param_init=_XAVIER_LINEAR,
            ),
            norm=QKNorm.Config(
                query_norm=RMSNorm.Config(
                    normalized_shape=head_dim,
                    param_init=_NORM_INIT,
                ),
                key_norm=RMSNorm.Config(
                    normalized_shape=head_dim,
                    param_init=_NORM_INIT,
                ),
            ),
        ),
        txt_attn=SelfAttention.Config(
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qkv=Linear.Config(
                in_features=hs,
                out_features=hs * 3,
                bias=qkv_bias,
                param_init=_XAVIER_LINEAR,
            ),
            proj=Linear.Config(
                in_features=hs,
                out_features=hs,
                bias=True,
                param_init=_XAVIER_LINEAR,
            ),
            norm=QKNorm.Config(
                query_norm=RMSNorm.Config(
                    normalized_shape=head_dim,
                    param_init=_NORM_INIT,
                ),
                key_norm=RMSNorm.Config(
                    normalized_shape=head_dim,
                    param_init=_NORM_INIT,
                ),
            ),
        ),
        img_mlp_in=Linear.Config(
            in_features=hs,
            out_features=mlp_hidden_dim,
            bias=True,
            param_init=_XAVIER_LINEAR,
        ),
        img_mlp_out=Linear.Config(
            in_features=mlp_hidden_dim,
            out_features=hs,
            bias=True,
            param_init=_XAVIER_LINEAR,
        ),
        txt_mlp_in=Linear.Config(
            in_features=hs,
            out_features=mlp_hidden_dim,
            bias=True,
            param_init=_XAVIER_LINEAR,
        ),
        txt_mlp_out=Linear.Config(
            in_features=mlp_hidden_dim,
            out_features=hs,
            bias=True,
            param_init=_XAVIER_LINEAR,
        ),
    )


def _make_single_block_config(
    hidden_size: int,
    num_heads: int,
    mlp_ratio: float,
) -> SingleStreamBlock.Config:
    hs = hidden_size
    head_dim = hs // num_heads
    mlp_hidden_dim = int(hs * mlp_ratio)
    return SingleStreamBlock.Config(
        hidden_size=hs,
        num_heads=num_heads,
        mlp_ratio=mlp_ratio,
        linear1=Linear.Config(
            in_features=hs,
            out_features=hs * 3 + mlp_hidden_dim,
            bias=True,
            param_init=_XAVIER_LINEAR,
        ),
        linear2=Linear.Config(
            in_features=hs + mlp_hidden_dim,
            out_features=hs,
            bias=True,
            param_init=_XAVIER_LINEAR,
        ),
        modulation=Modulation.Config(
            double=False,
            lin=Linear.Config(
                in_features=hs,
                out_features=3 * hs,
                bias=True,
                param_init=_ZERO_LINEAR,
            ),
        ),
        norm=QKNorm.Config(
            query_norm=RMSNorm.Config(
                normalized_shape=head_dim,
                param_init=_NORM_INIT,
            ),
            key_norm=RMSNorm.Config(
                normalized_shape=head_dim,
                param_init=_NORM_INIT,
            ),
        ),
    )


def _flux_dev() -> FluxModel.Config:
    hidden_size = 3072
    num_heads = 24
    mlp_ratio = 4.0
    qkv_bias = True
    vec_in_dim = 768
    in_channels = 64
    context_in_dim = 4096
    depth = 19
    depth_single_blocks = 38
    double_tmpl = _make_double_block_config(
        hidden_size=hidden_size,
        num_heads=num_heads,
        mlp_ratio=mlp_ratio,
        qkv_bias=qkv_bias,
    )
    single_tmpl = _make_single_block_config(
        hidden_size=hidden_size,
        num_heads=num_heads,
        mlp_ratio=mlp_ratio,
    )
    return FluxModel.Config(
        in_channels=in_channels,
        out_channels=64,
        vec_in_dim=vec_in_dim,
        context_in_dim=context_in_dim,
        hidden_size=hidden_size,
        mlp_ratio=mlp_ratio,
        num_heads=num_heads,
        depth=depth,
        depth_single_blocks=depth_single_blocks,
        axes_dim=(16, 56, 56),
        theta=10_000,
        qkv_bias=qkv_bias,
        img_in=Linear.Config(
            in_features=in_channels,
            out_features=hidden_size,
            bias=True,
            param_init=_XAVIER_LINEAR,
        ),
        txt_in=Linear.Config(
            in_features=context_in_dim,
            out_features=hidden_size,
            bias=True,
            param_init=_XAVIER_LINEAR,
        ),
        autoencoder=AutoEncoder.Config(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=(1, 2, 4, 4),
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
        clip_encoder=FluxEmbedder.Config(version=_CLIP_ENCODER_VERSION),
        t5_encoder=FluxEmbedder.Config(version=_T5_ENCODER_VERSION),
        pe_config=EmbedND.Config(dim=128, theta=10_000, axes_dim=(16, 56, 56)),
        time_in_config=MLPEmbedder.Config(
            in_dim=256,
            hidden_dim=hidden_size,
            in_layer=Linear.Config(
                in_features=256,
                out_features=hidden_size,
                bias=True,
                param_init=_NORMAL_02,
            ),
            out_layer=Linear.Config(
                in_features=hidden_size,
                out_features=hidden_size,
                bias=True,
                param_init=_NORMAL_02,
            ),
        ),
        vector_in_config=MLPEmbedder.Config(
            in_dim=vec_in_dim,
            hidden_dim=hidden_size,
            in_layer=Linear.Config(
                in_features=vec_in_dim,
                out_features=hidden_size,
                bias=True,
                param_init=_NORMAL_02,
            ),
            out_layer=Linear.Config(
                in_features=hidden_size,
                out_features=hidden_size,
                bias=True,
                param_init=_NORMAL_02,
            ),
        ),
        double_blocks=[deepcopy(double_tmpl) for _ in range(depth)],
        single_blocks=[deepcopy(single_tmpl) for _ in range(depth_single_blocks)],
        final_layer_config=LastLayer.Config(
            hidden_size=hidden_size,
            patch_size=1,
            out_channels=64,
            linear=Linear.Config(
                in_features=hidden_size,
                out_features=1 * 1 * 64,
                bias=True,
                param_init=_ZERO_LINEAR,
            ),
            adaln_linear=Linear.Config(
                in_features=hidden_size,
                out_features=2 * hidden_size,
                bias=True,
                param_init=_ZERO_LINEAR,
            ),
        ),
    )


def _flux_schnell() -> FluxModel.Config:
    hidden_size = 3072
    num_heads = 24
    mlp_ratio = 4.0
    qkv_bias = True
    vec_in_dim = 768
    in_channels = 64
    context_in_dim = 4096
    depth = 19
    depth_single_blocks = 38
    double_tmpl = _make_double_block_config(
        hidden_size=hidden_size,
        num_heads=num_heads,
        mlp_ratio=mlp_ratio,
        qkv_bias=qkv_bias,
    )
    single_tmpl = _make_single_block_config(
        hidden_size=hidden_size,
        num_heads=num_heads,
        mlp_ratio=mlp_ratio,
    )
    return FluxModel.Config(
        in_channels=in_channels,
        out_channels=64,
        vec_in_dim=vec_in_dim,
        context_in_dim=context_in_dim,
        hidden_size=hidden_size,
        mlp_ratio=mlp_ratio,
        num_heads=num_heads,
        depth=depth,
        depth_single_blocks=depth_single_blocks,
        axes_dim=(16, 56, 56),
        theta=10_000,
        qkv_bias=qkv_bias,
        img_in=Linear.Config(
            in_features=in_channels,
            out_features=hidden_size,
            bias=True,
            param_init=_XAVIER_LINEAR,
        ),
        txt_in=Linear.Config(
            in_features=context_in_dim,
            out_features=hidden_size,
            bias=True,
            param_init=_XAVIER_LINEAR,
        ),
        autoencoder=AutoEncoder.Config(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=(1, 2, 4, 4),
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
        clip_encoder=FluxEmbedder.Config(version=_CLIP_ENCODER_VERSION),
        t5_encoder=FluxEmbedder.Config(version=_T5_ENCODER_VERSION),
        pe_config=EmbedND.Config(dim=128, theta=10_000, axes_dim=(16, 56, 56)),
        time_in_config=MLPEmbedder.Config(
            in_dim=256,
            hidden_dim=hidden_size,
            in_layer=Linear.Config(
                in_features=256,
                out_features=hidden_size,
                bias=True,
                param_init=_NORMAL_02,
            ),
            out_layer=Linear.Config(
                in_features=hidden_size,
                out_features=hidden_size,
                bias=True,
                param_init=_NORMAL_02,
            ),
        ),
        vector_in_config=MLPEmbedder.Config(
            in_dim=vec_in_dim,
            hidden_dim=hidden_size,
            in_layer=Linear.Config(
                in_features=vec_in_dim,
                out_features=hidden_size,
                bias=True,
                param_init=_NORMAL_02,
            ),
            out_layer=Linear.Config(
                in_features=hidden_size,
                out_features=hidden_size,
                bias=True,
                param_init=_NORMAL_02,
            ),
        ),
        double_blocks=[deepcopy(double_tmpl) for _ in range(depth)],
        single_blocks=[deepcopy(single_tmpl) for _ in range(depth_single_blocks)],
        final_layer_config=LastLayer.Config(
            hidden_size=hidden_size,
            patch_size=1,
            out_channels=64,
            linear=Linear.Config(
                in_features=hidden_size,
                out_features=1 * 1 * 64,
                bias=True,
                param_init=_ZERO_LINEAR,
            ),
            adaln_linear=Linear.Config(
                in_features=hidden_size,
                out_features=2 * hidden_size,
                bias=True,
                param_init=_ZERO_LINEAR,
            ),
        ),
    )


def _flux_debug() -> FluxModel.Config:
    hidden_size = 1536
    num_heads = 12
    mlp_ratio = 4.0
    qkv_bias = True
    vec_in_dim = 768
    in_channels = 64
    context_in_dim = 4096
    depth = 2
    depth_single_blocks = 2
    double_tmpl = _make_double_block_config(
        hidden_size=hidden_size,
        num_heads=num_heads,
        mlp_ratio=mlp_ratio,
        qkv_bias=qkv_bias,
    )
    single_tmpl = _make_single_block_config(
        hidden_size=hidden_size,
        num_heads=num_heads,
        mlp_ratio=mlp_ratio,
    )
    return FluxModel.Config(
        in_channels=in_channels,
        out_channels=64,
        vec_in_dim=vec_in_dim,
        context_in_dim=context_in_dim,
        hidden_size=hidden_size,
        mlp_ratio=mlp_ratio,
        num_heads=num_heads,
        depth=depth,
        depth_single_blocks=depth_single_blocks,
        axes_dim=(16, 56, 56),
        theta=10_000,
        qkv_bias=qkv_bias,
        img_in=Linear.Config(
            in_features=in_channels,
            out_features=hidden_size,
            bias=True,
            param_init=_XAVIER_LINEAR,
        ),
        txt_in=Linear.Config(
            in_features=context_in_dim,
            out_features=hidden_size,
            bias=True,
            param_init=_XAVIER_LINEAR,
        ),
        autoencoder=AutoEncoder.Config(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=(1, 2, 4, 4),
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
        clip_encoder=FluxEmbedder.Config(version=_CLIP_ENCODER_VERSION),
        t5_encoder=FluxEmbedder.Config(version=_T5_ENCODER_VERSION),
        pe_config=EmbedND.Config(dim=128, theta=10_000, axes_dim=(16, 56, 56)),
        time_in_config=MLPEmbedder.Config(
            in_dim=256,
            hidden_dim=hidden_size,
            in_layer=Linear.Config(
                in_features=256,
                out_features=hidden_size,
                bias=True,
                param_init=_NORMAL_02,
            ),
            out_layer=Linear.Config(
                in_features=hidden_size,
                out_features=hidden_size,
                bias=True,
                param_init=_NORMAL_02,
            ),
        ),
        vector_in_config=MLPEmbedder.Config(
            in_dim=vec_in_dim,
            hidden_dim=hidden_size,
            in_layer=Linear.Config(
                in_features=vec_in_dim,
                out_features=hidden_size,
                bias=True,
                param_init=_NORMAL_02,
            ),
            out_layer=Linear.Config(
                in_features=hidden_size,
                out_features=hidden_size,
                bias=True,
                param_init=_NORMAL_02,
            ),
        ),
        double_blocks=[deepcopy(double_tmpl) for _ in range(depth)],
        single_blocks=[deepcopy(single_tmpl) for _ in range(depth_single_blocks)],
        final_layer_config=LastLayer.Config(
            hidden_size=hidden_size,
            patch_size=1,
            out_channels=64,
            linear=Linear.Config(
                in_features=hidden_size,
                out_features=1 * 1 * 64,
                bias=True,
                param_init=_ZERO_LINEAR,
            ),
            adaln_linear=Linear.Config(
                in_features=hidden_size,
                out_features=2 * hidden_size,
                bias=True,
                param_init=_ZERO_LINEAR,
            ),
        ),
    )


# The default lengths are ``_flux_seq_len(img_size, max_t5_encoding_len)`` for
# the img_size / T5 length each shipped trainer config uses.
MODEL_FLAVORS = {
    "flux-dev": _flux_dev,
    "flux-schnell": _flux_schnell,
    "flux-debug": _flux_debug,
}


def build_model_config(
    flavor: str,
    converters: list[ModelConfigConverter.Config] | None = None,
) -> FluxModel.Config:
    config = MODEL_FLAVORS[flavor]()
    if converters is not None:
        validate_converter_compatibility(converters)
        for c in converters:
            config = c.build().convert(config)
    return config
