Skip to content

vllm_omni.diffusion.layers.adalayernorm

logger module-attribute

logger = init_logger(__name__)

AdaLayerNorm

Bases: CustomOp

AdaLayerNorm

out = layernorm(x) * (1 + scale) + shift

elementwise_affine instance-attribute

elementwise_affine = elementwise_affine

eps instance-attribute

eps = eps

hidden_size instance-attribute

hidden_size = hidden_size

layernorm instance-attribute

layernorm = LayerNorm(
    hidden_size,
    elementwise_affine=elementwise_affine,
    eps=eps,
)

forward_cuda

forward_cuda(
    x: Tensor, scale: Tensor, shift: Tensor
) -> Tensor

forward_hip

forward_hip(
    x: Tensor, scale: Tensor, shift: Tensor
) -> Tensor

forward_native

forward_native(
    x: Tensor, scale: Tensor, shift: Tensor
) -> Tensor

forward_npu

forward_npu(
    x: Tensor, scale: Tensor, shift: Tensor
) -> Tensor

forward_xpu

forward_xpu(
    x: Tensor, scale: Tensor, shift: Tensor
) -> Tensor

AdaLayerNormContinuous

Bases: Module

linear instance-attribute

linear = ReplicatedLinear(
    conditioning_embedding_dim,
    embedding_dim * 2,
    bias=bias,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.linear",
)

norm instance-attribute

norm = LayerNorm(
    embedding_dim,
    eps=eps,
    elementwise_affine=elementwise_affine,
)

silu instance-attribute

silu = SiLU()

forward

forward(
    x: Tensor, conditioning_embedding: Tensor
) -> Tensor

AdaLayerNormZero

Bases: Module

emb instance-attribute

emb = None

linear instance-attribute

linear = ReplicatedLinear(
    embedding_dim,
    6 * embedding_dim,
    bias=bias,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.linear",
)

norm instance-attribute

norm = LayerNorm(
    embedding_dim, elementwise_affine=False, eps=1e-06
)

silu instance-attribute

silu = SiLU()

forward

forward(
    x: Tensor, emb: Tensor
) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]

AdaLayerNormZeroSingle

Bases: Module

linear instance-attribute

linear = ReplicatedLinear(
    embedding_dim,
    3 * embedding_dim,
    bias=bias,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.linear",
)

norm instance-attribute

norm = LayerNorm(
    embedding_dim, elementwise_affine=False, eps=1e-06
)

silu instance-attribute

silu = SiLU()

forward

forward(x: Tensor, emb: Tensor) -> tuple[Tensor, Tensor]