| | import numpy as np |
| | import torch |
| | import torch.nn.functional as F |
| | import torchvision.transforms.functional as TF |
| | from diffusers import DiffusionPipeline |
| |
|
| |
|
| | class LGMPipeline(DiffusionPipeline): |
| | def __init__(self, lgm): |
| | super().__init__() |
| |
|
| | self.imagenet_default_mean = (0.485, 0.456, 0.406) |
| | self.imagenet_default_std = (0.229, 0.224, 0.225) |
| |
|
| | lgm = lgm.half().cuda() |
| | self.register_modules(lgm=lgm) |
| |
|
| | def save_ply(self, gaussians, path): |
| | self.lgm.gs.save_ply(gaussians, path) |
| |
|
| | @torch.no_grad() |
| | def __call__(self, images): |
| | images = np.stack([images[1], images[2], images[3], images[0]], axis=0) |
| | images = torch.from_numpy(images).permute(0, 3, 1, 2).float().cuda() |
| | images = F.interpolate( |
| | images, |
| | size=(256, 256), |
| | mode="bilinear", |
| | align_corners=False, |
| | ) |
| | images = TF.normalize( |
| | images, self.imagenet_default_mean, self.imagenet_default_std |
| | ) |
| |
|
| | rays_embeddings = self.lgm.prepare_default_rays("cuda", elevation=0) |
| | images = torch.cat([images, rays_embeddings], dim=1).unsqueeze(0) |
| | images = images.half().cuda() |
| |
|
| | result = self.lgm(images) |
| | return result |
| |
|
