souging/TRELLIS_TextTo3D

1 2from typing import *3import torch4from . import _C5from . import pytorch6 7 8@torch.no_grad()9def encode(coords: torch.Tensor, permute: List[int] = [0, 1, 2], mode: Literal['z_order', 'hilbert'] = 'z_order') -> torch.Tensor:10    """11    Encodes 3D coordinates into a 30-bit code.12 13    Args:14        coords: a tensor of shape [N, 3] containing the 3D coordinates.15        permute: the permutation of the coordinates.16        mode: the encoding mode to use.17    """18    assert coords.shape[-1] == 3 and coords.ndim == 2, "Input coordinates must be of shape [N, 3]"19    x = coords[:, permute[0]].int()20    y = coords[:, permute[1]].int()21    z = coords[:, permute[2]].int()22    if mode == 'z_order':23        return _C.z_order_encode(x, y, z)24    elif mode == 'hilbert':25        return _C.hilbert_encode(x, y, z)26    else:27        raise ValueError(f"Unknown encoding mode: {mode}")28 29 30@torch.no_grad()31def decode(code: torch.Tensor, permute: List[int] = [0, 1, 2], mode: Literal['z_order', 'hilbert'] = 'z_order') -> torch.Tensor:32    """33    Decodes a 30-bit code into 3D coordinates.34 35    Args:36        code: a tensor of shape [N] containing the 30-bit code.37        permute: the permutation of the coordinates.38        mode: the decoding mode to use.39    """40    assert code.ndim == 1, "Input code must be of shape [N]"41    if mode == 'z_order':42        coords = _C.z_order_decode(code)43    elif mode == 'hilbert':44        coords = _C.hilbert_decode(code)45    else:46        raise ValueError(f"Unknown decoding mode: {mode}")47    x = coords[permute.index(0)]48    y = coords[permute.index(1)]49    z = coords[permute.index(2)]50    return torch.stack([x, y, z], dim=-1)51