souging/TRELLIS_TextTo3D

1import torch.nn as nn2from ..modules import sparse as sp3 4FP16_MODULES = (5    nn.Conv1d,6    nn.Conv2d,7    nn.Conv3d,8    nn.ConvTranspose1d,9    nn.ConvTranspose2d,10    nn.ConvTranspose3d,11    nn.Linear,12    sp.SparseConv3d,13    sp.SparseInverseConv3d,14    sp.SparseLinear,15)16 17def convert_module_to_f16(l):18    """19    Convert primitive modules to float16.20    """21    if isinstance(l, FP16_MODULES):22        for p in l.parameters():23            p.data = p.data.half()24 25 26def convert_module_to_f32(l):27    """28    Convert primitive modules to float32, undoing convert_module_to_f16().29    """30    if isinstance(l, FP16_MODULES):31        for p in l.parameters():32            p.data = p.data.float()33 34 35def zero_module(module):36    """37    Zero out the parameters of a module and return it.38    """39    for p in module.parameters():40        p.detach().zero_()41    return module42 43 44def scale_module(module, scale):45    """46    Scale the parameters of a module and return it.47    """48    for p in module.parameters():49        p.detach().mul_(scale)50    return module51 52 53def modulate(x, shift, scale):54    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)55