52 lines
1.4 KiB
Python
52 lines
1.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Small platform-agnostic helpers for UNet-style blocks (adapted from diffusers).
|
|
|
|
``conv_nd`` / ``normalization`` / ``linear`` / ``zero_module`` are used by the
|
|
``ResBlock`` / ``UNetDown`` / ``UNetUp`` implementations across HunyuanImage3's
|
|
platform-specific layer packages. They are plain ``nn.ConvNd`` / ``nn.GroupNorm``
|
|
/ ``nn.Linear`` wrappers with no hardware-specific behaviour, so they live here
|
|
instead of being duplicated per backend.
|
|
"""
|
|
|
|
from torch import nn
|
|
|
|
|
|
def conv_nd(dims, *args, **kwargs): # noqa: N802
|
|
"""
|
|
Create a 1D, 2D, or 3D convolution module.
|
|
"""
|
|
if dims == 1:
|
|
return nn.Conv1d(*args, **kwargs)
|
|
elif dims == 2:
|
|
return nn.Conv2d(*args, **kwargs)
|
|
elif dims == 3:
|
|
return nn.Conv3d(*args, **kwargs)
|
|
raise ValueError(f"unsupported dimensions: {dims}")
|
|
|
|
|
|
def normalization(channels, **kwargs):
|
|
"""
|
|
Make a standard normalization layer.
|
|
:param channels: number of input channels.
|
|
:return: a nn.Module for normalization.
|
|
"""
|
|
return nn.GroupNorm(32, channels, **kwargs)
|
|
|
|
|
|
def linear(*args, **kwargs):
|
|
"""
|
|
Create a linear module.
|
|
"""
|
|
return nn.Linear(*args, **kwargs)
|
|
|
|
|
|
def zero_module(module):
|
|
"""
|
|
Zero out the parameters of a module and return it.
|
|
"""
|
|
for p in module.parameters():
|
|
p.detach().zero_()
|
|
return module
|