souging/TRELLIS_TextTo3D

1import torch2 3 4def pixel_shuffle_3d(x: torch.Tensor, scale_factor: int) -> torch.Tensor:5    """6    3D pixel shuffle.7    """8    B, C, H, W, D = x.shape9    C_ = C // scale_factor**310    x = x.reshape(B, C_, scale_factor, scale_factor, scale_factor, H, W, D)11    x = x.permute(0, 1, 5, 2, 6, 3, 7, 4)12    x = x.reshape(B, C_, H*scale_factor, W*scale_factor, D*scale_factor)13    return x14 15 16def patchify(x: torch.Tensor, patch_size: int):17    """18    Patchify a tensor.19 20    Args:21        x (torch.Tensor): (N, C, *spatial) tensor22        patch_size (int): Patch size23    """24    DIM = x.dim() - 225    for d in range(2, DIM + 2):26        assert x.shape[d] % patch_size == 0, f"Dimension {d} of input tensor must be divisible by patch size, got {x.shape[d]} and {patch_size}"27 28    x = x.reshape(*x.shape[:2], *sum([[x.shape[d] // patch_size, patch_size] for d in range(2, DIM + 2)], []))29    x = x.permute(0, 1, *([2 * i + 3 for i in range(DIM)] + [2 * i + 2 for i in range(DIM)]))30    x = x.reshape(x.shape[0], x.shape[1] * (patch_size ** DIM), *(x.shape[-DIM:]))31    return x32 33 34def unpatchify(x: torch.Tensor, patch_size: int):35    """36    Unpatchify a tensor.37 38    Args:39        x (torch.Tensor): (N, C, *spatial) tensor40        patch_size (int): Patch size41    """42    DIM = x.dim() - 243    assert x.shape[1] % (patch_size ** DIM) == 0, f"Second dimension of input tensor must be divisible by patch size to unpatchify, got {x.shape[1]} and {patch_size ** DIM}"44 45    x = x.reshape(x.shape[0], x.shape[1] // (patch_size ** DIM), *([patch_size] * DIM), *(x.shape[-DIM:]))46    x = x.permute(0, 1, *(sum([[2 + DIM + i, 2 + i] for i in range(DIM)], [])))47    x = x.reshape(x.shape[0], x.shape[1], *[x.shape[2 + 2 * i] * patch_size for i in range(DIM)])48    return x49