Function get_normalization

src/diffusers/models/normalization.py:632–647

(
    norm_type: str = "batch_norm",
    num_features: int | None = None,
    eps: float = 1e-5,
    elementwise_affine: bool = True,
    bias: bool = True,
)

def get_normalization(
    norm_type: str = "batch_norm",
    num_features: int | None = None,
    eps: float = 1e-5,
    elementwise_affine: bool = True,
    bias: bool = True,
) -> nn.Module:
    if norm_type == "rms_norm":
        norm = RMSNorm(num_features, eps=eps, elementwise_affine=elementwise_affine, bias=bias)
    elif norm_type == "layer_norm":
        norm = nn.LayerNorm(num_features, eps=eps, elementwise_affine=elementwise_affine, bias=bias)
    elif norm_type == "batch_norm":
        norm = nn.BatchNorm2d(num_features, eps=eps, affine=elementwise_affine)
    else:
        raise ValueError(f"{norm_type=} is not supported.")
    return norm