88 lines
2.7 KiB
Python
88 lines
2.7 KiB
Python
import os
|
|
|
|
import slangtorch
|
|
import torch
|
|
import torch.nn as nn
|
|
from jaxtyping import Bool, Float
|
|
from torch import Tensor
|
|
|
|
|
|
class TextureBaker(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.baker = slangtorch.loadModule(
|
|
os.path.join(os.path.dirname(__file__), "texture_baker.slang")
|
|
)
|
|
|
|
def rasterize(
|
|
self,
|
|
uv: Float[Tensor, "Nv 2"],
|
|
face_indices: Float[Tensor, "Nf 3"],
|
|
bake_resolution: int,
|
|
) -> Float[Tensor, "bake_resolution bake_resolution 4"]:
|
|
if not face_indices.is_cuda or not uv.is_cuda:
|
|
raise ValueError("All input tensors must be on cuda")
|
|
|
|
face_indices = face_indices.to(torch.int32)
|
|
uv = uv.to(torch.float32)
|
|
|
|
rast_result = torch.empty(
|
|
bake_resolution, bake_resolution, 4, device=uv.device, dtype=torch.float32
|
|
)
|
|
|
|
block_size = 16
|
|
grid_size = bake_resolution // block_size
|
|
self.baker.bake_uv(uv=uv, indices=face_indices, output=rast_result).launchRaw(
|
|
blockSize=(block_size, block_size, 1), gridSize=(grid_size, grid_size, 1)
|
|
)
|
|
|
|
return rast_result
|
|
|
|
def get_mask(
|
|
self, rast: Float[Tensor, "bake_resolution bake_resolution 4"]
|
|
) -> Bool[Tensor, "bake_resolution bake_resolution"]:
|
|
return rast[..., -1] >= 0
|
|
|
|
def interpolate(
|
|
self,
|
|
attr: Float[Tensor, "Nv 3"],
|
|
rast: Float[Tensor, "bake_resolution bake_resolution 4"],
|
|
face_indices: Float[Tensor, "Nf 3"],
|
|
uv: Float[Tensor, "Nv 2"],
|
|
) -> Float[Tensor, "bake_resolution bake_resolution 3"]:
|
|
# Make sure all input tensors are on torch
|
|
if not attr.is_cuda or not face_indices.is_cuda or not rast.is_cuda:
|
|
raise ValueError("All input tensors must be on cuda")
|
|
|
|
attr = attr.to(torch.float32)
|
|
face_indices = face_indices.to(torch.int32)
|
|
uv = uv.to(torch.float32)
|
|
|
|
pos_bake = torch.zeros(
|
|
rast.shape[0],
|
|
rast.shape[1],
|
|
3,
|
|
device=attr.device,
|
|
dtype=attr.dtype,
|
|
)
|
|
|
|
block_size = 16
|
|
grid_size = rast.shape[0] // block_size
|
|
self.baker.interpolate(
|
|
attr=attr, indices=face_indices, rast=rast, output=pos_bake
|
|
).launchRaw(
|
|
blockSize=(block_size, block_size, 1), gridSize=(grid_size, grid_size, 1)
|
|
)
|
|
|
|
return pos_bake
|
|
|
|
def forward(
|
|
self,
|
|
attr: Float[Tensor, "Nv 3"],
|
|
uv: Float[Tensor, "Nv 2"],
|
|
face_indices: Float[Tensor, "Nf 3"],
|
|
bake_resolution: int,
|
|
) -> Float[Tensor, "bake_resolution bake_resolution 3"]:
|
|
rast = self.rasterize(uv, face_indices, bake_resolution)
|
|
return self.interpolate(attr, rast, face_indices, uv)
|