|
import numpy as np |
|
import rembg |
|
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.bg_remover = rembg.new_session() |
|
|
|
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): |
|
unstacked = [] |
|
for i in range(4): |
|
image = rembg.remove(images[i], session=self.bg_remover) |
|
image = images.astype(np.float32) / 255.0 |
|
image = image[..., :3] * image[..., -1:] + (1 - image[..., -1:]) |
|
unstacked.append(image) |
|
images = np.concatenate( |
|
[ |
|
np.concatenate([unstacked[1], unstacked[2]], axis=1), |
|
np.concatenate([unstacked[3], unstacked[0]], axis=1), |
|
], |
|
axis=0, |
|
) |
|
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 |
|
|
