souging/TRELLIS_TextTo3D

1import torch2import vox2seq3 4 5if __name__ == "__main__":6    RES = 2567    coords = torch.meshgrid(torch.arange(RES), torch.arange(RES), torch.arange(RES))8    coords = torch.stack(coords, dim=-1).reshape(-1, 3).int().cuda()9    code_z_cuda = vox2seq.encode(coords, mode='z_order')10    code_z_pytorch = vox2seq.pytorch.encode(coords, mode='z_order')11    code_h_cuda = vox2seq.encode(coords, mode='hilbert')12    code_h_pytorch = vox2seq.pytorch.encode(coords, mode='hilbert')13    assert torch.equal(code_z_cuda, code_z_pytorch)14    assert torch.equal(code_h_cuda, code_h_pytorch)15 16    code = torch.arange(RES**3).int().cuda()17    coords_z_cuda = vox2seq.decode(code, mode='z_order')18    coords_z_pytorch = vox2seq.pytorch.decode(code, mode='z_order')19    coords_h_cuda = vox2seq.decode(code, mode='hilbert')20    coords_h_pytorch = vox2seq.pytorch.decode(code, mode='hilbert')21    assert torch.equal(coords_z_cuda, coords_z_pytorch)22    assert torch.equal(coords_h_cuda, coords_h_pytorch)23 24    print("All tests passed.")25 26