nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) # FLUX-VAE GroupNorm

nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)

# 입력
[1, 2, 3, 4]
[4, 5, 6, 7]

# Norm 후
[-1.34, -0.45, 0.45, 1.34]
[-1.34, -0.45, 0.45, 1.34]

# QWEN-IMAGE Norm
class AdaLayerNormContinuous(nn.Module):
    r"""
    Adaptive normalization layer with a norm layer (layer_norm or rms_norm).

    Args:
        embedding_dim (`int`): Embedding dimension to use during projection.
        conditioning_embedding_dim (`int`): Dimension of the input condition.
        elementwise_affine (`bool`, defaults to `True`):
            Boolean flag to denote if affine transformation should be applied.
        eps (`float`, defaults to 1e-5): Epsilon factor.
        bias (`bias`, defaults to `True`): Boolean flag to denote if bias should be use.
        norm_type (`str`, defaults to `"layer_norm"`):
            Normalization layer to use. Values supported: "layer_norm", "rms_norm".
    """

    def __init__(
        self,
        embedding_dim: int,
        conditioning_embedding_dim: int,
        # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters
        # because the output is immediately scaled and shifted by the projected conditioning embeddings.
        # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters.
        # However, this is how it was implemented in the original code, and it's rather likely you should
        # set `elementwise_affine` to False.
        elementwise_affine=True,
        eps=1e-5,
        bias=True,
        norm_type="layer_norm",
    ):
        super().__init__()
        self.silu = nn.SiLU()
        self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
        if norm_type == "layer_norm":
            self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
        elif norm_type == "rms_norm":
            self.norm = RMSNorm(embedding_dim, eps, elementwise_affine)
        else:
            raise ValueError(f"unknown norm_type {norm_type}")

    def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor:
        # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT)
        emb = self.linear(self.silu(conditioning_embedding).to(x.dtype))
        scale, shift = torch.chunk(emb, 2, dim=1)
        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
        return x


# QWEN-Image 몸통
class QwenImageTransformer2DModel(nn.Module): 
    def __init__(
                self,
                params: QwenParams,
        ):
        ...
        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
        ...
    
    def forward(self, hidden_states, ..., timestep, ...):
        ...
        
        hidden_states = self.norm_out(hidden_states, temb)
        output = self.proj_out(hidden_states)

        return output

# FLUX RMSNorm
class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(dim))

    def forward(self, x: Tensor):
        x_dtype = x.dtype
        x = x.float()
        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
        # return (x * rrms).to(dtype=x_dtype) * self.scale
        return ((x * rrms) * self.scale.float()).to(dtype=x_dtype)
    
# QWEN-IMAGE RMSNorm
class RMSNorm(nn.Module):
    r"""
    RMS Norm as introduced in https://huggingface.co/papers/1910.07467 by Zhang et al.

    Args:
        dim (`int`): Number of dimensions to use for `weights`. Only effective when `elementwise_affine` is True.
        eps (`float`): Small value to use when calculating the reciprocal of the square-root.
        elementwise_affine (`bool`, defaults to `True`):
            Boolean flag to denote if affine transformation should be applied.
        bias (`bool`, defaults to False): If also training the `bias` param.
    """

    def __init__(self, dim, eps: float, elementwise_affine: bool = True, bias: bool = False):
        super().__init__()

        self.eps = eps
        self.elementwise_affine = elementwise_affine

        if isinstance(dim, numbers.Integral):
            dim = (dim,)

        self.dim = torch.Size(dim)

        self.weight = None
        self.bias = None

        if elementwise_affine:
            self.weight = nn.Parameter(torch.ones(dim))
            if bias:
                self.bias = nn.Parameter(torch.zeros(dim))

    def forward(self, hidden_states):
        # if is_torch_npu_available():
        #     import torch_npu
        #     if self.weight is not None:
        #         # convert into half-precision if necessary
        #         if self.weight.dtype in [torch.float16, torch.bfloat16]:
        #             hidden_states = hidden_states.to(self.weight.dtype)
        #     hidden_states = torch_npu.npu_rms_norm(hidden_states, self.weight, epsilon=self.eps)[0]
        #     if self.bias is not None:
        #         hidden_states = hidden_states + self.bias
        # else:
        input_dtype = hidden_states.dtype
        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)

        if self.weight is not None:
            # convert into half-precision if necessary
            if self.weight.dtype in [torch.float16, torch.bfloat16]:
                hidden_states = hidden_states.to(self.weight.dtype)
            elif self.weight.dtype == torch.float8_e4m3fn:  # fp8 support
                hidden_states = hidden_states * self.weight.to(hidden_states.dtype)
                hidden_states = hidden_states + (self.bias.to(hidden_states.dtype) if self.bias is not None else 0)
                hidden_states = hidden_states.to(input_dtype)
                return hidden_states

            hidden_states = hidden_states * self.weight
            if self.bias is not None:
                hidden_states = hidden_states + self.bias
        else:
            hidden_states = hidden_states.to(input_dtype)

        return hidden_states

sqrt(mean(x**2))

x_dtype = x.dtype
x = x.float()
rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
# return (x * rrms).to(dtype=x_dtype) * self.scale
return ((x * rrms) * self.scale.float()).to(dtype=x_dtype)

@dataclass
class ModulationOut:
    shift: Tensor
    scale: Tensor
    gate: Tensor


class Modulation(nn.Module):
    def __init__(self, dim: int, double: bool):
        super().__init__()
        self.is_double = double
        self.multiplier = 6 if double else 3
        self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)

    def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]:
        out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1)

        return (
            ModulationOut(*out[:3]),
            ModulationOut(*out[3:]) if self.is_double else None,
        )

def forward(~~) :
    img_mod1, img_mod2 = self.img_mod(vec)
    ...

    img_modulated = self.img_norm1(img)
    img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift

    ...
    # 각종 Layer, Attention 연산 거친 후

    img = img + img_mod1.gate * self.img_attn.proj(img_attn)
    img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)

def __init__(~~) :
    self.img_mod = nn.Sequential(
        nn.SiLU(),
        nn.Linear(dim, 6 * dim, bias=True),  # For scale, shift, gate for norm1 and norm2
    )
    ...

def _modulate(self, x, mod_params):
    """Apply modulation to input tensor"""
    shift, scale, gate = mod_params.chunk(3, dim=-1)
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1), gate.unsqueeze(1)

def forward(~~) :
    img_mod_params = self.img_mod(temb)  # [B, 6*dim]
    img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
    
    img_normed = self.img_norm1(hidden_states)
    img_modulated, img_gate1 = self._modulate(img_normed, img_mod1)
    
    ...
    
    hidden_states = hidden_states + img_gate1 * img_attn_output
    
    img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2)
    img_normed2 = self.img_norm2(hidden_states)
    img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2)
    del img_normed2, img_mod2
    img_mlp_output = self.img_mlp(img_modulated2)
    del img_modulated2
    hidden_states = hidden_states + img_gate2 * img_mlp_output