Spaces:
Runtime error
Runtime error
| import nvdiffrast.torch as dr | |
| import torch | |
| from .typing import * | |
| class NVDiffRasterizerContext: | |
| def __init__(self, context_type: str, device: torch.device) -> None: | |
| self.device = device | |
| self.ctx = self.initialize_context(context_type, device) | |
| def initialize_context( | |
| self, context_type: str, device: torch.device | |
| ) -> Union[dr.RasterizeGLContext, dr.RasterizeCudaContext]: | |
| context_type = "cuda" | |
| if context_type == "gl": | |
| return dr.RasterizeGLContext(device=device) | |
| elif context_type == "cuda": | |
| return dr.RasterizeCudaContext(device=device) | |
| else: | |
| raise ValueError(f"Unknown rasterizer context type: {context_type}") | |
| def vertex_transform( | |
| self, verts: Float[Tensor, "Nv 3"], mvp_mtx: Float[Tensor, "B 4 4"] | |
| ) -> Float[Tensor, "B Nv 4"]: | |
| with torch.cuda.amp.autocast(enabled=False): | |
| verts_homo = torch.cat( | |
| [verts, torch.ones([verts.shape[0], 1]).to(verts)], dim=-1 | |
| ) | |
| verts_clip = torch.matmul(verts_homo, mvp_mtx.permute(0, 2, 1)) | |
| return verts_clip | |
| def rasterize( | |
| self, | |
| pos: Float[Tensor, "B Nv 4"], | |
| tri: Integer[Tensor, "Nf 3"], | |
| resolution: Union[int, Tuple[int, int]], | |
| ): | |
| # rasterize in instance mode (single topology) | |
| return dr.rasterize(self.ctx, pos.float(), tri.int(), resolution, grad_db=True) | |
| def rasterize_one( | |
| self, | |
| pos: Float[Tensor, "Nv 4"], | |
| tri: Integer[Tensor, "Nf 3"], | |
| resolution: Union[int, Tuple[int, int]], | |
| ): | |
| # rasterize one single mesh under a single viewpoint | |
| rast, rast_db = self.rasterize(pos[None, ...], tri, resolution) | |
| return rast[0], rast_db[0] | |
| def antialias( | |
| self, | |
| color: Float[Tensor, "B H W C"], | |
| rast: Float[Tensor, "B H W 4"], | |
| pos: Float[Tensor, "B Nv 4"], | |
| tri: Integer[Tensor, "Nf 3"], | |
| ) -> Float[Tensor, "B H W C"]: | |
| return dr.antialias(color.float(), rast, pos.float(), tri.int()) | |
| def interpolate( | |
| self, | |
| attr: Float[Tensor, "B Nv C"], | |
| rast: Float[Tensor, "B H W 4"], | |
| tri: Integer[Tensor, "Nf 3"], | |
| rast_db=None, | |
| diff_attrs=None, | |
| ) -> Float[Tensor, "B H W C"]: | |
| return dr.interpolate( | |
| attr.float(), rast, tri.int(), rast_db=rast_db, diff_attrs=diff_attrs | |
| ) | |
| def interpolate_one( | |
| self, | |
| attr: Float[Tensor, "Nv C"], | |
| rast: Float[Tensor, "B H W 4"], | |
| tri: Integer[Tensor, "Nf 3"], | |
| rast_db=None, | |
| diff_attrs=None, | |
| ) -> Float[Tensor, "B H W C"]: | |
| return self.interpolate(attr[None, ...], rast, tri, rast_db, diff_attrs) | |