| 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 |
|
|
