Integrated Wonder3D & NeuS, Add vertex color support for mesh render
This commit is contained in:
MrForExample
@@ -0,0 +1,306 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import cv2
|
||||
import numpy as np
|
||||
import os
|
||||
from glob import glob
|
||||
from scipy.spatial.transform import Rotation as Rot
|
||||
from scipy.spatial.transform import Slerp
|
||||
from glob import glob
|
||||
|
||||
from shared_utils.image_utils import prepare_torch_img
|
||||
|
||||
def camNormal2worldNormal(rot_c2w, camNormal):
|
||||
H,W,_ = camNormal.shape
|
||||
normal_img = np.matmul(rot_c2w[None, :, :], camNormal.reshape(-1,3)[:, :, None]).reshape([H, W, 3])
|
||||
|
||||
return normal_img
|
||||
|
||||
def worldNormal2camNormal(rot_w2c, worldNormal):
|
||||
H,W,_ = worldNormal.shape
|
||||
normal_img = np.matmul(rot_w2c[None, :, :], worldNormal.reshape(-1,3)[:, :, None]).reshape([H, W, 3])
|
||||
|
||||
return normal_img
|
||||
|
||||
def trans_normal(normal, RT_w2c, RT_w2c_target):
|
||||
|
||||
normal_world = camNormal2worldNormal(np.linalg.inv(RT_w2c[:3,:3]), normal)
|
||||
normal_target_cam = worldNormal2camNormal(RT_w2c_target[:3,:3], normal_world)
|
||||
|
||||
return normal_target_cam
|
||||
|
||||
def img2normal(img):
|
||||
return (img/255.)*2-1
|
||||
|
||||
def normal2img(normal):
|
||||
return np.uint8((normal*0.5+0.5)*255)
|
||||
|
||||
def norm_normalize(normal, dim=-1):
|
||||
|
||||
normal = normal/(np.linalg.norm(normal, axis=dim, keepdims=True)+1e-6)
|
||||
|
||||
return normal
|
||||
|
||||
def RT_opengl2opencv(RT):
|
||||
# Build the coordinate transform matrix from world to computer vision camera
|
||||
# R_world2cv = R_bcam2cv@R_world2bcam
|
||||
# T_world2cv = R_bcam2cv@T_world2bcam
|
||||
|
||||
R = RT[:3, :3]
|
||||
t = RT[:3, 3]
|
||||
|
||||
R_bcam2cv = np.asarray([[1, 0, 0], [0, -1, 0], [0, 0, -1]], np.float32)
|
||||
|
||||
R_world2cv = R_bcam2cv @ R
|
||||
t_world2cv = R_bcam2cv @ t
|
||||
|
||||
RT = np.concatenate([R_world2cv,t_world2cv[:,None]],1)
|
||||
|
||||
return RT
|
||||
|
||||
def normal_opengl2opencv(normal):
|
||||
H,W,C = np.shape(normal)
|
||||
# normal_img = np.reshape(normal, (H*W,C))
|
||||
R_bcam2cv = np.array([1, -1, -1], np.float32)
|
||||
normal_cv = normal * R_bcam2cv[None, None, :]
|
||||
|
||||
return normal_cv
|
||||
|
||||
def inv_RT(RT):
|
||||
RT_h = np.concatenate([RT, np.array([[0,0,0,1]])], axis=0)
|
||||
RT_inv = np.linalg.inv(RT_h)
|
||||
|
||||
return RT_inv[:3, :]
|
||||
|
||||
def load_pose_and_preprocess(mv_images, mv_masks, mv_normals, imSize, view_types, cam_pose_dir=None, normal_system='front'):
|
||||
|
||||
all_images = []
|
||||
all_masks = []
|
||||
all_normals = []
|
||||
all_normals_world = []
|
||||
all_poses = []
|
||||
all_w2cs = []
|
||||
|
||||
RT_front = np.loadtxt(glob(os.path.join(cam_pose_dir, '*_%s_RT.txt'%( 'front')))[0]) # world2cam matrix
|
||||
RT_front_cv = RT_opengl2opencv(RT_front) # convert normal from opengl to opencv
|
||||
img_H, img_W = imSize[0], imSize[1]
|
||||
|
||||
# rgb to bgr
|
||||
mv_images = mv_images[:, :, :, [2, 1, 0]].contiguous()
|
||||
|
||||
for idx, view in enumerate(view_types):
|
||||
|
||||
image = prepare_torch_img(mv_images[idx].unsqueeze(0), img_H, img_W, keep_shape=True).squeeze(0).detach().cpu().numpy()
|
||||
mask = prepare_torch_img(mv_masks[idx].unsqueeze(2).unsqueeze(0), img_H, img_W, keep_shape=True).squeeze(0).squeeze(2).detach().cpu().numpy()
|
||||
normal = prepare_torch_img(mv_normals[idx].unsqueeze(0), img_H, img_W, keep_shape=True).detach().squeeze(0).cpu().numpy()
|
||||
|
||||
RT = np.loadtxt(os.path.join(cam_pose_dir, '000_%s_RT.txt'%( view))) # world2cam matrix
|
||||
|
||||
# [0, 1] -> [-1, 1]
|
||||
normal = normal * 2 - 1
|
||||
normal[mask==0] = [0,0,0]
|
||||
|
||||
RT_cv = RT_opengl2opencv(RT) # convert normal from opengl to opencv
|
||||
all_poses.append(inv_RT(RT_cv)) # cam2world
|
||||
all_w2cs.append(RT_cv)
|
||||
|
||||
# whether to
|
||||
normal_cam_cv = normal_opengl2opencv(normal)
|
||||
|
||||
if normal_system == 'front':
|
||||
normal_world = camNormal2worldNormal(inv_RT(RT_front_cv)[:3, :3], normal_cam_cv)
|
||||
elif normal_system == 'self':
|
||||
normal_world = camNormal2worldNormal(inv_RT(RT_cv)[:3, :3], normal_cam_cv)
|
||||
|
||||
all_images.append(image)
|
||||
all_masks.append(mask)
|
||||
all_normals.append(normal_cam_cv)
|
||||
all_normals_world.append(normal_world)
|
||||
|
||||
return np.stack(all_images), np.stack(all_masks), np.stack(all_normals), np.stack(all_normals_world), np.stack(all_poses), np.stack(all_w2cs)
|
||||
|
||||
|
||||
class Dataset:
|
||||
def __init__(self, mv_images, mv_masks, mv_normals, conf, cam_pose_dir, num_views=6):
|
||||
super(Dataset, self).__init__()
|
||||
self.device = torch.device('cuda')
|
||||
self.conf = conf
|
||||
|
||||
self.object_viewidx = conf.get_int('object_viewidx')
|
||||
self.imSize = conf['imSize']
|
||||
self.stage = conf['stage']
|
||||
self.mtype = conf['mtype']
|
||||
|
||||
if num_views == 4:
|
||||
view_types = ['front', 'right', 'back', 'left']
|
||||
elif num_views == 5:
|
||||
view_types = ['front', 'front_right', 'right', 'back', 'left']
|
||||
elif num_views == 6:
|
||||
view_types = ['front', 'front_right', 'right', 'back', 'left', 'front_left']
|
||||
|
||||
self.images_np, self.masks_np, self.normals_cam_np, self.normals_world_np ,self.pose_all_np, self.w2c_all_np = load_pose_and_preprocess(
|
||||
mv_images, mv_masks, mv_normals, self.imSize, view_types, cam_pose_dir
|
||||
)
|
||||
|
||||
self.n_images = num_views
|
||||
|
||||
self.images = torch.from_numpy(self.images_np.astype(np.float32)).cpu() # [n_images, H, W, 3]
|
||||
self.masks = torch.from_numpy(self.masks_np.astype(np.float32)).cpu() # [n_images, H, W]
|
||||
self.normals_cam = torch.from_numpy(self.normals_cam_np.astype(np.float32)).cpu() # [n_images, H, W, 3]
|
||||
self.normals_world = torch.from_numpy(self.normals_world_np.astype(np.float32)).cpu() # [n_images, H, W, 3]
|
||||
self.pose_all = torch.from_numpy(self.pose_all_np.astype(np.float32)).cpu() # [n_images,3, 4] cam2world
|
||||
# self.pose_all = torch.stack(self.pose_all).to(self.device) # [n_images, 4, 4]
|
||||
self.H, self.W = self.images.shape[1], self.images.shape[2]
|
||||
self.image_pixels = self.H * self.W
|
||||
|
||||
self.intrinsic = torch.from_numpy(np.array([
|
||||
[self.W/2.0, 0, self.W / 2.0, 0],
|
||||
[0, self.H/2.0, self.H/ 2.0, 0],
|
||||
[0, 0, 1, 0],
|
||||
[0, 0, 0, 1]
|
||||
]).astype(np.float32))
|
||||
|
||||
self.intrinsics_all = torch.stack([self.intrinsic]*num_views, dim=0).cpu()
|
||||
self.intrinsics_all_inv = torch.inverse(self.intrinsics_all).cpu() # [n_images, 4, 4]
|
||||
|
||||
object_bbox_min = np.array([-1.01, -1.01, -1.01, 1.0])
|
||||
object_bbox_max = np.array([ 1.01, 1.01, 1.01, 1.0])
|
||||
|
||||
self.object_bbox_min = object_bbox_min[:3]
|
||||
self.object_bbox_max = object_bbox_max[:3]
|
||||
|
||||
self.near = 0.2
|
||||
self.far = 2.4
|
||||
|
||||
self.cos = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
|
||||
self.all_rays = self.prepare_all_rays()
|
||||
|
||||
def gen_rays_at(self, img_idx, resolution_level=1):
|
||||
"""
|
||||
Generate rays at world space from one camera.
|
||||
"""
|
||||
l = resolution_level
|
||||
tx = torch.linspace(0, self.W - 1, self.W // l)
|
||||
ty = torch.linspace(0, self.H - 1, self.H // l)
|
||||
pixels_x, pixels_y = torch.meshgrid(tx, ty)
|
||||
q = torch.stack([(pixels_x/self.W-0.5)*2, (pixels_y/self.H-0.5)*2, torch.zeros_like(pixels_y)], dim=-1) # W, H, 3
|
||||
|
||||
v = torch.stack([torch.zeros_like(pixels_y), torch.zeros_like(pixels_y), torch.ones_like(pixels_y)], dim=-1) # W, H, 3
|
||||
|
||||
# orthogonal projection
|
||||
rays_v = v / torch.linalg.norm(v, ord=2, dim=-1, keepdim=True) # W, H, 3
|
||||
rays_v = torch.matmul(self.pose_all[img_idx, None, None, :3, :3].cuda(), rays_v[:, :, :, None].cuda()).squeeze() # W, H, 3
|
||||
|
||||
rays_o = torch.matmul(self.pose_all[img_idx, None, None, :3, :3].cuda(), q[:, :, :, None].cuda()).squeeze() # W, H, 3
|
||||
rays_o = self.pose_all[img_idx, None, None, :3, 3].expand(rays_v.shape).cuda() + rays_o # W, H, 3
|
||||
return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
|
||||
|
||||
def get_batch_of_random_rays(self, batch_size):
|
||||
"""
|
||||
Get batch of random rays at world space.
|
||||
"""
|
||||
batch_index = torch.randint(low=0, high=self.__len__(), size=[batch_size])
|
||||
return self.all_rays[batch_index, :] # batch_size, 14
|
||||
|
||||
def prepare_rays_a_view(self, img_idx):
|
||||
"""
|
||||
Generate random rays at world space from one camera.
|
||||
"""
|
||||
tx = torch.linspace(0, self.W - 1, self.W)
|
||||
ty = torch.linspace(0, self.H - 1, self.H)
|
||||
pixels_x, pixels_y = torch.meshgrid(tx, ty)
|
||||
|
||||
pixels_x = pixels_x.reshape(-1).long()
|
||||
pixels_y = pixels_y.reshape(-1).long()
|
||||
color = self.images[img_idx][(pixels_y, pixels_x)] # batch_size, 3
|
||||
mask = self.masks[img_idx][(pixels_y, pixels_x)] # batch_size, 3
|
||||
normal = self.normals_world[img_idx][(pixels_y, pixels_x)] # batch_size, 3
|
||||
|
||||
q = torch.stack([(pixels_x / self.W-0.5)*2, (pixels_y / self.H-0.5)*2, torch.zeros_like(pixels_y)], dim=-1).float() # batch_size, 3
|
||||
v = torch.stack([torch.zeros_like(pixels_y), torch.zeros_like(pixels_y), torch.ones_like(pixels_y)], dim=-1).float()
|
||||
|
||||
rays_v = v / torch.linalg.norm(v, ord=2, dim=-1, keepdim=True) # batch_size, 3
|
||||
rays_v = torch.matmul(self.pose_all[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze() # batch_size, 3
|
||||
|
||||
rays_o = torch.matmul(self.pose_all[img_idx, None, :3, :3], q[:, :, None]).squeeze() # batch_size, 3
|
||||
rays_o = self.pose_all[img_idx, None, :3, 3].expand(rays_v.shape) + rays_o # batch_size, 3
|
||||
|
||||
cosines = self.cos(rays_v, normal)
|
||||
|
||||
# pdb.set_trace()
|
||||
|
||||
return torch.cat([rays_o.cpu(), rays_v.cpu(), color, mask[:, None], normal, cosines[:, None]], dim=-1) # batch_size, 10
|
||||
|
||||
|
||||
def prepare_all_rays(self,):
|
||||
all_rays = []
|
||||
for idx in range(self.n_images):
|
||||
rays = self.prepare_rays_a_view(idx)
|
||||
all_rays.append(rays)
|
||||
all_rays = torch.concat(all_rays, dim=0)
|
||||
return all_rays
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return self.all_rays[idx]
|
||||
|
||||
def __len__(self):
|
||||
return self.all_rays.shape[0]
|
||||
|
||||
def gen_rays_between(self, idx_0, idx_1, ratio, resolution_level=1):
|
||||
"""
|
||||
Interpolate pose between two cameras.
|
||||
"""
|
||||
l = resolution_level
|
||||
tx = torch.linspace(0, self.W - 1, self.W // l)
|
||||
ty = torch.linspace(0, self.H - 1, self.H // l)
|
||||
pixels_x, pixels_y = torch.meshgrid(tx, ty)
|
||||
p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1) # W, H, 3
|
||||
p = torch.matmul(self.intrinsics_all_inv[0, None, None, :3, :3], p[:, :, :, None]).squeeze() # W, H, 3
|
||||
rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True) # W, H, 3
|
||||
trans = self.pose_all[idx_0, :3, 3] * (1.0 - ratio) + self.pose_all[idx_1, :3, 3] * ratio
|
||||
pose_0 = self.pose_all[idx_0].detach().cpu().numpy()
|
||||
pose_1 = self.pose_all[idx_1].detach().cpu().numpy()
|
||||
pose_0 = np.linalg.inv(pose_0)
|
||||
pose_1 = np.linalg.inv(pose_1)
|
||||
rot_0 = pose_0[:3, :3]
|
||||
rot_1 = pose_1[:3, :3]
|
||||
rots = Rot.from_matrix(np.stack([rot_0, rot_1]))
|
||||
key_times = [0, 1]
|
||||
slerp = Slerp(key_times, rots)
|
||||
rot = slerp(ratio)
|
||||
pose = np.diag([1.0, 1.0, 1.0, 1.0])
|
||||
pose = pose.astype(np.float32)
|
||||
pose[:3, :3] = rot.as_matrix()
|
||||
pose[:3, 3] = ((1.0 - ratio) * pose_0 + ratio * pose_1)[:3, 3]
|
||||
pose = np.linalg.inv(pose)
|
||||
rot = torch.from_numpy(pose[:3, :3]).cuda()
|
||||
trans = torch.from_numpy(pose[:3, 3]).cuda()
|
||||
rays_v = torch.matmul(rot[None, None, :3, :3], rays_v[:, :, :, None]).squeeze() # W, H, 3
|
||||
rays_o = trans[None, None, :3].expand(rays_v.shape) # W, H, 3
|
||||
return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
|
||||
|
||||
def near_far_from_sphere(self, rays_o, rays_d):
|
||||
a = torch.sum(rays_d**2, dim=-1, keepdim=True)
|
||||
b = 2.0 * torch.sum(rays_o * rays_d, dim=-1, keepdim=True)
|
||||
mid = 0.5 * (-b) / a
|
||||
near = mid - 1.0
|
||||
far = mid + 1.0
|
||||
return near, far
|
||||
|
||||
def get_near_far(self,):
|
||||
return self.near, self.far
|
||||
|
||||
def image_at(self, idx, resolution_level):
|
||||
img = self.images_np[idx]
|
||||
return (cv2.resize(img, (self.W // resolution_level, self.H // resolution_level))).clip(0, 255)
|
||||
|
||||
def normal_cam_at(self, idx, resolution_level):
|
||||
normal_cam = self.normals_cam_np[idx]
|
||||
img = normal2img(normal_cam)
|
||||
return (cv2.resize(img, (self.W // resolution_level, self.H // resolution_level))).clip(0, 255)
|
||||
|
||||
def mask_at(self, idx, resolution_level):
|
||||
mask = np.uint8(self.masks_np[idx]*255)[:, :, None]
|
||||
mask = np.concatenate([mask]*3, axis=-1)
|
||||
return (cv2.resize(mask, (self.W // resolution_level, self.H // resolution_level))).clip(0, 255)
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
# Positional encoding embedding. Code was taken from https://github.com/bmild/nerf.
|
||||
class Embedder:
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
self.create_embedding_fn()
|
||||
|
||||
def create_embedding_fn(self):
|
||||
embed_fns = []
|
||||
d = self.kwargs['input_dims']
|
||||
out_dim = 0
|
||||
if self.kwargs['include_input']:
|
||||
embed_fns.append(lambda x: x)
|
||||
out_dim += d
|
||||
|
||||
max_freq = self.kwargs['max_freq_log2']
|
||||
N_freqs = self.kwargs['num_freqs']
|
||||
|
||||
if self.kwargs['log_sampling']:
|
||||
freq_bands = 2. ** torch.linspace(0., max_freq, N_freqs)
|
||||
else:
|
||||
freq_bands = torch.linspace(2.**0., 2.**max_freq, N_freqs)
|
||||
|
||||
for freq in freq_bands:
|
||||
for p_fn in self.kwargs['periodic_fns']:
|
||||
embed_fns.append(lambda x, p_fn=p_fn, freq=freq: p_fn(x * freq))
|
||||
out_dim += d
|
||||
|
||||
self.embed_fns = embed_fns
|
||||
self.out_dim = out_dim
|
||||
|
||||
def embed(self, inputs):
|
||||
return torch.cat([fn(inputs) for fn in self.embed_fns], -1)
|
||||
|
||||
|
||||
def get_embedder(multires, input_dims=3):
|
||||
embed_kwargs = {
|
||||
'include_input': True,
|
||||
'input_dims': input_dims,
|
||||
'max_freq_log2': multires-1,
|
||||
'num_freqs': multires,
|
||||
'log_sampling': True,
|
||||
'periodic_fns': [torch.sin, torch.cos],
|
||||
}
|
||||
|
||||
embedder_obj = Embedder(**embed_kwargs)
|
||||
def embed(x, eo=embedder_obj): return eo.embed(x)
|
||||
return embed, embedder_obj.out_dim
|
||||
@@ -0,0 +1,264 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from .embedder import get_embedder
|
||||
|
||||
|
||||
# This implementation is borrowed from IDR: https://github.com/lioryariv/idr
|
||||
class SDFNetwork(nn.Module):
|
||||
def __init__(self,
|
||||
d_in,
|
||||
d_out,
|
||||
d_hidden,
|
||||
n_layers,
|
||||
skip_in=(4,),
|
||||
multires=0,
|
||||
bias=0.5,
|
||||
scale=1,
|
||||
geometric_init=True,
|
||||
weight_norm=True,
|
||||
inside_outside=False):
|
||||
super(SDFNetwork, self).__init__()
|
||||
|
||||
dims = [d_in] + [d_hidden for _ in range(n_layers)] + [d_out]
|
||||
|
||||
self.embed_fn_fine = None
|
||||
|
||||
if multires > 0:
|
||||
embed_fn, input_ch = get_embedder(multires, input_dims=d_in)
|
||||
self.embed_fn_fine = embed_fn
|
||||
dims[0] = input_ch
|
||||
|
||||
self.num_layers = len(dims)
|
||||
self.skip_in = skip_in
|
||||
self.scale = scale
|
||||
|
||||
for l in range(0, self.num_layers - 1):
|
||||
if l + 1 in self.skip_in:
|
||||
out_dim = dims[l + 1] - dims[0]
|
||||
else:
|
||||
out_dim = dims[l + 1]
|
||||
|
||||
lin = nn.Linear(dims[l], out_dim)
|
||||
|
||||
if geometric_init:
|
||||
if l == self.num_layers - 2:
|
||||
if not inside_outside:
|
||||
torch.nn.init.normal_(lin.weight, mean=np.sqrt(np.pi) / np.sqrt(dims[l]), std=0.0001)
|
||||
torch.nn.init.constant_(lin.bias, -bias)
|
||||
else:
|
||||
torch.nn.init.normal_(lin.weight, mean=-np.sqrt(np.pi) / np.sqrt(dims[l]), std=0.0001)
|
||||
torch.nn.init.constant_(lin.bias, bias)
|
||||
elif multires > 0 and l == 0:
|
||||
torch.nn.init.constant_(lin.bias, 0.0)
|
||||
torch.nn.init.constant_(lin.weight[:, 3:], 0.0)
|
||||
torch.nn.init.normal_(lin.weight[:, :3], 0.0, np.sqrt(2) / np.sqrt(out_dim))
|
||||
elif multires > 0 and l in self.skip_in:
|
||||
torch.nn.init.constant_(lin.bias, 0.0)
|
||||
torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim))
|
||||
torch.nn.init.constant_(lin.weight[:, -(dims[0] - 3):], 0.0)
|
||||
else:
|
||||
torch.nn.init.constant_(lin.bias, 0.0)
|
||||
torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim))
|
||||
|
||||
if weight_norm:
|
||||
lin = nn.utils.weight_norm(lin)
|
||||
|
||||
setattr(self, "lin" + str(l), lin)
|
||||
|
||||
self.activation = nn.Softplus(beta=100)
|
||||
|
||||
def forward(self, inputs):
|
||||
inputs = inputs * self.scale
|
||||
if self.embed_fn_fine is not None:
|
||||
inputs = self.embed_fn_fine(inputs)
|
||||
|
||||
x = inputs
|
||||
for l in range(0, self.num_layers - 1):
|
||||
lin = getattr(self, "lin" + str(l))
|
||||
|
||||
if l in self.skip_in:
|
||||
x = torch.cat([x, inputs], 1) / np.sqrt(2)
|
||||
|
||||
x = lin(x)
|
||||
|
||||
if l < self.num_layers - 2:
|
||||
x = self.activation(x)
|
||||
return torch.cat([x[:, :1] / self.scale, x[:, 1:]], dim=-1)
|
||||
|
||||
def sdf(self, x):
|
||||
return self.forward(x)[:, :1]
|
||||
|
||||
def sdf_hidden_appearance(self, x):
|
||||
return self.forward(x)
|
||||
|
||||
def gradient(self, x):
|
||||
x.requires_grad_(True)
|
||||
with torch.enable_grad():
|
||||
y = self.sdf(x)
|
||||
d_output = torch.ones_like(y, requires_grad=False, device=y.device)
|
||||
gradients = torch.autograd.grad(
|
||||
outputs=y,
|
||||
inputs=x,
|
||||
grad_outputs=d_output,
|
||||
create_graph=True,
|
||||
retain_graph=True,
|
||||
only_inputs=True)[0]
|
||||
return gradients.unsqueeze(1)
|
||||
|
||||
|
||||
# This implementation is borrowed from IDR: https://github.com/lioryariv/idr
|
||||
class RenderingNetwork(nn.Module):
|
||||
def __init__(self,
|
||||
d_feature,
|
||||
mode,
|
||||
d_in,
|
||||
d_out,
|
||||
d_hidden,
|
||||
n_layers,
|
||||
weight_norm=True,
|
||||
multires_view=0,
|
||||
squeeze_out=True):
|
||||
super().__init__()
|
||||
|
||||
self.mode = mode
|
||||
self.squeeze_out = squeeze_out
|
||||
dims = [d_in + d_feature] + [d_hidden for _ in range(n_layers)] + [d_out]
|
||||
|
||||
self.embedview_fn = None
|
||||
if multires_view > 0:
|
||||
embedview_fn, input_ch = get_embedder(multires_view)
|
||||
self.embedview_fn = embedview_fn
|
||||
dims[0] += (input_ch - 3)
|
||||
|
||||
self.num_layers = len(dims)
|
||||
|
||||
for l in range(0, self.num_layers - 1):
|
||||
out_dim = dims[l + 1]
|
||||
lin = nn.Linear(dims[l], out_dim)
|
||||
|
||||
if weight_norm:
|
||||
lin = nn.utils.weight_norm(lin)
|
||||
|
||||
setattr(self, "lin" + str(l), lin)
|
||||
|
||||
self.relu = nn.ReLU()
|
||||
|
||||
def forward(self, points, normals, view_dirs, feature_vectors):
|
||||
if self.embedview_fn is not None:
|
||||
view_dirs = self.embedview_fn(view_dirs)
|
||||
|
||||
rendering_input = None
|
||||
|
||||
if self.mode == 'idr':
|
||||
rendering_input = torch.cat([points, view_dirs, normals, feature_vectors], dim=-1)
|
||||
elif self.mode == 'no_view_dir':
|
||||
rendering_input = torch.cat([points, normals, feature_vectors], dim=-1)
|
||||
elif self.mode == 'no_normal':
|
||||
rendering_input = torch.cat([points, view_dirs, feature_vectors], dim=-1)
|
||||
|
||||
x = rendering_input
|
||||
|
||||
for l in range(0, self.num_layers - 1):
|
||||
lin = getattr(self, "lin" + str(l))
|
||||
|
||||
x = lin(x)
|
||||
|
||||
if l < self.num_layers - 2:
|
||||
x = self.relu(x)
|
||||
|
||||
if self.squeeze_out:
|
||||
x = torch.sigmoid(x)
|
||||
return x
|
||||
|
||||
|
||||
# This implementation is borrowed from nerf-pytorch: https://github.com/yenchenlin/nerf-pytorch
|
||||
class NeRF(nn.Module):
|
||||
def __init__(self,
|
||||
D=8,
|
||||
W=256,
|
||||
d_in=3,
|
||||
d_in_view=3,
|
||||
multires=0,
|
||||
multires_view=0,
|
||||
output_ch=4,
|
||||
skips=[4],
|
||||
use_viewdirs=False):
|
||||
super(NeRF, self).__init__()
|
||||
self.D = D
|
||||
self.W = W
|
||||
self.d_in = d_in
|
||||
self.d_in_view = d_in_view
|
||||
self.input_ch = 3
|
||||
self.input_ch_view = 3
|
||||
self.embed_fn = None
|
||||
self.embed_fn_view = None
|
||||
|
||||
if multires > 0:
|
||||
embed_fn, input_ch = get_embedder(multires, input_dims=d_in)
|
||||
self.embed_fn = embed_fn
|
||||
self.input_ch = input_ch
|
||||
|
||||
if multires_view > 0:
|
||||
embed_fn_view, input_ch_view = get_embedder(multires_view, input_dims=d_in_view)
|
||||
self.embed_fn_view = embed_fn_view
|
||||
self.input_ch_view = input_ch_view
|
||||
|
||||
self.skips = skips
|
||||
self.use_viewdirs = use_viewdirs
|
||||
|
||||
self.pts_linears = nn.ModuleList(
|
||||
[nn.Linear(self.input_ch, W)] +
|
||||
[nn.Linear(W, W) if i not in self.skips else nn.Linear(W + self.input_ch, W) for i in range(D - 1)])
|
||||
|
||||
### Implementation according to the official code release
|
||||
### (https://github.com/bmild/nerf/blob/master/run_nerf_helpers.py#L104-L105)
|
||||
self.views_linears = nn.ModuleList([nn.Linear(self.input_ch_view + W, W // 2)])
|
||||
|
||||
### Implementation according to the paper
|
||||
# self.views_linears = nn.ModuleList(
|
||||
# [nn.Linear(input_ch_views + W, W//2)] + [nn.Linear(W//2, W//2) for i in range(D//2)])
|
||||
|
||||
if use_viewdirs:
|
||||
self.feature_linear = nn.Linear(W, W)
|
||||
self.alpha_linear = nn.Linear(W, 1)
|
||||
self.rgb_linear = nn.Linear(W // 2, 3)
|
||||
else:
|
||||
self.output_linear = nn.Linear(W, output_ch)
|
||||
|
||||
def forward(self, input_pts, input_views):
|
||||
if self.embed_fn is not None:
|
||||
input_pts = self.embed_fn(input_pts)
|
||||
if self.embed_fn_view is not None:
|
||||
input_views = self.embed_fn_view(input_views)
|
||||
|
||||
h = input_pts
|
||||
for i, l in enumerate(self.pts_linears):
|
||||
h = self.pts_linears[i](h)
|
||||
h = F.relu(h)
|
||||
if i in self.skips:
|
||||
h = torch.cat([input_pts, h], -1)
|
||||
|
||||
if self.use_viewdirs:
|
||||
alpha = self.alpha_linear(h)
|
||||
feature = self.feature_linear(h)
|
||||
h = torch.cat([feature, input_views], -1)
|
||||
|
||||
for i, l in enumerate(self.views_linears):
|
||||
h = self.views_linears[i](h)
|
||||
h = F.relu(h)
|
||||
|
||||
rgb = self.rgb_linear(h)
|
||||
return alpha, rgb
|
||||
else:
|
||||
assert False
|
||||
|
||||
|
||||
class SingleVarianceNetwork(nn.Module):
|
||||
def __init__(self, init_val):
|
||||
super(SingleVarianceNetwork, self).__init__()
|
||||
self.register_parameter('variance', nn.Parameter(torch.tensor(init_val)))
|
||||
|
||||
def forward(self, x):
|
||||
return torch.ones([len(x), 1]).to(x.device) * torch.exp(self.variance * 10.0)
|
||||
@@ -0,0 +1,3 @@
|
||||
-5.266582965850830078e-01 7.410295009613037109e-01 -4.165407419204711914e-01 -5.960464477539062500e-08
|
||||
5.865638996738198330e-08 4.900035560131072998e-01 8.717204332351684570e-01 -9.462351613365171943e-08
|
||||
8.500770330429077148e-01 4.590988159179687500e-01 -2.580644786357879639e-01 -1.300000071525573730e+00
|
||||
@@ -0,0 +1,3 @@
|
||||
-9.734988808631896973e-01 1.993551850318908691e-01 -1.120596975088119507e-01 -1.713633537292480469e-07
|
||||
3.790224578636980368e-09 4.900034964084625244e-01 8.717204928398132324e-01 1.772203575001185527e-07
|
||||
2.286916375160217285e-01 8.486189246177673340e-01 -4.770178496837615967e-01 -1.838477611541748047e+00
|
||||
@@ -0,0 +1,3 @@
|
||||
2.286914736032485962e-01 8.486190438270568848e-01 -4.770178198814392090e-01 1.564621925354003906e-07
|
||||
-3.417914484771245043e-08 4.900034070014953613e-01 8.717205524444580078e-01 -7.293811421504869941e-08
|
||||
9.734990000724792480e-01 -1.993550658226013184e-01 1.120596155524253845e-01 -1.838477969169616699e+00
|
||||
@@ -0,0 +1,3 @@
|
||||
5.266583561897277832e-01 -7.410295009613037109e-01 4.165407419204711914e-01 0.000000000000000000e+00
|
||||
5.865638996738198330e-08 4.900035560131072998e-01 8.717204332351684570e-01 9.462351613365171943e-08
|
||||
-8.500770330429077148e-01 -4.590988159179687500e-01 2.580645382404327393e-01 -1.300000071525573730e+00
|
||||
@@ -0,0 +1,3 @@
|
||||
-2.286916971206665039e-01 -8.486189842224121094e-01 4.770179092884063721e-01 -2.458691596984863281e-07
|
||||
9.085837859856837895e-09 4.900034666061401367e-01 8.717205524444580078e-01 1.205695667749751010e-07
|
||||
-9.734990000724792480e-01 1.993551701307296753e-01 -1.120597645640373230e-01 -1.838477969169616699e+00
|
||||
@@ -0,0 +1,3 @@
|
||||
9.734989404678344727e-01 -1.993551850318908691e-01 1.120596975088119507e-01 -1.415610313415527344e-07
|
||||
3.790224578636980368e-09 4.900034964084625244e-01 8.717204928398132324e-01 -1.772203575001185527e-07
|
||||
-2.286916375160217285e-01 -8.486189246177673340e-01 4.770178794860839844e-01 -1.838477611541748047e+00
|
||||
@@ -0,0 +1,3 @@
|
||||
-8.500771522521972656e-01 -4.590989053249359131e-01 2.580644488334655762e-01 0.000000000000000000e+00
|
||||
-4.257411134744870651e-08 4.900034964084625244e-01 8.717204928398132324e-01 9.006067358541258727e-08
|
||||
-5.266583561897277832e-01 7.410295605659484863e-01 -4.165408313274383545e-01 -1.300000071525573730e+00
|
||||
@@ -0,0 +1,3 @@
|
||||
8.500770330429077148e-01 4.590989053249359131e-01 -2.580644488334655762e-01 5.960464477539062500e-08
|
||||
-4.257411134744870651e-08 4.900034964084625244e-01 8.717204928398132324e-01 -9.006067358541258727e-08
|
||||
5.266583561897277832e-01 -7.410295605659484863e-01 4.165407419204711914e-01 -1.300000071525573730e+00
|
||||
@@ -0,0 +1,3 @@
|
||||
9.958608150482177734e-01 7.923202216625213623e-02 -4.453715682029724121e-02 -3.098167056236889039e-09
|
||||
-9.089154005050659180e-02 8.681122064590454102e-01 -4.879753291606903076e-01 5.784738377201392723e-08
|
||||
-2.028124157504862524e-08 4.900035560131072998e-01 8.717204332351684570e-01 -1.300000071525573730e+00
|
||||
@@ -0,0 +1,45 @@
|
||||
import numpy as np
|
||||
|
||||
def camNormal2worldNormal(rot_c2w, camNormal):
|
||||
H,W,_ = camNormal.shape
|
||||
normal_img = np.matmul(rot_c2w[None, :, :], camNormal.reshape(-1,3)[:, :, None]).reshape([H, W, 3])
|
||||
|
||||
return normal_img
|
||||
|
||||
def worldNormal2camNormal(rot_w2c, normal_map_world):
|
||||
H,W,_ = normal_map_world.shape
|
||||
# normal_img = np.matmul(rot_w2c[None, :, :], worldNormal.reshape(-1,3)[:, :, None]).reshape([H, W, 3])
|
||||
|
||||
# faster version
|
||||
# Reshape the normal map into a 2D array where each row represents a normal vector
|
||||
normal_map_flat = normal_map_world.reshape(-1, 3)
|
||||
|
||||
# Transform the normal vectors using the transformation matrix
|
||||
normal_map_camera_flat = np.dot(normal_map_flat, rot_w2c.T)
|
||||
|
||||
# Reshape the transformed normal map back to its original shape
|
||||
normal_map_camera = normal_map_camera_flat.reshape(normal_map_world.shape)
|
||||
|
||||
return normal_map_camera
|
||||
|
||||
def trans_normal(normal, RT_w2c, RT_w2c_target):
|
||||
|
||||
# normal_world = camNormal2worldNormal(np.linalg.inv(RT_w2c[:3,:3]), normal)
|
||||
# normal_target_cam = worldNormal2camNormal(RT_w2c_target[:3,:3], normal_world)
|
||||
|
||||
relative_RT = np.matmul(RT_w2c_target[:3,:3], np.linalg.inv(RT_w2c[:3,:3]))
|
||||
normal_target_cam = worldNormal2camNormal(relative_RT[:3,:3], normal)
|
||||
|
||||
return normal_target_cam
|
||||
|
||||
def img2normal(img):
|
||||
return (img/255.)*2-1
|
||||
|
||||
def normal2img(normal):
|
||||
return np.uint8((normal*0.5+0.5)*255)
|
||||
|
||||
def norm_normalize(normal, dim=-1):
|
||||
|
||||
normal = normal/(np.linalg.norm(normal, axis=dim, keepdims=True)+1e-6)
|
||||
|
||||
return normal
|
||||
@@ -0,0 +1,19 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
def visualize_depth_numpy(depth, minmax=None, cmap=cv2.COLORMAP_JET):
|
||||
"""
|
||||
depth: (H, W)
|
||||
"""
|
||||
|
||||
x = np.nan_to_num(depth) # change nan to 0
|
||||
if minmax is None:
|
||||
mi = np.min(x[x > 0]) # get minimum positive depth (ignore background)
|
||||
ma = np.max(x)
|
||||
else:
|
||||
mi, ma = minmax
|
||||
|
||||
x = (x - mi) / (ma - mi + 1e-8) # normalize to 0~1
|
||||
x = (255 * x).astype(np.uint8)
|
||||
x_ = cv2.applyColorMap(x, cmap)
|
||||
return x_, [mi, ma]
|
||||
@@ -0,0 +1,431 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import mcubes
|
||||
|
||||
def extract_fields(bound_min, bound_max, resolution, query_func):
|
||||
N = 64
|
||||
X = torch.linspace(bound_min[0], bound_max[0], resolution).split(N)
|
||||
Y = torch.linspace(bound_min[1], bound_max[1], resolution).split(N)
|
||||
Z = torch.linspace(bound_min[2], bound_max[2], resolution).split(N)
|
||||
|
||||
u = np.zeros([resolution, resolution, resolution], dtype=np.float32)
|
||||
with torch.no_grad():
|
||||
for xi, xs in enumerate(X):
|
||||
for yi, ys in enumerate(Y):
|
||||
for zi, zs in enumerate(Z):
|
||||
xx, yy, zz = torch.meshgrid(xs, ys, zs)
|
||||
pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1)
|
||||
val = query_func(pts.cuda()).reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy()
|
||||
u[xi * N: xi * N + len(xs), yi * N: yi * N + len(ys), zi * N: zi * N + len(zs)] = val
|
||||
return u
|
||||
|
||||
|
||||
def extract_geometry(bound_min, bound_max, resolution, threshold, query_func, color_func):
|
||||
u = extract_fields(bound_min, bound_max, resolution, query_func)
|
||||
vertices, triangles = mcubes.marching_cubes(u, threshold)
|
||||
b_max_np = bound_max.detach().cpu().numpy()
|
||||
b_min_np = bound_min.detach().cpu().numpy()
|
||||
|
||||
vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
|
||||
|
||||
vertices_color = color_func(vertices)
|
||||
|
||||
return vertices, triangles, vertices_color
|
||||
|
||||
|
||||
def sample_pdf(bins, weights, n_samples, det=False):
|
||||
# This implementation is from NeRF
|
||||
# Get pdf
|
||||
device = weights.device
|
||||
weights = weights + 1e-5 # prevent nans
|
||||
pdf = weights / torch.sum(weights, -1, keepdim=True)
|
||||
cdf = torch.cumsum(pdf, -1)
|
||||
cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1)
|
||||
# Take uniform samples
|
||||
if det:
|
||||
u = torch.linspace(0. + 0.5 / n_samples, 1. - 0.5 / n_samples, steps=n_samples).to(device)
|
||||
u = u.expand(list(cdf.shape[:-1]) + [n_samples])
|
||||
else:
|
||||
u = torch.rand(list(cdf.shape[:-1]) + [n_samples]).to(device)
|
||||
|
||||
# Invert CDF
|
||||
u = u.contiguous()
|
||||
inds = torch.searchsorted(cdf, u, right=True)
|
||||
below = torch.max(torch.zeros_like(inds - 1).to(device), inds - 1)
|
||||
above = torch.min((cdf.shape[-1] - 1) * torch.ones_like(inds).to(device), inds)
|
||||
inds_g = torch.stack([below, above], -1) # (batch, N_samples, 2)
|
||||
|
||||
matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]]
|
||||
cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g)
|
||||
bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g)
|
||||
|
||||
denom = (cdf_g[..., 1] - cdf_g[..., 0])
|
||||
denom = torch.where(denom < 1e-5, torch.ones_like(denom).to(device), denom)
|
||||
t = (u - cdf_g[..., 0]) / denom
|
||||
samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])
|
||||
|
||||
return samples
|
||||
|
||||
|
||||
class NeuSRenderer:
|
||||
def __init__(self,
|
||||
nerf,
|
||||
sdf_network,
|
||||
deviation_network,
|
||||
color_network,
|
||||
n_samples,
|
||||
n_importance,
|
||||
n_outside,
|
||||
up_sample_steps,
|
||||
perturb,
|
||||
sdf_decay_param):
|
||||
self.nerf = nerf
|
||||
self.sdf_network = sdf_network
|
||||
self.deviation_network = deviation_network
|
||||
self.color_network = color_network
|
||||
self.n_samples = n_samples
|
||||
self.n_importance = n_importance
|
||||
self.n_outside = n_outside
|
||||
self.up_sample_steps = up_sample_steps
|
||||
self.perturb = perturb
|
||||
|
||||
self.sdf_decay_param = sdf_decay_param
|
||||
|
||||
def render_core_outside(self, rays_o, rays_d, z_vals, sample_dist, nerf, background_rgb=None):
|
||||
"""
|
||||
Render background
|
||||
"""
|
||||
batch_size, n_samples = z_vals.shape
|
||||
|
||||
# Section length
|
||||
dists = z_vals[..., 1:] - z_vals[..., :-1]
|
||||
dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)
|
||||
mid_z_vals = z_vals + dists * 0.5
|
||||
|
||||
# Section midpoints
|
||||
pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None] # batch_size, n_samples, 3
|
||||
|
||||
dis_to_center = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).clip(1.0, 1e10)
|
||||
pts = torch.cat([pts / dis_to_center, 1.0 / dis_to_center], dim=-1) # batch_size, n_samples, 4
|
||||
|
||||
dirs = rays_d[:, None, :].expand(batch_size, n_samples, 3)
|
||||
|
||||
pts = pts.reshape(-1, 3 + int(self.n_outside > 0))
|
||||
dirs = dirs.reshape(-1, 3)
|
||||
|
||||
density, sampled_color = nerf(pts, dirs)
|
||||
sampled_color = torch.sigmoid(sampled_color)
|
||||
alpha = 1.0 - torch.exp(-F.softplus(density.reshape(batch_size, n_samples)) * dists)
|
||||
alpha = alpha.reshape(batch_size, n_samples)
|
||||
weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
|
||||
sampled_color = sampled_color.reshape(batch_size, n_samples, 3)
|
||||
color = (weights[:, :, None] * sampled_color).sum(dim=1)
|
||||
if background_rgb is not None:
|
||||
color = color + background_rgb * (1.0 - weights.sum(dim=-1, keepdim=True))
|
||||
|
||||
return {
|
||||
'color': color,
|
||||
'sampled_color': sampled_color,
|
||||
'alpha': alpha,
|
||||
'weights': weights,
|
||||
}
|
||||
|
||||
def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
|
||||
"""
|
||||
Up sampling give a fixed inv_s
|
||||
"""
|
||||
device = rays_o.device
|
||||
batch_size, n_samples = z_vals.shape
|
||||
pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None] # n_rays, n_samples, 3
|
||||
radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
|
||||
inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
|
||||
sdf = sdf.reshape(batch_size, n_samples)
|
||||
prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
|
||||
prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
|
||||
mid_sdf = (prev_sdf + next_sdf) * 0.5
|
||||
cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)
|
||||
|
||||
# ----------------------------------------------------------------------------------------------------------
|
||||
# Use min value of [ cos, prev_cos ]
|
||||
# Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
|
||||
# robust when meeting situations like below:
|
||||
#
|
||||
# SDF
|
||||
# ^
|
||||
# |\ -----x----...
|
||||
# | \ /
|
||||
# | x x
|
||||
# |---\----/-------------> 0 level
|
||||
# | \ /
|
||||
# | \/
|
||||
# |
|
||||
# ----------------------------------------------------------------------------------------------------------
|
||||
prev_cos_val = torch.cat([torch.zeros([batch_size, 1]).to(device), cos_val[:, :-1]], dim=-1)
|
||||
cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
|
||||
cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
|
||||
cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere
|
||||
|
||||
dist = (next_z_vals - prev_z_vals)
|
||||
prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
|
||||
next_esti_sdf = mid_sdf + cos_val * dist * 0.5
|
||||
prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
|
||||
next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
|
||||
alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
|
||||
weights = alpha * torch.cumprod(
|
||||
torch.cat([torch.ones([batch_size, 1]).to(device), 1. - alpha + 1e-7], -1), -1)[:, :-1]
|
||||
|
||||
z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
|
||||
return z_samples
|
||||
|
||||
def cat_z_vals(self, rays_o, rays_d, z_vals, new_z_vals, sdf, last=False):
|
||||
batch_size, n_samples = z_vals.shape
|
||||
_, n_importance = new_z_vals.shape
|
||||
pts = rays_o[:, None, :] + rays_d[:, None, :] * new_z_vals[..., :, None]
|
||||
z_vals = torch.cat([z_vals, new_z_vals], dim=-1)
|
||||
z_vals, index = torch.sort(z_vals, dim=-1)
|
||||
|
||||
if not last:
|
||||
new_sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, n_importance)
|
||||
sdf = torch.cat([sdf, new_sdf], dim=-1)
|
||||
xx = torch.arange(batch_size)[:, None].expand(batch_size, n_samples + n_importance).reshape(-1)
|
||||
index = index.reshape(-1)
|
||||
sdf = sdf[(xx, index)].reshape(batch_size, n_samples + n_importance)
|
||||
|
||||
return z_vals, sdf
|
||||
|
||||
def render_core(self,
|
||||
rays_o,
|
||||
rays_d,
|
||||
z_vals,
|
||||
sample_dist,
|
||||
sdf_network,
|
||||
deviation_network,
|
||||
color_network,
|
||||
background_alpha=None,
|
||||
background_sampled_color=None,
|
||||
background_rgb=None,
|
||||
cos_anneal_ratio=0.0):
|
||||
batch_size, n_samples = z_vals.shape
|
||||
device = rays_o.device
|
||||
# Section length
|
||||
dists = z_vals[..., 1:] - z_vals[..., :-1]
|
||||
dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape).to(device)], -1)
|
||||
mid_z_vals = z_vals + dists * 0.5
|
||||
|
||||
# Section midpoints
|
||||
pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None] # n_rays, n_samples, 3
|
||||
dirs = rays_d[:, None, :].expand(pts.shape)
|
||||
|
||||
pts = pts.reshape(-1, 3)
|
||||
dirs = dirs.reshape(-1, 3)
|
||||
|
||||
sdf_nn_output = sdf_network(pts)
|
||||
sdf = sdf_nn_output[:, :1]
|
||||
feature_vector = sdf_nn_output[:, 1:]
|
||||
|
||||
gradients = sdf_network.gradient(pts).squeeze()
|
||||
sampled_color = color_network(pts, gradients, dirs, feature_vector).reshape(batch_size, n_samples, 3)
|
||||
|
||||
inv_s = deviation_network(torch.zeros([1, 3]).to(device))[:, :1].clip(1e-6, 1e6) # Single parameter
|
||||
inv_s = inv_s.expand(batch_size * n_samples, 1)
|
||||
|
||||
true_cos = (dirs * gradients).sum(-1, keepdim=True)
|
||||
|
||||
# "cos_anneal_ratio" grows from 0 to 1 in the beginning training iterations. The anneal strategy below makes
|
||||
# the cos value "not dead" at the beginning training iterations, for better convergence.
|
||||
iter_cos = -(F.relu(-true_cos * 0.5 + 0.5) * (1.0 - cos_anneal_ratio) +
|
||||
F.relu(-true_cos) * cos_anneal_ratio) # always non-positive
|
||||
|
||||
# Estimate signed distances at section points
|
||||
estimated_next_sdf = sdf + iter_cos * dists.reshape(-1, 1) * 0.5
|
||||
estimated_prev_sdf = sdf - iter_cos * dists.reshape(-1, 1) * 0.5
|
||||
|
||||
prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_s)
|
||||
next_cdf = torch.sigmoid(estimated_next_sdf * inv_s)
|
||||
|
||||
p = prev_cdf - next_cdf
|
||||
c = prev_cdf
|
||||
|
||||
alpha = ((p + 1e-5) / (c + 1e-5)).reshape(batch_size, n_samples).clip(0.0, 1.0)
|
||||
|
||||
pts_norm = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).reshape(batch_size, n_samples)
|
||||
inside_sphere = (pts_norm < 1.0).float().detach()
|
||||
relax_inside_sphere = (pts_norm < 1.2).float().detach()
|
||||
|
||||
# Render with background
|
||||
if background_alpha is not None:
|
||||
alpha = alpha * inside_sphere + background_alpha[:, :n_samples] * (1.0 - inside_sphere)
|
||||
alpha = torch.cat([alpha, background_alpha[:, n_samples:]], dim=-1)
|
||||
sampled_color = sampled_color * inside_sphere[:, :, None] +\
|
||||
background_sampled_color[:, :n_samples] * (1.0 - inside_sphere)[:, :, None]
|
||||
sampled_color = torch.cat([sampled_color, background_sampled_color[:, n_samples:]], dim=1)
|
||||
|
||||
weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]).to(device), 1. - alpha + 1e-7], -1), -1)[:, :-1]
|
||||
weights_sum = weights.sum(dim=-1, keepdim=True)
|
||||
|
||||
color = (sampled_color * weights[:, :, None]).sum(dim=1)
|
||||
if background_rgb is not None: # Fixed background, usually black
|
||||
color = color + background_rgb.to(device) * (1.0 - weights_sum)
|
||||
|
||||
# Eikonal loss
|
||||
gradient_error = (torch.linalg.norm(gradients.reshape(batch_size, n_samples, 3), ord=2,
|
||||
dim=-1) - 1.0) ** 2
|
||||
gradient_error = (relax_inside_sphere * gradient_error).sum() / (relax_inside_sphere.sum() + 1e-5)
|
||||
|
||||
depth = (mid_z_vals * weights[:, :n_samples]).sum(dim=1, keepdim=True)
|
||||
|
||||
return {
|
||||
'color': color,
|
||||
'sdf': sdf,
|
||||
'dists': dists,
|
||||
'gradients': gradients.reshape(batch_size, n_samples, 3),
|
||||
's_val': 1.0 / inv_s,
|
||||
'mid_z_vals': mid_z_vals,
|
||||
'weights': weights,
|
||||
'cdf': c.reshape(batch_size, n_samples),
|
||||
'gradient_error': gradient_error,
|
||||
'inside_sphere': inside_sphere,
|
||||
'depth': depth,
|
||||
}
|
||||
|
||||
def render(self, rays_o, rays_d, near, far, perturb_overwrite=-1, background_rgb=None, cos_anneal_ratio=0.0):
|
||||
batch_size = len(rays_o)
|
||||
device = rays_o.device
|
||||
sample_dist = 2.0 / self.n_samples # Assuming the region of interest is a unit sphere
|
||||
z_vals = torch.linspace(0.0, 1.0, self.n_samples)
|
||||
z_vals = near + (far - near) * z_vals[None, :]
|
||||
|
||||
z_vals_outside = None
|
||||
if self.n_outside > 0:
|
||||
z_vals_outside = torch.linspace(1e-3, 1.0 - 1.0 / (self.n_outside + 1.0), self.n_outside)
|
||||
|
||||
n_samples = self.n_samples
|
||||
perturb = self.perturb
|
||||
|
||||
if perturb_overwrite >= 0:
|
||||
perturb = perturb_overwrite
|
||||
if perturb > 0:
|
||||
t_rand = (torch.rand([batch_size, 1]) - 0.5)
|
||||
z_vals = z_vals + t_rand * 2.0 / self.n_samples
|
||||
z_vals = z_vals.to(device)
|
||||
|
||||
if self.n_outside > 0:
|
||||
mids = .5 * (z_vals_outside[..., 1:] + z_vals_outside[..., :-1])
|
||||
upper = torch.cat([mids, z_vals_outside[..., -1:]], -1)
|
||||
lower = torch.cat([z_vals_outside[..., :1], mids], -1)
|
||||
t_rand = torch.rand([batch_size, z_vals_outside.shape[-1]])
|
||||
z_vals_outside = lower[None, :] + (upper - lower)[None, :] * t_rand
|
||||
|
||||
if self.n_outside > 0:
|
||||
z_vals_outside = far / torch.flip(z_vals_outside, dims=[-1]) + 1.0 / self.n_samples
|
||||
|
||||
background_alpha = None
|
||||
background_sampled_color = None
|
||||
|
||||
# Up sample
|
||||
if self.n_importance > 0:
|
||||
with torch.no_grad():
|
||||
pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
|
||||
sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, self.n_samples)
|
||||
|
||||
for i in range(self.up_sample_steps):
|
||||
new_z_vals = self.up_sample(rays_o,
|
||||
rays_d,
|
||||
z_vals,
|
||||
sdf,
|
||||
self.n_importance // self.up_sample_steps,
|
||||
64 * 2**i)
|
||||
z_vals, sdf = self.cat_z_vals(rays_o,
|
||||
rays_d,
|
||||
z_vals,
|
||||
new_z_vals,
|
||||
sdf,
|
||||
last=(i + 1 == self.up_sample_steps))
|
||||
|
||||
n_samples = self.n_samples + self.n_importance
|
||||
|
||||
# Background model
|
||||
if self.n_outside > 0:
|
||||
z_vals_feed = torch.cat([z_vals, z_vals_outside], dim=-1)
|
||||
z_vals_feed, _ = torch.sort(z_vals_feed, dim=-1)
|
||||
ret_outside = self.render_core_outside(rays_o, rays_d, z_vals_feed, sample_dist, self.nerf)
|
||||
|
||||
background_sampled_color = ret_outside['sampled_color']
|
||||
background_alpha = ret_outside['alpha']
|
||||
|
||||
# Render core
|
||||
ret_fine = self.render_core(rays_o,
|
||||
rays_d,
|
||||
z_vals,
|
||||
sample_dist,
|
||||
self.sdf_network,
|
||||
self.deviation_network,
|
||||
self.color_network,
|
||||
background_rgb=background_rgb,
|
||||
background_alpha=background_alpha,
|
||||
background_sampled_color=background_sampled_color,
|
||||
cos_anneal_ratio=cos_anneal_ratio)
|
||||
|
||||
color_fine = ret_fine['color']
|
||||
weights = ret_fine['weights']
|
||||
weights_sum = weights.sum(dim=-1, keepdim=True)
|
||||
gradients = ret_fine['gradients']
|
||||
s_val = ret_fine['s_val'].reshape(batch_size, n_samples).mean(dim=-1, keepdim=True)
|
||||
|
||||
|
||||
# - randomly sample points from the volume, and maximize the sdf
|
||||
pts_random = torch.rand([1024, 3]).float().cuda() * 2 - 1 # normalized to (-1, 1)
|
||||
sdf_random = self.sdf_network(pts_random)[:, :1]
|
||||
|
||||
sparse_loss_1 = torch.exp(
|
||||
-1 * torch.abs(sdf_random) * self.sdf_decay_param).mean() # - should equal
|
||||
sparse_loss_2 = torch.exp(-1 * torch.abs(ret_fine['sdf']) * self.sdf_decay_param).mean()
|
||||
sparse_loss = (sparse_loss_1 + sparse_loss_2) / 2
|
||||
|
||||
return {
|
||||
'color_fine': color_fine,
|
||||
'depth': ret_fine['depth'],
|
||||
's_val': s_val,
|
||||
'sparse_loss': sparse_loss,
|
||||
'cdf_fine': ret_fine['cdf'],
|
||||
'weight_sum': weights_sum,
|
||||
'weight_max': torch.max(weights, dim=-1, keepdim=True)[0],
|
||||
'gradients': gradients,
|
||||
'weights': weights,
|
||||
'gradient_error': ret_fine['gradient_error'],
|
||||
'inside_sphere': ret_fine['inside_sphere']
|
||||
}
|
||||
|
||||
|
||||
def get_vertex_colors(self, vertices):
|
||||
"""
|
||||
@param vertices: n,3
|
||||
@return:
|
||||
"""
|
||||
V = vertices.shape[0]
|
||||
bn = 20480
|
||||
verts_colors = []
|
||||
with torch.no_grad():
|
||||
for vi in range(0, V, bn):
|
||||
verts = torch.from_numpy(vertices[vi:vi+bn].astype(np.float32)).cuda()
|
||||
feats = self.sdf_network(verts)[..., 1:]
|
||||
# pdb.set_trace()
|
||||
gradients = self.sdf_network.gradient(verts) # ...,3
|
||||
gradients = F.normalize(gradients, dim=-1).squeeze(1)
|
||||
# pdb.set_trace()
|
||||
colors = self.color_network(verts, gradients, gradients, feats)
|
||||
colors = torch.clamp(colors,min=0,max=1).cpu().numpy()
|
||||
verts_colors.append(colors)
|
||||
|
||||
verts_colors = (np.concatenate(verts_colors, 0)*255).astype(np.uint8)
|
||||
return verts_colors[:, ::-1]
|
||||
|
||||
def extract_geometry(self, bound_min, bound_max, resolution, threshold=0.0):
|
||||
return extract_geometry(bound_min,
|
||||
bound_max,
|
||||
resolution=resolution,
|
||||
threshold=threshold,
|
||||
query_func=lambda pts: -self.sdf_network.sdf(pts),
|
||||
color_func=lambda pts: self.get_vertex_colors(pts))
|
||||
@@ -12,9 +12,13 @@
|
||||
### Note: this project is still a WIP
|
||||
|
||||
## Currently support:
|
||||
- For use case please check [Example Workflows](./_Example_Workflows/). [**Last update: 11/02/2024**]
|
||||
- For use case please check [Example Workflows](./_Example_Workflows/). [**Last update: 02/03/2024**]
|
||||
- **Note:** you need to put [Example Inputs Files & Folders](_Example_Workflows/_Example_Inputs_Files/) under ComfyUI Root Directory\ComfyUI\input folder before you can run the example workflow
|
||||
|
||||
- **Wonder3D**: [xxlong0/Wonder3D](https://github.com/xxlong0/Wonder3D)
|
||||
- Generate spatial consistent 6 views images & normal maps from a single image
|
||||
|
||||
|
||||
- **Large Multiview Gaussian Model**: [3DTopia/LGM](https://github.com/3DTopia/LGM)
|
||||
- Enable single image to 3D Gaussian in less than 30 seconds on a RTX3080 GPU, later you can also convert 3D Gaussian to mesh
|
||||
|
||||
@@ -56,6 +60,9 @@
|
||||
- Bake Multi-View images into UVTexture of given 3D mesh using [Nvdiffrast](https://github.com/NVlabs/nvdiffrast), supports:
|
||||
- Export to .obj, .ply, .glb
|
||||
|
||||
- **NeuS**
|
||||
- Fit a coarse mesh from sparse multi-view images & normal maps, as little as 4 to 6 views, pretty good at reconstruct the shape from reference images but texture lacking details.
|
||||
|
||||
- **Deep Marching Tetrahedrons**
|
||||
- Allow convert 3DGS .ply file to 3D mesh <br>
|
||||
*Note: I didn't spent time to turn the hyperprameters yet, the result will be improved in the future!*
|
||||
@@ -180,7 +187,7 @@ Copy the files inside folder [__New_ComfyUI_Bats](./_New_ComfyUI_Bats/) to your
|
||||
- Alternatively you can just activate the Conda env: `python_miniconda_env\ComfyUI`, and go to your ComfyUI root directory then run command `python ./ComfyUI/main.py`
|
||||
|
||||
## Tips
|
||||
* The world & camera coordinate system is the same as OpenGL:
|
||||
* OpenGL world & camera coordinate system:
|
||||
```
|
||||
World Camera
|
||||
|
||||
@@ -197,5 +204,13 @@ elevation: in (-90, 90), from +y to -y is (-90, 90)
|
||||
azimuth: in (-180, 180), from +z to +x is (0, 90)
|
||||
```
|
||||
|
||||
* Wonder3D world & camera coordinate system:
|
||||
|
||||
|
||||
|
||||
* Three.js coordinate system: (z-axis is pointing towards you and is coming out of the screen)
|
||||
|
||||
|
||||
|
||||
* If you encounter OpenGL errors (e.g., `[F glutil.cpp:338] eglInitialize() failed`), then set `force_cuda_rasterize` to true on corresponding node
|
||||
* If after the installation, your ComfyUI get stucked at starting or running, you could following the instruction in following link to solve the problem: [Code Hangs Indefinitely When Evaluating Neuron Models on GPU](https://github.com/lava-nc/lava-dl/discussions/211)
|
||||
|
||||
+782
@@ -0,0 +1,782 @@
|
||||
{
|
||||
"last_node_id": 23,
|
||||
"last_link_id": 36,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 17,
|
||||
"type": "[Comfy3D] Large Multiview Gaussian Model",
|
||||
"pos": [
|
||||
1270,
|
||||
461
|
||||
],
|
||||
"size": {
|
||||
"0": 252,
|
||||
"1": 46
|
||||
},
|
||||
"flags": {},
|
||||
"order": 8,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "multiview_images",
|
||||
"type": "IMAGE",
|
||||
"link": 35
|
||||
},
|
||||
{
|
||||
"name": "lgm_model",
|
||||
"type": "LGM_MODEL",
|
||||
"link": 22
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "gs_ply",
|
||||
"type": "GS_PLY",
|
||||
"links": [
|
||||
19,
|
||||
20
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Large Multiview Gaussian Model"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 15,
|
||||
"type": "[Comfy3D] Preview 3DMesh",
|
||||
"pos": [
|
||||
2326.1025658000017,
|
||||
214.90587429004222
|
||||
],
|
||||
"size": [
|
||||
600,
|
||||
500
|
||||
],
|
||||
"flags": {},
|
||||
"order": 15,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mesh_file_path",
|
||||
"type": "STRING",
|
||||
"link": 16,
|
||||
"widget": {
|
||||
"name": "mesh_file_path"
|
||||
}
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Preview 3DMesh"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
null
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 12,
|
||||
"type": "[Comfy3D] Switch 3DGS Axis",
|
||||
"pos": [
|
||||
1558,
|
||||
460
|
||||
],
|
||||
"size": {
|
||||
"0": 210,
|
||||
"1": 106
|
||||
},
|
||||
"flags": {},
|
||||
"order": 10,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "gs_ply",
|
||||
"type": "GS_PLY",
|
||||
"link": 19
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "switched_gs_ply",
|
||||
"type": "GS_PLY",
|
||||
"links": [
|
||||
13
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Switch 3DGS Axis"
|
||||
},
|
||||
"widgets_values": [
|
||||
"+x",
|
||||
"-y",
|
||||
"-z"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 10,
|
||||
"type": "[Comfy3D] Preview 3DGS",
|
||||
"pos": [
|
||||
2331,
|
||||
776
|
||||
],
|
||||
"size": [
|
||||
600,
|
||||
500
|
||||
],
|
||||
"flags": {},
|
||||
"order": 14,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "gs_file_path",
|
||||
"type": "STRING",
|
||||
"link": 11,
|
||||
"widget": {
|
||||
"name": "gs_file_path"
|
||||
}
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Preview 3DGS"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
null
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 16,
|
||||
"type": "[Comfy3D] Convert 3DGS to Mesh with NeRF and Marching Cubes",
|
||||
"pos": [
|
||||
1566,
|
||||
215
|
||||
],
|
||||
"size": {
|
||||
"0": 411.6000061035156,
|
||||
"1": 82
|
||||
},
|
||||
"flags": {},
|
||||
"order": 11,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "gs_ply",
|
||||
"type": "GS_PLY",
|
||||
"link": 20
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "mesh",
|
||||
"type": "MESH",
|
||||
"links": [
|
||||
17
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Convert 3DGS to Mesh with NeRF and Marching Cubes"
|
||||
},
|
||||
"widgets_values": [
|
||||
"big",
|
||||
false
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 5,
|
||||
"type": "InvertMask",
|
||||
"pos": [
|
||||
170,
|
||||
540
|
||||
],
|
||||
"size": {
|
||||
"0": 210,
|
||||
"1": 26
|
||||
},
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mask",
|
||||
"type": "MASK",
|
||||
"link": 3
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
25,
|
||||
32
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "InvertMask"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 20,
|
||||
"type": "[Comfy3D] Load Diffusers Pipeline",
|
||||
"pos": [
|
||||
-300,
|
||||
330
|
||||
],
|
||||
"size": {
|
||||
"0": 392.1646728515625,
|
||||
"1": 86.24864196777344
|
||||
},
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "pipe",
|
||||
"type": "DIFFUSERS_PIPE",
|
||||
"links": [
|
||||
29
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Load Diffusers Pipeline"
|
||||
},
|
||||
"widgets_values": [
|
||||
"Wonder3DMVDiffusionPipeline",
|
||||
"flamehaze1115/wonder3d-v1.0"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 22,
|
||||
"type": "[Comfy3D] Wonder3D MVDiffusion Model",
|
||||
"pos": [
|
||||
440,
|
||||
460
|
||||
],
|
||||
"size": {
|
||||
"0": 430.8000183105469,
|
||||
"1": 170
|
||||
},
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mvdiffusion_pipe",
|
||||
"type": "DIFFUSERS_PIPE",
|
||||
"link": 29
|
||||
},
|
||||
{
|
||||
"name": "reference_image",
|
||||
"type": "IMAGE",
|
||||
"link": 33
|
||||
},
|
||||
{
|
||||
"name": "reference_mask",
|
||||
"type": "MASK",
|
||||
"link": 32
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "multiview_images",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
34
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "multiview_normals",
|
||||
"type": "IMAGE",
|
||||
"links": [],
|
||||
"shape": 3,
|
||||
"slot_index": 1
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Wonder3D MVDiffusion Model"
|
||||
},
|
||||
"widgets_values": [
|
||||
0,
|
||||
"fixed",
|
||||
1,
|
||||
30
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 23,
|
||||
"type": "GetImagesFromBatchIndexed",
|
||||
"pos": [
|
||||
918,
|
||||
461
|
||||
],
|
||||
"size": {
|
||||
"0": 236.4545135498047,
|
||||
"1": 76
|
||||
},
|
||||
"flags": {},
|
||||
"order": 7,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 34
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
35,
|
||||
36
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "GetImagesFromBatchIndexed"
|
||||
},
|
||||
"widgets_values": [
|
||||
"0, 2, 3, 4"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 1,
|
||||
"type": "[Comfy3D] Load Large Multiview Gaussian Model",
|
||||
"pos": [
|
||||
917,
|
||||
598
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 82
|
||||
},
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "lgm_model",
|
||||
"type": "LGM_MODEL",
|
||||
"links": [
|
||||
22
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Load Large Multiview Gaussian Model"
|
||||
},
|
||||
"widgets_values": [
|
||||
"model_fp16.safetensors",
|
||||
"big"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"type": "[Comfy3D] Load Diffusers Pipeline",
|
||||
"pos": [
|
||||
-138.35285154113785,
|
||||
-248.56920693359376
|
||||
],
|
||||
"size": {
|
||||
"0": 375.3948669433594,
|
||||
"1": 83.33004760742188
|
||||
},
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 4,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "pipe",
|
||||
"type": "DIFFUSERS_PIPE",
|
||||
"links": [
|
||||
23
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Load Diffusers Pipeline"
|
||||
},
|
||||
"widgets_values": [
|
||||
"MVDreamPipeline",
|
||||
"ashawkey/imagedream-ipmv-diffusers"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 18,
|
||||
"type": "[Comfy3D] MVDream Model",
|
||||
"pos": [
|
||||
451.6471484588624,
|
||||
-248.56920693359376
|
||||
],
|
||||
"size": {
|
||||
"0": 405.5999755859375,
|
||||
"1": 286
|
||||
},
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 4,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mvdream_pipe",
|
||||
"type": "DIFFUSERS_PIPE",
|
||||
"link": 23
|
||||
},
|
||||
{
|
||||
"name": "reference_image",
|
||||
"type": "IMAGE",
|
||||
"link": 24
|
||||
},
|
||||
{
|
||||
"name": "reference_mask",
|
||||
"type": "MASK",
|
||||
"link": 25
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "multiview_images",
|
||||
"type": "IMAGE",
|
||||
"links": [],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] MVDream Model"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
"ugly, blurry, pixelated obscure, unnatural colors, poor lighting, dull, unclear, cropped, lowres, low quality, artifacts, duplicate",
|
||||
0,
|
||||
"fixed",
|
||||
4.99,
|
||||
30,
|
||||
0
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 14,
|
||||
"type": "[Comfy3D] Save 3D Mesh",
|
||||
"pos": [
|
||||
1996.1025658,
|
||||
214.90587429004222
|
||||
],
|
||||
"size": {
|
||||
"0": 307.8253173828125,
|
||||
"1": 60.248291015625
|
||||
},
|
||||
"flags": {},
|
||||
"order": 13,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mesh",
|
||||
"type": "MESH",
|
||||
"link": 17
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "save_path",
|
||||
"type": "STRING",
|
||||
"links": [
|
||||
16
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Save 3D Mesh"
|
||||
},
|
||||
"widgets_values": [
|
||||
"LGMTest/KungFuCat.obj"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 9,
|
||||
"type": "[Comfy3D] Save 3DGS",
|
||||
"pos": [
|
||||
1798,
|
||||
459
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 58
|
||||
},
|
||||
"flags": {},
|
||||
"order": 12,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "gs_ply",
|
||||
"type": "GS_PLY",
|
||||
"link": 13
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "save_path",
|
||||
"type": "STRING",
|
||||
"links": [
|
||||
11
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Save 3DGS"
|
||||
},
|
||||
"widgets_values": [
|
||||
"LGMTest/KungFuCat.ply"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
-190,
|
||||
480
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 314
|
||||
},
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
24,
|
||||
33
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
3
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 1
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"KungFuCat.png",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 19,
|
||||
"type": "PreviewImage",
|
||||
"pos": [
|
||||
1278,
|
||||
138
|
||||
],
|
||||
"size": {
|
||||
"0": 210,
|
||||
"1": 246
|
||||
},
|
||||
"flags": {},
|
||||
"order": 9,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 36
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "PreviewImage"
|
||||
}
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
3,
|
||||
4,
|
||||
1,
|
||||
5,
|
||||
0,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
11,
|
||||
9,
|
||||
0,
|
||||
10,
|
||||
0,
|
||||
"STRING"
|
||||
],
|
||||
[
|
||||
13,
|
||||
12,
|
||||
0,
|
||||
9,
|
||||
0,
|
||||
"GS_PLY"
|
||||
],
|
||||
[
|
||||
16,
|
||||
14,
|
||||
0,
|
||||
15,
|
||||
0,
|
||||
"STRING"
|
||||
],
|
||||
[
|
||||
17,
|
||||
16,
|
||||
0,
|
||||
14,
|
||||
0,
|
||||
"MESH"
|
||||
],
|
||||
[
|
||||
19,
|
||||
17,
|
||||
0,
|
||||
12,
|
||||
0,
|
||||
"GS_PLY"
|
||||
],
|
||||
[
|
||||
20,
|
||||
17,
|
||||
0,
|
||||
16,
|
||||
0,
|
||||
"GS_PLY"
|
||||
],
|
||||
[
|
||||
22,
|
||||
1,
|
||||
0,
|
||||
17,
|
||||
1,
|
||||
"LGM_MODEL"
|
||||
],
|
||||
[
|
||||
23,
|
||||
2,
|
||||
0,
|
||||
18,
|
||||
0,
|
||||
"DIFFUSERS_PIPE"
|
||||
],
|
||||
[
|
||||
24,
|
||||
4,
|
||||
0,
|
||||
18,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
25,
|
||||
5,
|
||||
0,
|
||||
18,
|
||||
2,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
29,
|
||||
20,
|
||||
0,
|
||||
22,
|
||||
0,
|
||||
"DIFFUSERS_PIPE"
|
||||
],
|
||||
[
|
||||
32,
|
||||
5,
|
||||
0,
|
||||
22,
|
||||
2,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
33,
|
||||
4,
|
||||
0,
|
||||
22,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
34,
|
||||
22,
|
||||
0,
|
||||
23,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
35,
|
||||
23,
|
||||
0,
|
||||
17,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
36,
|
||||
23,
|
||||
0,
|
||||
19,
|
||||
0,
|
||||
"IMAGE"
|
||||
]
|
||||
],
|
||||
"groups": [
|
||||
{
|
||||
"title": "MVDream: 4 multi-views",
|
||||
"bounding": [
|
||||
-172,
|
||||
-366,
|
||||
1139,
|
||||
436
|
||||
],
|
||||
"color": "#3f789e",
|
||||
"font_size": 24,
|
||||
"locked": false
|
||||
}
|
||||
],
|
||||
"config": {},
|
||||
"extra": {},
|
||||
"version": 0.4
|
||||
}
|
||||
@@ -0,0 +1,644 @@
|
||||
{
|
||||
"last_node_id": 47,
|
||||
"last_link_id": 68,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 21,
|
||||
"type": "PreviewImage",
|
||||
"pos": [
|
||||
1265,
|
||||
880
|
||||
],
|
||||
"size": {
|
||||
"0": 199.94778442382812,
|
||||
"1": 274.4532775878906
|
||||
},
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 32
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "PreviewImage"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 28,
|
||||
"type": "Image Rembg (Remove Background)",
|
||||
"pos": [
|
||||
1487,
|
||||
475
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 250
|
||||
},
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 37
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
38,
|
||||
39
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "Image Rembg (Remove Background)"
|
||||
},
|
||||
"widgets_values": [
|
||||
false,
|
||||
"u2net",
|
||||
false,
|
||||
true,
|
||||
true,
|
||||
240,
|
||||
10,
|
||||
10,
|
||||
"none"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 30,
|
||||
"type": "ImageToMask",
|
||||
"pos": [
|
||||
1859,
|
||||
475
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 58
|
||||
},
|
||||
"flags": {},
|
||||
"order": 8,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 39
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
46
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "ImageToMask"
|
||||
},
|
||||
"widgets_values": [
|
||||
"red"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 32,
|
||||
"type": "[Comfy3D] Preview 3DMesh",
|
||||
"pos": [
|
||||
3339,
|
||||
369
|
||||
],
|
||||
"size": [
|
||||
600,
|
||||
500
|
||||
],
|
||||
"flags": {},
|
||||
"order": 12,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mesh_file_path",
|
||||
"type": "STRING",
|
||||
"link": 44,
|
||||
"widget": {
|
||||
"name": "mesh_file_path"
|
||||
}
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Preview 3DMesh"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
null
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 29,
|
||||
"type": "PreviewImage",
|
||||
"pos": [
|
||||
1857,
|
||||
609
|
||||
],
|
||||
"size": {
|
||||
"0": 210,
|
||||
"1": 246
|
||||
},
|
||||
"flags": {},
|
||||
"order": 7,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 38
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "PreviewImage"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 33,
|
||||
"type": "[Comfy3D] Save 3D Mesh",
|
||||
"pos": [
|
||||
2980,
|
||||
370
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 58
|
||||
},
|
||||
"flags": {},
|
||||
"order": 11,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mesh",
|
||||
"type": "MESH",
|
||||
"link": 65
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "save_path",
|
||||
"type": "STRING",
|
||||
"links": [
|
||||
44
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Save 3D Mesh"
|
||||
},
|
||||
"widgets_values": [
|
||||
"NeuS/Cat.obj"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 22,
|
||||
"type": "InvertMask",
|
||||
"pos": [
|
||||
562,
|
||||
468
|
||||
],
|
||||
"size": {
|
||||
"0": 210,
|
||||
"1": 26
|
||||
},
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mask",
|
||||
"type": "MASK",
|
||||
"link": 33
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
34
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "InvertMask"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 19,
|
||||
"type": "PreviewImage",
|
||||
"pos": [
|
||||
1293,
|
||||
538
|
||||
],
|
||||
"size": {
|
||||
"0": 172.59315490722656,
|
||||
"1": 299.3858947753906
|
||||
},
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 31
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "PreviewImage"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
186,
|
||||
391
|
||||
],
|
||||
"size": {
|
||||
"0": 318.80224609375,
|
||||
"1": 348.95855712890625
|
||||
},
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
30
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
33
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 1
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"round_cat.png",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"type": "[Comfy3D] Load Diffusers Pipeline",
|
||||
"pos": [
|
||||
307,
|
||||
216
|
||||
],
|
||||
"size": {
|
||||
"0": 392.1646728515625,
|
||||
"1": 86.24864196777344
|
||||
},
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "pipe",
|
||||
"type": "DIFFUSERS_PIPE",
|
||||
"links": [
|
||||
29
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Load Diffusers Pipeline"
|
||||
},
|
||||
"widgets_values": [
|
||||
"Wonder3DMVDiffusionPipeline",
|
||||
"flamehaze1115/wonder3d-v1.0"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 20,
|
||||
"type": "[Comfy3D] Wonder3D MVDiffusion Model",
|
||||
"pos": [
|
||||
806,
|
||||
370
|
||||
],
|
||||
"size": {
|
||||
"0": 430.8000183105469,
|
||||
"1": 170
|
||||
},
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mvdiffusion_pipe",
|
||||
"type": "DIFFUSERS_PIPE",
|
||||
"link": 29
|
||||
},
|
||||
{
|
||||
"name": "reference_image",
|
||||
"type": "IMAGE",
|
||||
"link": 30
|
||||
},
|
||||
{
|
||||
"name": "reference_mask",
|
||||
"type": "MASK",
|
||||
"link": 34
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "multiview_images",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
31,
|
||||
37,
|
||||
47
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "multiview_normals",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
32,
|
||||
48
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 1
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Wonder3D MVDiffusion Model"
|
||||
},
|
||||
"widgets_values": [
|
||||
999,
|
||||
"fixed",
|
||||
1,
|
||||
30
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 34,
|
||||
"type": "[Comfy3D] NeuS",
|
||||
"pos": [
|
||||
2270,
|
||||
370
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 410
|
||||
},
|
||||
"flags": {},
|
||||
"order": 9,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "reference_image",
|
||||
"type": "IMAGE",
|
||||
"link": 47
|
||||
},
|
||||
{
|
||||
"name": "reference_mask",
|
||||
"type": "MASK",
|
||||
"link": 46
|
||||
},
|
||||
{
|
||||
"name": "reference_normals",
|
||||
"type": "IMAGE",
|
||||
"link": 48
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "mesh",
|
||||
"type": "MESH",
|
||||
"links": [
|
||||
64
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] NeuS"
|
||||
},
|
||||
"widgets_values": [
|
||||
1500,
|
||||
512,
|
||||
0.0005,
|
||||
0.0005,
|
||||
1,
|
||||
0.1,
|
||||
1,
|
||||
1,
|
||||
0.1,
|
||||
500,
|
||||
0,
|
||||
true,
|
||||
512,
|
||||
0
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 44,
|
||||
"type": "[Comfy3D] Switch Mesh Axis",
|
||||
"pos": [
|
||||
2630,
|
||||
370
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 130
|
||||
},
|
||||
"flags": {},
|
||||
"order": 10,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mesh",
|
||||
"type": "MESH",
|
||||
"link": 64
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "switched_mesh",
|
||||
"type": "MESH",
|
||||
"links": [
|
||||
65
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Switch Mesh Axis"
|
||||
},
|
||||
"widgets_values": [
|
||||
"-x",
|
||||
"-z",
|
||||
"+y",
|
||||
false
|
||||
]
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
29,
|
||||
2,
|
||||
0,
|
||||
20,
|
||||
0,
|
||||
"DIFFUSERS_PIPE"
|
||||
],
|
||||
[
|
||||
30,
|
||||
4,
|
||||
0,
|
||||
20,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
31,
|
||||
20,
|
||||
0,
|
||||
19,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
32,
|
||||
20,
|
||||
1,
|
||||
21,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
33,
|
||||
4,
|
||||
1,
|
||||
22,
|
||||
0,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
34,
|
||||
22,
|
||||
0,
|
||||
20,
|
||||
2,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
37,
|
||||
20,
|
||||
0,
|
||||
28,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
38,
|
||||
28,
|
||||
0,
|
||||
29,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
39,
|
||||
28,
|
||||
0,
|
||||
30,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
44,
|
||||
33,
|
||||
0,
|
||||
32,
|
||||
0,
|
||||
"STRING"
|
||||
],
|
||||
[
|
||||
46,
|
||||
30,
|
||||
0,
|
||||
34,
|
||||
1,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
47,
|
||||
20,
|
||||
0,
|
||||
34,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
48,
|
||||
20,
|
||||
1,
|
||||
34,
|
||||
2,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
64,
|
||||
34,
|
||||
0,
|
||||
44,
|
||||
0,
|
||||
"MESH"
|
||||
],
|
||||
[
|
||||
65,
|
||||
44,
|
||||
0,
|
||||
33,
|
||||
0,
|
||||
"MESH"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {},
|
||||
"version": 0.4
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 484 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.0 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 280 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 18 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 131 KiB |
@@ -0,0 +1,504 @@
|
||||
import os
|
||||
import math
|
||||
import numpy as np
|
||||
import cv2 as cv
|
||||
import trimesh
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import tqdm
|
||||
import comfy.utils
|
||||
from pyhocon import ConfigFactory
|
||||
|
||||
from NeuS.models.dataset_mvdiff import Dataset
|
||||
from NeuS.models.fields import RenderingNetwork, SDFNetwork, SingleVarianceNetwork, NeRF
|
||||
from NeuS.models.renderer import NeuSRenderer
|
||||
|
||||
from ..mesh_processer.mesh import Mesh
|
||||
|
||||
|
||||
def ranking_loss(error, penalize_ratio=0.7, type='mean'):
|
||||
error, indices = torch.sort(error)
|
||||
# only sum relatively small errors
|
||||
s_error = torch.index_select(error, 0, index=indices[: int(penalize_ratio * indices.shape[0])])
|
||||
if type == 'mean':
|
||||
return torch.mean(s_error)
|
||||
elif type == 'sum':
|
||||
return torch.sum(s_error)
|
||||
|
||||
class NeuSParams:
|
||||
def __init__(
|
||||
self,
|
||||
end_iter = 1000, # longer time, better result. 1w will be ok for most cases
|
||||
batch_size = 512,
|
||||
learning_rate = 5e-4,
|
||||
learning_rate_alpha = 0.05,
|
||||
color_weight = 1.0,
|
||||
igr_weight = 0.1,
|
||||
mask_weight = 1.0,
|
||||
normal_weight = 1.0,
|
||||
sparse_weight = 0.1,
|
||||
warm_up_end = 500,
|
||||
anneal_end = 0,
|
||||
use_white_bkgd = True,
|
||||
):
|
||||
self.end_iter = end_iter
|
||||
self.batch_size = batch_size
|
||||
self.learning_rate = learning_rate
|
||||
self.learning_rate_alpha = learning_rate_alpha
|
||||
self.color_weight = color_weight
|
||||
self.igr_weight = igr_weight
|
||||
self.mask_weight = mask_weight
|
||||
self.normal_weight = normal_weight
|
||||
self.sparse_weight = sparse_weight
|
||||
self.warm_up_end = warm_up_end
|
||||
self.anneal_end = anneal_end
|
||||
self.use_white_bkgd = use_white_bkgd
|
||||
|
||||
|
||||
class NeuSRunner:
|
||||
def __init__(self, mv_images, mv_masks, mv_normals, conf_path, cam_pose_dir, num_views, NeuS_params=None, mode='train', is_continue=False, device='cuda'):
|
||||
self.device = torch.device(device)
|
||||
|
||||
if NeuS_params is None:
|
||||
NeuS_params = NeuSParams()
|
||||
|
||||
# Configuration
|
||||
self.conf_path = conf_path
|
||||
f = open(self.conf_path)
|
||||
conf_text = f.read()
|
||||
f.close()
|
||||
self.conf = ConfigFactory.parse_string(conf_text)
|
||||
|
||||
self.base_exp_dir = self.conf['general.base_exp_dir']
|
||||
|
||||
# Training parameters
|
||||
self.end_iter = NeuS_params.end_iter
|
||||
self.batch_size = NeuS_params.batch_size
|
||||
self.learning_rate = NeuS_params.learning_rate
|
||||
self.learning_rate_alpha = NeuS_params.learning_rate_alpha
|
||||
self.color_weight = NeuS_params.color_weight
|
||||
self.igr_weight = NeuS_params.igr_weight
|
||||
self.mask_weight = NeuS_params.mask_weight
|
||||
self.normal_weight = NeuS_params.normal_weight
|
||||
self.sparse_weight = NeuS_params.sparse_weight
|
||||
self.warm_up_end = NeuS_params.warm_up_end
|
||||
self.anneal_end = NeuS_params.anneal_end
|
||||
self.use_white_bkgd = NeuS_params.use_white_bkgd
|
||||
|
||||
# Validation parameters
|
||||
self.validate_resolution_level = self.conf.get_int('validate.validate_resolution_level')
|
||||
self.save_freq = self.conf.get_int('validate.save_freq')
|
||||
self.report_freq = self.conf.get_int('validate.report_freq')
|
||||
self.val_freq = self.conf.get_int('validate.val_freq')
|
||||
self.val_mesh_freq = self.conf.get_int('validate.val_mesh_freq')
|
||||
|
||||
self.is_continue = is_continue
|
||||
self.mode = mode
|
||||
self.model_list = []
|
||||
|
||||
# Networks
|
||||
params_to_train_slow = []
|
||||
self.nerf_outside = NeRF(**self.conf['model.nerf']).to(self.device)
|
||||
self.sdf_network = SDFNetwork(**self.conf['model.sdf_network']).to(self.device)
|
||||
self.deviation_network = SingleVarianceNetwork(**self.conf['model.variance_network']).to(self.device)
|
||||
self.color_network = RenderingNetwork(**self.conf['model.rendering_network']).to(self.device)
|
||||
# params_to_train += list(self.nerf_outside.parameters())
|
||||
params_to_train_slow += list(self.sdf_network.parameters())
|
||||
params_to_train_slow += list(self.deviation_network.parameters())
|
||||
# params_to_train += list(self.color_network.parameters())
|
||||
|
||||
self.optimizer = torch.optim.Adam(
|
||||
[{'params': params_to_train_slow}, {'params': self.color_network.parameters(), 'lr': self.learning_rate * 2}], lr=self.learning_rate
|
||||
)
|
||||
|
||||
self.renderer = NeuSRenderer(
|
||||
self.nerf_outside, self.sdf_network, self.deviation_network, self.color_network, **self.conf['model.neus_renderer']
|
||||
)
|
||||
|
||||
# Prepare dataset
|
||||
self.dataset = Dataset(mv_images, mv_masks, mv_normals, self.conf['dataset'], cam_pose_dir, num_views)
|
||||
self.iter_step = 1
|
||||
|
||||
# Load checkpoint
|
||||
latest_model_name = None
|
||||
if is_continue:
|
||||
model_list_raw = os.listdir(os.path.join(self.base_exp_dir, 'checkpoints'))
|
||||
model_list = []
|
||||
for model_name in model_list_raw:
|
||||
if model_name[-3:] == 'pth' and int(model_name[5:-4]) <= self.end_iter:
|
||||
model_list.append(model_name)
|
||||
model_list.sort()
|
||||
latest_model_name = model_list[-1]
|
||||
|
||||
if latest_model_name is not None:
|
||||
self.load_checkpoint(latest_model_name)
|
||||
|
||||
def train(self):
|
||||
self.update_learning_rate()
|
||||
res_step = self.end_iter - self.iter_step
|
||||
|
||||
comfy_pbar = comfy.utils.ProgressBar(self.end_iter)
|
||||
|
||||
for epoch in tqdm.trange(res_step):
|
||||
data = self.dataset.get_batch_of_random_rays(self.batch_size).to(self.device)
|
||||
|
||||
rays_o, rays_d, true_rgb, mask, true_normal, cosines = (
|
||||
data[:, :3],
|
||||
data[:, 3:6],
|
||||
data[:, 6:9],
|
||||
data[:, 9:10],
|
||||
data[:, 10:13],
|
||||
data[:, 13:],
|
||||
)
|
||||
|
||||
# near, far = self.dataset.near_far_from_sphere(rays_o, rays_d)
|
||||
near, far = self.dataset.get_near_far()
|
||||
|
||||
background_rgb = None
|
||||
if self.use_white_bkgd:
|
||||
background_rgb = torch.ones([1, 3])
|
||||
|
||||
if self.mask_weight > 0.0:
|
||||
mask = (mask > 0.5).float()
|
||||
else:
|
||||
mask = torch.ones_like(mask)
|
||||
|
||||
cosines[cosines > -0.1] = 0
|
||||
mask = ((mask > 0) & (cosines < -0.1)).to(torch.float32)
|
||||
#mask_sum = mask.sum() + 1e-5
|
||||
render_out = self.renderer.render(
|
||||
rays_o, rays_d, near, far, background_rgb=background_rgb, cos_anneal_ratio=self.get_cos_anneal_ratio()
|
||||
)
|
||||
|
||||
color_fine = render_out['color_fine']
|
||||
s_val = render_out['s_val']
|
||||
#cdf_fine = render_out['cdf_fine']
|
||||
gradient_error = render_out['gradient_error']
|
||||
#weight_max = render_out['weight_max']
|
||||
weight_sum = render_out['weight_sum']
|
||||
|
||||
# Loss
|
||||
# color_error = (color_fine - true_rgb) * mask
|
||||
# color_fine_loss = F.l1_loss(color_error, torch.zeros_like(color_error), reduction='sum') / mask_sum
|
||||
|
||||
color_errors = (color_fine - true_rgb).abs().sum(dim=1)
|
||||
color_fine_loss = ranking_loss(color_errors[mask[:, 0] > 0])
|
||||
|
||||
#psnr = 20.0 * torch.log10(1.0 / (((color_fine - true_rgb) ** 2 * mask).sum() / (mask_sum * 3.0)).sqrt())
|
||||
|
||||
eikonal_loss = gradient_error
|
||||
|
||||
# pdb.set_trace()
|
||||
mask_errors = F.binary_cross_entropy(weight_sum.clip(1e-3, 1.0 - 1e-3), mask, reduction='none')
|
||||
mask_loss = ranking_loss(mask_errors[:, 0], penalize_ratio=0.8)
|
||||
|
||||
def feasible(key):
|
||||
return (key in render_out) and (render_out[key] is not None)
|
||||
|
||||
# calculate normal loss
|
||||
n_samples = self.renderer.n_samples + self.renderer.n_importance
|
||||
normals = render_out['gradients'] * render_out['weights'][:, :n_samples, None]
|
||||
if feasible('inside_sphere'):
|
||||
normals = normals * render_out['inside_sphere'][..., None]
|
||||
normals = normals.sum(dim=1)
|
||||
|
||||
# pdb.set_trace()
|
||||
normal_errors = 1 - F.cosine_similarity(normals, true_normal, dim=1)
|
||||
# normal_error = normal_error * mask[:, 0]
|
||||
# normal_loss = F.l1_loss(normal_error, torch.zeros_like(normal_error), reduction='sum') / mask_sum
|
||||
normal_errors = normal_errors * torch.exp(cosines.abs()[:, 0]) / torch.exp(cosines.abs()).sum()
|
||||
normal_loss = ranking_loss(normal_errors[mask[:, 0] > 0], penalize_ratio=0.9, type='sum')
|
||||
|
||||
sparse_loss = render_out['sparse_loss']
|
||||
|
||||
loss = (
|
||||
color_fine_loss * self.color_weight
|
||||
+ eikonal_loss * self.igr_weight
|
||||
+ sparse_loss * self.sparse_weight
|
||||
+ mask_loss * self.mask_weight
|
||||
+ normal_loss * self.normal_weight
|
||||
)
|
||||
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
self.optimizer.step()
|
||||
|
||||
if self.iter_step % self.report_freq == 0:
|
||||
print(
|
||||
'iter:{:8>d} loss = {:4>f} color_ls = {:4>f} eik_ls = {:4>f} normal_ls = {:4>f} mask_ls = {:4>f} sparse_ls = {:4>f} lr={:5>f}'.format(
|
||||
self.iter_step,
|
||||
loss,
|
||||
color_fine_loss,
|
||||
eikonal_loss,
|
||||
normal_loss,
|
||||
mask_loss,
|
||||
sparse_loss,
|
||||
self.optimizer.param_groups[0]['lr'],
|
||||
)
|
||||
)
|
||||
print('iter:{:8>d} s_val = {:4>f}'.format(self.iter_step, s_val.mean()))
|
||||
|
||||
if self.iter_step % self.val_mesh_freq == 0:
|
||||
self.extract_mesh(resolution=256)
|
||||
|
||||
self.update_learning_rate()
|
||||
|
||||
self.iter_step += 1
|
||||
|
||||
if self.iter_step % self.val_freq == 0:
|
||||
self.validate_image(idx=0)
|
||||
self.validate_image(idx=1)
|
||||
self.validate_image(idx=2)
|
||||
self.validate_image(idx=3)
|
||||
|
||||
if self.iter_step % self.save_freq == 0:
|
||||
self.save_checkpoint()
|
||||
|
||||
comfy_pbar.update_absolute(self.iter_step + 1)
|
||||
|
||||
def get_cos_anneal_ratio(self):
|
||||
if self.anneal_end == 0.0:
|
||||
return 1.0
|
||||
else:
|
||||
return np.min([1.0, self.iter_step / self.anneal_end])
|
||||
|
||||
def update_learning_rate(self):
|
||||
if self.iter_step < self.warm_up_end:
|
||||
learning_factor = self.iter_step / self.warm_up_end
|
||||
else:
|
||||
alpha = self.learning_rate_alpha
|
||||
progress = (self.iter_step - self.warm_up_end) / (self.end_iter - self.warm_up_end)
|
||||
learning_factor = (np.cos(np.pi * progress) + 1.0) * 0.5 * (1 - alpha) + alpha
|
||||
|
||||
for g in self.optimizer.param_groups:
|
||||
g['lr'] = self.learning_rate * learning_factor
|
||||
|
||||
def extract_mesh(self, world_space=False, resolution=64, threshold=0.0):
|
||||
bound_min = torch.tensor(self.dataset.object_bbox_min, dtype=torch.float32)
|
||||
bound_max = torch.tensor(self.dataset.object_bbox_max, dtype=torch.float32)
|
||||
|
||||
vertices, triangles, vertex_colors = self.renderer.extract_geometry(bound_min, bound_max, resolution=resolution, threshold=threshold)
|
||||
|
||||
if world_space:
|
||||
vertices = vertices * self.dataset.scale_mats_np[0][0, 0] + self.dataset.scale_mats_np[0][:3, 3][None]
|
||||
|
||||
print(f"vertex_colors: {vertex_colors[123]}, {vertex_colors[213]}, {vertex_colors[321]}")
|
||||
|
||||
mesh = trimesh.Trimesh(vertices, triangles, vertex_colors=vertex_colors)
|
||||
mesh = Mesh.load_trimesh(given_mesh=mesh)
|
||||
|
||||
# export as glb
|
||||
# os.makedirs(os.path.join(self.base_exp_dir, 'meshes'), exist_ok=True)
|
||||
# mesh.export(os.path.join(self.base_exp_dir, 'meshes', '{:0>8d}.ply'.format(self.iter_step)))
|
||||
|
||||
return mesh
|
||||
|
||||
def load_checkpoint(self, checkpoint_name):
|
||||
checkpoint = torch.load(os.path.join(self.base_exp_dir, 'checkpoints', checkpoint_name), map_location=self.device)
|
||||
self.nerf_outside.load_state_dict(checkpoint['nerf'])
|
||||
self.sdf_network.load_state_dict(checkpoint['sdf_network_fine'])
|
||||
self.deviation_network.load_state_dict(checkpoint['variance_network_fine'])
|
||||
self.color_network.load_state_dict(checkpoint['color_network_fine'])
|
||||
self.optimizer.load_state_dict(checkpoint['optimizer'])
|
||||
self.iter_step = checkpoint['iter_step']
|
||||
|
||||
def save_checkpoint(self):
|
||||
checkpoint = {
|
||||
'nerf': self.nerf_outside.state_dict(),
|
||||
'sdf_network_fine': self.sdf_network.state_dict(),
|
||||
'variance_network_fine': self.deviation_network.state_dict(),
|
||||
'color_network_fine': self.color_network.state_dict(),
|
||||
'optimizer': self.optimizer.state_dict(),
|
||||
'iter_step': self.iter_step,
|
||||
}
|
||||
|
||||
os.makedirs(os.path.join(self.base_exp_dir, 'checkpoints'), exist_ok=True)
|
||||
torch.save(checkpoint, os.path.join(self.base_exp_dir, 'checkpoints', 'ckpt_{:0>6d}.pth'.format(self.iter_step)))
|
||||
|
||||
def validate_image(self, idx=-1, resolution_level=-1):
|
||||
if idx < 0:
|
||||
idx = np.random.randint(self.dataset.n_images)
|
||||
|
||||
print('Validate: iter: {}, camera: {}'.format(self.iter_step, idx))
|
||||
|
||||
if resolution_level < 0:
|
||||
resolution_level = self.validate_resolution_level
|
||||
rays_o, rays_d = self.dataset.gen_rays_at(idx, resolution_level=resolution_level)
|
||||
H, W, _ = rays_o.shape
|
||||
rays_o = rays_o.reshape(-1, 3).split(self.batch_size)
|
||||
rays_d = rays_d.reshape(-1, 3).split(self.batch_size)
|
||||
|
||||
out_rgb_fine = []
|
||||
out_normal_fine = []
|
||||
out_mask = []
|
||||
|
||||
for rays_o_batch, rays_d_batch in zip(rays_o, rays_d):
|
||||
# near, far = self.dataset.near_far_from_sphere(rays_o_batch, rays_d_batch)
|
||||
near, far = self.dataset.get_near_far()
|
||||
background_rgb = torch.ones([1, 3]) if self.use_white_bkgd else None
|
||||
|
||||
render_out = self.renderer.render(
|
||||
rays_o_batch, rays_d_batch, near, far, cos_anneal_ratio=self.get_cos_anneal_ratio(), background_rgb=background_rgb
|
||||
)
|
||||
|
||||
def feasible(key):
|
||||
return (key in render_out) and (render_out[key] is not None)
|
||||
|
||||
if feasible('color_fine'):
|
||||
out_rgb_fine.append(render_out['color_fine'].detach().cpu().numpy())
|
||||
if feasible('gradients') and feasible('weights'):
|
||||
n_samples = self.renderer.n_samples + self.renderer.n_importance
|
||||
normals = render_out['gradients'] * render_out['weights'][:, :n_samples, None]
|
||||
if feasible('inside_sphere'):
|
||||
normals = normals * render_out['inside_sphere'][..., None]
|
||||
normals = normals.sum(dim=1).detach().cpu().numpy()
|
||||
out_normal_fine.append(normals)
|
||||
|
||||
if feasible('weight_sum'):
|
||||
out_mask.append(render_out['weight_sum'].detach().clip(0, 1).cpu().numpy())
|
||||
|
||||
del render_out
|
||||
|
||||
img_fine = None
|
||||
if len(out_rgb_fine) > 0:
|
||||
img_fine = (np.concatenate(out_rgb_fine, axis=0).reshape([H, W, 3, -1]) * 256).clip(0, 255)
|
||||
|
||||
mask_map = None
|
||||
if len(out_mask) > 0:
|
||||
mask_map = (np.concatenate(out_mask, axis=0).reshape([H, W, -1]) * 256).clip(0, 255)
|
||||
|
||||
normal_img = None
|
||||
if len(out_normal_fine) > 0:
|
||||
normal_img = np.concatenate(out_normal_fine, axis=0)
|
||||
rot = np.linalg.inv(self.dataset.pose_all[idx, :3, :3].detach().cpu().numpy())
|
||||
normal_img = (np.matmul(rot[None, :, :], normal_img[:, :, None]).reshape([H, W, 3, -1]) * 128 + 128).clip(0, 255)
|
||||
|
||||
os.makedirs(os.path.join(self.base_exp_dir, 'validations_fine'), exist_ok=True)
|
||||
os.makedirs(os.path.join(self.base_exp_dir, 'normals'), exist_ok=True)
|
||||
|
||||
for i in range(img_fine.shape[-1]):
|
||||
if len(out_rgb_fine) > 0:
|
||||
cv.imwrite(
|
||||
os.path.join(self.base_exp_dir, 'validations_fine', '{:0>8d}_{}_{}.png'.format(self.iter_step, i, idx)),
|
||||
np.concatenate(
|
||||
[
|
||||
img_fine[..., i],
|
||||
self.dataset.image_at(idx, resolution_level=resolution_level),
|
||||
self.dataset.mask_at(idx, resolution_level=resolution_level),
|
||||
]
|
||||
),
|
||||
)
|
||||
if len(out_normal_fine) > 0:
|
||||
cv.imwrite(
|
||||
os.path.join(self.base_exp_dir, 'normals', '{:0>8d}_{}_{}.png'.format(self.iter_step, i, idx)),
|
||||
np.concatenate([normal_img[..., i], self.dataset.normal_cam_at(idx, resolution_level=resolution_level)])[:, :, ::-1],
|
||||
)
|
||||
if len(out_mask) > 0:
|
||||
cv.imwrite(os.path.join(self.base_exp_dir, 'normals', '{:0>8d}_{}_{}_mask.png'.format(self.iter_step, i, idx)), mask_map[..., i])
|
||||
|
||||
def save_maps(self, idx, img_idx, resolution_level=1):
|
||||
view_types = ['front', 'back', 'left', 'right']
|
||||
print('Validate: iter: {}, camera: {}'.format(self.iter_step, idx))
|
||||
|
||||
rays_o, rays_d = self.dataset.gen_rays_at(idx, resolution_level=resolution_level)
|
||||
H, W, _ = rays_o.shape
|
||||
rays_o = rays_o.reshape(-1, 3).split(self.batch_size)
|
||||
rays_d = rays_d.reshape(-1, 3).split(self.batch_size)
|
||||
|
||||
out_rgb_fine = []
|
||||
out_normal_fine = []
|
||||
out_mask = []
|
||||
|
||||
for rays_o_batch, rays_d_batch in zip(rays_o, rays_d):
|
||||
# near, far = self.dataset.near_far_from_sphere(rays_o_batch, rays_d_batch)
|
||||
near, far = self.dataset.get_near_far()
|
||||
background_rgb = torch.ones([1, 3]) if self.use_white_bkgd else None
|
||||
|
||||
render_out = self.renderer.render(
|
||||
rays_o_batch, rays_d_batch, near, far, cos_anneal_ratio=self.get_cos_anneal_ratio(), background_rgb=background_rgb
|
||||
)
|
||||
|
||||
def feasible(key):
|
||||
return (key in render_out) and (render_out[key] is not None)
|
||||
|
||||
if feasible('color_fine'):
|
||||
out_rgb_fine.append(render_out['color_fine'].detach().cpu().numpy())
|
||||
if feasible('gradients') and feasible('weights'):
|
||||
n_samples = self.renderer.n_samples + self.renderer.n_importance
|
||||
normals = render_out['gradients'] * render_out['weights'][:, :n_samples, None]
|
||||
if feasible('inside_sphere'):
|
||||
normals = normals * render_out['inside_sphere'][..., None]
|
||||
normals = normals.sum(dim=1).detach().cpu().numpy()
|
||||
out_normal_fine.append(normals)
|
||||
|
||||
if feasible('weight_sum'):
|
||||
out_mask.append(render_out['weight_sum'].detach().clip(0, 1).cpu().numpy())
|
||||
|
||||
del render_out
|
||||
|
||||
img_fine = None
|
||||
if len(out_rgb_fine) > 0:
|
||||
img_fine = (np.concatenate(out_rgb_fine, axis=0).reshape([H, W, 3]) * 256).clip(0, 255)
|
||||
|
||||
mask_map = None
|
||||
if len(out_mask) > 0:
|
||||
mask_map = (np.concatenate(out_mask, axis=0).reshape([H, W, 1]) * 256).clip(0, 255)
|
||||
|
||||
normal_img = None
|
||||
if len(out_normal_fine) > 0:
|
||||
normal_img = np.concatenate(out_normal_fine, axis=0)
|
||||
# rot = np.linalg.inv(self.dataset.pose_all[idx, :3, :3].detach().cpu().numpy())
|
||||
world_normal_img = (normal_img.reshape([H, W, 3]) * 128 + 128).clip(0, 255)
|
||||
|
||||
os.makedirs(os.path.join(self.base_exp_dir, 'coarse_maps'), exist_ok=True)
|
||||
img_rgba = np.concatenate([img_fine[:, :, ::-1], mask_map], axis=-1)
|
||||
normal_rgba = np.concatenate([world_normal_img[:, :, ::-1], mask_map], axis=-1)
|
||||
|
||||
cv.imwrite(os.path.join(self.base_exp_dir, 'coarse_maps', "normals_mlp_%03d_%s.png" % (img_idx, view_types[idx])), img_rgba)
|
||||
cv.imwrite(os.path.join(self.base_exp_dir, 'coarse_maps', "normals_grad_%03d_%s.png" % (img_idx, view_types[idx])), normal_rgba)
|
||||
|
||||
def render_novel_image(self, idx_0, idx_1, ratio, resolution_level):
|
||||
"""
|
||||
Interpolate view between two cameras.
|
||||
"""
|
||||
rays_o, rays_d = self.dataset.gen_rays_between(idx_0, idx_1, ratio, resolution_level=resolution_level)
|
||||
H, W, _ = rays_o.shape
|
||||
rays_o = rays_o.reshape(-1, 3).split(self.batch_size)
|
||||
rays_d = rays_d.reshape(-1, 3).split(self.batch_size)
|
||||
|
||||
out_rgb_fine = []
|
||||
for rays_o_batch, rays_d_batch in zip(rays_o, rays_d):
|
||||
# near, far = self.dataset.near_far_from_sphere(rays_o_batch, rays_d_batch)
|
||||
near, far = self.dataset.get_near_far()
|
||||
background_rgb = torch.ones([1, 3]) if self.use_white_bkgd else None
|
||||
|
||||
render_out = self.renderer.render(
|
||||
rays_o_batch, rays_d_batch, near, far, cos_anneal_ratio=self.get_cos_anneal_ratio(), background_rgb=background_rgb
|
||||
)
|
||||
|
||||
out_rgb_fine.append(render_out['color_fine'].detach().cpu().numpy())
|
||||
|
||||
del render_out
|
||||
|
||||
img_fine = (np.concatenate(out_rgb_fine, axis=0).reshape([H, W, 3]) * 256).clip(0, 255).astype(np.uint8)
|
||||
return img_fine
|
||||
|
||||
def interpolate_view(self, img_idx_0, img_idx_1):
|
||||
images = []
|
||||
n_frames = 60
|
||||
for i in range(n_frames):
|
||||
print(i)
|
||||
images.append(self.render_novel_image(img_idx_0, img_idx_1, np.sin(((i / n_frames) - 0.5) * np.pi) * 0.5 + 0.5, resolution_level=4))
|
||||
for i in range(n_frames):
|
||||
images.append(images[n_frames - i - 1])
|
||||
|
||||
fourcc = cv.VideoWriter_fourcc(*'mp4v')
|
||||
video_dir = os.path.join(self.base_exp_dir, 'render')
|
||||
os.makedirs(video_dir, exist_ok=True)
|
||||
h, w, _ = images[0].shape
|
||||
writer = cv.VideoWriter(os.path.join(video_dir, '{:0>8d}_{}_{}.mp4'.format(self.iter_step, img_idx_0, img_idx_1)), fourcc, 30, (w, h))
|
||||
|
||||
for image in images:
|
||||
writer.write(image)
|
||||
|
||||
writer.release()
|
||||
+27
-25
@@ -14,31 +14,33 @@ from ..shared_utils.camera_utils import OrbitCamera, MiniCam, calculate_fovX, ge
|
||||
from ..shared_utils.image_utils import prepare_torch_img
|
||||
|
||||
class GSParams:
|
||||
def __init__(self,
|
||||
training_iterations=30_000,
|
||||
batch_size=1,
|
||||
lambda_ssim=0.2,
|
||||
lambda_alpha=3,
|
||||
lambda_offset=0,
|
||||
lambda_offset_opacity=0,
|
||||
invert_bg_prob = 0.5,
|
||||
feature_lr=0.0025,
|
||||
opacity_lr=0.05,
|
||||
scaling_lr=0.005,
|
||||
rotation_lr=0.001,
|
||||
position_lr_init=0.00016,
|
||||
position_lr_final=0.0000016,
|
||||
position_lr_delay_mult=0.01,
|
||||
position_lr_max_steps=30_000,
|
||||
num_pts=10_000,
|
||||
K=3,
|
||||
percent_dense=0.01,
|
||||
density_start_iter=500,
|
||||
density_end_iter=15_000,
|
||||
densification_interval=100,
|
||||
opacity_reset_interval=3000,
|
||||
densify_grad_threshold=0.0002,
|
||||
sh_degree=3):
|
||||
def __init__(
|
||||
self,
|
||||
training_iterations=30_000,
|
||||
batch_size=1,
|
||||
lambda_ssim=0.2,
|
||||
lambda_alpha=3,
|
||||
lambda_offset=0,
|
||||
lambda_offset_opacity=0,
|
||||
invert_bg_prob = 0.5,
|
||||
feature_lr=0.0025,
|
||||
opacity_lr=0.05,
|
||||
scaling_lr=0.005,
|
||||
rotation_lr=0.001,
|
||||
position_lr_init=0.00016,
|
||||
position_lr_final=0.0000016,
|
||||
position_lr_delay_mult=0.01,
|
||||
position_lr_max_steps=30_000,
|
||||
num_pts=10_000,
|
||||
K=3,
|
||||
percent_dense=0.01,
|
||||
density_start_iter=500,
|
||||
density_end_iter=15_000,
|
||||
densification_interval=100,
|
||||
opacity_reset_interval=3000,
|
||||
densify_grad_threshold=0.0002,
|
||||
sh_degree=3
|
||||
):
|
||||
|
||||
# training params
|
||||
self.training_iterations = training_iterations
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
*.7z filter=lfs diff=lfs merge=lfs -text
|
||||
*.arrow filter=lfs diff=lfs merge=lfs -text
|
||||
*.bin filter=lfs diff=lfs merge=lfs -text
|
||||
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
||||
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
||||
*.ftz filter=lfs diff=lfs merge=lfs -text
|
||||
*.gz filter=lfs diff=lfs merge=lfs -text
|
||||
*.h5 filter=lfs diff=lfs merge=lfs -text
|
||||
*.joblib filter=lfs diff=lfs merge=lfs -text
|
||||
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
||||
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
||||
*.model filter=lfs diff=lfs merge=lfs -text
|
||||
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
||||
*.npy filter=lfs diff=lfs merge=lfs -text
|
||||
*.npz filter=lfs diff=lfs merge=lfs -text
|
||||
*.onnx filter=lfs diff=lfs merge=lfs -text
|
||||
*.ot filter=lfs diff=lfs merge=lfs -text
|
||||
*.parquet filter=lfs diff=lfs merge=lfs -text
|
||||
*.pb filter=lfs diff=lfs merge=lfs -text
|
||||
*.pickle filter=lfs diff=lfs merge=lfs -text
|
||||
*.pkl filter=lfs diff=lfs merge=lfs -text
|
||||
*.pt filter=lfs diff=lfs merge=lfs -text
|
||||
*.pth filter=lfs diff=lfs merge=lfs -text
|
||||
*.rar filter=lfs diff=lfs merge=lfs -text
|
||||
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
||||
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
||||
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
||||
*.tar filter=lfs diff=lfs merge=lfs -text
|
||||
*.tflite filter=lfs diff=lfs merge=lfs -text
|
||||
*.tgz filter=lfs diff=lfs merge=lfs -text
|
||||
*.wasm filter=lfs diff=lfs merge=lfs -text
|
||||
*.xz filter=lfs diff=lfs merge=lfs -text
|
||||
*.zip filter=lfs diff=lfs merge=lfs -text
|
||||
*.zst filter=lfs diff=lfs merge=lfs -text
|
||||
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
||||
@@ -0,0 +1,18 @@
|
||||
---
|
||||
license: openrail
|
||||
---
|
||||
# MVDream-diffusers Model Card
|
||||
This is a port of https://huggingface.co/Peng-Wang/ImageDream into diffusers.
|
||||
|
||||
For usage, please check: https://github.com/ashawkey/mvdream_diffusers
|
||||
## Citation
|
||||
```
|
||||
@article{wang2023imagedream,
|
||||
title={ImageDream: Image-Prompt Multi-view Diffusion for 3D Generation},
|
||||
author={Wang, Peng and Shi, Yichun},
|
||||
journal={arXiv preprint arXiv:2312.02201},
|
||||
year={2023}
|
||||
}
|
||||
```
|
||||
## Misuse, Malicious Use, and Out-of-Scope Use
|
||||
The model should not be used to intentionally create or disseminate images that create hostile or alienating environments for people. This includes generating images that people would foreseeably find disturbing, distressing, or offensive; or content that propagates historical or current stereotypes.
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
{
|
||||
"crop_size": {
|
||||
"height": 224,
|
||||
"width": 224
|
||||
},
|
||||
"do_center_crop": true,
|
||||
"do_convert_rgb": true,
|
||||
"do_normalize": true,
|
||||
"do_rescale": true,
|
||||
"do_resize": true,
|
||||
"feature_extractor_type": "CLIPFeatureExtractor",
|
||||
"image_mean": [
|
||||
0.48145466,
|
||||
0.4578275,
|
||||
0.40821073
|
||||
],
|
||||
"image_processor_type": "CLIPImageProcessor",
|
||||
"image_std": [
|
||||
0.26862954,
|
||||
0.26130258,
|
||||
0.27577711
|
||||
],
|
||||
"resample": 3,
|
||||
"rescale_factor": 0.00392156862745098,
|
||||
"size": {
|
||||
"shortest_edge": 224
|
||||
},
|
||||
"use_square_size": false
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
{
|
||||
"_name_or_path": "laion/CLIP-ViT-H-14-laion2B-s32B-b79K",
|
||||
"architectures": [
|
||||
"CLIPVisionModel"
|
||||
],
|
||||
"attention_dropout": 0.0,
|
||||
"dropout": 0.0,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_size": 1280,
|
||||
"image_size": 224,
|
||||
"initializer_factor": 1.0,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 5120,
|
||||
"layer_norm_eps": 1e-05,
|
||||
"model_type": "clip_vision_model",
|
||||
"num_attention_heads": 16,
|
||||
"num_channels": 3,
|
||||
"num_hidden_layers": 32,
|
||||
"patch_size": 14,
|
||||
"projection_dim": 1024,
|
||||
"torch_dtype": "float16",
|
||||
"transformers_version": "4.35.2"
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
{
|
||||
"_class_name": "MVDreamPipeline",
|
||||
"_diffusers_version": "0.25.0",
|
||||
"feature_extractor": [
|
||||
"transformers",
|
||||
"CLIPImageProcessor"
|
||||
],
|
||||
"image_encoder": [
|
||||
"transformers",
|
||||
"CLIPVisionModel"
|
||||
],
|
||||
"requires_safety_checker": false,
|
||||
"scheduler": [
|
||||
"diffusers",
|
||||
"DDIMScheduler"
|
||||
],
|
||||
"text_encoder": [
|
||||
"transformers",
|
||||
"CLIPTextModel"
|
||||
],
|
||||
"tokenizer": [
|
||||
"transformers",
|
||||
"CLIPTokenizer"
|
||||
],
|
||||
"unet": [
|
||||
"mv_unet",
|
||||
"MultiViewUNetModel"
|
||||
],
|
||||
"vae": [
|
||||
"diffusers",
|
||||
"AutoencoderKL"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,558 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import inspect
|
||||
import numpy as np
|
||||
from typing import Callable, List, Optional, Union
|
||||
from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModel, CLIPImageProcessor
|
||||
from diffusers import AutoencoderKL, DiffusionPipeline
|
||||
from diffusers.utils import (
|
||||
deprecate,
|
||||
is_accelerate_available,
|
||||
is_accelerate_version,
|
||||
logging,
|
||||
)
|
||||
from diffusers.configuration_utils import FrozenDict
|
||||
from diffusers.schedulers import DDIMScheduler
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
|
||||
from mv_unet import MultiViewUNetModel, get_camera
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
class MVDreamPipeline(DiffusionPipeline):
|
||||
|
||||
_optional_components = ["feature_extractor", "image_encoder"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vae: AutoencoderKL,
|
||||
unet: MultiViewUNetModel,
|
||||
tokenizer: CLIPTokenizer,
|
||||
text_encoder: CLIPTextModel,
|
||||
scheduler: DDIMScheduler,
|
||||
# imagedream variant
|
||||
feature_extractor: CLIPImageProcessor,
|
||||
image_encoder: CLIPVisionModel,
|
||||
requires_safety_checker: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1: # type: ignore
|
||||
deprecation_message = (
|
||||
f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
|
||||
f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " # type: ignore
|
||||
"to update the config accordingly as leaving `steps_offset` might led to incorrect results"
|
||||
" in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
|
||||
" it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
|
||||
" file"
|
||||
)
|
||||
deprecate(
|
||||
"steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False
|
||||
)
|
||||
new_config = dict(scheduler.config)
|
||||
new_config["steps_offset"] = 1
|
||||
scheduler._internal_dict = FrozenDict(new_config)
|
||||
|
||||
if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True: # type: ignore
|
||||
deprecation_message = (
|
||||
f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
|
||||
" `clip_sample` should be set to False in the configuration file. Please make sure to update the"
|
||||
" config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
|
||||
" future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
|
||||
" nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
|
||||
)
|
||||
deprecate(
|
||||
"clip_sample not set", "1.0.0", deprecation_message, standard_warn=False
|
||||
)
|
||||
new_config = dict(scheduler.config)
|
||||
new_config["clip_sample"] = False
|
||||
scheduler._internal_dict = FrozenDict(new_config)
|
||||
|
||||
self.register_modules(
|
||||
vae=vae,
|
||||
unet=unet,
|
||||
scheduler=scheduler,
|
||||
tokenizer=tokenizer,
|
||||
text_encoder=text_encoder,
|
||||
feature_extractor=feature_extractor,
|
||||
image_encoder=image_encoder,
|
||||
)
|
||||
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
|
||||
self.register_to_config(requires_safety_checker=requires_safety_checker)
|
||||
|
||||
def enable_vae_slicing(self):
|
||||
r"""
|
||||
Enable sliced VAE decoding.
|
||||
|
||||
When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several
|
||||
steps. This is useful to save some memory and allow larger batch sizes.
|
||||
"""
|
||||
self.vae.enable_slicing()
|
||||
|
||||
def disable_vae_slicing(self):
|
||||
r"""
|
||||
Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to
|
||||
computing decoding in one step.
|
||||
"""
|
||||
self.vae.disable_slicing()
|
||||
|
||||
def enable_vae_tiling(self):
|
||||
r"""
|
||||
Enable tiled VAE decoding.
|
||||
|
||||
When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in
|
||||
several steps. This is useful to save a large amount of memory and to allow the processing of larger images.
|
||||
"""
|
||||
self.vae.enable_tiling()
|
||||
|
||||
def disable_vae_tiling(self):
|
||||
r"""
|
||||
Disable tiled VAE decoding. If `enable_vae_tiling` was previously invoked, this method will go back to
|
||||
computing decoding in one step.
|
||||
"""
|
||||
self.vae.disable_tiling()
|
||||
|
||||
def enable_sequential_cpu_offload(self, gpu_id=0):
|
||||
r"""
|
||||
Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
|
||||
text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
|
||||
`torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
|
||||
Note that offloading happens on a submodule basis. Memory savings are higher than with
|
||||
`enable_model_cpu_offload`, but performance is lower.
|
||||
"""
|
||||
if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):
|
||||
from accelerate import cpu_offload
|
||||
else:
|
||||
raise ImportError(
|
||||
"`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher"
|
||||
)
|
||||
|
||||
device = torch.device(f"cuda:{gpu_id}")
|
||||
|
||||
if self.device.type != "cpu":
|
||||
self.to("cpu", silence_dtype_warnings=True)
|
||||
torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
|
||||
|
||||
for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
|
||||
cpu_offload(cpu_offloaded_model, device)
|
||||
|
||||
def enable_model_cpu_offload(self, gpu_id=0):
|
||||
r"""
|
||||
Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
|
||||
to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
|
||||
method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
|
||||
`enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
|
||||
"""
|
||||
if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
|
||||
from accelerate import cpu_offload_with_hook
|
||||
else:
|
||||
raise ImportError(
|
||||
"`enable_model_offload` requires `accelerate v0.17.0` or higher."
|
||||
)
|
||||
|
||||
device = torch.device(f"cuda:{gpu_id}")
|
||||
|
||||
if self.device.type != "cpu":
|
||||
self.to("cpu", silence_dtype_warnings=True)
|
||||
torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
|
||||
|
||||
hook = None
|
||||
for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:
|
||||
_, hook = cpu_offload_with_hook(
|
||||
cpu_offloaded_model, device, prev_module_hook=hook
|
||||
)
|
||||
|
||||
# We'll offload the last model manually.
|
||||
self.final_offload_hook = hook
|
||||
|
||||
@property
|
||||
def _execution_device(self):
|
||||
r"""
|
||||
Returns the device on which the pipeline's models will be executed. After calling
|
||||
`pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
|
||||
hooks.
|
||||
"""
|
||||
if not hasattr(self.unet, "_hf_hook"):
|
||||
return self.device
|
||||
for module in self.unet.modules():
|
||||
if (
|
||||
hasattr(module, "_hf_hook")
|
||||
and hasattr(module._hf_hook, "execution_device")
|
||||
and module._hf_hook.execution_device is not None
|
||||
):
|
||||
return torch.device(module._hf_hook.execution_device)
|
||||
return self.device
|
||||
|
||||
def _encode_prompt(
|
||||
self,
|
||||
prompt,
|
||||
device,
|
||||
num_images_per_prompt,
|
||||
do_classifier_free_guidance: bool,
|
||||
negative_prompt=None,
|
||||
):
|
||||
r"""
|
||||
Encodes the prompt into text encoder hidden states.
|
||||
|
||||
Args:
|
||||
prompt (`str` or `List[str]`, *optional*):
|
||||
prompt to be encoded
|
||||
device: (`torch.device`):
|
||||
torch device
|
||||
num_images_per_prompt (`int`):
|
||||
number of images that should be generated per prompt
|
||||
do_classifier_free_guidance (`bool`):
|
||||
whether to use classifier free guidance or not
|
||||
negative_prompt (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts not to guide the image generation. If not defined, one has to pass
|
||||
`negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
|
||||
Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
|
||||
prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
|
||||
provided, text embeddings will be generated from `prompt` input argument.
|
||||
negative_prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
|
||||
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
|
||||
argument.
|
||||
"""
|
||||
if prompt is not None and isinstance(prompt, str):
|
||||
batch_size = 1
|
||||
elif prompt is not None and isinstance(prompt, list):
|
||||
batch_size = len(prompt)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"`prompt` should be either a string or a list of strings, but got {type(prompt)}."
|
||||
)
|
||||
|
||||
text_inputs = self.tokenizer(
|
||||
prompt,
|
||||
padding="max_length",
|
||||
max_length=self.tokenizer.model_max_length,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
text_input_ids = text_inputs.input_ids
|
||||
untruncated_ids = self.tokenizer(
|
||||
prompt, padding="longest", return_tensors="pt"
|
||||
).input_ids
|
||||
|
||||
if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
|
||||
text_input_ids, untruncated_ids
|
||||
):
|
||||
removed_text = self.tokenizer.batch_decode(
|
||||
untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
|
||||
)
|
||||
logger.warning(
|
||||
"The following part of your input was truncated because CLIP can only handle sequences up to"
|
||||
f" {self.tokenizer.model_max_length} tokens: {removed_text}"
|
||||
)
|
||||
|
||||
if (
|
||||
hasattr(self.text_encoder.config, "use_attention_mask")
|
||||
and self.text_encoder.config.use_attention_mask
|
||||
):
|
||||
attention_mask = text_inputs.attention_mask.to(device)
|
||||
else:
|
||||
attention_mask = None
|
||||
|
||||
prompt_embeds = self.text_encoder(
|
||||
text_input_ids.to(device),
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
prompt_embeds = prompt_embeds[0]
|
||||
|
||||
prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
|
||||
|
||||
bs_embed, seq_len, _ = prompt_embeds.shape
|
||||
# duplicate text embeddings for each generation per prompt, using mps friendly method
|
||||
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
|
||||
prompt_embeds = prompt_embeds.view(
|
||||
bs_embed * num_images_per_prompt, seq_len, -1
|
||||
)
|
||||
|
||||
# get unconditional embeddings for classifier free guidance
|
||||
if do_classifier_free_guidance:
|
||||
uncond_tokens: List[str]
|
||||
if negative_prompt is None:
|
||||
uncond_tokens = [""] * batch_size
|
||||
elif type(prompt) is not type(negative_prompt):
|
||||
raise TypeError(
|
||||
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
|
||||
f" {type(prompt)}."
|
||||
)
|
||||
elif isinstance(negative_prompt, str):
|
||||
uncond_tokens = [negative_prompt]
|
||||
elif batch_size != len(negative_prompt):
|
||||
raise ValueError(
|
||||
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
|
||||
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
|
||||
" the batch size of `prompt`."
|
||||
)
|
||||
else:
|
||||
uncond_tokens = negative_prompt
|
||||
|
||||
max_length = prompt_embeds.shape[1]
|
||||
uncond_input = self.tokenizer(
|
||||
uncond_tokens,
|
||||
padding="max_length",
|
||||
max_length=max_length,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
if (
|
||||
hasattr(self.text_encoder.config, "use_attention_mask")
|
||||
and self.text_encoder.config.use_attention_mask
|
||||
):
|
||||
attention_mask = uncond_input.attention_mask.to(device)
|
||||
else:
|
||||
attention_mask = None
|
||||
|
||||
negative_prompt_embeds = self.text_encoder(
|
||||
uncond_input.input_ids.to(device),
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
negative_prompt_embeds = negative_prompt_embeds[0]
|
||||
|
||||
# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
|
||||
seq_len = negative_prompt_embeds.shape[1]
|
||||
|
||||
negative_prompt_embeds = negative_prompt_embeds.to(
|
||||
dtype=self.text_encoder.dtype, device=device
|
||||
)
|
||||
|
||||
negative_prompt_embeds = negative_prompt_embeds.repeat(
|
||||
1, num_images_per_prompt, 1
|
||||
)
|
||||
negative_prompt_embeds = negative_prompt_embeds.view(
|
||||
batch_size * num_images_per_prompt, seq_len, -1
|
||||
)
|
||||
|
||||
# For classifier free guidance, we need to do two forward passes.
|
||||
# Here we concatenate the unconditional and text embeddings into a single batch
|
||||
# to avoid doing two forward passes
|
||||
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
|
||||
|
||||
return prompt_embeds
|
||||
|
||||
def decode_latents(self, latents):
|
||||
latents = 1 / self.vae.config.scaling_factor * latents
|
||||
image = self.vae.decode(latents).sample
|
||||
image = (image / 2 + 0.5).clamp(0, 1)
|
||||
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
|
||||
image = image.cpu().permute(0, 2, 3, 1).float().numpy()
|
||||
return image
|
||||
|
||||
def prepare_extra_step_kwargs(self, generator, eta):
|
||||
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
|
||||
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
|
||||
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
|
||||
# and should be between [0, 1]
|
||||
|
||||
accepts_eta = "eta" in set(
|
||||
inspect.signature(self.scheduler.step).parameters.keys()
|
||||
)
|
||||
extra_step_kwargs = {}
|
||||
if accepts_eta:
|
||||
extra_step_kwargs["eta"] = eta
|
||||
|
||||
# check if the scheduler accepts generator
|
||||
accepts_generator = "generator" in set(
|
||||
inspect.signature(self.scheduler.step).parameters.keys()
|
||||
)
|
||||
if accepts_generator:
|
||||
extra_step_kwargs["generator"] = generator
|
||||
return extra_step_kwargs
|
||||
|
||||
def prepare_latents(
|
||||
self,
|
||||
batch_size,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
dtype,
|
||||
device,
|
||||
generator,
|
||||
latents=None,
|
||||
):
|
||||
shape = (
|
||||
batch_size,
|
||||
num_channels_latents,
|
||||
height // self.vae_scale_factor,
|
||||
width // self.vae_scale_factor,
|
||||
)
|
||||
if isinstance(generator, list) and len(generator) != batch_size:
|
||||
raise ValueError(
|
||||
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
|
||||
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
|
||||
)
|
||||
|
||||
if latents is None:
|
||||
latents = randn_tensor(
|
||||
shape, generator=generator, device=device, dtype=dtype
|
||||
)
|
||||
else:
|
||||
latents = latents.to(device)
|
||||
|
||||
# scale the initial noise by the standard deviation required by the scheduler
|
||||
latents = latents * self.scheduler.init_noise_sigma
|
||||
return latents
|
||||
|
||||
def encode_image(self, image, device, num_images_per_prompt):
|
||||
dtype = next(self.image_encoder.parameters()).dtype
|
||||
|
||||
if image.dtype == np.float32:
|
||||
image = (image * 255).astype(np.uint8)
|
||||
|
||||
image = self.feature_extractor(image, return_tensors="pt").pixel_values
|
||||
image = image.to(device=device, dtype=dtype)
|
||||
|
||||
image_embeds = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
|
||||
image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
|
||||
|
||||
return torch.zeros_like(image_embeds), image_embeds
|
||||
|
||||
def encode_image_latents(self, image, device, num_images_per_prompt):
|
||||
|
||||
dtype = next(self.image_encoder.parameters()).dtype
|
||||
|
||||
image = torch.from_numpy(image).unsqueeze(0).permute(0, 3, 1, 2).to(device=device) # [1, 3, H, W]
|
||||
image = 2 * image - 1
|
||||
image = F.interpolate(image, (256, 256), mode='bilinear', align_corners=False)
|
||||
image = image.to(dtype=dtype)
|
||||
|
||||
posterior = self.vae.encode(image).latent_dist
|
||||
latents = posterior.sample() * self.vae.config.scaling_factor # [B, C, H, W]
|
||||
latents = latents.repeat_interleave(num_images_per_prompt, dim=0)
|
||||
|
||||
return torch.zeros_like(latents), latents
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
prompt: str = "",
|
||||
image: Optional[np.ndarray] = None,
|
||||
height: int = 256,
|
||||
width: int = 256,
|
||||
num_inference_steps: int = 50,
|
||||
guidance_scale: float = 7.0,
|
||||
negative_prompt: str = "",
|
||||
num_images_per_prompt: int = 1,
|
||||
eta: float = 0.0,
|
||||
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
||||
output_type: Optional[str] = "numpy", # pil, numpy, latents
|
||||
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
|
||||
callback_steps: int = 1,
|
||||
num_frames: int = 4,
|
||||
device=torch.device("cuda:0"),
|
||||
):
|
||||
self.unet = self.unet.to(device=device)
|
||||
self.vae = self.vae.to(device=device)
|
||||
self.text_encoder = self.text_encoder.to(device=device)
|
||||
|
||||
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
|
||||
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
|
||||
# corresponds to doing no classifier free guidance.
|
||||
do_classifier_free_guidance = guidance_scale > 1.0
|
||||
|
||||
# Prepare timesteps
|
||||
self.scheduler.set_timesteps(num_inference_steps, device=device)
|
||||
timesteps = self.scheduler.timesteps
|
||||
|
||||
# imagedream variant
|
||||
if image is not None:
|
||||
assert isinstance(image, np.ndarray) and image.dtype == np.float32
|
||||
self.image_encoder = self.image_encoder.to(device=device)
|
||||
image_embeds_neg, image_embeds_pos = self.encode_image(image, device, num_images_per_prompt)
|
||||
image_latents_neg, image_latents_pos = self.encode_image_latents(image, device, num_images_per_prompt)
|
||||
|
||||
_prompt_embeds = self._encode_prompt(
|
||||
prompt=prompt,
|
||||
device=device,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
do_classifier_free_guidance=do_classifier_free_guidance,
|
||||
negative_prompt=negative_prompt,
|
||||
) # type: ignore
|
||||
prompt_embeds_neg, prompt_embeds_pos = _prompt_embeds.chunk(2)
|
||||
|
||||
# Prepare latent variables
|
||||
actual_num_frames = num_frames if image is None else num_frames + 1
|
||||
latents: torch.Tensor = self.prepare_latents(
|
||||
actual_num_frames * num_images_per_prompt,
|
||||
4,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds_pos.dtype,
|
||||
device,
|
||||
generator,
|
||||
None,
|
||||
)
|
||||
|
||||
if image is not None:
|
||||
camera = get_camera(num_frames, elevation=5, extra_view=True).to(dtype=latents.dtype, device=device)
|
||||
else:
|
||||
camera = get_camera(num_frames, elevation=15, extra_view=False).to(dtype=latents.dtype, device=device)
|
||||
camera = camera.repeat_interleave(num_images_per_prompt, dim=0)
|
||||
|
||||
# Prepare extra step kwargs.
|
||||
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
|
||||
|
||||
# Denoising loop
|
||||
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
|
||||
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
||||
for i, t in enumerate(timesteps):
|
||||
# expand the latents if we are doing classifier free guidance
|
||||
multiplier = 2 if do_classifier_free_guidance else 1
|
||||
latent_model_input = torch.cat([latents] * multiplier)
|
||||
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
|
||||
|
||||
unet_inputs = {
|
||||
'x': latent_model_input,
|
||||
'timesteps': torch.tensor([t] * actual_num_frames * multiplier, dtype=latent_model_input.dtype, device=device),
|
||||
'context': torch.cat([prompt_embeds_neg] * actual_num_frames + [prompt_embeds_pos] * actual_num_frames),
|
||||
'num_frames': actual_num_frames,
|
||||
'camera': torch.cat([camera] * multiplier),
|
||||
}
|
||||
|
||||
if image is not None:
|
||||
unet_inputs['ip'] = torch.cat([image_embeds_neg] * actual_num_frames + [image_embeds_pos] * actual_num_frames)
|
||||
unet_inputs['ip_img'] = torch.cat([image_latents_neg] + [image_latents_pos]) # no repeat
|
||||
|
||||
# predict the noise residual
|
||||
noise_pred = self.unet.forward(**unet_inputs)
|
||||
|
||||
# perform guidance
|
||||
if do_classifier_free_guidance:
|
||||
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
|
||||
noise_pred = noise_pred_uncond + guidance_scale * (
|
||||
noise_pred_text - noise_pred_uncond
|
||||
)
|
||||
|
||||
# compute the previous noisy sample x_t -> x_t-1
|
||||
latents: torch.Tensor = self.scheduler.step(
|
||||
noise_pred, t, latents, **extra_step_kwargs, return_dict=False
|
||||
)[0]
|
||||
|
||||
# call the callback, if provided
|
||||
if i == len(timesteps) - 1 or (
|
||||
(i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
|
||||
):
|
||||
progress_bar.update()
|
||||
if callback is not None and i % callback_steps == 0:
|
||||
callback(i, t, latents) # type: ignore
|
||||
|
||||
# Post-processing
|
||||
if output_type == "latent":
|
||||
image = latents
|
||||
elif output_type == "pil":
|
||||
image = self.decode_latents(latents)
|
||||
image = self.numpy_to_pil(image)
|
||||
else: # numpy
|
||||
image = self.decode_latents(latents)
|
||||
|
||||
# Offload last model to CPU
|
||||
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
|
||||
self.final_offload_hook.offload()
|
||||
|
||||
return image
|
||||
@@ -0,0 +1,19 @@
|
||||
{
|
||||
"_class_name": "DDIMScheduler",
|
||||
"_diffusers_version": "0.25.0",
|
||||
"beta_end": 0.012,
|
||||
"beta_schedule": "scaled_linear",
|
||||
"beta_start": 0.00085,
|
||||
"clip_sample": false,
|
||||
"clip_sample_range": 1.0,
|
||||
"dynamic_thresholding_ratio": 0.995,
|
||||
"num_train_timesteps": 1000,
|
||||
"prediction_type": "epsilon",
|
||||
"rescale_betas_zero_snr": false,
|
||||
"sample_max_value": 1.0,
|
||||
"set_alpha_to_one": false,
|
||||
"steps_offset": 1,
|
||||
"thresholding": false,
|
||||
"timestep_spacing": "leading",
|
||||
"trained_betas": null
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
{
|
||||
"_name_or_path": "stabilityai/stable-diffusion-2-1",
|
||||
"architectures": [
|
||||
"CLIPTextModel"
|
||||
],
|
||||
"attention_dropout": 0.0,
|
||||
"bos_token_id": 0,
|
||||
"dropout": 0.0,
|
||||
"eos_token_id": 2,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_size": 1024,
|
||||
"initializer_factor": 1.0,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 4096,
|
||||
"layer_norm_eps": 1e-05,
|
||||
"max_position_embeddings": 77,
|
||||
"model_type": "clip_text_model",
|
||||
"num_attention_heads": 16,
|
||||
"num_hidden_layers": 23,
|
||||
"pad_token_id": 1,
|
||||
"projection_dim": 512,
|
||||
"torch_dtype": "float16",
|
||||
"transformers_version": "4.35.2",
|
||||
"vocab_size": 49408
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,24 @@
|
||||
{
|
||||
"bos_token": {
|
||||
"content": "<|startoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"eos_token": {
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"pad_token": "!",
|
||||
"unk_token": {
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
{
|
||||
"add_prefix_space": false,
|
||||
"added_tokens_decoder": {
|
||||
"0": {
|
||||
"content": "!",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"49406": {
|
||||
"content": "<|startoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"49407": {
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
}
|
||||
},
|
||||
"bos_token": "<|startoftext|>",
|
||||
"clean_up_tokenization_spaces": true,
|
||||
"do_lower_case": true,
|
||||
"eos_token": "<|endoftext|>",
|
||||
"errors": "replace",
|
||||
"model_max_length": 77,
|
||||
"pad_token": "!",
|
||||
"tokenizer_class": "CLIPTokenizer",
|
||||
"unk_token": "<|endoftext|>"
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,25 @@
|
||||
{
|
||||
"_class_name": "MultiViewUNetModel",
|
||||
"_diffusers_version": "0.25.0",
|
||||
"attention_resolutions": [
|
||||
4,
|
||||
2,
|
||||
1
|
||||
],
|
||||
"camera_dim": 16,
|
||||
"channel_mult": [
|
||||
1,
|
||||
2,
|
||||
4,
|
||||
4
|
||||
],
|
||||
"context_dim": 1024,
|
||||
"image_size": 32,
|
||||
"in_channels": 4,
|
||||
"ip_dim": 16,
|
||||
"model_channels": 320,
|
||||
"num_head_channels": 64,
|
||||
"num_res_blocks": 2,
|
||||
"out_channels": 4,
|
||||
"transformer_depth": 1
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,31 @@
|
||||
{
|
||||
"_class_name": "AutoencoderKL",
|
||||
"_diffusers_version": "0.25.0",
|
||||
"act_fn": "silu",
|
||||
"block_out_channels": [
|
||||
128,
|
||||
256,
|
||||
512,
|
||||
512
|
||||
],
|
||||
"down_block_types": [
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D"
|
||||
],
|
||||
"force_upcast": true,
|
||||
"in_channels": 3,
|
||||
"latent_channels": 4,
|
||||
"layers_per_block": 2,
|
||||
"norm_num_groups": 32,
|
||||
"out_channels": 3,
|
||||
"sample_size": 256,
|
||||
"scaling_factor": 0.18215,
|
||||
"up_block_types": [
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
{
|
||||
"crop_size": {
|
||||
"height": 224,
|
||||
"width": 224
|
||||
},
|
||||
"do_center_crop": true,
|
||||
"do_convert_rgb": true,
|
||||
"do_normalize": true,
|
||||
"do_rescale": true,
|
||||
"do_resize": true,
|
||||
"feature_extractor_type": "CLIPFeatureExtractor",
|
||||
"image_mean": [
|
||||
0.48145466,
|
||||
0.4578275,
|
||||
0.40821073
|
||||
],
|
||||
"image_processor_type": "CLIPImageProcessor",
|
||||
"image_std": [
|
||||
0.26862954,
|
||||
0.26130258,
|
||||
0.27577711
|
||||
],
|
||||
"resample": 3,
|
||||
"rescale_factor": 0.00392156862745098,
|
||||
"size": {
|
||||
"shortest_edge": 224
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
{
|
||||
"_name_or_path": "lambdalabs/sd-image-variations-diffusers",
|
||||
"architectures": [
|
||||
"CLIPVisionModelWithProjection"
|
||||
],
|
||||
"attention_dropout": 0.0,
|
||||
"dropout": 0.0,
|
||||
"hidden_act": "quick_gelu",
|
||||
"hidden_size": 1024,
|
||||
"image_size": 224,
|
||||
"initializer_factor": 1.0,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 4096,
|
||||
"layer_norm_eps": 1e-05,
|
||||
"model_type": "clip_vision_model",
|
||||
"num_attention_heads": 16,
|
||||
"num_channels": 3,
|
||||
"num_hidden_layers": 24,
|
||||
"patch_size": 14,
|
||||
"projection_dim": 768,
|
||||
"torch_dtype": "float16",
|
||||
"transformers_version": "4.34.1"
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
{
|
||||
"_class_name": "MVDiffusionImagePipeline",
|
||||
"_diffusers_version": "0.19.3",
|
||||
"feature_extractor": [
|
||||
"transformers",
|
||||
"CLIPImageProcessor"
|
||||
],
|
||||
"image_encoder": [
|
||||
"transformers",
|
||||
"CLIPVisionModelWithProjection"
|
||||
],
|
||||
"requires_safety_checker": true,
|
||||
"safety_checker": [
|
||||
null,
|
||||
null
|
||||
],
|
||||
"scheduler": [
|
||||
"diffusers",
|
||||
"DDIMScheduler"
|
||||
],
|
||||
"unet": [
|
||||
"mvdiffusion.models.unet_mv2d_condition",
|
||||
"UNetMV2DConditionModel"
|
||||
],
|
||||
"vae": [
|
||||
"diffusers",
|
||||
"AutoencoderKL"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
{
|
||||
"_class_name": "DDIMScheduler",
|
||||
"_diffusers_version": "0.19.3",
|
||||
"beta_end": 0.012,
|
||||
"beta_schedule": "scaled_linear",
|
||||
"beta_start": 0.00085,
|
||||
"clip_sample": false,
|
||||
"clip_sample_range": 1.0,
|
||||
"dynamic_thresholding_ratio": 0.995,
|
||||
"num_train_timesteps": 1000,
|
||||
"prediction_type": "epsilon",
|
||||
"rescale_betas_zero_snr": false,
|
||||
"sample_max_value": 1.0,
|
||||
"set_alpha_to_one": false,
|
||||
"skip_prk_steps": true,
|
||||
"steps_offset": 1,
|
||||
"thresholding": false,
|
||||
"timestep_spacing": "leading",
|
||||
"trained_betas": null
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
{
|
||||
"_class_name": "UNetMV2DConditionModel",
|
||||
"_diffusers_version": "0.19.3",
|
||||
"_name_or_path": "/home/xiaoxiao/Workplace/wonder3Dplus/ckpts",
|
||||
"act_fn": "silu",
|
||||
"addition_embed_type": null,
|
||||
"addition_embed_type_num_heads": 64,
|
||||
"addition_time_embed_dim": null,
|
||||
"attention_head_dim": 8,
|
||||
"block_out_channels": [
|
||||
320,
|
||||
640,
|
||||
1280,
|
||||
1280
|
||||
],
|
||||
"cd_attention_last": false,
|
||||
"cd_attention_mid": true,
|
||||
"center_input_sample": false,
|
||||
"class_embed_type": "projection",
|
||||
"class_embeddings_concat": false,
|
||||
"conv_in_kernel": 3,
|
||||
"conv_out_kernel": 3,
|
||||
"cross_attention_dim": 768,
|
||||
"cross_attention_norm": null,
|
||||
"cross_domain_attention": false,
|
||||
"down_block_types": [
|
||||
"CrossAttnDownBlockMV2D",
|
||||
"CrossAttnDownBlockMV2D",
|
||||
"CrossAttnDownBlockMV2D",
|
||||
"DownBlock2D"
|
||||
],
|
||||
"downsample_padding": 1,
|
||||
"dual_cross_attention": false,
|
||||
"encoder_hid_dim": null,
|
||||
"encoder_hid_dim_type": null,
|
||||
"flip_sin_to_cos": true,
|
||||
"freq_shift": 0,
|
||||
"in_channels": 8,
|
||||
"joint_attention": false,
|
||||
"joint_attention_twice": true,
|
||||
"layers_per_block": 2,
|
||||
"mid_block_only_cross_attention": null,
|
||||
"mid_block_scale_factor": 1,
|
||||
"mid_block_type": "UNetMidBlockMV2DCrossAttn",
|
||||
"multiview_attention": true,
|
||||
"mvcd_attention": false,
|
||||
"norm_eps": 1e-05,
|
||||
"norm_num_groups": 32,
|
||||
"num_attention_heads": null,
|
||||
"num_class_embeds": null,
|
||||
"num_views": 6,
|
||||
"only_cross_attention": false,
|
||||
"out_channels": 4,
|
||||
"projection_class_embeddings_input_dim": 10,
|
||||
"resnet_out_scale_factor": 1.0,
|
||||
"resnet_skip_time_act": false,
|
||||
"resnet_time_scale_shift": "default",
|
||||
"sample_size": 32,
|
||||
"sparse_mv_attention": false,
|
||||
"time_cond_proj_dim": null,
|
||||
"time_embedding_act_fn": null,
|
||||
"time_embedding_dim": null,
|
||||
"time_embedding_type": "positional",
|
||||
"timestep_post_act": null,
|
||||
"transformer_layers_per_block": 1,
|
||||
"up_block_types": [
|
||||
"UpBlock2D",
|
||||
"CrossAttnUpBlockMV2D",
|
||||
"CrossAttnUpBlockMV2D",
|
||||
"CrossAttnUpBlockMV2D"
|
||||
],
|
||||
"upcast_attention": false,
|
||||
"use_linear_projection": false
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
{
|
||||
"_class_name": "AutoencoderKL",
|
||||
"_diffusers_version": "0.19.3",
|
||||
"_name_or_path": "lambdalabs/sd-image-variations-diffusers",
|
||||
"act_fn": "silu",
|
||||
"block_out_channels": [
|
||||
128,
|
||||
256,
|
||||
512,
|
||||
512
|
||||
],
|
||||
"down_block_types": [
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D"
|
||||
],
|
||||
"force_upcast": true,
|
||||
"in_channels": 3,
|
||||
"latent_channels": 4,
|
||||
"layers_per_block": 2,
|
||||
"norm_num_groups": 32,
|
||||
"out_channels": 3,
|
||||
"sample_size": 256,
|
||||
"scaling_factor": 0.18215,
|
||||
"up_block_types": [
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D"
|
||||
]
|
||||
}
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
{
|
||||
"_class_name": "MVDiffusionImagePipeline",
|
||||
"_diffusers_version": "0.19.3",
|
||||
"feature_extractor": [
|
||||
"transformers",
|
||||
"CLIPImageProcessor"
|
||||
],
|
||||
"image_encoder": [
|
||||
"transformers",
|
||||
"CLIPVisionModelWithProjection"
|
||||
],
|
||||
"requires_safety_checker": true,
|
||||
"safety_checker": [
|
||||
null,
|
||||
null
|
||||
],
|
||||
"scheduler": [
|
||||
"diffusers",
|
||||
"DDIMScheduler"
|
||||
],
|
||||
"unet": [
|
||||
"mvdiffusion.models.unet_mv2d_condition",
|
||||
"UNetMV2DConditionModel"
|
||||
],
|
||||
"vae": [
|
||||
"diffusers",
|
||||
"AutoencoderKL"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
d6d2efc033a06a74d3761268de7295c97e6935d2
|
||||
+1
@@ -0,0 +1 @@
|
||||
../../blobs/6d7b44061a6f8f191002c2af48e47bb887816288
|
||||
@@ -0,0 +1,70 @@
|
||||
general {
|
||||
base_exp_dir = ./outputs/NeuS/
|
||||
}
|
||||
|
||||
dataset {
|
||||
object_viewidx = 1,
|
||||
imSize = [256, 256],
|
||||
stage = coarse,
|
||||
mtype = mlp
|
||||
}
|
||||
|
||||
validate {
|
||||
validate_resolution_level = 1,
|
||||
save_freq = 5000,
|
||||
val_freq = 5000,
|
||||
val_mesh_freq =5000,
|
||||
report_freq = 200,
|
||||
}
|
||||
|
||||
model {
|
||||
nerf {
|
||||
D = 8,
|
||||
d_in = 4,
|
||||
d_in_view = 3,
|
||||
W = 256,
|
||||
multires = 10,
|
||||
multires_view = 4,
|
||||
output_ch = 4,
|
||||
skips=[4],
|
||||
use_viewdirs=True
|
||||
}
|
||||
|
||||
sdf_network {
|
||||
d_out = 257,
|
||||
d_in = 3,
|
||||
d_hidden = 256,
|
||||
n_layers = 8,
|
||||
skip_in = [4],
|
||||
multires = 6,
|
||||
bias = 0.5,
|
||||
scale = 1.0,
|
||||
geometric_init = True,
|
||||
weight_norm = True
|
||||
}
|
||||
|
||||
variance_network {
|
||||
init_val = 0.3
|
||||
}
|
||||
|
||||
rendering_network {
|
||||
d_feature = 256,
|
||||
mode = no_view_dir,
|
||||
d_in = 6,
|
||||
d_out = 3,
|
||||
d_hidden = 256,
|
||||
n_layers = 4,
|
||||
weight_norm = True,
|
||||
multires_view = 0,
|
||||
squeeze_out = True
|
||||
}
|
||||
|
||||
neus_renderer {
|
||||
n_samples = 64,
|
||||
n_importance = 64,
|
||||
n_outside = 0,
|
||||
up_sample_steps = 4, # 1 for simple coarse-to-fine sampling
|
||||
perturb = 1.0,
|
||||
sdf_decay_param = 100
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
#pretrained_model_name_or_path: 'flamehaze1115/wonder3d-v1.0' # or './ckpts'
|
||||
revision: null
|
||||
validation_dataset:
|
||||
root_dir: "./example_images" # the folder path stores testing images
|
||||
num_views: 6
|
||||
bg_color: 'white'
|
||||
img_wh: [256, 256]
|
||||
num_validation_samples: 1000
|
||||
crop_size: 192
|
||||
filepaths: ['owl.png'] # the test image names. leave it empty, test all images in the folder
|
||||
|
||||
save_dir: 'outputs/'
|
||||
|
||||
pred_type: 'joint'
|
||||
seed: 42
|
||||
validation_batch_size: 1
|
||||
dataloader_num_workers: 64
|
||||
|
||||
local_rank: -1
|
||||
|
||||
pipe_kwargs:
|
||||
camera_embedding_type: 'e_de_da_sincos'
|
||||
num_views: 6
|
||||
|
||||
validation_guidance_scales: [1.0]
|
||||
pipe_validation_kwargs:
|
||||
eta: 1.0
|
||||
validation_grid_nrow: 6
|
||||
|
||||
unet_from_pretrained_kwargs:
|
||||
camera_embedding_type: 'e_de_da_sincos'
|
||||
projection_class_embeddings_input_dim: 10
|
||||
num_views: 6
|
||||
sample_size: 32
|
||||
cd_attention_mid: true
|
||||
zero_init_conv_in: false
|
||||
zero_init_camera_projection: false
|
||||
|
||||
num_views: 6
|
||||
camera_embedding_type: 'e_de_da_sincos'
|
||||
|
||||
enable_xformers_memory_efficient_attention: true
|
||||
+349
-127
@@ -5,34 +5,47 @@ import torch
|
||||
import trimesh
|
||||
import numpy as np
|
||||
|
||||
from kiui.op import safe_normalize, dot
|
||||
from kiui.typing import *
|
||||
|
||||
from shared_utils.sh_utils import SH2RGB
|
||||
from shared_utils.image_utils import prepare_torch_img
|
||||
|
||||
def dot(x, y):
|
||||
return torch.sum(x * y, -1, keepdim=True)
|
||||
|
||||
|
||||
def length(x, eps=1e-20):
|
||||
return torch.sqrt(torch.clamp(dot(x, x), min=eps))
|
||||
|
||||
|
||||
def safe_normalize(x, eps=1e-20):
|
||||
return x / length(x, eps)
|
||||
|
||||
class Mesh:
|
||||
"""
|
||||
A torch-native trimesh class, with support for ``ply/obj/glb`` formats.
|
||||
|
||||
Note:
|
||||
This class only supports one mesh with a single texture image (an albedo texture and a metallic-roughness texture).
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
v=None,
|
||||
f=None,
|
||||
vn=None,
|
||||
fn=None,
|
||||
vt=None,
|
||||
ft=None,
|
||||
albedo=None,
|
||||
vc=None, # vertex color
|
||||
device=None,
|
||||
v: Optional[Tensor] = None,
|
||||
f: Optional[Tensor] = None,
|
||||
vn: Optional[Tensor] = None,
|
||||
fn: Optional[Tensor] = None,
|
||||
vt: Optional[Tensor] = None,
|
||||
ft: Optional[Tensor] = None,
|
||||
vc: Optional[Tensor] = None, # vertex color
|
||||
albedo: Optional[Tensor] = None,
|
||||
metallicRoughness: Optional[Tensor] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
):
|
||||
"""Init a mesh directly using all attributes.
|
||||
|
||||
Args:
|
||||
v (Optional[Tensor]): vertices, float [N, 3]. Defaults to None.
|
||||
f (Optional[Tensor]): faces, int [M, 3]. Defaults to None.
|
||||
vn (Optional[Tensor]): vertex normals, float [N, 3]. Defaults to None.
|
||||
fn (Optional[Tensor]): faces for normals, int [M, 3]. Defaults to None.
|
||||
vt (Optional[Tensor]): vertex uv coordinates, float [N, 2]. Defaults to None.
|
||||
ft (Optional[Tensor]): faces for uvs, int [M, 3]. Defaults to None.
|
||||
vc (Optional[Tensor]): vertex colors, float [N, 3]. Defaults to None.
|
||||
albedo (Optional[Tensor]): albedo texture, float [H, W, 3], RGB format. Defaults to None.
|
||||
metallicRoughness (Optional[Tensor]): metallic-roughness texture, float [H, W, 3], metallic(Blue) = metallicRoughness[..., 2], roughness(Green) = metallicRoughness[..., 1]. Defaults to None.
|
||||
device (Optional[torch.device]): torch device. Defaults to None.
|
||||
"""
|
||||
self.device = device
|
||||
self.v = v
|
||||
self.vn = vn
|
||||
@@ -40,32 +53,58 @@ class Mesh:
|
||||
self.f = f
|
||||
self.fn = fn
|
||||
self.ft = ft
|
||||
# only support a single albedo
|
||||
self.albedo = albedo
|
||||
# support vertex color is no albedo
|
||||
# will first see if there is vertex color to use
|
||||
self.vc = vc
|
||||
# only support a single albedo image
|
||||
self.albedo = albedo
|
||||
# pbr extension, metallic(Blue) = metallicRoughness[..., 2], roughness(Green) = metallicRoughness[..., 1]
|
||||
# ref: https://registry.khronos.org/glTF/specs/2.0/glTF-2.0.html
|
||||
self.metallicRoughness = metallicRoughness
|
||||
|
||||
self.ori_center = 0
|
||||
self.ori_scale = 1
|
||||
|
||||
|
||||
|
||||
@classmethod
|
||||
def load(cls, path=None, resize=True, renormal=True, retex=False, front_dir='+z', **kwargs):
|
||||
# assume init with kwargs
|
||||
if path is None:
|
||||
mesh = cls(**kwargs)
|
||||
def load(cls, path, resize=True, clean=False, renormal=True, retex=False, bound=0.9, front_dir='+z', **kwargs):
|
||||
"""load mesh from path.
|
||||
|
||||
Args:
|
||||
path (str): path to mesh file, supports ply, obj, glb.
|
||||
clean (bool, optional): perform mesh cleaning at load (e.g., merge close vertices). Defaults to False.
|
||||
resize (bool, optional): auto resize the mesh using ``bound`` into [-bound, bound]^3. Defaults to True.
|
||||
renormal (bool, optional): re-calc the vertex normals. Defaults to True.
|
||||
retex (bool, optional): re-calc the uv coordinates, will overwrite the existing uv coordinates. Defaults to False.
|
||||
bound (float, optional): bound to resize. Defaults to 0.9.
|
||||
front_dir (str, optional): front-view direction of the mesh, should be [+-][xyz][ 123]. Defaults to '+z'.
|
||||
device (torch.device, optional): torch device. Defaults to None.
|
||||
|
||||
Note:
|
||||
a ``device`` keyword argument can be provided to specify the torch device.
|
||||
If it's not provided, we will try to use ``'cuda'`` as the device if it's available.
|
||||
|
||||
Returns:
|
||||
Mesh: the loaded Mesh object.
|
||||
"""
|
||||
# obj supports face uv
|
||||
elif path.lower().endswith(".obj"):
|
||||
if path.endswith(".obj"):
|
||||
mesh = cls.load_obj(path, **kwargs)
|
||||
# trimesh only supports vertex uv, but can load more formats
|
||||
else:
|
||||
mesh = cls.load_trimesh(path, **kwargs)
|
||||
|
||||
# clean
|
||||
if clean:
|
||||
from kiui.mesh_utils import clean_mesh
|
||||
vertices = mesh.v.detach().cpu().numpy()
|
||||
triangles = mesh.f.detach().cpu().numpy()
|
||||
vertices, triangles = clean_mesh(vertices, triangles, remesh=False)
|
||||
mesh.v = torch.from_numpy(vertices).contiguous().float().to(mesh.device)
|
||||
mesh.f = torch.from_numpy(triangles).contiguous().int().to(mesh.device)
|
||||
|
||||
print(f"[Mesh loading] v: {mesh.v.shape}, f: {mesh.f.shape}")
|
||||
# auto-normalize
|
||||
if resize:
|
||||
mesh.auto_size()
|
||||
mesh.auto_size(bound=bound)
|
||||
# auto-fix normal
|
||||
if renormal or mesh.vn is None:
|
||||
mesh.auto_normal()
|
||||
@@ -105,6 +144,20 @@ class Mesh:
|
||||
# load from obj file
|
||||
@classmethod
|
||||
def load_obj(cls, path, albedo_path=None, device=None):
|
||||
"""load an ``obj`` mesh.
|
||||
|
||||
Args:
|
||||
path (str): path to mesh.
|
||||
albedo_path (str, optional): path to the albedo texture image, will overwrite the existing texture path if specified in mtl. Defaults to None.
|
||||
device (torch.device, optional): torch device. Defaults to None.
|
||||
|
||||
Note:
|
||||
We will try to read `mtl` path from `obj`, else we assume the file name is the same as `obj` but with `mtl` extension.
|
||||
The `usemtl` statement is ignored, and we only use the last material path in `mtl` file.
|
||||
|
||||
Returns:
|
||||
Mesh: the loaded Mesh object.
|
||||
"""
|
||||
assert os.path.splitext(path)[-1] == ".obj"
|
||||
|
||||
mesh = cls()
|
||||
@@ -121,7 +174,7 @@ class Mesh:
|
||||
|
||||
def parse_f_v(fv):
|
||||
# pass in a vertex term of a face, return {v, vt, vn} (-1 if not provided)
|
||||
# supported forms: (Only supports triangulated face)
|
||||
# supported forms:
|
||||
# f v1 v2 v3
|
||||
# f v1/vt1 v2/vt2 v3/vt3
|
||||
# f v1/vt1/vn1 v2/vt2/vn2 v3/vt3/vn3
|
||||
@@ -130,7 +183,6 @@ class Mesh:
|
||||
xs.extend([-1] * (3 - len(xs)))
|
||||
return xs[0], xs[1], xs[2]
|
||||
|
||||
# NOTE: we ignore usemtl, and assume the mesh ONLY uses one material (first in mtl)
|
||||
vertices, texcoords, normals = [], [], []
|
||||
faces, tfaces, nfaces = [], [], []
|
||||
mtl_path = None
|
||||
@@ -214,21 +266,28 @@ class Mesh:
|
||||
break
|
||||
|
||||
# if albedo_path is not provided, try retrieve it from mtl
|
||||
metallic_path = None
|
||||
roughness_path = None
|
||||
if mtl_path is not None and albedo_path is None:
|
||||
with open(mtl_path, "r") as f:
|
||||
lines = f.readlines()
|
||||
|
||||
for line in lines:
|
||||
split_line = line.split()
|
||||
# empty line
|
||||
if len(split_line) == 0:
|
||||
continue
|
||||
prefix = split_line[0]
|
||||
# NOTE: simply use the first map_Kd as albedo!
|
||||
|
||||
if "map_Kd" in prefix:
|
||||
# assume relative path!
|
||||
albedo_path = os.path.join(os.path.dirname(path), split_line[1])
|
||||
print(f"[load_obj] use texture from: {albedo_path}")
|
||||
break
|
||||
|
||||
elif "map_Pm" in prefix:
|
||||
metallic_path = os.path.join(os.path.dirname(path), split_line[1])
|
||||
elif "map_Pr" in prefix:
|
||||
roughness_path = os.path.join(os.path.dirname(path), split_line[1])
|
||||
|
||||
# still not found albedo_path, or the path doesn't exist
|
||||
if albedo_path is None or not os.path.exists(albedo_path):
|
||||
# init an empty texture
|
||||
@@ -240,13 +299,40 @@ class Mesh:
|
||||
albedo = cv2.cvtColor(albedo, cv2.COLOR_BGR2RGB)
|
||||
albedo = albedo.astype(np.float32) / 255
|
||||
print(f"[load_obj] load texture: {albedo.shape}")
|
||||
|
||||
|
||||
mesh.albedo = torch.tensor(albedo, dtype=torch.float32, device=device)
|
||||
|
||||
# try to load metallic and roughness
|
||||
if metallic_path is not None and roughness_path is not None:
|
||||
print(f"[load_obj] load metallicRoughness from: {metallic_path}, {roughness_path}")
|
||||
metallic = cv2.imread(metallic_path, cv2.IMREAD_UNCHANGED)
|
||||
metallic = metallic.astype(np.float32) / 255
|
||||
roughness = cv2.imread(roughness_path, cv2.IMREAD_UNCHANGED)
|
||||
roughness = roughness.astype(np.float32) / 255
|
||||
metallicRoughness = np.stack([np.zeros_like(metallic), roughness, metallic], axis=-1)
|
||||
|
||||
mesh.metallicRoughness = torch.tensor(metallicRoughness, dtype=torch.float32, device=device).contiguous()
|
||||
|
||||
return mesh
|
||||
|
||||
@classmethod
|
||||
def load_trimesh(cls, path, device=None):
|
||||
def load_trimesh(cls, path=None, given_mesh=None, device=None):
|
||||
"""load a mesh using ``trimesh.load()``.
|
||||
|
||||
Can load various formats like ``glb`` and serves as a fallback.
|
||||
|
||||
Note:
|
||||
We will try to merge all meshes if the glb contains more than one,
|
||||
but **this may cause the texture to lose**, since we only support one texture image!
|
||||
|
||||
Args:
|
||||
path (str): path to the mesh file.
|
||||
given_mesh (trimesh): trimesh instance
|
||||
device (torch.device, optional): torch device. Defaults to None.
|
||||
|
||||
Returns:
|
||||
Mesh: the loaded Mesh object.
|
||||
"""
|
||||
mesh = cls()
|
||||
|
||||
# device
|
||||
@@ -256,19 +342,29 @@ class Mesh:
|
||||
mesh.device = device
|
||||
|
||||
# use trimesh to load ply/glb, assume only has one single RootMesh...
|
||||
_data = trimesh.load(path)
|
||||
if isinstance(_data, trimesh.Scene):
|
||||
if len(_data.geometry) == 1:
|
||||
_mesh = list(_data.geometry.values())[0]
|
||||
if given_mesh is not None:
|
||||
_mesh = given_mesh
|
||||
elif path is not None:
|
||||
_data = trimesh.load(path)
|
||||
if isinstance(_data, trimesh.Scene):
|
||||
if len(_data.geometry) == 1:
|
||||
_mesh = list(_data.geometry.values())[0]
|
||||
else:
|
||||
print(f"[load_trimesh] concatenating {len(_data.geometry)} meshes.")
|
||||
_concat = []
|
||||
# loop the scene graph and apply transform to each mesh
|
||||
scene_graph = _data.graph.to_flattened() # dict {name: {transform: 4x4 mat, geometry: str}}
|
||||
for k, v in scene_graph.items():
|
||||
name = v['geometry']
|
||||
if name in _data.geometry and isinstance(_data.geometry[name], trimesh.Trimesh):
|
||||
transform = v['transform']
|
||||
_concat.append(_data.geometry[name].apply_transform(transform))
|
||||
_mesh = trimesh.util.concatenate(_concat)
|
||||
else:
|
||||
# manual concat, will lose texture
|
||||
_concat = []
|
||||
for g in _data.geometry.values():
|
||||
if isinstance(g, trimesh.Trimesh):
|
||||
_concat.append(g)
|
||||
_mesh = trimesh.util.concatenate(_concat)
|
||||
_mesh = _data
|
||||
else:
|
||||
_mesh = _data
|
||||
print(f"[load_trimesh] failed to load mesh, either path or given_mesh must be given")
|
||||
return None
|
||||
|
||||
if _mesh.visual.kind == 'vertex':
|
||||
vertex_colors = _mesh.visual.vertex_colors
|
||||
@@ -279,11 +375,15 @@ class Mesh:
|
||||
_material = _mesh.visual.material
|
||||
if isinstance(_material, trimesh.visual.material.PBRMaterial):
|
||||
texture = np.array(_material.baseColorTexture).astype(np.float32) / 255
|
||||
# load metallicRoughness if present
|
||||
if _material.metallicRoughnessTexture is not None:
|
||||
metallicRoughness = np.array(_material.metallicRoughnessTexture).astype(np.float32) / 255
|
||||
mesh.metallicRoughness = torch.tensor(metallicRoughness, dtype=torch.float32, device=device).contiguous()
|
||||
elif isinstance(_material, trimesh.visual.material.SimpleMaterial):
|
||||
texture = np.array(_material.to_pbr().baseColorTexture).astype(np.float32) / 255
|
||||
else:
|
||||
raise NotImplementedError(f"material type {type(_material)} not supported!")
|
||||
mesh.albedo = torch.tensor(texture, dtype=torch.float32, device=device)
|
||||
mesh.albedo = torch.tensor(texture[..., :3], dtype=torch.float32, device=device).contiguous()
|
||||
print(f"[load_trimesh] load texture: {texture.shape}")
|
||||
else:
|
||||
mesh.set_new_albedo(1024, 1024)
|
||||
@@ -340,17 +440,29 @@ class Mesh:
|
||||
|
||||
# aabb
|
||||
def aabb(self):
|
||||
"""get the axis-aligned bounding box of the mesh.
|
||||
|
||||
Returns:
|
||||
Tuple[torch.Tensor]: the min xyz and max xyz of the mesh.
|
||||
"""
|
||||
return torch.min(self.v, dim=0).values, torch.max(self.v, dim=0).values
|
||||
|
||||
# unit size
|
||||
@torch.no_grad()
|
||||
def auto_size(self):
|
||||
def auto_size(self, bound=0.9):
|
||||
"""auto resize the mesh.
|
||||
|
||||
Args:
|
||||
bound (float, optional): resizing into ``[-bound, bound]^3``. Defaults to 0.9.
|
||||
"""
|
||||
vmin, vmax = self.aabb()
|
||||
self.ori_center = (vmax + vmin) / 2
|
||||
self.ori_scale = 1.2 / torch.max(vmax - vmin).item()
|
||||
self.ori_scale = 2 * bound / torch.max(vmax - vmin).item()
|
||||
self.v = (self.v - self.ori_center) * self.ori_scale
|
||||
|
||||
def auto_normal(self):
|
||||
"""auto calculate the vertex normals.
|
||||
"""
|
||||
i0, i1, i2 = self.f[:, 0].long(), self.f[:, 1].long(), self.f[:, 2].long()
|
||||
v0, v1, v2 = self.v[i0, :], self.v[i1, :], self.v[i2, :]
|
||||
|
||||
@@ -374,6 +486,13 @@ class Mesh:
|
||||
self.fn = self.f
|
||||
|
||||
def auto_uv(self, cache_path=None, vmap=True):
|
||||
"""auto calculate the uv coordinates.
|
||||
|
||||
Args:
|
||||
cache_path (str, optional): path to save/load the uv cache as a npz file, this can avoid calculating uv every time when loading the same mesh, which is time-consuming. Defaults to None.
|
||||
vmap (bool, optional): remap vertices based on uv coordinates, so each v correspond to a unique vt (necessary for formats like gltf).
|
||||
Usually this will duplicate the vertices on the edge of uv atlas. Defaults to True.
|
||||
"""
|
||||
# try to load cache
|
||||
if cache_path is not None:
|
||||
cache_path = os.path.splitext(cache_path)[0] + "_uv.npz"
|
||||
@@ -407,7 +526,11 @@ class Mesh:
|
||||
self.align_v_to_vt(vmapping)
|
||||
|
||||
def align_v_to_vt(self, vmapping=None):
|
||||
# remap v/f and vn/vn to vt/ft.
|
||||
""" remap v/f and vn/fn to vt/ft.
|
||||
|
||||
Args:
|
||||
vmapping (np.ndarray, optional): the mapping relationship from f to ft. Defaults to None.
|
||||
"""
|
||||
if vmapping is None:
|
||||
ft = self.ft.view(-1).long()
|
||||
f = self.f.view(-1).long()
|
||||
@@ -422,25 +545,45 @@ class Mesh:
|
||||
self.fn = self.ft
|
||||
|
||||
def to(self, device):
|
||||
"""move all tensor attributes to device.
|
||||
|
||||
Args:
|
||||
device (torch.device): target device.
|
||||
|
||||
Returns:
|
||||
Mesh: self.
|
||||
"""
|
||||
self.device = device
|
||||
for name in ["v", "f", "vn", "fn", "vt", "ft", "albedo"]:
|
||||
for name in ["v", "f", "vn", "fn", "vt", "ft", "albedo", "vc", "metallicRoughness"]:
|
||||
tensor = getattr(self, name)
|
||||
if tensor is not None:
|
||||
setattr(self, name, tensor.to(device))
|
||||
return self
|
||||
|
||||
def write(self, path):
|
||||
if path.lower().endswith(".ply"):
|
||||
"""write the mesh to a path.
|
||||
|
||||
Args:
|
||||
path (str): path to write, supports ply, obj and glb.
|
||||
"""
|
||||
if path.endswith(".ply"):
|
||||
self.write_ply(path)
|
||||
elif path.lower().endswith(".obj"):
|
||||
elif path.endswith(".obj"):
|
||||
self.write_obj(path)
|
||||
elif path.lower().endswith(".glb") or path.lower().endswith(".gltf"):
|
||||
elif path.endswith(".glb") or path.endswith(".gltf"):
|
||||
self.write_glb(path)
|
||||
else:
|
||||
raise NotImplementedError(f"format {path} not supported!")
|
||||
|
||||
# write to ply file (only geom)
|
||||
def write_ply(self, path):
|
||||
"""write the mesh in ply format. Only for geometry!
|
||||
|
||||
Args:
|
||||
path (str): path to write.
|
||||
"""
|
||||
|
||||
if self.albedo is not None:
|
||||
print(f'[WARN] ply format does not support exporting texture, will ignore!')
|
||||
|
||||
v_np = self.v.detach().cpu().numpy()
|
||||
f_np = self.f.detach().cpu().numpy()
|
||||
@@ -448,70 +591,44 @@ class Mesh:
|
||||
_mesh = trimesh.Trimesh(vertices=v_np, faces=f_np)
|
||||
_mesh.export(path)
|
||||
|
||||
# write to gltf/glb file (geom + texture)
|
||||
def write_glb(self, path):
|
||||
|
||||
assert self.vn is not None and self.vt is not None # should be improved to support export without texture...
|
||||
def write_glb(self, path):
|
||||
"""write the mesh in glb/gltf format.
|
||||
This will create a scene with a single mesh.
|
||||
|
||||
Args:
|
||||
path (str): path to write.
|
||||
"""
|
||||
|
||||
# assert self.v.shape[0] == self.vn.shape[0] and self.v.shape[0] == self.vt.shape[0]
|
||||
if self.v.shape[0] != self.vt.shape[0]:
|
||||
if self.vt is not None and self.v.shape[0] != self.vt.shape[0]:
|
||||
self.align_v_to_vt()
|
||||
|
||||
# assume f == fn == ft
|
||||
|
||||
import pygltflib
|
||||
|
||||
f_np = self.f.detach().cpu().numpy().astype(np.uint32)
|
||||
v_np = self.v.detach().cpu().numpy().astype(np.float32)
|
||||
# vn_np = self.vn.detach().cpu().numpy().astype(np.float32)
|
||||
vt_np = self.vt.detach().cpu().numpy().astype(np.float32)
|
||||
|
||||
albedo = self.albedo.detach().cpu().numpy()
|
||||
albedo = (albedo * 255).astype(np.uint8)
|
||||
albedo = cv2.cvtColor(albedo, cv2.COLOR_RGB2BGR)
|
||||
|
||||
f_np_blob = f_np.flatten().tobytes()
|
||||
v_np_blob = v_np.tobytes()
|
||||
# vn_np_blob = vn_np.tobytes()
|
||||
vt_np_blob = vt_np.tobytes()
|
||||
albedo_blob = cv2.imencode('.png', albedo)[1].tobytes()
|
||||
|
||||
v_np = self.v.detach().cpu().numpy().astype(np.float32)
|
||||
v_np_blob = v_np.tobytes()
|
||||
|
||||
blob = f_np_blob + v_np_blob
|
||||
byteOffset = len(blob)
|
||||
|
||||
# base mesh
|
||||
gltf = pygltflib.GLTF2(
|
||||
scene=0,
|
||||
scenes=[pygltflib.Scene(nodes=[0])],
|
||||
nodes=[pygltflib.Node(mesh=0)],
|
||||
meshes=[pygltflib.Mesh(primitives=[
|
||||
pygltflib.Primitive(
|
||||
# indices to accessors (0 is triangles)
|
||||
attributes=pygltflib.Attributes(
|
||||
POSITION=1, TEXCOORD_0=2,
|
||||
),
|
||||
indices=0, material=0,
|
||||
)
|
||||
])],
|
||||
materials=[
|
||||
pygltflib.Material(
|
||||
pbrMetallicRoughness=pygltflib.PbrMetallicRoughness(
|
||||
baseColorTexture=pygltflib.TextureInfo(index=0, texCoord=0),
|
||||
metallicFactor=0.0,
|
||||
roughnessFactor=1.0,
|
||||
),
|
||||
alphaCutoff=0,
|
||||
doubleSided=True,
|
||||
)
|
||||
],
|
||||
textures=[
|
||||
pygltflib.Texture(sampler=0, source=0),
|
||||
],
|
||||
samplers=[
|
||||
pygltflib.Sampler(magFilter=pygltflib.LINEAR, minFilter=pygltflib.LINEAR_MIPMAP_LINEAR, wrapS=pygltflib.REPEAT, wrapT=pygltflib.REPEAT),
|
||||
],
|
||||
images=[
|
||||
# use embedded (buffer) image
|
||||
pygltflib.Image(bufferView=3, mimeType="image/png"),
|
||||
],
|
||||
meshes=[pygltflib.Mesh(primitives=[pygltflib.Primitive(
|
||||
# indices to accessors (0 is triangles)
|
||||
attributes=pygltflib.Attributes(
|
||||
POSITION=1,
|
||||
),
|
||||
indices=0,
|
||||
)])],
|
||||
buffers=[
|
||||
pygltflib.Buffer(byteLength=len(f_np_blob) + len(v_np_blob) + len(vt_np_blob) + len(albedo_blob))
|
||||
pygltflib.Buffer(byteLength=len(f_np_blob) + len(v_np_blob))
|
||||
],
|
||||
# buffer view (based on dtype)
|
||||
bufferViews=[
|
||||
@@ -529,20 +646,6 @@ class Mesh:
|
||||
byteStride=12, # vec3
|
||||
target=pygltflib.ARRAY_BUFFER, # GL_ARRAY_BUFFER (34962)
|
||||
),
|
||||
# texcoords; as vec2 array
|
||||
pygltflib.BufferView(
|
||||
buffer=0,
|
||||
byteOffset=len(f_np_blob) + len(v_np_blob),
|
||||
byteLength=len(vt_np_blob),
|
||||
byteStride=8, # vec2
|
||||
target=pygltflib.ARRAY_BUFFER,
|
||||
),
|
||||
# texture; as none target
|
||||
pygltflib.BufferView(
|
||||
buffer=0,
|
||||
byteOffset=len(f_np_blob) + len(v_np_blob) + len(vt_np_blob),
|
||||
byteLength=len(albedo_blob),
|
||||
),
|
||||
],
|
||||
accessors=[
|
||||
# 0 = triangles
|
||||
@@ -563,6 +666,53 @@ class Mesh:
|
||||
max=v_np.max(axis=0).tolist(),
|
||||
min=v_np.min(axis=0).tolist(),
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
# append texture info
|
||||
if self.vt is not None:
|
||||
|
||||
vt_np = self.vt.detach().cpu().numpy().astype(np.float32)
|
||||
vt_np_blob = vt_np.tobytes()
|
||||
|
||||
albedo = self.albedo.detach().cpu().numpy()
|
||||
albedo = (albedo * 255).astype(np.uint8)
|
||||
albedo = cv2.cvtColor(albedo, cv2.COLOR_RGB2BGR)
|
||||
albedo_blob = cv2.imencode('.png', albedo)[1].tobytes()
|
||||
|
||||
# update primitive
|
||||
gltf.meshes[0].primitives[0].attributes.TEXCOORD_0 = 2
|
||||
gltf.meshes[0].primitives[0].material = 0
|
||||
|
||||
# update materials
|
||||
gltf.materials.append(pygltflib.Material(
|
||||
pbrMetallicRoughness=pygltflib.PbrMetallicRoughness(
|
||||
baseColorTexture=pygltflib.TextureInfo(index=0, texCoord=0),
|
||||
metallicFactor=0.0,
|
||||
roughnessFactor=1.0,
|
||||
),
|
||||
alphaMode=pygltflib.OPAQUE,
|
||||
alphaCutoff=None,
|
||||
doubleSided=True,
|
||||
))
|
||||
|
||||
gltf.textures.append(pygltflib.Texture(sampler=0, source=0))
|
||||
gltf.samplers.append(pygltflib.Sampler(magFilter=pygltflib.LINEAR, minFilter=pygltflib.LINEAR_MIPMAP_LINEAR, wrapS=pygltflib.REPEAT, wrapT=pygltflib.REPEAT))
|
||||
gltf.images.append(pygltflib.Image(bufferView=3, mimeType="image/png"))
|
||||
|
||||
# update buffers
|
||||
gltf.bufferViews.append(
|
||||
# index = 2, texcoords; as vec2 array
|
||||
pygltflib.BufferView(
|
||||
buffer=0,
|
||||
byteOffset=byteOffset,
|
||||
byteLength=len(vt_np_blob),
|
||||
byteStride=8, # vec2
|
||||
target=pygltflib.ARRAY_BUFFER,
|
||||
)
|
||||
)
|
||||
|
||||
gltf.accessors.append(
|
||||
# 2 = texcoords
|
||||
pygltflib.Accessor(
|
||||
bufferView=2,
|
||||
@@ -571,25 +721,81 @@ class Mesh:
|
||||
type=pygltflib.VEC2,
|
||||
max=vt_np.max(axis=0).tolist(),
|
||||
min=vt_np.min(axis=0).tolist(),
|
||||
),
|
||||
],
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
blob += vt_np_blob
|
||||
byteOffset += len(vt_np_blob)
|
||||
|
||||
gltf.bufferViews.append(
|
||||
# index = 3, albedo texture; as none target
|
||||
pygltflib.BufferView(
|
||||
buffer=0,
|
||||
byteOffset=byteOffset,
|
||||
byteLength=len(albedo_blob),
|
||||
)
|
||||
)
|
||||
|
||||
blob += albedo_blob
|
||||
byteOffset += len(albedo_blob)
|
||||
|
||||
gltf.buffers[0].byteLength = byteOffset
|
||||
|
||||
# append metllic roughness
|
||||
if self.metallicRoughness is not None:
|
||||
metallicRoughness = self.metallicRoughness.detach().cpu().numpy()
|
||||
metallicRoughness = (metallicRoughness * 255).astype(np.uint8)
|
||||
metallicRoughness = cv2.cvtColor(metallicRoughness, cv2.COLOR_RGB2BGR)
|
||||
metallicRoughness_blob = cv2.imencode('.png', metallicRoughness)[1].tobytes()
|
||||
|
||||
# update texture definition
|
||||
gltf.materials[0].pbrMetallicRoughness.metallicFactor = 1.0
|
||||
gltf.materials[0].pbrMetallicRoughness.roughnessFactor = 1.0
|
||||
gltf.materials[0].pbrMetallicRoughness.metallicRoughnessTexture = pygltflib.TextureInfo(index=1, texCoord=0)
|
||||
|
||||
gltf.textures.append(pygltflib.Texture(sampler=1, source=1))
|
||||
gltf.samplers.append(pygltflib.Sampler(magFilter=pygltflib.LINEAR, minFilter=pygltflib.LINEAR_MIPMAP_LINEAR, wrapS=pygltflib.REPEAT, wrapT=pygltflib.REPEAT))
|
||||
gltf.images.append(pygltflib.Image(bufferView=4, mimeType="image/png"))
|
||||
|
||||
# update buffers
|
||||
gltf.bufferViews.append(
|
||||
# index = 4, metallicRoughness texture; as none target
|
||||
pygltflib.BufferView(
|
||||
buffer=0,
|
||||
byteOffset=byteOffset,
|
||||
byteLength=len(metallicRoughness_blob),
|
||||
)
|
||||
)
|
||||
|
||||
blob += metallicRoughness_blob
|
||||
byteOffset += len(metallicRoughness_blob)
|
||||
|
||||
gltf.buffers[0].byteLength = byteOffset
|
||||
|
||||
|
||||
# set actual data
|
||||
gltf.set_binary_blob(f_np_blob + v_np_blob + vt_np_blob + albedo_blob)
|
||||
gltf.set_binary_blob(blob)
|
||||
|
||||
# glb = b"".join(gltf.save_to_bytes())
|
||||
gltf.save(path)
|
||||
|
||||
# write to obj file (geom + texture)
|
||||
|
||||
def write_obj(self, path):
|
||||
"""write the mesh in obj format. Will also write the texture and mtl files.
|
||||
|
||||
Args:
|
||||
path (str): path to write.
|
||||
"""
|
||||
|
||||
mtl_path = path.replace(".obj", ".mtl")
|
||||
albedo_path = path.replace(".obj", "_albedo.png")
|
||||
metallic_path = path.replace(".obj", "_metallic.png")
|
||||
roughness_path = path.replace(".obj", "_roughness.png")
|
||||
|
||||
v_np = self.v.detach().cpu().numpy()
|
||||
vt_np = self.vt.detach().cpu().numpy() if self.vt is not None else None
|
||||
vn_np = self.vn.detach().cpu().numpy() if self.vn is not None else None
|
||||
vc_np = self.vc.detach().cpu().numpy() if self.vc is not None else None
|
||||
f_np = self.f.detach().cpu().numpy()
|
||||
ft_np = self.ft.detach().cpu().numpy() if self.ft is not None else None
|
||||
fn_np = self.fn.detach().cpu().numpy() if self.fn is not None else None
|
||||
@@ -597,8 +803,14 @@ class Mesh:
|
||||
with open(path, "w") as fp:
|
||||
fp.write(f"mtllib {os.path.basename(mtl_path)} \n")
|
||||
|
||||
for v in v_np:
|
||||
fp.write(f"v {v[0]} {v[1]} {v[2]} \n")
|
||||
if vc_np is None:
|
||||
for v in v_np:
|
||||
fp.write(f"v {v[0]} {v[1]} {v[2]} \n")
|
||||
else:
|
||||
for v_i in range(v_np.shape[0]):
|
||||
v = v_np[v_i]
|
||||
vc = vc_np[v_i]
|
||||
fp.write(f"v {v[0]} {v[1]} {v[2]} {vc[0]} {vc[1]} {vc[2]}\n")
|
||||
|
||||
if vt_np is not None:
|
||||
for v in vt_np:
|
||||
@@ -626,12 +838,22 @@ class Mesh:
|
||||
fp.write(f"Ns 0 \n")
|
||||
if self.albedo is not None:
|
||||
fp.write(f"map_Kd {os.path.basename(albedo_path)} \n")
|
||||
if self.metallicRoughness is not None:
|
||||
# ref: https://en.wikipedia.org/wiki/Wavefront_.obj_file#Physically-based_Rendering
|
||||
fp.write(f"map_Pm {os.path.basename(metallic_path)} \n")
|
||||
fp.write(f"map_Pr {os.path.basename(roughness_path)} \n")
|
||||
|
||||
if self.albedo is not None:
|
||||
albedo = self.albedo.detach().cpu().numpy()
|
||||
albedo = (albedo * 255).astype(np.uint8)
|
||||
cv2.imwrite(albedo_path, cv2.cvtColor(albedo, cv2.COLOR_RGB2BGR))
|
||||
|
||||
if self.metallicRoughness is not None:
|
||||
metallicRoughness = self.metallicRoughness.detach().cpu().numpy()
|
||||
metallicRoughness = (metallicRoughness * 255).astype(np.uint8)
|
||||
cv2.imwrite(metallic_path, metallicRoughness[..., 2])
|
||||
cv2.imwrite(roughness_path, metallicRoughness[..., 1])
|
||||
|
||||
def convert_to_pointcloud(self):
|
||||
xyz = self.v.detach().cpu().numpy()
|
||||
num_pts = self.v.shape[0]
|
||||
|
||||
@@ -447,7 +447,7 @@ def switch_ply_axis_and_scale(plydata, target_axis, target_scale, coordinate_inv
|
||||
|
||||
return write_gs_ply(xyz, normals, features_dc_2d, features_extra_2d, opacities, scales, rots, construct_list_of_gs_attributes(features_dc, features_extra, scales, rots))
|
||||
|
||||
def switch_mesh_axis_and_scale(mesh, target_axis, target_scale):
|
||||
def switch_mesh_axis_and_scale(mesh, target_axis, target_scale, flip_normal=False):
|
||||
"""
|
||||
Args:
|
||||
target_axis (array): shape (3)
|
||||
@@ -456,5 +456,7 @@ def switch_mesh_axis_and_scale(mesh, target_axis, target_scale):
|
||||
target_scale = torch.tensor(target_scale).float().cuda()
|
||||
mesh.v = switch_vector_axis(mesh.v * target_scale, target_axis)
|
||||
mesh.vn = switch_vector_axis(mesh.vn * target_scale, target_axis)
|
||||
if flip_normal:
|
||||
mesh.vn *= -1
|
||||
return mesh
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
-5.266582965850830078e-01 7.410295009613037109e-01 -4.165407419204711914e-01 -5.960464477539062500e-08
|
||||
5.865638996738198330e-08 4.900035560131072998e-01 8.717204332351684570e-01 -9.462351613365171943e-08
|
||||
8.500770330429077148e-01 4.590988159179687500e-01 -2.580644786357879639e-01 -1.300000071525573730e+00
|
||||
@@ -0,0 +1,3 @@
|
||||
-9.734988808631896973e-01 1.993551850318908691e-01 -1.120596975088119507e-01 -1.713633537292480469e-07
|
||||
3.790224578636980368e-09 4.900034964084625244e-01 8.717204928398132324e-01 1.772203575001185527e-07
|
||||
2.286916375160217285e-01 8.486189246177673340e-01 -4.770178496837615967e-01 -1.838477611541748047e+00
|
||||
@@ -0,0 +1,3 @@
|
||||
2.286914736032485962e-01 8.486190438270568848e-01 -4.770178198814392090e-01 1.564621925354003906e-07
|
||||
-3.417914484771245043e-08 4.900034070014953613e-01 8.717205524444580078e-01 -7.293811421504869941e-08
|
||||
9.734990000724792480e-01 -1.993550658226013184e-01 1.120596155524253845e-01 -1.838477969169616699e+00
|
||||
@@ -0,0 +1,3 @@
|
||||
5.266583561897277832e-01 -7.410295009613037109e-01 4.165407419204711914e-01 0.000000000000000000e+00
|
||||
5.865638996738198330e-08 4.900035560131072998e-01 8.717204332351684570e-01 9.462351613365171943e-08
|
||||
-8.500770330429077148e-01 -4.590988159179687500e-01 2.580645382404327393e-01 -1.300000071525573730e+00
|
||||
@@ -0,0 +1,3 @@
|
||||
-2.286916971206665039e-01 -8.486189842224121094e-01 4.770179092884063721e-01 -2.458691596984863281e-07
|
||||
9.085837859856837895e-09 4.900034666061401367e-01 8.717205524444580078e-01 1.205695667749751010e-07
|
||||
-9.734990000724792480e-01 1.993551701307296753e-01 -1.120597645640373230e-01 -1.838477969169616699e+00
|
||||
@@ -0,0 +1,3 @@
|
||||
9.734989404678344727e-01 -1.993551850318908691e-01 1.120596975088119507e-01 -1.415610313415527344e-07
|
||||
3.790224578636980368e-09 4.900034964084625244e-01 8.717204928398132324e-01 -1.772203575001185527e-07
|
||||
-2.286916375160217285e-01 -8.486189246177673340e-01 4.770178794860839844e-01 -1.838477611541748047e+00
|
||||
@@ -0,0 +1,3 @@
|
||||
-8.500771522521972656e-01 -4.590989053249359131e-01 2.580644488334655762e-01 0.000000000000000000e+00
|
||||
-4.257411134744870651e-08 4.900034964084625244e-01 8.717204928398132324e-01 9.006067358541258727e-08
|
||||
-5.266583561897277832e-01 7.410295605659484863e-01 -4.165408313274383545e-01 -1.300000071525573730e+00
|
||||
@@ -0,0 +1,3 @@
|
||||
8.500770330429077148e-01 4.590989053249359131e-01 -2.580644488334655762e-01 5.960464477539062500e-08
|
||||
-4.257411134744870651e-08 4.900034964084625244e-01 8.717204928398132324e-01 -9.006067358541258727e-08
|
||||
5.266583561897277832e-01 -7.410295605659484863e-01 4.165407419204711914e-01 -1.300000071525573730e+00
|
||||
@@ -0,0 +1,3 @@
|
||||
9.958608150482177734e-01 7.923202216625213623e-02 -4.453715682029724121e-02 -3.098167056236889039e-09
|
||||
-9.089154005050659180e-02 8.681122064590454102e-01 -4.879753291606903076e-01 5.784738377201392723e-08
|
||||
-2.028124157504862524e-08 4.900035560131072998e-01 8.717204332351684570e-01 -1.300000071525573730e+00
|
||||
@@ -0,0 +1,45 @@
|
||||
import numpy as np
|
||||
|
||||
def camNormal2worldNormal(rot_c2w, camNormal):
|
||||
H,W,_ = camNormal.shape
|
||||
normal_img = np.matmul(rot_c2w[None, :, :], camNormal.reshape(-1,3)[:, :, None]).reshape([H, W, 3])
|
||||
|
||||
return normal_img
|
||||
|
||||
def worldNormal2camNormal(rot_w2c, normal_map_world):
|
||||
H,W,_ = normal_map_world.shape
|
||||
# normal_img = np.matmul(rot_w2c[None, :, :], worldNormal.reshape(-1,3)[:, :, None]).reshape([H, W, 3])
|
||||
|
||||
# faster version
|
||||
# Reshape the normal map into a 2D array where each row represents a normal vector
|
||||
normal_map_flat = normal_map_world.reshape(-1, 3)
|
||||
|
||||
# Transform the normal vectors using the transformation matrix
|
||||
normal_map_camera_flat = np.dot(normal_map_flat, rot_w2c.T)
|
||||
|
||||
# Reshape the transformed normal map back to its original shape
|
||||
normal_map_camera = normal_map_camera_flat.reshape(normal_map_world.shape)
|
||||
|
||||
return normal_map_camera
|
||||
|
||||
def trans_normal(normal, RT_w2c, RT_w2c_target):
|
||||
|
||||
# normal_world = camNormal2worldNormal(np.linalg.inv(RT_w2c[:3,:3]), normal)
|
||||
# normal_target_cam = worldNormal2camNormal(RT_w2c_target[:3,:3], normal_world)
|
||||
|
||||
relative_RT = np.matmul(RT_w2c_target[:3,:3], np.linalg.inv(RT_w2c[:3,:3]))
|
||||
normal_target_cam = worldNormal2camNormal(relative_RT[:3,:3], normal)
|
||||
|
||||
return normal_target_cam
|
||||
|
||||
def img2normal(img):
|
||||
return (img/255.)*2-1
|
||||
|
||||
def normal2img(normal):
|
||||
return np.uint8((normal*0.5+0.5)*255)
|
||||
|
||||
def norm_normalize(normal, dim=-1):
|
||||
|
||||
normal = normal/(np.linalg.norm(normal, axis=dim, keepdims=True)+1e-6)
|
||||
|
||||
return normal
|
||||
@@ -0,0 +1,510 @@
|
||||
from typing import Dict
|
||||
import numpy as np
|
||||
from omegaconf import DictConfig, ListConfig
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from pathlib import Path
|
||||
import json
|
||||
from PIL import Image
|
||||
from torchvision import transforms
|
||||
from einops import rearrange
|
||||
from typing import Literal, Tuple, Optional, Any
|
||||
import random
|
||||
|
||||
import json
|
||||
import os, sys
|
||||
import math
|
||||
|
||||
from .normal_utils import trans_normal, normal2img, img2normal
|
||||
|
||||
|
||||
class ObjaverseDataset(Dataset):
|
||||
def __init__(self,
|
||||
root_dir: str,
|
||||
num_views: int,
|
||||
bg_color: Any,
|
||||
img_wh: Tuple[int, int],
|
||||
object_list: str,
|
||||
groups_num: int=1,
|
||||
validation: bool = False,
|
||||
data_view_num: int = 6,
|
||||
num_validation_samples: int = 64,
|
||||
num_samples: Optional[int] = None,
|
||||
invalid_list: Optional[str] = None,
|
||||
trans_norm_system: bool = True, # if True, transform all normals map into the cam system of front view
|
||||
augment_data: bool = False,
|
||||
read_normal: bool = True,
|
||||
read_color: bool = False,
|
||||
read_depth: bool = False,
|
||||
read_mask: bool = True,
|
||||
mix_color_normal: bool = False,
|
||||
suffix: str = 'png',
|
||||
subscene_tag: int = 3,
|
||||
backup_scene: str = "9438abf986c7453a9f4df7c34aa2e65b"
|
||||
) -> None:
|
||||
"""Create a dataset from a folder of images.
|
||||
If you pass in a root directory it will be searched for images
|
||||
ending in ext (ext can be a list)
|
||||
"""
|
||||
self.root_dir = Path(root_dir)
|
||||
self.num_views = num_views
|
||||
self.bg_color = bg_color
|
||||
self.validation = validation
|
||||
self.num_samples = num_samples
|
||||
self.trans_norm_system = trans_norm_system
|
||||
self.augment_data = augment_data
|
||||
self.invalid_list = invalid_list
|
||||
self.groups_num = groups_num
|
||||
print("augment data: ", self.augment_data)
|
||||
self.img_wh = img_wh
|
||||
self.read_normal = read_normal
|
||||
self.read_color = read_color
|
||||
self.read_depth = read_depth
|
||||
self.read_mask = read_mask
|
||||
self.mix_color_normal = mix_color_normal # mix load color and normal maps
|
||||
self.suffix = suffix
|
||||
self.subscene_tag = subscene_tag
|
||||
|
||||
self.view_types = ['front', 'front_right', 'right', 'back', 'left', 'front_left']
|
||||
self.fix_cam_pose_dir = "./mvdiffusion/data/fixed_poses/nine_views"
|
||||
|
||||
self.fix_cam_poses = self.load_fixed_poses() # world2cam matrix
|
||||
|
||||
if object_list is not None:
|
||||
with open(object_list) as f:
|
||||
self.objects = json.load(f)
|
||||
self.objects = [os.path.basename(o).replace(".glb", "") for o in self.objects]
|
||||
else:
|
||||
self.objects = os.listdir(self.root_dir)
|
||||
self.objects = sorted(self.objects)
|
||||
|
||||
if self.invalid_list is not None:
|
||||
with open(self.invalid_list) as f:
|
||||
self.invalid_objects = json.load(f)
|
||||
self.invalid_objects = [os.path.basename(o).replace(".glb", "") for o in self.invalid_objects]
|
||||
else:
|
||||
self.invalid_objects = []
|
||||
|
||||
|
||||
self.all_objects = set(self.objects) - (set(self.invalid_objects) & set(self.objects))
|
||||
self.all_objects = list(self.all_objects)
|
||||
|
||||
if not validation:
|
||||
self.all_objects = self.all_objects[:-num_validation_samples]
|
||||
else:
|
||||
self.all_objects = self.all_objects[-num_validation_samples:]
|
||||
if num_samples is not None:
|
||||
self.all_objects = self.all_objects[:num_samples]
|
||||
|
||||
print("loading ", len(self.all_objects), " objects in the dataset")
|
||||
|
||||
if self.mix_color_normal:
|
||||
self.backup_data = self.__getitem_mix__(0, backup_scene)
|
||||
else:
|
||||
self.backup_data = self.__getitem_joint__(0, backup_scene)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.objects)*self.total_view
|
||||
|
||||
def load_fixed_poses(self):
|
||||
poses = {}
|
||||
for face in self.view_types:
|
||||
RT = np.loadtxt(os.path.join(self.fix_cam_pose_dir,'%03d_%s_RT.txt'%(0, face)))
|
||||
poses[face] = RT
|
||||
|
||||
return poses
|
||||
|
||||
def cartesian_to_spherical(self, xyz):
|
||||
ptsnew = np.hstack((xyz, np.zeros(xyz.shape)))
|
||||
xy = xyz[:,0]**2 + xyz[:,1]**2
|
||||
z = np.sqrt(xy + xyz[:,2]**2)
|
||||
theta = np.arctan2(np.sqrt(xy), xyz[:,2]) # for elevation angle defined from Z-axis down
|
||||
#ptsnew[:,4] = np.arctan2(xyz[:,2], np.sqrt(xy)) # for elevation angle defined from XY-plane up
|
||||
azimuth = np.arctan2(xyz[:,1], xyz[:,0])
|
||||
return np.array([theta, azimuth, z])
|
||||
|
||||
def get_T(self, target_RT, cond_RT):
|
||||
R, T = target_RT[:3, :3], target_RT[:, -1]
|
||||
T_target = -R.T @ T # change to cam2world
|
||||
|
||||
R, T = cond_RT[:3, :3], cond_RT[:, -1]
|
||||
T_cond = -R.T @ T
|
||||
|
||||
theta_cond, azimuth_cond, z_cond = self.cartesian_to_spherical(T_cond[None, :])
|
||||
theta_target, azimuth_target, z_target = self.cartesian_to_spherical(T_target[None, :])
|
||||
|
||||
d_theta = theta_target - theta_cond
|
||||
d_azimuth = (azimuth_target - azimuth_cond) % (2 * math.pi)
|
||||
d_z = z_target - z_cond
|
||||
|
||||
# d_T = torch.tensor([d_theta.item(), math.sin(d_azimuth.item()), math.cos(d_azimuth.item()), d_z.item()])
|
||||
return d_theta, d_azimuth
|
||||
|
||||
def get_bg_color(self):
|
||||
if self.bg_color == 'white':
|
||||
bg_color = np.array([1., 1., 1.], dtype=np.float32)
|
||||
elif self.bg_color == 'black':
|
||||
bg_color = np.array([0., 0., 0.], dtype=np.float32)
|
||||
elif self.bg_color == 'gray':
|
||||
bg_color = np.array([0.5, 0.5, 0.5], dtype=np.float32)
|
||||
elif self.bg_color == 'random':
|
||||
bg_color = np.random.rand(3)
|
||||
elif self.bg_color == 'three_choices':
|
||||
white = np.array([1., 1., 1.], dtype=np.float32)
|
||||
black = np.array([0., 0., 0.], dtype=np.float32)
|
||||
gray = np.array([0.5, 0.5, 0.5], dtype=np.float32)
|
||||
bg_color = random.choice([white, black, gray])
|
||||
elif isinstance(self.bg_color, float):
|
||||
bg_color = np.array([self.bg_color] * 3, dtype=np.float32)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
return bg_color
|
||||
|
||||
|
||||
|
||||
def load_mask(self, img_path, return_type='np'):
|
||||
# not using cv2 as may load in uint16 format
|
||||
# img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) # [0, 255]
|
||||
# img = cv2.resize(img, self.img_wh, interpolation=cv2.INTER_CUBIC)
|
||||
# pil always returns uint8
|
||||
img = np.array(Image.open(img_path).resize(self.img_wh))
|
||||
img = np.float32(img > 0)
|
||||
|
||||
assert len(np.shape(img)) == 2
|
||||
|
||||
if return_type == "np":
|
||||
pass
|
||||
elif return_type == "pt":
|
||||
img = torch.from_numpy(img)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
return img
|
||||
|
||||
def load_image(self, img_path, bg_color, alpha, return_type='np'):
|
||||
# not using cv2 as may load in uint16 format
|
||||
# img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) # [0, 255]
|
||||
# img = cv2.resize(img, self.img_wh, interpolation=cv2.INTER_CUBIC)
|
||||
# pil always returns uint8
|
||||
img = np.array(Image.open(img_path).resize(self.img_wh))
|
||||
img = img.astype(np.float32) / 255. # [0, 1]
|
||||
assert img.shape[-1] == 3 or img.shape[-1] == 4 # RGB or RGBA
|
||||
|
||||
if alpha is None and img.shape[-1] == 4:
|
||||
alpha = img[:, :, 3:]
|
||||
img = img[:, :, :3]
|
||||
|
||||
if alpha.shape[-1] != 1:
|
||||
alpha = alpha[:, :, None]
|
||||
|
||||
img = img[...,:3] * alpha + bg_color * (1 - alpha)
|
||||
|
||||
if return_type == "np":
|
||||
pass
|
||||
elif return_type == "pt":
|
||||
img = torch.from_numpy(img)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
return img
|
||||
|
||||
def load_depth(self, img_path, bg_color, alpha, return_type='np'):
|
||||
# not using cv2 as may load in uint16 format
|
||||
# img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) # [0, 255]
|
||||
# img = cv2.resize(img, self.img_wh, interpolation=cv2.INTER_CUBIC)
|
||||
# pil always returns uint8
|
||||
img = np.array(Image.open(img_path).resize(self.img_wh))
|
||||
img = img.astype(np.float32) / 65535. # [0, 1]
|
||||
|
||||
img[img > 0.4] = 0
|
||||
img = img / 0.4
|
||||
|
||||
assert img.ndim == 2 # depth
|
||||
img = np.stack([img]*3, axis=-1)
|
||||
|
||||
if alpha.shape[-1] != 1:
|
||||
alpha = alpha[:, :, None]
|
||||
|
||||
# print(np.max(img[:, :, 0]))
|
||||
|
||||
img = img[...,:3] * alpha + bg_color * (1 - alpha)
|
||||
|
||||
if return_type == "np":
|
||||
pass
|
||||
elif return_type == "pt":
|
||||
img = torch.from_numpy(img)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
return img
|
||||
|
||||
def load_normal(self, img_path, bg_color, alpha, RT_w2c=None, RT_w2c_cond=None, return_type='np'):
|
||||
# not using cv2 as may load in uint16 format
|
||||
# img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) # [0, 255]
|
||||
# img = cv2.resize(img, self.img_wh, interpolation=cv2.INTER_CUBIC)
|
||||
# pil always returns uint8
|
||||
normal = np.array(Image.open(img_path).resize(self.img_wh))
|
||||
|
||||
assert normal.shape[-1] == 3 or normal.shape[-1] == 4 # RGB or RGBA
|
||||
|
||||
if alpha is None and normal.shape[-1] == 4:
|
||||
alpha = normal[:, :, 3:] / 255.
|
||||
normal = normal[:, :, :3]
|
||||
|
||||
normal = trans_normal(img2normal(normal), RT_w2c, RT_w2c_cond)
|
||||
|
||||
img = (normal*0.5 + 0.5).astype(np.float32) # [0, 1]
|
||||
|
||||
if alpha.shape[-1] != 1:
|
||||
alpha = alpha[:, :, None]
|
||||
|
||||
img = img[...,:3] * alpha + bg_color * (1 - alpha)
|
||||
|
||||
if return_type == "np":
|
||||
pass
|
||||
elif return_type == "pt":
|
||||
img = torch.from_numpy(img)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
return img
|
||||
|
||||
def __len__(self):
|
||||
return len(self.all_objects)
|
||||
|
||||
def __getitem_mix__(self, index, debug_object=None):
|
||||
if debug_object is not None:
|
||||
object_name = debug_object #
|
||||
set_idx = random.sample(range(0, self.groups_num), 1)[0] # without replacement
|
||||
else:
|
||||
object_name = self.all_objects[index%len(self.all_objects)]
|
||||
set_idx = 0
|
||||
|
||||
if self.augment_data:
|
||||
cond_view = random.sample(self.view_types, k=1)[0]
|
||||
else:
|
||||
cond_view = 'front'
|
||||
|
||||
|
||||
# ! if you would like predict depth; modify here
|
||||
if random.random() < 0.5:
|
||||
read_color, read_normal, read_depth = True, False, False
|
||||
else:
|
||||
read_color, read_normal, read_depth = False, True, False
|
||||
|
||||
read_normal = read_normal & self.read_normal
|
||||
read_depth = read_depth & self.read_depth
|
||||
|
||||
assert (read_color and (read_normal or read_depth)) is False
|
||||
|
||||
view_types = self.view_types
|
||||
|
||||
cond_w2c = self.fix_cam_poses[cond_view]
|
||||
|
||||
tgt_w2cs = [self.fix_cam_poses[view] for view in view_types]
|
||||
|
||||
elevations = []
|
||||
azimuths = []
|
||||
|
||||
# get the bg color
|
||||
bg_color = self.get_bg_color()
|
||||
|
||||
if self.read_mask:
|
||||
cond_alpha = self.load_mask(os.path.join(self.root_dir, object_name[:self.subscene_tag], object_name, "mask_%03d_%s.%s" % (set_idx, cond_view, self.suffix)), return_type='np')
|
||||
else:
|
||||
cond_alpha = None
|
||||
img_tensors_in = [
|
||||
self.load_image(os.path.join(self.root_dir, object_name[:self.subscene_tag], object_name, "rgb_%03d_%s.%s" % (set_idx, cond_view, self.suffix)), bg_color, cond_alpha, return_type='pt').permute(2, 0, 1)
|
||||
] * self.num_views
|
||||
img_tensors_out = []
|
||||
|
||||
for view, tgt_w2c in zip(view_types, tgt_w2cs):
|
||||
img_path = os.path.join(self.root_dir, object_name[:self.subscene_tag], object_name, "rgb_%03d_%s.%s" % (set_idx, view, self.suffix))
|
||||
mask_path = os.path.join(self.root_dir, object_name[:self.subscene_tag], object_name, "mask_%03d_%s.%s" % (set_idx, view, self.suffix))
|
||||
normal_path = os.path.join(self.root_dir, object_name[:self.subscene_tag], object_name, "normals_%03d_%s.%s" % (set_idx, view, self.suffix))
|
||||
depth_path = os.path.join(self.root_dir, object_name[:self.subscene_tag], object_name, "depth_%03d_%s.%s" % (set_idx, view, self.suffix))
|
||||
if self.read_mask:
|
||||
alpha = self.load_mask(mask_path, return_type='np')
|
||||
else:
|
||||
alpha = None
|
||||
|
||||
if read_color:
|
||||
img_tensor = self.load_image(img_path, bg_color, alpha, return_type="pt")
|
||||
img_tensor = img_tensor.permute(2, 0, 1)
|
||||
img_tensors_out.append(img_tensor)
|
||||
|
||||
if read_normal:
|
||||
normal_tensor = self.load_normal(normal_path, bg_color, alpha, RT_w2c=tgt_w2c, RT_w2c_cond=cond_w2c, return_type="pt").permute(2, 0, 1)
|
||||
img_tensors_out.append(normal_tensor)
|
||||
if read_depth:
|
||||
depth_tensor = self.load_depth(depth_path, bg_color, alpha, return_type="pt").permute(2, 0, 1)
|
||||
img_tensors_out.append(depth_tensor)
|
||||
|
||||
# evelations, azimuths
|
||||
elevation, azimuth = self.get_T(tgt_w2c, cond_w2c)
|
||||
elevations.append(elevation)
|
||||
azimuths.append(azimuth)
|
||||
|
||||
img_tensors_in = torch.stack(img_tensors_in, dim=0).float() # (Nv, 3, H, W)
|
||||
img_tensors_out = torch.stack(img_tensors_out, dim=0).float() # (Nv, 3, H, W)
|
||||
|
||||
|
||||
elevations = torch.as_tensor(elevations).float().squeeze(1)
|
||||
azimuths = torch.as_tensor(azimuths).float().squeeze(1)
|
||||
elevations_cond = torch.as_tensor([0] * self.num_views).float() # fixed only use 4 views to train
|
||||
camera_embeddings = torch.stack([elevations_cond, elevations, azimuths], dim=-1) # (Nv, 3)
|
||||
|
||||
normal_class = torch.tensor([1, 0]).float()
|
||||
normal_task_embeddings = torch.stack([normal_class]*self.num_views, dim=0) # (Nv, 2)
|
||||
color_class = torch.tensor([0, 1]).float()
|
||||
color_task_embeddings = torch.stack([color_class]*self.num_views, dim=0) # (Nv, 2)
|
||||
if read_normal or read_depth:
|
||||
task_embeddings = normal_task_embeddings
|
||||
if read_color:
|
||||
task_embeddings = color_task_embeddings
|
||||
# print(elevations)
|
||||
# print(azimuths)
|
||||
return {
|
||||
'elevations_cond': elevations_cond,
|
||||
'elevations_cond_deg': torch.rad2deg(elevations_cond),
|
||||
'elevations': elevations,
|
||||
'azimuths': azimuths,
|
||||
'elevations_deg': torch.rad2deg(elevations),
|
||||
'azimuths_deg': torch.rad2deg(azimuths),
|
||||
'imgs_in': img_tensors_in,
|
||||
'imgs_out': img_tensors_out,
|
||||
'camera_embeddings': camera_embeddings,
|
||||
'task_embeddings': task_embeddings
|
||||
}
|
||||
|
||||
|
||||
def __getitem_joint__(self, index, debug_object=None):
|
||||
if debug_object is not None:
|
||||
object_name = debug_object #
|
||||
set_idx = random.sample(range(0, self.groups_num), 1)[0] # without replacement
|
||||
else:
|
||||
object_name = self.all_objects[index%len(self.all_objects)]
|
||||
set_idx = 0
|
||||
|
||||
if self.augment_data:
|
||||
cond_view = random.sample(self.view_types, k=1)[0]
|
||||
else:
|
||||
cond_view = 'front'
|
||||
|
||||
view_types = self.view_types
|
||||
|
||||
cond_w2c = self.fix_cam_poses[cond_view]
|
||||
|
||||
tgt_w2cs = [self.fix_cam_poses[view] for view in view_types]
|
||||
|
||||
elevations = []
|
||||
azimuths = []
|
||||
|
||||
# get the bg color
|
||||
bg_color = self.get_bg_color()
|
||||
|
||||
if self.read_mask:
|
||||
cond_alpha = self.load_mask(os.path.join(self.root_dir, object_name[:self.subscene_tag], object_name, "mask_%03d_%s.%s" % (set_idx, cond_view, self.suffix)), return_type='np')
|
||||
else:
|
||||
cond_alpha = None
|
||||
img_tensors_in = [
|
||||
self.load_image(os.path.join(self.root_dir, object_name[:self.subscene_tag], object_name, "rgb_%03d_%s.%s" % (set_idx, cond_view, self.suffix)), bg_color, cond_alpha, return_type='pt').permute(2, 0, 1)
|
||||
] * self.num_views
|
||||
img_tensors_out = []
|
||||
normal_tensors_out = []
|
||||
for view, tgt_w2c in zip(view_types, tgt_w2cs):
|
||||
img_path = os.path.join(self.root_dir, object_name[:self.subscene_tag], object_name, "rgb_%03d_%s.%s" % (set_idx, view, self.suffix))
|
||||
mask_path = os.path.join(self.root_dir, object_name[:self.subscene_tag], object_name, "mask_%03d_%s.%s" % (set_idx, view, self.suffix))
|
||||
if self.read_mask:
|
||||
alpha = self.load_mask(mask_path, return_type='np')
|
||||
else:
|
||||
alpha = None
|
||||
|
||||
if self.read_color:
|
||||
img_tensor = self.load_image(img_path, bg_color, alpha, return_type="pt")
|
||||
img_tensor = img_tensor.permute(2, 0, 1)
|
||||
img_tensors_out.append(img_tensor)
|
||||
|
||||
if self.read_normal:
|
||||
normal_path = os.path.join(self.root_dir, object_name[:self.subscene_tag], object_name, "normals_%03d_%s.%s" % (set_idx, view, self.suffix))
|
||||
normal_tensor = self.load_normal(normal_path, bg_color, alpha, RT_w2c=tgt_w2c, RT_w2c_cond=cond_w2c, return_type="pt").permute(2, 0, 1)
|
||||
normal_tensors_out.append(normal_tensor)
|
||||
|
||||
# evelations, azimuths
|
||||
elevation, azimuth = self.get_T(tgt_w2c, cond_w2c)
|
||||
elevations.append(elevation)
|
||||
azimuths.append(azimuth)
|
||||
|
||||
img_tensors_in = torch.stack(img_tensors_in, dim=0).float() # (Nv, 3, H, W)
|
||||
if self.read_color:
|
||||
img_tensors_out = torch.stack(img_tensors_out, dim=0).float() # (Nv, 3, H, W)
|
||||
if self.read_normal:
|
||||
normal_tensors_out = torch.stack(normal_tensors_out, dim=0).float() # (Nv, 3, H, W)
|
||||
|
||||
elevations = torch.as_tensor(elevations).float().squeeze(1)
|
||||
azimuths = torch.as_tensor(azimuths).float().squeeze(1)
|
||||
elevations_cond = torch.as_tensor([0] * self.num_views).float() # fixed only use 4 views to train
|
||||
|
||||
camera_embeddings = torch.stack([elevations_cond, elevations, azimuths], dim=-1) # (Nv, 3)
|
||||
|
||||
normal_class = torch.tensor([1, 0]).float()
|
||||
normal_task_embeddings = torch.stack([normal_class]*self.num_views, dim=0) # (Nv, 2)
|
||||
color_class = torch.tensor([0, 1]).float()
|
||||
color_task_embeddings = torch.stack([color_class]*self.num_views, dim=0) # (Nv, 2)
|
||||
|
||||
return {
|
||||
'elevations_cond': elevations_cond,
|
||||
'elevations_cond_deg': torch.rad2deg(elevations_cond),
|
||||
'elevations': elevations,
|
||||
'azimuths': azimuths,
|
||||
'elevations_deg': torch.rad2deg(elevations),
|
||||
'azimuths_deg': torch.rad2deg(azimuths),
|
||||
'imgs_in': img_tensors_in,
|
||||
'imgs_out': img_tensors_out,
|
||||
'normals_out': normal_tensors_out,
|
||||
'camera_embeddings': camera_embeddings,
|
||||
'normal_task_embeddings': normal_task_embeddings,
|
||||
'color_task_embeddings': color_task_embeddings
|
||||
}
|
||||
|
||||
def __getitem__(self, index):
|
||||
try:
|
||||
if self.mix_color_normal:
|
||||
data = self.__getitem_mix__(index)
|
||||
else:
|
||||
data = self.__getitem_joint__(index)
|
||||
return data
|
||||
except:
|
||||
print("load error ", self.all_objects[index%len(self.all_objects)] )
|
||||
return self.backup_data
|
||||
|
||||
|
||||
class ConcatDataset(torch.utils.data.Dataset):
|
||||
def __init__(self, datasets, weights):
|
||||
self.datasets = datasets
|
||||
self.weights = weights
|
||||
self.num_datasets = len(datasets)
|
||||
|
||||
def __getitem__(self, i):
|
||||
|
||||
chosen = random.choices(self.datasets, self.weights, k=1)[0]
|
||||
return chosen[i]
|
||||
|
||||
def __len__(self):
|
||||
return max(len(d) for d in self.datasets)
|
||||
|
||||
if __name__ == "__main__":
|
||||
train_dataset = ObjaverseDataset(
|
||||
root_dir="/ghome/l5/xxlong/.objaverse/hf-objaverse-v1/renderings",
|
||||
size=(128, 128),
|
||||
ext="hdf5",
|
||||
default_trans=torch.zeros(3),
|
||||
return_paths=False,
|
||||
total_view=8,
|
||||
validation=False,
|
||||
object_list=None,
|
||||
views_mode='fourviews'
|
||||
)
|
||||
data0 = train_dataset[0]
|
||||
data1 = train_dataset[50]
|
||||
# print(data)
|
||||
@@ -0,0 +1,292 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from PIL import Image
|
||||
from typing import Literal, Tuple, Optional, Any
|
||||
import cv2
|
||||
|
||||
import os
|
||||
import math
|
||||
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
def add_margin(pil_img, color=0, size=256):
|
||||
width, height = pil_img.size
|
||||
result = Image.new(pil_img.mode, (size, size), color)
|
||||
result.paste(pil_img, ((size - width) // 2, (size - height) // 2))
|
||||
return result
|
||||
|
||||
def scale_and_place_object(image, scale_factor):
|
||||
assert np.shape(image)[-1]==4 # RGBA
|
||||
|
||||
# Extract the alpha channel (transparency) and the object (RGB channels)
|
||||
alpha_channel = image[:, :, 3]
|
||||
|
||||
# Find the bounding box coordinates of the object
|
||||
coords = cv2.findNonZero(alpha_channel)
|
||||
x, y, width, height = cv2.boundingRect(coords)
|
||||
|
||||
# Calculate the scale factor for resizing
|
||||
original_height, original_width = image.shape[:2]
|
||||
|
||||
if width > height:
|
||||
size = width
|
||||
original_size = original_width
|
||||
else:
|
||||
size = height
|
||||
original_size = original_height
|
||||
|
||||
scale_factor = min(scale_factor, size / (original_size+0.0))
|
||||
|
||||
new_size = scale_factor * original_size
|
||||
scale_factor = new_size / size
|
||||
|
||||
# Calculate the new size based on the scale factor
|
||||
new_width = int(width * scale_factor)
|
||||
new_height = int(height * scale_factor)
|
||||
|
||||
center_x = original_width // 2
|
||||
center_y = original_height // 2
|
||||
|
||||
paste_x = center_x - (new_width // 2)
|
||||
paste_y = center_y - (new_height // 2)
|
||||
|
||||
# Resize the object (RGB channels) to the new size
|
||||
rescaled_object = cv2.resize(image[y:y+height, x:x+width], (new_width, new_height))
|
||||
|
||||
# Create a new RGBA image with the resized image
|
||||
new_image = np.zeros((original_height, original_width, 4), dtype=np.uint8)
|
||||
|
||||
new_image[paste_y:paste_y + new_height, paste_x:paste_x + new_width] = rescaled_object
|
||||
|
||||
return new_image
|
||||
|
||||
class SingleImageDataset(Dataset):
|
||||
def __init__(self,
|
||||
fix_cam_pose_dir: str,
|
||||
num_views: int,
|
||||
img_wh: Tuple[int, int],
|
||||
bg_color: str,
|
||||
crop_size: int = 224,
|
||||
single_image: Optional[Image.Image] = None,
|
||||
num_validation_samples: Optional[int] = None,
|
||||
root_dir: Optional[str] = '',
|
||||
filepaths: Optional[list] = None,
|
||||
cond_type: Optional[str] = None
|
||||
) -> None:
|
||||
"""Create a dataset from a folder of images.
|
||||
If you pass in a root directory it will be searched for images
|
||||
ending in ext (ext can be a list)
|
||||
"""
|
||||
self.root_dir = root_dir
|
||||
self.num_views = num_views
|
||||
self.img_wh = img_wh
|
||||
self.crop_size = crop_size
|
||||
self.bg_color = bg_color
|
||||
self.cond_type = cond_type
|
||||
|
||||
if self.num_views == 4:
|
||||
self.view_types = ['front', 'right', 'back', 'left']
|
||||
elif self.num_views == 5:
|
||||
self.view_types = ['front', 'front_right', 'right', 'back', 'left']
|
||||
elif self.num_views == 6:
|
||||
self.view_types = ['front', 'front_right', 'right', 'back', 'left', 'front_left']
|
||||
|
||||
self.fix_cam_pose_dir = fix_cam_pose_dir
|
||||
|
||||
self.fix_cam_poses = self.load_fixed_poses() # world2cam matrix
|
||||
|
||||
if single_image is None:
|
||||
if filepaths is None:
|
||||
# Get a list of all files in the directory
|
||||
file_list = os.listdir(self.root_dir)
|
||||
else:
|
||||
file_list = filepaths
|
||||
|
||||
# Filter the files that end with .png or .jpg
|
||||
self.file_list = [file for file in file_list if file.endswith(('.png', '.jpg'))]
|
||||
else:
|
||||
self.file_list = None
|
||||
|
||||
# load all images
|
||||
self.all_images = []
|
||||
self.all_alphas = []
|
||||
bg_color = self.get_bg_color()
|
||||
|
||||
if single_image is not None:
|
||||
image, alpha = self.load_image(None, bg_color, return_type='pt', Imagefile=single_image)
|
||||
self.all_images.append(image)
|
||||
self.all_alphas.append(alpha)
|
||||
else:
|
||||
for file in self.file_list:
|
||||
print(os.path.join(self.root_dir, file))
|
||||
image, alpha = self.load_image(os.path.join(self.root_dir, file), bg_color, return_type='pt')
|
||||
self.all_images.append(image)
|
||||
self.all_alphas.append(alpha)
|
||||
|
||||
self.all_images = self.all_images[:num_validation_samples]
|
||||
self.all_alphas = self.all_alphas[:num_validation_samples]
|
||||
|
||||
|
||||
def __len__(self):
|
||||
return len(self.all_images)
|
||||
|
||||
def load_fixed_poses(self):
|
||||
poses = {}
|
||||
for face in self.view_types:
|
||||
RT = np.loadtxt(os.path.join(self.fix_cam_pose_dir,'%03d_%s_RT.txt'%(0, face)))
|
||||
poses[face] = RT
|
||||
|
||||
return poses
|
||||
|
||||
def cartesian_to_spherical(self, xyz):
|
||||
ptsnew = np.hstack((xyz, np.zeros(xyz.shape)))
|
||||
xy = xyz[:,0]**2 + xyz[:,1]**2
|
||||
z = np.sqrt(xy + xyz[:,2]**2)
|
||||
theta = np.arctan2(np.sqrt(xy), xyz[:,2]) # for elevation angle defined from Z-axis down
|
||||
#ptsnew[:,4] = np.arctan2(xyz[:,2], np.sqrt(xy)) # for elevation angle defined from XY-plane up
|
||||
azimuth = np.arctan2(xyz[:,1], xyz[:,0])
|
||||
return np.array([theta, azimuth, z])
|
||||
|
||||
def get_T(self, target_RT, cond_RT):
|
||||
R, T = target_RT[:3, :3], target_RT[:, -1]
|
||||
T_target = -R.T @ T # change to cam2world
|
||||
|
||||
R, T = cond_RT[:3, :3], cond_RT[:, -1]
|
||||
T_cond = -R.T @ T
|
||||
|
||||
theta_cond, azimuth_cond, z_cond = self.cartesian_to_spherical(T_cond[None, :])
|
||||
theta_target, azimuth_target, z_target = self.cartesian_to_spherical(T_target[None, :])
|
||||
|
||||
d_theta = theta_target - theta_cond
|
||||
d_azimuth = (azimuth_target - azimuth_cond) % (2 * math.pi)
|
||||
d_z = z_target - z_cond
|
||||
|
||||
# d_T = torch.tensor([d_theta.item(), math.sin(d_azimuth.item()), math.cos(d_azimuth.item()), d_z.item()])
|
||||
return d_theta, d_azimuth
|
||||
|
||||
def get_bg_color(self):
|
||||
if self.bg_color == 'white':
|
||||
bg_color = np.array([1., 1., 1.], dtype=np.float32)
|
||||
elif self.bg_color == 'black':
|
||||
bg_color = np.array([0., 0., 0.], dtype=np.float32)
|
||||
elif self.bg_color == 'gray':
|
||||
bg_color = np.array([0.5, 0.5, 0.5], dtype=np.float32)
|
||||
elif self.bg_color == 'random':
|
||||
bg_color = np.random.rand(3)
|
||||
elif isinstance(self.bg_color, float):
|
||||
bg_color = np.array([self.bg_color] * 3, dtype=np.float32)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
return bg_color
|
||||
|
||||
|
||||
def load_image(self, img_path, bg_color, return_type='np', Imagefile=None):
|
||||
# pil always returns uint8
|
||||
if Imagefile is None:
|
||||
image_input = Image.open(img_path)
|
||||
else:
|
||||
image_input = Imagefile
|
||||
image_size = self.img_wh[0]
|
||||
|
||||
if self.crop_size!=-1:
|
||||
alpha_np = np.asarray(image_input)[:, :, 3]
|
||||
coords = np.stack(np.nonzero(alpha_np), 1)[:, (1, 0)]
|
||||
min_x, min_y = np.min(coords, 0)
|
||||
max_x, max_y = np.max(coords, 0)
|
||||
ref_img_ = image_input.crop((min_x, min_y, max_x, max_y))
|
||||
h, w = ref_img_.height, ref_img_.width
|
||||
scale = self.crop_size / max(h, w)
|
||||
h_, w_ = int(scale * h), int(scale * w)
|
||||
ref_img_ = ref_img_.resize((w_, h_))
|
||||
image_input = add_margin(ref_img_, size=image_size)
|
||||
else:
|
||||
image_input = add_margin(image_input, size=max(image_input.height, image_input.width))
|
||||
image_input = image_input.resize((image_size, image_size))
|
||||
|
||||
# img = scale_and_place_object(img, self.scale_ratio)
|
||||
img = np.array(image_input)
|
||||
img = img.astype(np.float32) / 255. # [0, 1]
|
||||
assert img.shape[-1] == 4 # RGBA
|
||||
|
||||
alpha = img[...,3:4]
|
||||
img = img[...,:3] * alpha + bg_color * (1 - alpha)
|
||||
|
||||
if return_type == "np":
|
||||
pass
|
||||
elif return_type == "pt":
|
||||
img = torch.from_numpy(img)
|
||||
alpha = torch.from_numpy(alpha)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
return img, alpha
|
||||
|
||||
|
||||
def __len__(self):
|
||||
return len(self.all_images)
|
||||
|
||||
def __getitem__(self, index):
|
||||
|
||||
image = self.all_images[index%len(self.all_images)]
|
||||
alpha = self.all_alphas[index%len(self.all_images)]
|
||||
if self.file_list is not None:
|
||||
filename = self.file_list[index%len(self.all_images)].replace(".png", "")
|
||||
else:
|
||||
filename = 'null'
|
||||
|
||||
cond_w2c = self.fix_cam_poses['front']
|
||||
|
||||
tgt_w2cs = [self.fix_cam_poses[view] for view in self.view_types]
|
||||
|
||||
elevations = []
|
||||
azimuths = []
|
||||
|
||||
img_tensors_in = [
|
||||
image.permute(2, 0, 1)
|
||||
] * self.num_views
|
||||
|
||||
alpha_tensors_in = [
|
||||
alpha.permute(2, 0, 1)
|
||||
] * self.num_views
|
||||
|
||||
for view, tgt_w2c in zip(self.view_types, tgt_w2cs):
|
||||
# evelations, azimuths
|
||||
elevation, azimuth = self.get_T(tgt_w2c, cond_w2c)
|
||||
elevations.append(elevation)
|
||||
azimuths.append(azimuth)
|
||||
|
||||
img_tensors_in = torch.stack(img_tensors_in, dim=0).float() # (Nv, 3, H, W)
|
||||
alpha_tensors_in = torch.stack(alpha_tensors_in, dim=0).float() # (Nv, 3, H, W)
|
||||
|
||||
elevations = torch.as_tensor(elevations).float().squeeze(1)
|
||||
azimuths = torch.as_tensor(azimuths).float().squeeze(1)
|
||||
elevations_cond = torch.as_tensor([0] * self.num_views).float()
|
||||
|
||||
normal_class = torch.tensor([1, 0]).float()
|
||||
normal_task_embeddings = torch.stack([normal_class]*self.num_views, dim=0) # (Nv, 2)
|
||||
color_class = torch.tensor([0, 1]).float()
|
||||
color_task_embeddings = torch.stack([color_class]*self.num_views, dim=0) # (Nv, 2)
|
||||
|
||||
camera_embeddings = torch.stack([elevations_cond, elevations, azimuths], dim=-1) # (Nv, 3)
|
||||
|
||||
out = {
|
||||
'elevations_cond': elevations_cond,
|
||||
'elevations_cond_deg': torch.rad2deg(elevations_cond),
|
||||
'elevations': elevations,
|
||||
'azimuths': azimuths,
|
||||
'elevations_deg': torch.rad2deg(elevations),
|
||||
'azimuths_deg': torch.rad2deg(azimuths),
|
||||
'imgs_in': img_tensors_in,
|
||||
'alphas': alpha_tensors_in,
|
||||
'camera_embeddings': camera_embeddings,
|
||||
'normal_task_embeddings': normal_task_embeddings,
|
||||
'color_task_embeddings': color_task_embeddings,
|
||||
'filename': filename,
|
||||
}
|
||||
|
||||
return out
|
||||
|
||||
|
||||
@@ -0,0 +1,985 @@
|
||||
# Copyright 2023 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models.embeddings import ImagePositionalEmbeddings
|
||||
from diffusers.utils import BaseOutput, deprecate
|
||||
from diffusers.utils.torch_utils import maybe_allow_in_graph
|
||||
from diffusers.models.attention import FeedForward, AdaLayerNorm, AdaLayerNormZero, Attention
|
||||
from diffusers.models.embeddings import PatchEmbed
|
||||
from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.utils.import_utils import is_xformers_available
|
||||
|
||||
from einops import rearrange, repeat
|
||||
|
||||
|
||||
if is_xformers_available():
|
||||
import xformers
|
||||
import xformers.ops
|
||||
else:
|
||||
xformers = None
|
||||
|
||||
def my_repeat(tensor, num_repeats):
|
||||
"""
|
||||
Repeat a tensor along a given dimension
|
||||
"""
|
||||
if len(tensor.shape) == 3:
|
||||
return repeat(tensor, "b d c -> (b v) d c", v=num_repeats)
|
||||
elif len(tensor.shape) == 4:
|
||||
return repeat(tensor, "a b d c -> (a v) b d c", v=num_repeats)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TransformerMV2DModelOutput(BaseOutput):
|
||||
"""
|
||||
The output of [`Transformer2DModel`].
|
||||
|
||||
Args:
|
||||
sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete):
|
||||
The hidden states output conditioned on the `encoder_hidden_states` input. If discrete, returns probability
|
||||
distributions for the unnoised latent pixels.
|
||||
"""
|
||||
|
||||
sample: torch.FloatTensor
|
||||
|
||||
|
||||
class TransformerMV2DModel(ModelMixin, ConfigMixin):
|
||||
"""
|
||||
A 2D Transformer model for image-like data.
|
||||
|
||||
Parameters:
|
||||
num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.
|
||||
attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.
|
||||
in_channels (`int`, *optional*):
|
||||
The number of channels in the input and output (specify if the input is **continuous**).
|
||||
num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.
|
||||
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
|
||||
cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
|
||||
sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**).
|
||||
This is fixed during training since it is used to learn a number of position embeddings.
|
||||
num_vector_embeds (`int`, *optional*):
|
||||
The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**).
|
||||
Includes the class for the masked latent pixel.
|
||||
activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward.
|
||||
num_embeds_ada_norm ( `int`, *optional*):
|
||||
The number of diffusion steps used during training. Pass if at least one of the norm_layers is
|
||||
`AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are
|
||||
added to the hidden states.
|
||||
|
||||
During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`.
|
||||
attention_bias (`bool`, *optional*):
|
||||
Configure if the `TransformerBlocks` attention should contain a bias parameter.
|
||||
"""
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
num_attention_heads: int = 16,
|
||||
attention_head_dim: int = 88,
|
||||
in_channels: Optional[int] = None,
|
||||
out_channels: Optional[int] = None,
|
||||
num_layers: int = 1,
|
||||
dropout: float = 0.0,
|
||||
norm_num_groups: int = 32,
|
||||
cross_attention_dim: Optional[int] = None,
|
||||
attention_bias: bool = False,
|
||||
sample_size: Optional[int] = None,
|
||||
num_vector_embeds: Optional[int] = None,
|
||||
patch_size: Optional[int] = None,
|
||||
activation_fn: str = "geglu",
|
||||
num_embeds_ada_norm: Optional[int] = None,
|
||||
use_linear_projection: bool = False,
|
||||
only_cross_attention: bool = False,
|
||||
upcast_attention: bool = False,
|
||||
norm_type: str = "layer_norm",
|
||||
norm_elementwise_affine: bool = True,
|
||||
num_views: int = 1,
|
||||
cd_attention_last: bool=False,
|
||||
cd_attention_mid: bool=False,
|
||||
multiview_attention: bool=True,
|
||||
sparse_mv_attention: bool = False,
|
||||
mvcd_attention: bool=False
|
||||
):
|
||||
super().__init__()
|
||||
self.use_linear_projection = use_linear_projection
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.attention_head_dim = attention_head_dim
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
|
||||
# 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
|
||||
# Define whether input is continuous or discrete depending on configuration
|
||||
self.is_input_continuous = (in_channels is not None) and (patch_size is None)
|
||||
self.is_input_vectorized = num_vector_embeds is not None
|
||||
self.is_input_patches = in_channels is not None and patch_size is not None
|
||||
|
||||
if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
|
||||
deprecation_message = (
|
||||
f"The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"
|
||||
" incorrectly set to `'layer_norm'`.Make sure to set `norm_type` to `'ada_norm'` in the config."
|
||||
" Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect"
|
||||
" results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it"
|
||||
" would be very nice if you could open a Pull request for the `transformer/config.json` file"
|
||||
)
|
||||
deprecate("norm_type!=num_embeds_ada_norm", "1.0.0", deprecation_message, standard_warn=False)
|
||||
norm_type = "ada_norm"
|
||||
|
||||
if self.is_input_continuous and self.is_input_vectorized:
|
||||
raise ValueError(
|
||||
f"Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"
|
||||
" sure that either `in_channels` or `num_vector_embeds` is None."
|
||||
)
|
||||
elif self.is_input_vectorized and self.is_input_patches:
|
||||
raise ValueError(
|
||||
f"Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make"
|
||||
" sure that either `num_vector_embeds` or `num_patches` is None."
|
||||
)
|
||||
elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
|
||||
raise ValueError(
|
||||
f"Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:"
|
||||
f" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None."
|
||||
)
|
||||
|
||||
# 2. Define input layers
|
||||
if self.is_input_continuous:
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)
|
||||
if use_linear_projection:
|
||||
self.proj_in = LoRACompatibleLinear(in_channels, inner_dim)
|
||||
else:
|
||||
self.proj_in = LoRACompatibleConv(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
|
||||
elif self.is_input_vectorized:
|
||||
assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
|
||||
assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"
|
||||
|
||||
self.height = sample_size
|
||||
self.width = sample_size
|
||||
self.num_vector_embeds = num_vector_embeds
|
||||
self.num_latent_pixels = self.height * self.width
|
||||
|
||||
self.latent_image_embedding = ImagePositionalEmbeddings(
|
||||
num_embed=num_vector_embeds, embed_dim=inner_dim, height=self.height, width=self.width
|
||||
)
|
||||
elif self.is_input_patches:
|
||||
assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"
|
||||
|
||||
self.height = sample_size
|
||||
self.width = sample_size
|
||||
|
||||
self.patch_size = patch_size
|
||||
self.pos_embed = PatchEmbed(
|
||||
height=sample_size,
|
||||
width=sample_size,
|
||||
patch_size=patch_size,
|
||||
in_channels=in_channels,
|
||||
embed_dim=inner_dim,
|
||||
)
|
||||
|
||||
# 3. Define transformers blocks
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
BasicMVTransformerBlock(
|
||||
inner_dim,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
dropout=dropout,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
activation_fn=activation_fn,
|
||||
num_embeds_ada_norm=num_embeds_ada_norm,
|
||||
attention_bias=attention_bias,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
norm_type=norm_type,
|
||||
norm_elementwise_affine=norm_elementwise_affine,
|
||||
num_views=num_views,
|
||||
cd_attention_last=cd_attention_last,
|
||||
cd_attention_mid=cd_attention_mid,
|
||||
multiview_attention=multiview_attention,
|
||||
sparse_mv_attention=sparse_mv_attention,
|
||||
mvcd_attention=mvcd_attention
|
||||
)
|
||||
for d in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
# 4. Define output layers
|
||||
self.out_channels = in_channels if out_channels is None else out_channels
|
||||
if self.is_input_continuous:
|
||||
# TODO: should use out_channels for continuous projections
|
||||
if use_linear_projection:
|
||||
self.proj_out = LoRACompatibleLinear(inner_dim, in_channels)
|
||||
else:
|
||||
self.proj_out = LoRACompatibleConv(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
elif self.is_input_vectorized:
|
||||
self.norm_out = nn.LayerNorm(inner_dim)
|
||||
self.out = nn.Linear(inner_dim, self.num_vector_embeds - 1)
|
||||
elif self.is_input_patches:
|
||||
self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6)
|
||||
self.proj_out_1 = nn.Linear(inner_dim, 2 * inner_dim)
|
||||
self.proj_out_2 = nn.Linear(inner_dim, patch_size * patch_size * self.out_channels)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: Optional[torch.Tensor] = None,
|
||||
timestep: Optional[torch.LongTensor] = None,
|
||||
class_labels: Optional[torch.LongTensor] = None,
|
||||
cross_attention_kwargs: Dict[str, Any] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
encoder_attention_mask: Optional[torch.Tensor] = None,
|
||||
return_dict: bool = True,
|
||||
):
|
||||
"""
|
||||
The [`Transformer2DModel`] forward method.
|
||||
|
||||
Args:
|
||||
hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
|
||||
Input `hidden_states`.
|
||||
encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
|
||||
Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
|
||||
self-attention.
|
||||
timestep ( `torch.LongTensor`, *optional*):
|
||||
Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
|
||||
class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
|
||||
Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
|
||||
`AdaLayerZeroNorm`.
|
||||
encoder_attention_mask ( `torch.Tensor`, *optional*):
|
||||
Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:
|
||||
|
||||
* Mask `(batch, sequence_length)` True = keep, False = discard.
|
||||
* Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.
|
||||
|
||||
If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
|
||||
above. This bias will be added to the cross-attention scores.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
|
||||
tuple.
|
||||
|
||||
Returns:
|
||||
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
|
||||
`tuple` where the first element is the sample tensor.
|
||||
"""
|
||||
# ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
|
||||
# we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
|
||||
# we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
|
||||
# expects mask of shape:
|
||||
# [batch, key_tokens]
|
||||
# adds singleton query_tokens dimension:
|
||||
# [batch, 1, key_tokens]
|
||||
# this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
|
||||
# [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)
|
||||
# [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
|
||||
if attention_mask is not None and attention_mask.ndim == 2:
|
||||
# assume that mask is expressed as:
|
||||
# (1 = keep, 0 = discard)
|
||||
# convert mask into a bias that can be added to attention scores:
|
||||
# (keep = +0, discard = -10000.0)
|
||||
attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
|
||||
# convert encoder_attention_mask to a bias the same way we do for attention_mask
|
||||
if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
|
||||
encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
|
||||
encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
|
||||
|
||||
# 1. Input
|
||||
if self.is_input_continuous:
|
||||
batch, _, height, width = hidden_states.shape
|
||||
residual = hidden_states
|
||||
|
||||
hidden_states = self.norm(hidden_states)
|
||||
if not self.use_linear_projection:
|
||||
hidden_states = self.proj_in(hidden_states)
|
||||
inner_dim = hidden_states.shape[1]
|
||||
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim)
|
||||
else:
|
||||
inner_dim = hidden_states.shape[1]
|
||||
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim)
|
||||
hidden_states = self.proj_in(hidden_states)
|
||||
elif self.is_input_vectorized:
|
||||
hidden_states = self.latent_image_embedding(hidden_states)
|
||||
elif self.is_input_patches:
|
||||
hidden_states = self.pos_embed(hidden_states)
|
||||
|
||||
# 2. Blocks
|
||||
for block in self.transformer_blocks:
|
||||
hidden_states = block(
|
||||
hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
timestep=timestep,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
class_labels=class_labels,
|
||||
)
|
||||
|
||||
# 3. Output
|
||||
if self.is_input_continuous:
|
||||
if not self.use_linear_projection:
|
||||
hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
else:
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()
|
||||
|
||||
output = hidden_states + residual
|
||||
elif self.is_input_vectorized:
|
||||
hidden_states = self.norm_out(hidden_states)
|
||||
logits = self.out(hidden_states)
|
||||
# (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
|
||||
logits = logits.permute(0, 2, 1)
|
||||
|
||||
# log(p(x_0))
|
||||
output = F.log_softmax(logits.double(), dim=1).float()
|
||||
elif self.is_input_patches:
|
||||
# TODO: cleanup!
|
||||
conditioning = self.transformer_blocks[0].norm1.emb(
|
||||
timestep, class_labels, hidden_dtype=hidden_states.dtype
|
||||
)
|
||||
shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)
|
||||
hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
|
||||
hidden_states = self.proj_out_2(hidden_states)
|
||||
|
||||
# unpatchify
|
||||
height = width = int(hidden_states.shape[1] ** 0.5)
|
||||
hidden_states = hidden_states.reshape(
|
||||
shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)
|
||||
)
|
||||
hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
|
||||
output = hidden_states.reshape(
|
||||
shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)
|
||||
)
|
||||
|
||||
if not return_dict:
|
||||
return (output,)
|
||||
|
||||
return TransformerMV2DModelOutput(sample=output)
|
||||
|
||||
|
||||
@maybe_allow_in_graph
|
||||
class BasicMVTransformerBlock(nn.Module):
|
||||
r"""
|
||||
A basic Transformer block.
|
||||
|
||||
Parameters:
|
||||
dim (`int`): The number of channels in the input and output.
|
||||
num_attention_heads (`int`): The number of heads to use for multi-head attention.
|
||||
attention_head_dim (`int`): The number of channels in each head.
|
||||
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
|
||||
cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
|
||||
only_cross_attention (`bool`, *optional*):
|
||||
Whether to use only cross-attention layers. In this case two cross attention layers are used.
|
||||
double_self_attention (`bool`, *optional*):
|
||||
Whether to use two self-attention layers. In this case no cross attention layers are used.
|
||||
activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
|
||||
num_embeds_ada_norm (:
|
||||
obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
|
||||
attention_bias (:
|
||||
obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_attention_heads: int,
|
||||
attention_head_dim: int,
|
||||
dropout=0.0,
|
||||
cross_attention_dim: Optional[int] = None,
|
||||
activation_fn: str = "geglu",
|
||||
num_embeds_ada_norm: Optional[int] = None,
|
||||
attention_bias: bool = False,
|
||||
only_cross_attention: bool = False,
|
||||
double_self_attention: bool = False,
|
||||
upcast_attention: bool = False,
|
||||
norm_elementwise_affine: bool = True,
|
||||
norm_type: str = "layer_norm",
|
||||
final_dropout: bool = False,
|
||||
num_views: int = 1,
|
||||
cd_attention_last: bool = False,
|
||||
cd_attention_mid: bool = False,
|
||||
multiview_attention: bool = True,
|
||||
sparse_mv_attention: bool = False,
|
||||
mvcd_attention: bool = False
|
||||
):
|
||||
super().__init__()
|
||||
self.only_cross_attention = only_cross_attention
|
||||
|
||||
self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
|
||||
self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm"
|
||||
|
||||
if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
|
||||
raise ValueError(
|
||||
f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"
|
||||
f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."
|
||||
)
|
||||
|
||||
# Define 3 blocks. Each block has its own normalization layer.
|
||||
# 1. Self-Attn
|
||||
if self.use_ada_layer_norm:
|
||||
self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm)
|
||||
elif self.use_ada_layer_norm_zero:
|
||||
self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)
|
||||
else:
|
||||
self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
|
||||
|
||||
self.multiview_attention = multiview_attention
|
||||
self.sparse_mv_attention = sparse_mv_attention
|
||||
self.mvcd_attention = mvcd_attention
|
||||
|
||||
self.attn1 = CustomAttention(
|
||||
query_dim=dim,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
dropout=dropout,
|
||||
bias=attention_bias,
|
||||
cross_attention_dim=cross_attention_dim if only_cross_attention else None,
|
||||
upcast_attention=upcast_attention,
|
||||
processor=MVAttnProcessor()
|
||||
)
|
||||
|
||||
# 2. Cross-Attn
|
||||
if cross_attention_dim is not None or double_self_attention:
|
||||
# We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
|
||||
# I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
|
||||
# the second cross attention block.
|
||||
self.norm2 = (
|
||||
AdaLayerNorm(dim, num_embeds_ada_norm)
|
||||
if self.use_ada_layer_norm
|
||||
else nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
|
||||
)
|
||||
self.attn2 = Attention(
|
||||
query_dim=dim,
|
||||
cross_attention_dim=cross_attention_dim if not double_self_attention else None,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
dropout=dropout,
|
||||
bias=attention_bias,
|
||||
upcast_attention=upcast_attention,
|
||||
) # is self-attn if encoder_hidden_states is none
|
||||
else:
|
||||
self.norm2 = None
|
||||
self.attn2 = None
|
||||
|
||||
# 3. Feed-forward
|
||||
self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
|
||||
self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout)
|
||||
|
||||
# let chunk size default to None
|
||||
self._chunk_size = None
|
||||
self._chunk_dim = 0
|
||||
|
||||
self.num_views = num_views
|
||||
|
||||
self.cd_attention_last = cd_attention_last
|
||||
|
||||
if self.cd_attention_last:
|
||||
# Joint task -Attn
|
||||
self.attn_joint_last = CustomJointAttention(
|
||||
query_dim=dim,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
dropout=dropout,
|
||||
bias=attention_bias,
|
||||
cross_attention_dim=cross_attention_dim if only_cross_attention else None,
|
||||
upcast_attention=upcast_attention,
|
||||
processor=JointAttnProcessor()
|
||||
)
|
||||
nn.init.zeros_(self.attn_joint_last.to_out[0].weight.data)
|
||||
self.norm_joint_last = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
|
||||
|
||||
|
||||
self.cd_attention_mid = cd_attention_mid
|
||||
|
||||
if self.cd_attention_mid:
|
||||
# print("cross-domain attn in the middle")
|
||||
# Joint task -Attn
|
||||
self.attn_joint_mid = CustomJointAttention(
|
||||
query_dim=dim,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
dropout=dropout,
|
||||
bias=attention_bias,
|
||||
cross_attention_dim=cross_attention_dim if only_cross_attention else None,
|
||||
upcast_attention=upcast_attention,
|
||||
processor=JointAttnProcessor()
|
||||
)
|
||||
nn.init.zeros_(self.attn_joint_mid.to_out[0].weight.data)
|
||||
self.norm_joint_mid = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
|
||||
|
||||
def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int):
|
||||
# Sets chunk feed-forward
|
||||
self._chunk_size = chunk_size
|
||||
self._chunk_dim = dim
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor,
|
||||
attention_mask: Optional[torch.FloatTensor] = None,
|
||||
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
||||
encoder_attention_mask: Optional[torch.FloatTensor] = None,
|
||||
timestep: Optional[torch.LongTensor] = None,
|
||||
cross_attention_kwargs: Dict[str, Any] = None,
|
||||
class_labels: Optional[torch.LongTensor] = None,
|
||||
):
|
||||
assert attention_mask is None # not supported yet
|
||||
# Notice that normalization is always applied before the real computation in the following blocks.
|
||||
# 1. Self-Attention
|
||||
if self.use_ada_layer_norm:
|
||||
norm_hidden_states = self.norm1(hidden_states, timestep)
|
||||
elif self.use_ada_layer_norm_zero:
|
||||
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
|
||||
hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
|
||||
)
|
||||
else:
|
||||
norm_hidden_states = self.norm1(hidden_states)
|
||||
|
||||
cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
|
||||
|
||||
attn_output = self.attn1(
|
||||
norm_hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
|
||||
attention_mask=attention_mask,
|
||||
num_views=self.num_views,
|
||||
multiview_attention=self.multiview_attention,
|
||||
sparse_mv_attention=self.sparse_mv_attention,
|
||||
mvcd_attention=self.mvcd_attention,
|
||||
**cross_attention_kwargs,
|
||||
)
|
||||
|
||||
|
||||
if self.use_ada_layer_norm_zero:
|
||||
attn_output = gate_msa.unsqueeze(1) * attn_output
|
||||
hidden_states = attn_output + hidden_states
|
||||
|
||||
# joint attention twice
|
||||
if self.cd_attention_mid:
|
||||
norm_hidden_states = (
|
||||
self.norm_joint_mid(hidden_states, timestep) if self.use_ada_layer_norm else self.norm_joint_mid(hidden_states)
|
||||
)
|
||||
hidden_states = self.attn_joint_mid(norm_hidden_states) + hidden_states
|
||||
|
||||
# 2. Cross-Attention
|
||||
if self.attn2 is not None:
|
||||
norm_hidden_states = (
|
||||
self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)
|
||||
)
|
||||
|
||||
attn_output = self.attn2(
|
||||
norm_hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
attention_mask=encoder_attention_mask,
|
||||
**cross_attention_kwargs,
|
||||
)
|
||||
hidden_states = attn_output + hidden_states
|
||||
|
||||
# 3. Feed-forward
|
||||
norm_hidden_states = self.norm3(hidden_states)
|
||||
|
||||
if self.use_ada_layer_norm_zero:
|
||||
norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
|
||||
|
||||
if self._chunk_size is not None:
|
||||
# "feed_forward_chunk_size" can be used to save memory
|
||||
if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
|
||||
raise ValueError(
|
||||
f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
|
||||
)
|
||||
|
||||
num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
|
||||
ff_output = torch.cat(
|
||||
[self.ff(hid_slice) for hid_slice in norm_hidden_states.chunk(num_chunks, dim=self._chunk_dim)],
|
||||
dim=self._chunk_dim,
|
||||
)
|
||||
else:
|
||||
ff_output = self.ff(norm_hidden_states)
|
||||
|
||||
if self.use_ada_layer_norm_zero:
|
||||
ff_output = gate_mlp.unsqueeze(1) * ff_output
|
||||
|
||||
hidden_states = ff_output + hidden_states
|
||||
|
||||
if self.cd_attention_last:
|
||||
norm_hidden_states = (
|
||||
self.norm_joint_last(hidden_states, timestep) if self.use_ada_layer_norm else self.norm_joint_last(hidden_states)
|
||||
)
|
||||
hidden_states = self.attn_joint_last(norm_hidden_states) + hidden_states
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class CustomAttention(Attention):
|
||||
def set_use_memory_efficient_attention_xformers(
|
||||
self, use_memory_efficient_attention_xformers: bool, *args, **kwargs
|
||||
):
|
||||
processor = XFormersMVAttnProcessor()
|
||||
self.set_processor(processor)
|
||||
# print("using xformers attention processor")
|
||||
|
||||
|
||||
class CustomJointAttention(Attention):
|
||||
def set_use_memory_efficient_attention_xformers(
|
||||
self, use_memory_efficient_attention_xformers: bool, *args, **kwargs
|
||||
):
|
||||
processor = XFormersJointAttnProcessor()
|
||||
self.set_processor(processor)
|
||||
# print("using xformers attention processor")
|
||||
|
||||
class MVAttnProcessor:
|
||||
r"""
|
||||
Default processor for performing attention-related computations.
|
||||
"""
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
num_views=1,
|
||||
multiview_attention=True
|
||||
):
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
# print('query', query.shape, 'key', key.shape, 'value', value.shape)
|
||||
#([bx4, 1024, 320]) key torch.Size([bx4, 1024, 320]) value torch.Size([bx4, 1024, 320])
|
||||
# pdb.set_trace()
|
||||
# multi-view self-attention
|
||||
if multiview_attention:
|
||||
key = rearrange(key, "(b t) d c -> b (t d) c", t=num_views).repeat_interleave(num_views, dim=0)
|
||||
value = rearrange(value, "(b t) d c -> b (t d) c", t=num_views).repeat_interleave(num_views, dim=0)
|
||||
|
||||
query = attn.head_to_batch_dim(query).contiguous()
|
||||
key = attn.head_to_batch_dim(key).contiguous()
|
||||
value = attn.head_to_batch_dim(value).contiguous()
|
||||
|
||||
attention_probs = attn.get_attention_scores(query, key, attention_mask)
|
||||
hidden_states = torch.bmm(attention_probs, value)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class XFormersMVAttnProcessor:
|
||||
r"""
|
||||
Default processor for performing attention-related computations.
|
||||
"""
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
num_views=1.,
|
||||
multiview_attention=True,
|
||||
sparse_mv_attention=False,
|
||||
mvcd_attention=False,
|
||||
):
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
|
||||
# from yuancheng; here attention_mask is None
|
||||
if attention_mask is not None:
|
||||
# expand our mask's singleton query_tokens dimension:
|
||||
# [batch*heads, 1, key_tokens] ->
|
||||
# [batch*heads, query_tokens, key_tokens]
|
||||
# so that it can be added as a bias onto the attention scores that xformers computes:
|
||||
# [batch*heads, query_tokens, key_tokens]
|
||||
# we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
|
||||
_, query_tokens, _ = hidden_states.shape
|
||||
attention_mask = attention_mask.expand(-1, query_tokens, -1)
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
key_raw = attn.to_k(encoder_hidden_states)
|
||||
value_raw = attn.to_v(encoder_hidden_states)
|
||||
|
||||
# print('query', query.shape, 'key', key.shape, 'value', value.shape)
|
||||
#([bx4, 1024, 320]) key torch.Size([bx4, 1024, 320]) value torch.Size([bx4, 1024, 320])
|
||||
# pdb.set_trace()
|
||||
# multi-view self-attention
|
||||
if multiview_attention:
|
||||
if not sparse_mv_attention:
|
||||
key = my_repeat(rearrange(key_raw, "(b t) d c -> b (t d) c", t=num_views), num_views)
|
||||
value = my_repeat(rearrange(value_raw, "(b t) d c -> b (t d) c", t=num_views), num_views)
|
||||
else:
|
||||
key_front = my_repeat(rearrange(key_raw, "(b t) d c -> b t d c", t=num_views)[:, 0, :, :], num_views) # [(b t), d, c]
|
||||
value_front = my_repeat(rearrange(value_raw, "(b t) d c -> b t d c", t=num_views)[:, 0, :, :], num_views)
|
||||
key = torch.cat([key_front, key_raw], dim=1) # shape (b t) (2 d) c
|
||||
value = torch.cat([value_front, value_raw], dim=1)
|
||||
|
||||
else:
|
||||
# print("don't use multiview attention.")
|
||||
key = key_raw
|
||||
value = value_raw
|
||||
|
||||
query = attn.head_to_batch_dim(query)
|
||||
key = attn.head_to_batch_dim(key)
|
||||
value = attn.head_to_batch_dim(value)
|
||||
|
||||
hidden_states = xformers.ops.memory_efficient_attention(query, key, value, attn_bias=attention_mask)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
|
||||
class XFormersJointAttnProcessor:
|
||||
r"""
|
||||
Default processor for performing attention-related computations.
|
||||
"""
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
num_tasks=2
|
||||
):
|
||||
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
|
||||
# from yuancheng; here attention_mask is None
|
||||
if attention_mask is not None:
|
||||
# expand our mask's singleton query_tokens dimension:
|
||||
# [batch*heads, 1, key_tokens] ->
|
||||
# [batch*heads, query_tokens, key_tokens]
|
||||
# so that it can be added as a bias onto the attention scores that xformers computes:
|
||||
# [batch*heads, query_tokens, key_tokens]
|
||||
# we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
|
||||
_, query_tokens, _ = hidden_states.shape
|
||||
attention_mask = attention_mask.expand(-1, query_tokens, -1)
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
assert num_tasks == 2 # only support two tasks now
|
||||
|
||||
key_0, key_1 = torch.chunk(key, dim=0, chunks=2) # keys shape (b t) d c
|
||||
value_0, value_1 = torch.chunk(value, dim=0, chunks=2)
|
||||
key = torch.cat([key_0, key_1], dim=1) # (b t) 2d c
|
||||
value = torch.cat([value_0, value_1], dim=1) # (b t) 2d c
|
||||
key = torch.cat([key]*2, dim=0) # ( 2 b t) 2d c
|
||||
value = torch.cat([value]*2, dim=0) # (2 b t) 2d c
|
||||
|
||||
|
||||
query = attn.head_to_batch_dim(query).contiguous()
|
||||
key = attn.head_to_batch_dim(key).contiguous()
|
||||
value = attn.head_to_batch_dim(value).contiguous()
|
||||
|
||||
hidden_states = xformers.ops.memory_efficient_attention(query, key, value, attn_bias=attention_mask)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class JointAttnProcessor:
|
||||
r"""
|
||||
Default processor for performing attention-related computations.
|
||||
"""
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
num_tasks=2
|
||||
):
|
||||
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
assert num_tasks == 2 # only support two tasks now
|
||||
|
||||
key_0, key_1 = torch.chunk(key, dim=0, chunks=2) # keys shape (b t) d c
|
||||
value_0, value_1 = torch.chunk(value, dim=0, chunks=2)
|
||||
key = torch.cat([key_0, key_1], dim=1) # (b t) 2d c
|
||||
value = torch.cat([value_0, value_1], dim=1) # (b t) 2d c
|
||||
key = torch.cat([key]*2, dim=0) # ( 2 b t) 2d c
|
||||
value = torch.cat([value]*2, dim=0) # (2 b t) 2d c
|
||||
|
||||
|
||||
query = attn.head_to_batch_dim(query).contiguous()
|
||||
key = attn.head_to_batch_dim(key).contiguous()
|
||||
value = attn.head_to_batch_dim(value).contiguous()
|
||||
|
||||
attention_probs = attn.get_attention_scores(query, key, attention_mask)
|
||||
hidden_states = torch.bmm(attention_probs, value)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
@@ -0,0 +1,922 @@
|
||||
# Copyright 2023 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from diffusers.utils import is_torch_version, logging
|
||||
from diffusers.models.normalization import AdaGroupNorm
|
||||
from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor, AttnAddedKVProcessor2_0
|
||||
from diffusers.models.transformers.dual_transformer_2d import DualTransformer2DModel
|
||||
from diffusers.models.resnet import Downsample2D, FirDownsample2D, FirUpsample2D, KDownsample2D, KUpsample2D, ResnetBlock2D, Upsample2D
|
||||
from mvdiffusion.models.transformer_mv2d import TransformerMV2DModel
|
||||
|
||||
from diffusers.models.unets.unet_2d_blocks import DownBlock2D, ResnetDownsampleBlock2D, AttnDownBlock2D, CrossAttnDownBlock2D, SimpleCrossAttnDownBlock2D, SkipDownBlock2D, AttnSkipDownBlock2D, DownEncoderBlock2D, AttnDownEncoderBlock2D, KDownBlock2D, KCrossAttnDownBlock2D
|
||||
from diffusers.models.unets.unet_2d_blocks import UpBlock2D, ResnetUpsampleBlock2D, CrossAttnUpBlock2D, SimpleCrossAttnUpBlock2D, AttnUpBlock2D, SkipUpBlock2D, AttnSkipUpBlock2D, UpDecoderBlock2D, AttnUpDecoderBlock2D, KUpBlock2D, KCrossAttnUpBlock2D
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
def get_down_block(
|
||||
down_block_type,
|
||||
num_layers,
|
||||
in_channels,
|
||||
out_channels,
|
||||
temb_channels,
|
||||
add_downsample,
|
||||
resnet_eps,
|
||||
resnet_act_fn,
|
||||
transformer_layers_per_block=1,
|
||||
num_attention_heads=None,
|
||||
resnet_groups=None,
|
||||
cross_attention_dim=None,
|
||||
downsample_padding=None,
|
||||
dual_cross_attention=False,
|
||||
use_linear_projection=False,
|
||||
only_cross_attention=False,
|
||||
upcast_attention=False,
|
||||
resnet_time_scale_shift="default",
|
||||
resnet_skip_time_act=False,
|
||||
resnet_out_scale_factor=1.0,
|
||||
cross_attention_norm=None,
|
||||
attention_head_dim=None,
|
||||
downsample_type=None,
|
||||
num_views=1,
|
||||
cd_attention_last: bool = False,
|
||||
cd_attention_mid: bool = False,
|
||||
multiview_attention: bool = True,
|
||||
sparse_mv_attention: bool = False,
|
||||
mvcd_attention: bool=False
|
||||
):
|
||||
# If attn head dim is not defined, we default it to the number of heads
|
||||
if attention_head_dim is None:
|
||||
logger.warn(
|
||||
f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
|
||||
)
|
||||
attention_head_dim = num_attention_heads
|
||||
|
||||
down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type
|
||||
if down_block_type == "DownBlock2D":
|
||||
return DownBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
downsample_padding=downsample_padding,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif down_block_type == "ResnetDownsampleBlock2D":
|
||||
return ResnetDownsampleBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
skip_time_act=resnet_skip_time_act,
|
||||
output_scale_factor=resnet_out_scale_factor,
|
||||
)
|
||||
elif down_block_type == "AttnDownBlock2D":
|
||||
if add_downsample is False:
|
||||
downsample_type = None
|
||||
else:
|
||||
downsample_type = downsample_type or "conv" # default to 'conv'
|
||||
return AttnDownBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
downsample_padding=downsample_padding,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
downsample_type=downsample_type,
|
||||
)
|
||||
elif down_block_type == "CrossAttnDownBlock2D":
|
||||
if cross_attention_dim is None:
|
||||
raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock2D")
|
||||
return CrossAttnDownBlock2D(
|
||||
num_layers=num_layers,
|
||||
transformer_layers_per_block=transformer_layers_per_block,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
downsample_padding=downsample_padding,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
# custom MV2D attention block
|
||||
elif down_block_type == "CrossAttnDownBlockMV2D":
|
||||
if cross_attention_dim is None:
|
||||
raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlockMV2D")
|
||||
return CrossAttnDownBlockMV2D(
|
||||
num_layers=num_layers,
|
||||
transformer_layers_per_block=transformer_layers_per_block,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
downsample_padding=downsample_padding,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
num_views=num_views,
|
||||
cd_attention_last=cd_attention_last,
|
||||
cd_attention_mid=cd_attention_mid,
|
||||
multiview_attention=multiview_attention,
|
||||
sparse_mv_attention=sparse_mv_attention,
|
||||
mvcd_attention=mvcd_attention
|
||||
)
|
||||
elif down_block_type == "SimpleCrossAttnDownBlock2D":
|
||||
if cross_attention_dim is None:
|
||||
raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnDownBlock2D")
|
||||
return SimpleCrossAttnDownBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
skip_time_act=resnet_skip_time_act,
|
||||
output_scale_factor=resnet_out_scale_factor,
|
||||
only_cross_attention=only_cross_attention,
|
||||
cross_attention_norm=cross_attention_norm,
|
||||
)
|
||||
elif down_block_type == "SkipDownBlock2D":
|
||||
return SkipDownBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
downsample_padding=downsample_padding,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif down_block_type == "AttnSkipDownBlock2D":
|
||||
return AttnSkipDownBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif down_block_type == "DownEncoderBlock2D":
|
||||
return DownEncoderBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
downsample_padding=downsample_padding,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif down_block_type == "AttnDownEncoderBlock2D":
|
||||
return AttnDownEncoderBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
downsample_padding=downsample_padding,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif down_block_type == "KDownBlock2D":
|
||||
return KDownBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
)
|
||||
elif down_block_type == "KCrossAttnDownBlock2D":
|
||||
return KCrossAttnDownBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attention_head_dim=attention_head_dim,
|
||||
add_self_attention=True if not add_downsample else False,
|
||||
)
|
||||
raise ValueError(f"{down_block_type} does not exist.")
|
||||
|
||||
|
||||
def get_up_block(
|
||||
up_block_type,
|
||||
num_layers,
|
||||
in_channels,
|
||||
out_channels,
|
||||
prev_output_channel,
|
||||
temb_channels,
|
||||
add_upsample,
|
||||
resnet_eps,
|
||||
resnet_act_fn,
|
||||
transformer_layers_per_block=1,
|
||||
num_attention_heads=None,
|
||||
resnet_groups=None,
|
||||
cross_attention_dim=None,
|
||||
dual_cross_attention=False,
|
||||
use_linear_projection=False,
|
||||
only_cross_attention=False,
|
||||
upcast_attention=False,
|
||||
resnet_time_scale_shift="default",
|
||||
resnet_skip_time_act=False,
|
||||
resnet_out_scale_factor=1.0,
|
||||
cross_attention_norm=None,
|
||||
attention_head_dim=None,
|
||||
upsample_type=None,
|
||||
num_views=1,
|
||||
cd_attention_last: bool = False,
|
||||
cd_attention_mid: bool = False,
|
||||
multiview_attention: bool = True,
|
||||
sparse_mv_attention: bool = False,
|
||||
mvcd_attention: bool=False
|
||||
):
|
||||
# If attn head dim is not defined, we default it to the number of heads
|
||||
if attention_head_dim is None:
|
||||
logger.warn(
|
||||
f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
|
||||
)
|
||||
attention_head_dim = num_attention_heads
|
||||
|
||||
up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
|
||||
if up_block_type == "UpBlock2D":
|
||||
return UpBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif up_block_type == "ResnetUpsampleBlock2D":
|
||||
return ResnetUpsampleBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
skip_time_act=resnet_skip_time_act,
|
||||
output_scale_factor=resnet_out_scale_factor,
|
||||
)
|
||||
elif up_block_type == "CrossAttnUpBlock2D":
|
||||
if cross_attention_dim is None:
|
||||
raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock2D")
|
||||
return CrossAttnUpBlock2D(
|
||||
num_layers=num_layers,
|
||||
transformer_layers_per_block=transformer_layers_per_block,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
# custom MV2D attention block
|
||||
elif up_block_type == "CrossAttnUpBlockMV2D":
|
||||
if cross_attention_dim is None:
|
||||
raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlockMV2D")
|
||||
return CrossAttnUpBlockMV2D(
|
||||
num_layers=num_layers,
|
||||
transformer_layers_per_block=transformer_layers_per_block,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
num_views=num_views,
|
||||
cd_attention_last=cd_attention_last,
|
||||
cd_attention_mid=cd_attention_mid,
|
||||
multiview_attention=multiview_attention,
|
||||
sparse_mv_attention=sparse_mv_attention,
|
||||
mvcd_attention=mvcd_attention
|
||||
)
|
||||
elif up_block_type == "SimpleCrossAttnUpBlock2D":
|
||||
if cross_attention_dim is None:
|
||||
raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnUpBlock2D")
|
||||
return SimpleCrossAttnUpBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
skip_time_act=resnet_skip_time_act,
|
||||
output_scale_factor=resnet_out_scale_factor,
|
||||
only_cross_attention=only_cross_attention,
|
||||
cross_attention_norm=cross_attention_norm,
|
||||
)
|
||||
elif up_block_type == "AttnUpBlock2D":
|
||||
if add_upsample is False:
|
||||
upsample_type = None
|
||||
else:
|
||||
upsample_type = upsample_type or "conv" # default to 'conv'
|
||||
|
||||
return AttnUpBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
upsample_type=upsample_type,
|
||||
)
|
||||
elif up_block_type == "SkipUpBlock2D":
|
||||
return SkipUpBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif up_block_type == "AttnSkipUpBlock2D":
|
||||
return AttnSkipUpBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif up_block_type == "UpDecoderBlock2D":
|
||||
return UpDecoderBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
temb_channels=temb_channels,
|
||||
)
|
||||
elif up_block_type == "AttnUpDecoderBlock2D":
|
||||
return AttnUpDecoderBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
temb_channels=temb_channels,
|
||||
)
|
||||
elif up_block_type == "KUpBlock2D":
|
||||
return KUpBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
)
|
||||
elif up_block_type == "KCrossAttnUpBlock2D":
|
||||
return KCrossAttnUpBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attention_head_dim=attention_head_dim,
|
||||
)
|
||||
|
||||
raise ValueError(f"{up_block_type} does not exist.")
|
||||
|
||||
|
||||
class UNetMidBlockMV2DCrossAttn(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
temb_channels: int,
|
||||
dropout: float = 0.0,
|
||||
num_layers: int = 1,
|
||||
transformer_layers_per_block: int = 1,
|
||||
resnet_eps: float = 1e-6,
|
||||
resnet_time_scale_shift: str = "default",
|
||||
resnet_act_fn: str = "swish",
|
||||
resnet_groups: int = 32,
|
||||
resnet_pre_norm: bool = True,
|
||||
num_attention_heads=1,
|
||||
output_scale_factor=1.0,
|
||||
cross_attention_dim=1280,
|
||||
dual_cross_attention=False,
|
||||
use_linear_projection=False,
|
||||
upcast_attention=False,
|
||||
num_views: int = 1,
|
||||
cd_attention_last: bool = False,
|
||||
cd_attention_mid: bool = False,
|
||||
multiview_attention: bool = True,
|
||||
sparse_mv_attention: bool = False,
|
||||
mvcd_attention: bool=False
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.has_cross_attention = True
|
||||
self.num_attention_heads = num_attention_heads
|
||||
resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
|
||||
|
||||
# there is always at least one resnet
|
||||
resnets = [
|
||||
ResnetBlock2D(
|
||||
in_channels=in_channels,
|
||||
out_channels=in_channels,
|
||||
temb_channels=temb_channels,
|
||||
eps=resnet_eps,
|
||||
groups=resnet_groups,
|
||||
dropout=dropout,
|
||||
time_embedding_norm=resnet_time_scale_shift,
|
||||
non_linearity=resnet_act_fn,
|
||||
output_scale_factor=output_scale_factor,
|
||||
pre_norm=resnet_pre_norm,
|
||||
)
|
||||
]
|
||||
attentions = []
|
||||
|
||||
for _ in range(num_layers):
|
||||
if not dual_cross_attention:
|
||||
attentions.append(
|
||||
TransformerMV2DModel(
|
||||
num_attention_heads,
|
||||
in_channels // num_attention_heads,
|
||||
in_channels=in_channels,
|
||||
num_layers=transformer_layers_per_block,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
norm_num_groups=resnet_groups,
|
||||
use_linear_projection=use_linear_projection,
|
||||
upcast_attention=upcast_attention,
|
||||
num_views=num_views,
|
||||
cd_attention_last=cd_attention_last,
|
||||
cd_attention_mid=cd_attention_mid,
|
||||
multiview_attention=multiview_attention,
|
||||
sparse_mv_attention=sparse_mv_attention,
|
||||
mvcd_attention=mvcd_attention
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
resnets.append(
|
||||
ResnetBlock2D(
|
||||
in_channels=in_channels,
|
||||
out_channels=in_channels,
|
||||
temb_channels=temb_channels,
|
||||
eps=resnet_eps,
|
||||
groups=resnet_groups,
|
||||
dropout=dropout,
|
||||
time_embedding_norm=resnet_time_scale_shift,
|
||||
non_linearity=resnet_act_fn,
|
||||
output_scale_factor=output_scale_factor,
|
||||
pre_norm=resnet_pre_norm,
|
||||
)
|
||||
)
|
||||
|
||||
self.attentions = nn.ModuleList(attentions)
|
||||
self.resnets = nn.ModuleList(resnets)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor,
|
||||
temb: Optional[torch.FloatTensor] = None,
|
||||
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
||||
attention_mask: Optional[torch.FloatTensor] = None,
|
||||
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
encoder_attention_mask: Optional[torch.FloatTensor] = None,
|
||||
) -> torch.FloatTensor:
|
||||
hidden_states = self.resnets[0](hidden_states, temb)
|
||||
for attn, resnet in zip(self.attentions, self.resnets[1:]):
|
||||
hidden_states = attn(
|
||||
hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
attention_mask=attention_mask,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
hidden_states = resnet(hidden_states, temb)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class CrossAttnUpBlockMV2D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
prev_output_channel: int,
|
||||
temb_channels: int,
|
||||
dropout: float = 0.0,
|
||||
num_layers: int = 1,
|
||||
transformer_layers_per_block: int = 1,
|
||||
resnet_eps: float = 1e-6,
|
||||
resnet_time_scale_shift: str = "default",
|
||||
resnet_act_fn: str = "swish",
|
||||
resnet_groups: int = 32,
|
||||
resnet_pre_norm: bool = True,
|
||||
num_attention_heads=1,
|
||||
cross_attention_dim=1280,
|
||||
output_scale_factor=1.0,
|
||||
add_upsample=True,
|
||||
dual_cross_attention=False,
|
||||
use_linear_projection=False,
|
||||
only_cross_attention=False,
|
||||
upcast_attention=False,
|
||||
num_views: int = 1,
|
||||
cd_attention_last: bool = False,
|
||||
cd_attention_mid: bool = False,
|
||||
multiview_attention: bool = True,
|
||||
sparse_mv_attention: bool = False,
|
||||
mvcd_attention: bool=False
|
||||
):
|
||||
super().__init__()
|
||||
resnets = []
|
||||
attentions = []
|
||||
|
||||
self.has_cross_attention = True
|
||||
self.num_attention_heads = num_attention_heads
|
||||
|
||||
for i in range(num_layers):
|
||||
res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
|
||||
resnet_in_channels = prev_output_channel if i == 0 else out_channels
|
||||
|
||||
resnets.append(
|
||||
ResnetBlock2D(
|
||||
in_channels=resnet_in_channels + res_skip_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
eps=resnet_eps,
|
||||
groups=resnet_groups,
|
||||
dropout=dropout,
|
||||
time_embedding_norm=resnet_time_scale_shift,
|
||||
non_linearity=resnet_act_fn,
|
||||
output_scale_factor=output_scale_factor,
|
||||
pre_norm=resnet_pre_norm,
|
||||
)
|
||||
)
|
||||
if not dual_cross_attention:
|
||||
attentions.append(
|
||||
TransformerMV2DModel(
|
||||
num_attention_heads,
|
||||
out_channels // num_attention_heads,
|
||||
in_channels=out_channels,
|
||||
num_layers=transformer_layers_per_block,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
norm_num_groups=resnet_groups,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
num_views=num_views,
|
||||
cd_attention_last=cd_attention_last,
|
||||
cd_attention_mid=cd_attention_mid,
|
||||
multiview_attention=multiview_attention,
|
||||
sparse_mv_attention=sparse_mv_attention,
|
||||
mvcd_attention=mvcd_attention
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
self.attentions = nn.ModuleList(attentions)
|
||||
self.resnets = nn.ModuleList(resnets)
|
||||
|
||||
if add_upsample:
|
||||
self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
|
||||
else:
|
||||
self.upsamplers = None
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor,
|
||||
res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
|
||||
temb: Optional[torch.FloatTensor] = None,
|
||||
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
||||
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
upsample_size: Optional[int] = None,
|
||||
attention_mask: Optional[torch.FloatTensor] = None,
|
||||
encoder_attention_mask: Optional[torch.FloatTensor] = None,
|
||||
):
|
||||
for resnet, attn in zip(self.resnets, self.attentions):
|
||||
# pop res hidden states
|
||||
res_hidden_states = res_hidden_states_tuple[-1]
|
||||
res_hidden_states_tuple = res_hidden_states_tuple[:-1]
|
||||
hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
|
||||
|
||||
if self.training and self.gradient_checkpointing:
|
||||
|
||||
def create_custom_forward(module, return_dict=None):
|
||||
def custom_forward(*inputs):
|
||||
if return_dict is not None:
|
||||
return module(*inputs, return_dict=return_dict)
|
||||
else:
|
||||
return module(*inputs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(resnet),
|
||||
hidden_states,
|
||||
temb,
|
||||
**ckpt_kwargs,
|
||||
)
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(attn, return_dict=False),
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
None, # timestep
|
||||
None, # class_labels
|
||||
cross_attention_kwargs,
|
||||
attention_mask,
|
||||
encoder_attention_mask,
|
||||
**ckpt_kwargs,
|
||||
)[0]
|
||||
else:
|
||||
hidden_states = resnet(hidden_states, temb)
|
||||
hidden_states = attn(
|
||||
hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
attention_mask=attention_mask,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
if self.upsamplers is not None:
|
||||
for upsampler in self.upsamplers:
|
||||
hidden_states = upsampler(hidden_states, upsample_size)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class CrossAttnDownBlockMV2D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
temb_channels: int,
|
||||
dropout: float = 0.0,
|
||||
num_layers: int = 1,
|
||||
transformer_layers_per_block: int = 1,
|
||||
resnet_eps: float = 1e-6,
|
||||
resnet_time_scale_shift: str = "default",
|
||||
resnet_act_fn: str = "swish",
|
||||
resnet_groups: int = 32,
|
||||
resnet_pre_norm: bool = True,
|
||||
num_attention_heads=1,
|
||||
cross_attention_dim=1280,
|
||||
output_scale_factor=1.0,
|
||||
downsample_padding=1,
|
||||
add_downsample=True,
|
||||
dual_cross_attention=False,
|
||||
use_linear_projection=False,
|
||||
only_cross_attention=False,
|
||||
upcast_attention=False,
|
||||
num_views: int = 1,
|
||||
cd_attention_last: bool = False,
|
||||
cd_attention_mid: bool = False,
|
||||
multiview_attention: bool = True,
|
||||
sparse_mv_attention: bool = False,
|
||||
mvcd_attention: bool=False
|
||||
):
|
||||
super().__init__()
|
||||
resnets = []
|
||||
attentions = []
|
||||
|
||||
self.has_cross_attention = True
|
||||
self.num_attention_heads = num_attention_heads
|
||||
|
||||
for i in range(num_layers):
|
||||
in_channels = in_channels if i == 0 else out_channels
|
||||
resnets.append(
|
||||
ResnetBlock2D(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
eps=resnet_eps,
|
||||
groups=resnet_groups,
|
||||
dropout=dropout,
|
||||
time_embedding_norm=resnet_time_scale_shift,
|
||||
non_linearity=resnet_act_fn,
|
||||
output_scale_factor=output_scale_factor,
|
||||
pre_norm=resnet_pre_norm,
|
||||
)
|
||||
)
|
||||
if not dual_cross_attention:
|
||||
attentions.append(
|
||||
TransformerMV2DModel(
|
||||
num_attention_heads,
|
||||
out_channels // num_attention_heads,
|
||||
in_channels=out_channels,
|
||||
num_layers=transformer_layers_per_block,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
norm_num_groups=resnet_groups,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
num_views=num_views,
|
||||
cd_attention_last=cd_attention_last,
|
||||
cd_attention_mid=cd_attention_mid,
|
||||
multiview_attention=multiview_attention,
|
||||
sparse_mv_attention=sparse_mv_attention,
|
||||
mvcd_attention=mvcd_attention
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
self.attentions = nn.ModuleList(attentions)
|
||||
self.resnets = nn.ModuleList(resnets)
|
||||
|
||||
if add_downsample:
|
||||
self.downsamplers = nn.ModuleList(
|
||||
[
|
||||
Downsample2D(
|
||||
out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
|
||||
)
|
||||
]
|
||||
)
|
||||
else:
|
||||
self.downsamplers = None
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor,
|
||||
temb: Optional[torch.FloatTensor] = None,
|
||||
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
||||
attention_mask: Optional[torch.FloatTensor] = None,
|
||||
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
encoder_attention_mask: Optional[torch.FloatTensor] = None,
|
||||
additional_residuals=None,
|
||||
):
|
||||
output_states = ()
|
||||
|
||||
blocks = list(zip(self.resnets, self.attentions))
|
||||
|
||||
for i, (resnet, attn) in enumerate(blocks):
|
||||
if self.training and self.gradient_checkpointing:
|
||||
|
||||
def create_custom_forward(module, return_dict=None):
|
||||
def custom_forward(*inputs):
|
||||
if return_dict is not None:
|
||||
return module(*inputs, return_dict=return_dict)
|
||||
else:
|
||||
return module(*inputs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(resnet),
|
||||
hidden_states,
|
||||
temb,
|
||||
**ckpt_kwargs,
|
||||
)
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(attn, return_dict=False),
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
None, # timestep
|
||||
None, # class_labels
|
||||
cross_attention_kwargs,
|
||||
attention_mask,
|
||||
encoder_attention_mask,
|
||||
**ckpt_kwargs,
|
||||
)[0]
|
||||
else:
|
||||
hidden_states = resnet(hidden_states, temb)
|
||||
hidden_states = attn(
|
||||
hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
attention_mask=attention_mask,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
# apply additional residuals to the output of the last pair of resnet and attention blocks
|
||||
if i == len(blocks) - 1 and additional_residuals is not None:
|
||||
hidden_states = hidden_states + additional_residuals
|
||||
|
||||
output_states = output_states + (hidden_states,)
|
||||
|
||||
if self.downsamplers is not None:
|
||||
for downsampler in self.downsamplers:
|
||||
hidden_states = downsampler(hidden_states)
|
||||
|
||||
output_states = output_states + (hidden_states,)
|
||||
|
||||
return hidden_states, output_states
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,510 @@
|
||||
# Copyright 2023 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import inspect
|
||||
import warnings
|
||||
from typing import Callable, List, Optional, Union
|
||||
|
||||
import PIL
|
||||
import torch
|
||||
import torchvision.transforms.functional as TF
|
||||
from packaging import version
|
||||
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
|
||||
|
||||
from diffusers.configuration_utils import FrozenDict
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
from diffusers.models import AutoencoderKL, UNet2DConditionModel
|
||||
from diffusers.schedulers import KarrasDiffusionSchedulers
|
||||
from diffusers.utils import deprecate, logging
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
|
||||
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
|
||||
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
|
||||
from einops import rearrange, repeat
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
class MVDiffusionImagePipeline(DiffusionPipeline):
|
||||
r"""
|
||||
Pipeline to generate image variations from an input image using Stable Diffusion.
|
||||
|
||||
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
|
||||
implemented for all pipelines (downloading, saving, running on a particular device, etc.).
|
||||
|
||||
Args:
|
||||
vae ([`AutoencoderKL`]):
|
||||
Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
|
||||
image_encoder ([`~transformers.CLIPVisionModelWithProjection`]):
|
||||
Frozen CLIP image-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
|
||||
text_encoder ([`~transformers.CLIPTextModel`]):
|
||||
Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
|
||||
tokenizer ([`~transformers.CLIPTokenizer`]):
|
||||
A `CLIPTokenizer` to tokenize text.
|
||||
unet ([`UNet2DConditionModel`]):
|
||||
A `UNet2DConditionModel` to denoise the encoded image latents.
|
||||
scheduler ([`SchedulerMixin`]):
|
||||
A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
|
||||
[`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
|
||||
safety_checker ([`StableDiffusionSafetyChecker`]):
|
||||
Classification module that estimates whether generated images could be considered offensive or harmful.
|
||||
Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details
|
||||
about a model's potential harms.
|
||||
feature_extractor ([`~transformers.CLIPImageProcessor`]):
|
||||
A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`.
|
||||
"""
|
||||
# TODO: feature_extractor is required to encode images (if they are in PIL format),
|
||||
# we should give a descriptive message if the pipeline doesn't have one.
|
||||
_optional_components = ["safety_checker"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vae: AutoencoderKL,
|
||||
image_encoder: CLIPVisionModelWithProjection,
|
||||
unet: UNet2DConditionModel,
|
||||
scheduler: KarrasDiffusionSchedulers,
|
||||
safety_checker: StableDiffusionSafetyChecker,
|
||||
feature_extractor: CLIPImageProcessor,
|
||||
requires_safety_checker: bool = False,
|
||||
camera_embedding_type: str = 'e_de_da_sincos',
|
||||
num_views: int = 6
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
if safety_checker is None and requires_safety_checker:
|
||||
logger.warn(
|
||||
f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
|
||||
" that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
|
||||
" results in services or applications open to the public. Both the diffusers team and Hugging Face"
|
||||
" strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
|
||||
" it only for use-cases that involve analyzing network behavior or auditing its results. For more"
|
||||
" information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
|
||||
)
|
||||
|
||||
if safety_checker is not None and feature_extractor is None:
|
||||
raise ValueError(
|
||||
"Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
|
||||
" checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
|
||||
)
|
||||
|
||||
is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(
|
||||
version.parse(unet.config._diffusers_version).base_version
|
||||
) < version.parse("0.9.0.dev0")
|
||||
is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64
|
||||
if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:
|
||||
deprecation_message = (
|
||||
"The configuration file of the unet has set the default `sample_size` to smaller than"
|
||||
" 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the"
|
||||
" following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"
|
||||
" CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"
|
||||
" \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"
|
||||
" configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"
|
||||
" in the config might lead to incorrect results in future versions. If you have downloaded this"
|
||||
" checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"
|
||||
" the `unet/config.json` file"
|
||||
)
|
||||
deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)
|
||||
new_config = dict(unet.config)
|
||||
new_config["sample_size"] = 64
|
||||
unet._internal_dict = FrozenDict(new_config)
|
||||
|
||||
self.register_modules(
|
||||
vae=vae,
|
||||
image_encoder=image_encoder,
|
||||
unet=unet,
|
||||
scheduler=scheduler,
|
||||
safety_checker=safety_checker,
|
||||
feature_extractor=feature_extractor,
|
||||
)
|
||||
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
|
||||
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
|
||||
self.register_to_config(requires_safety_checker=requires_safety_checker)
|
||||
|
||||
self.camera_embedding_type: str = camera_embedding_type
|
||||
self.num_views: int = num_views
|
||||
|
||||
self.camera_embedding = torch.tensor(
|
||||
[[ 0.0000, 0.0000, 0.0000, 1.0000, 0.0000],
|
||||
[ 0.0000, -0.2362, 0.8125, 1.0000, 0.0000],
|
||||
[ 0.0000, -0.1686, 1.6934, 1.0000, 0.0000],
|
||||
[ 0.0000, 0.5220, 3.1406, 1.0000, 0.0000],
|
||||
[ 0.0000, 0.6904, 4.8359, 1.0000, 0.0000],
|
||||
[ 0.0000, 0.3733, 5.5859, 1.0000, 0.0000],
|
||||
[ 0.0000, 0.0000, 0.0000, 0.0000, 1.0000],
|
||||
[ 0.0000, -0.2362, 0.8125, 0.0000, 1.0000],
|
||||
[ 0.0000, -0.1686, 1.6934, 0.0000, 1.0000],
|
||||
[ 0.0000, 0.5220, 3.1406, 0.0000, 1.0000],
|
||||
[ 0.0000, 0.6904, 4.8359, 0.0000, 1.0000],
|
||||
[ 0.0000, 0.3733, 5.5859, 0.0000, 1.0000]], dtype=torch.float16)
|
||||
|
||||
def _encode_image(self, image_pil, device, num_images_per_prompt, do_classifier_free_guidance):
|
||||
dtype = next(self.image_encoder.parameters()).dtype
|
||||
|
||||
image_pt = self.feature_extractor(images=image_pil, return_tensors="pt").pixel_values
|
||||
image_pt = image_pt.to(device=device, dtype=dtype)
|
||||
image_embeddings = self.image_encoder(image_pt).image_embeds
|
||||
image_embeddings = image_embeddings.unsqueeze(1)
|
||||
|
||||
# duplicate image embeddings for each generation per prompt, using mps friendly method
|
||||
# Note: repeat differently from official pipelines
|
||||
# B1B2B3B4 -> B1B2B3B4B1B2B3B4
|
||||
bs_embed, seq_len, _ = image_embeddings.shape
|
||||
image_embeddings = image_embeddings.repeat(num_images_per_prompt, 1, 1)
|
||||
|
||||
if do_classifier_free_guidance:
|
||||
negative_prompt_embeds = torch.zeros_like(image_embeddings)
|
||||
|
||||
# For classifier free guidance, we need to do two forward passes.
|
||||
# Here we concatenate the unconditional and text embeddings into a single batch
|
||||
# to avoid doing two forward passes
|
||||
image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings])
|
||||
|
||||
image_pt = torch.stack([TF.to_tensor(img) for img in image_pil], dim=0).to(device).to(dtype)
|
||||
image_pt = image_pt * 2.0 - 1.0
|
||||
image_latents = self.vae.encode(image_pt).latent_dist.mode() * self.vae.config.scaling_factor
|
||||
# Note: repeat differently from official pipelines
|
||||
# B1B2B3B4 -> B1B2B3B4B1B2B3B4
|
||||
image_latents = image_latents.repeat(num_images_per_prompt, 1, 1, 1)
|
||||
|
||||
if do_classifier_free_guidance:
|
||||
image_latents = torch.cat([torch.zeros_like(image_latents), image_latents])
|
||||
|
||||
return image_embeddings, image_latents
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker
|
||||
def run_safety_checker(self, image, device, dtype):
|
||||
if self.safety_checker is None:
|
||||
has_nsfw_concept = None
|
||||
else:
|
||||
if torch.is_tensor(image):
|
||||
feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")
|
||||
else:
|
||||
feature_extractor_input = self.image_processor.numpy_to_pil(image)
|
||||
safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)
|
||||
image, has_nsfw_concept = self.safety_checker(
|
||||
images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
|
||||
)
|
||||
return image, has_nsfw_concept
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents
|
||||
def decode_latents(self, latents):
|
||||
warnings.warn(
|
||||
"The decode_latents method is deprecated and will be removed in a future version. Please"
|
||||
" use VaeImageProcessor instead",
|
||||
FutureWarning,
|
||||
)
|
||||
latents = 1 / self.vae.config.scaling_factor * latents
|
||||
image = self.vae.decode(latents, return_dict=False)[0]
|
||||
image = (image / 2 + 0.5).clamp(0, 1)
|
||||
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
|
||||
image = image.cpu().permute(0, 2, 3, 1).float().numpy()
|
||||
return image
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
|
||||
def prepare_extra_step_kwargs(self, generator, eta):
|
||||
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
|
||||
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
|
||||
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
|
||||
# and should be between [0, 1]
|
||||
|
||||
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
|
||||
extra_step_kwargs = {}
|
||||
if accepts_eta:
|
||||
extra_step_kwargs["eta"] = eta
|
||||
|
||||
# check if the scheduler accepts generator
|
||||
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
|
||||
if accepts_generator:
|
||||
extra_step_kwargs["generator"] = generator
|
||||
return extra_step_kwargs
|
||||
|
||||
def check_inputs(self, image, height, width, callback_steps):
|
||||
if (
|
||||
not isinstance(image, torch.Tensor)
|
||||
and not isinstance(image, PIL.Image.Image)
|
||||
and not isinstance(image, list)
|
||||
):
|
||||
raise ValueError(
|
||||
"`image` has to be of type `torch.FloatTensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is"
|
||||
f" {type(image)}"
|
||||
)
|
||||
|
||||
if height % 8 != 0 or width % 8 != 0:
|
||||
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
|
||||
|
||||
if (callback_steps is None) or (
|
||||
callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
|
||||
):
|
||||
raise ValueError(
|
||||
f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
|
||||
f" {type(callback_steps)}."
|
||||
)
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
|
||||
def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
|
||||
shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
|
||||
if isinstance(generator, list) and len(generator) != batch_size:
|
||||
raise ValueError(
|
||||
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
|
||||
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
|
||||
)
|
||||
|
||||
if latents is None:
|
||||
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
||||
else:
|
||||
latents = latents.to(device)
|
||||
|
||||
# scale the initial noise by the standard deviation required by the scheduler
|
||||
latents = latents * self.scheduler.init_noise_sigma
|
||||
return latents
|
||||
|
||||
def prepare_camera_embedding(self, camera_embedding: Union[float, torch.Tensor], do_classifier_free_guidance, num_images_per_prompt=1):
|
||||
# (B, 3)
|
||||
camera_embedding = camera_embedding.to(dtype=self.unet.dtype, device=self.unet.device)
|
||||
|
||||
if self.camera_embedding_type == 'e_de_da_sincos':
|
||||
# (B, 6)
|
||||
camera_embedding = torch.cat([
|
||||
torch.sin(camera_embedding),
|
||||
torch.cos(camera_embedding)
|
||||
], dim=-1)
|
||||
assert self.unet.config.class_embed_type == 'projection'
|
||||
assert self.unet.config.projection_class_embeddings_input_dim == 6 or self.unet.config.projection_class_embeddings_input_dim == 10
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
# Note: repeat differently from official pipelines
|
||||
# B1B2B3B4 -> B1B2B3B4B1B2B3B4
|
||||
camera_embedding = camera_embedding.repeat(num_images_per_prompt, 1)
|
||||
|
||||
if do_classifier_free_guidance:
|
||||
camera_embedding = torch.cat([
|
||||
camera_embedding,
|
||||
camera_embedding
|
||||
], dim=0)
|
||||
|
||||
return camera_embedding
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
image: Union[List[PIL.Image.Image], torch.FloatTensor],
|
||||
# elevation_cond: torch.FloatTensor,
|
||||
# elevation: torch.FloatTensor,
|
||||
# azimuth: torch.FloatTensor,
|
||||
camera_embedding: Optional[torch.FloatTensor]=None,
|
||||
height: Optional[int] = None,
|
||||
width: Optional[int] = None,
|
||||
num_inference_steps: int = 50,
|
||||
guidance_scale: float = 7.5,
|
||||
num_images_per_prompt: Optional[int] = 1,
|
||||
eta: float = 0.0,
|
||||
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
||||
latents: Optional[torch.FloatTensor] = None,
|
||||
output_type: Optional[str] = "pil",
|
||||
return_dict: bool = True,
|
||||
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
|
||||
callback_steps: int = 1,
|
||||
normal_cond: Optional[Union[List[PIL.Image.Image], torch.FloatTensor]] = None,
|
||||
):
|
||||
r"""
|
||||
The call function to the pipeline for generation.
|
||||
|
||||
Args:
|
||||
image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.FloatTensor`):
|
||||
Image or images to guide image generation. If you provide a tensor, it needs to be compatible with
|
||||
[`CLIPImageProcessor`](https://huggingface.co/lambdalabs/sd-image-variations-diffusers/blob/main/feature_extractor/preprocessor_config.json).
|
||||
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
|
||||
The height in pixels of the generated image.
|
||||
width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
|
||||
The width in pixels of the generated image.
|
||||
num_inference_steps (`int`, *optional*, defaults to 50):
|
||||
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
|
||||
expense of slower inference. This parameter is modulated by `strength`.
|
||||
guidance_scale (`float`, *optional*, defaults to 7.5):
|
||||
A higher guidance scale value encourages the model to generate images closely linked to the text
|
||||
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
|
||||
num_images_per_prompt (`int`, *optional*, defaults to 1):
|
||||
The number of images to generate per prompt.
|
||||
eta (`float`, *optional*, defaults to 0.0):
|
||||
Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
|
||||
to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
|
||||
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
|
||||
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
|
||||
generation deterministic.
|
||||
latents (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
|
||||
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
|
||||
tensor is generated by sampling using the supplied random `generator`.
|
||||
output_type (`str`, *optional*, defaults to `"pil"`):
|
||||
The output format of the generated image. Choose between `PIL.Image` or `np.array`.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
|
||||
plain tuple.
|
||||
callback (`Callable`, *optional*):
|
||||
A function that calls every `callback_steps` steps during inference. The function is called with the
|
||||
following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
|
||||
callback_steps (`int`, *optional*, defaults to 1):
|
||||
The frequency at which the `callback` function is called. If not specified, the callback is called at
|
||||
every step.
|
||||
|
||||
Returns:
|
||||
[`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
|
||||
If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,
|
||||
otherwise a `tuple` is returned where the first element is a list with the generated images and the
|
||||
second element is a list of `bool`s indicating whether the corresponding generated image contains
|
||||
"not-safe-for-work" (nsfw) content.
|
||||
|
||||
Examples:
|
||||
|
||||
```py
|
||||
from diffusers import StableDiffusionImageVariationPipeline
|
||||
from PIL import Image
|
||||
from io import BytesIO
|
||||
import requests
|
||||
|
||||
pipe = StableDiffusionImageVariationPipeline.from_pretrained(
|
||||
"lambdalabs/sd-image-variations-diffusers", revision="v2.0"
|
||||
)
|
||||
pipe = pipe.to("cuda")
|
||||
|
||||
url = "https://lh3.googleusercontent.com/y-iFOHfLTwkuQSUegpwDdgKmOjRSTvPxat63dQLB25xkTs4lhIbRUFeNBWZzYf370g=s1200"
|
||||
|
||||
response = requests.get(url)
|
||||
image = Image.open(BytesIO(response.content)).convert("RGB")
|
||||
|
||||
out = pipe(image, num_images_per_prompt=3, guidance_scale=15)
|
||||
out["images"][0].save("result.jpg")
|
||||
```
|
||||
"""
|
||||
# 0. Default height and width to unet
|
||||
height = height or self.unet.config.sample_size * self.vae_scale_factor
|
||||
width = width or self.unet.config.sample_size * self.vae_scale_factor
|
||||
|
||||
# 1. Check inputs. Raise error if not correct
|
||||
self.check_inputs(image, height, width, callback_steps)
|
||||
|
||||
|
||||
# 2. Define call parameters
|
||||
if isinstance(image, list):
|
||||
batch_size = len(image)
|
||||
elif isinstance(image, torch.Tensor):
|
||||
batch_size = image.shape[0]
|
||||
assert batch_size >= self.num_views and batch_size % self.num_views == 0
|
||||
elif isinstance(image, PIL.Image.Image):
|
||||
image = [image]*self.num_views*2
|
||||
batch_size = self.num_views*2
|
||||
|
||||
device = self._execution_device
|
||||
dtype = self.vae.dtype
|
||||
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
|
||||
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
|
||||
# corresponds to doing no classifier free guidance.
|
||||
do_classifier_free_guidance = guidance_scale != 1.0
|
||||
|
||||
# 3. Encode input image
|
||||
if isinstance(image, list):
|
||||
image_pil = image
|
||||
elif isinstance(image, torch.Tensor):
|
||||
image_pil = [TF.to_pil_image(image[i]) for i in range(image.shape[0])]
|
||||
image_embeddings, image_latents = self._encode_image(image_pil, device, num_images_per_prompt, do_classifier_free_guidance)
|
||||
|
||||
if normal_cond is not None:
|
||||
if isinstance(normal_cond, list):
|
||||
normal_cond_pil = normal_cond
|
||||
elif isinstance(normal_cond, torch.Tensor):
|
||||
normal_cond_pil = [TF.to_pil_image(normal_cond[i]) for i in range(normal_cond.shape[0])]
|
||||
_, image_latents = self._encode_image(normal_cond_pil, device, num_images_per_prompt, do_classifier_free_guidance)
|
||||
|
||||
|
||||
# assert len(elevation_cond) == batch_size and len(elevation) == batch_size and len(azimuth) == batch_size
|
||||
# camera_embeddings = self.prepare_camera_condition(elevation_cond, elevation, azimuth, do_classifier_free_guidance=do_classifier_free_guidance, num_images_per_prompt=num_images_per_prompt)
|
||||
|
||||
if camera_embedding is not None:
|
||||
assert len(camera_embedding) == batch_size
|
||||
else:
|
||||
camera_embedding = self.camera_embedding.to(dtype)
|
||||
camera_embedding = repeat(camera_embedding, "Nv Nce -> (B Nv) Nce", B=batch_size//len(camera_embedding))
|
||||
camera_embeddings = self.prepare_camera_embedding(camera_embedding, do_classifier_free_guidance=do_classifier_free_guidance, num_images_per_prompt=num_images_per_prompt)
|
||||
|
||||
# 4. Prepare timesteps
|
||||
self.scheduler.set_timesteps(num_inference_steps, device=device)
|
||||
timesteps = self.scheduler.timesteps
|
||||
|
||||
# 5. Prepare latent variables
|
||||
num_channels_latents = self.unet.config.out_channels
|
||||
latents = self.prepare_latents(
|
||||
batch_size * num_images_per_prompt,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
image_embeddings.dtype,
|
||||
device,
|
||||
generator,
|
||||
latents,
|
||||
)
|
||||
|
||||
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
|
||||
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
|
||||
|
||||
# 7. Denoising loop
|
||||
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
|
||||
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
||||
for i, t in enumerate(timesteps):
|
||||
# expand the latents if we are doing classifier free guidance
|
||||
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
|
||||
latent_model_input = torch.cat([
|
||||
latent_model_input, image_latents
|
||||
], dim=1)
|
||||
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
|
||||
|
||||
# predict the noise residual
|
||||
noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=image_embeddings, class_labels=camera_embeddings).sample
|
||||
|
||||
# perform guidance
|
||||
if do_classifier_free_guidance:
|
||||
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
|
||||
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
|
||||
|
||||
# compute the previous noisy sample x_t -> x_t-1
|
||||
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
|
||||
|
||||
# call the callback, if provided
|
||||
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
|
||||
progress_bar.update()
|
||||
if callback is not None and i % callback_steps == 0:
|
||||
callback(i, t, latents)
|
||||
|
||||
if not output_type == "latent":
|
||||
if num_channels_latents == 8:
|
||||
latents = torch.cat([latents[:, :4], latents[:, 4:]], dim=0)
|
||||
|
||||
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
|
||||
image, has_nsfw_concept = self.run_safety_checker(image, device, image_embeddings.dtype)
|
||||
else:
|
||||
image = latents
|
||||
has_nsfw_concept = None
|
||||
|
||||
if has_nsfw_concept is None:
|
||||
do_denormalize = [True] * image.shape[0]
|
||||
else:
|
||||
do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
|
||||
|
||||
image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
|
||||
|
||||
if not return_dict:
|
||||
return (image, has_nsfw_concept)
|
||||
|
||||
return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
import os
|
||||
from omegaconf import OmegaConf
|
||||
from packaging import version
|
||||
from typing import Dict, Optional, Tuple, List
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
# ============ Register OmegaConf Recolvers ============= #
|
||||
OmegaConf.register_new_resolver('calc_exp_lr_decay_rate', lambda factor, n: factor**(1./n), replace=True)
|
||||
OmegaConf.register_new_resolver('add', lambda a, b: a + b, replace=True)
|
||||
OmegaConf.register_new_resolver('sub', lambda a, b: a - b, replace=True)
|
||||
OmegaConf.register_new_resolver('mul', lambda a, b: a * b, replace=True)
|
||||
OmegaConf.register_new_resolver('div', lambda a, b: a / b, replace=True)
|
||||
OmegaConf.register_new_resolver('idiv', lambda a, b: a // b, replace=True)
|
||||
OmegaConf.register_new_resolver('basename', lambda p: os.path.basename(p), replace=True)
|
||||
# ======================================================= #
|
||||
|
||||
|
||||
def prompt(question):
|
||||
inp = input(f"{question} (y/n)").lower().strip()
|
||||
if inp and inp == 'y':
|
||||
return True
|
||||
if inp and inp == 'n':
|
||||
return False
|
||||
return prompt(question)
|
||||
|
||||
@dataclass
|
||||
class MVConfig:
|
||||
pretrained_unet_path: str
|
||||
revision: Optional[str]
|
||||
validation_dataset: Dict
|
||||
save_dir: str
|
||||
seed: Optional[int]
|
||||
validation_batch_size: int
|
||||
dataloader_num_workers: int
|
||||
|
||||
local_rank: int
|
||||
|
||||
pipe_kwargs: Dict
|
||||
pipe_validation_kwargs: Dict
|
||||
unet_from_pretrained_kwargs: Dict
|
||||
validation_guidance_scales: List[float]
|
||||
validation_grid_nrow: int
|
||||
camera_embedding_lr_mult: float
|
||||
|
||||
num_views: int
|
||||
camera_embedding_type: str
|
||||
|
||||
pred_type: str # joint, or ablation
|
||||
|
||||
enable_xformers_memory_efficient_attention: bool
|
||||
|
||||
cond_on_normals: bool
|
||||
cond_on_colors: bool
|
||||
|
||||
def load_config(*yaml_files, cli_args=[]):
|
||||
yaml_confs = [OmegaConf.load(f) for f in yaml_files]
|
||||
cli_conf = OmegaConf.from_cli(cli_args)
|
||||
conf = OmegaConf.merge(*yaml_confs, cli_conf)
|
||||
OmegaConf.resolve(conf)
|
||||
|
||||
schema = OmegaConf.structured(MVConfig)
|
||||
conf = OmegaConf.merge(schema, conf)
|
||||
return conf
|
||||
|
||||
def config_to_primitive(config, resolve=True):
|
||||
return OmegaConf.to_container(config, resolve=resolve)
|
||||
|
||||
|
||||
def dump_config(path, config):
|
||||
with open(path, 'w') as fp:
|
||||
OmegaConf.save(config=config, f=fp)
|
||||
|
||||
def get_rank():
|
||||
# SLURM_PROCID can be set even if SLURM is not managing the multiprocessing,
|
||||
# therefore LOCAL_RANK needs to be checked first
|
||||
rank_keys = ("RANK", "LOCAL_RANK", "SLURM_PROCID", "JSM_NAMESPACE_RANK")
|
||||
for key in rank_keys:
|
||||
rank = os.environ.get(key)
|
||||
if rank is not None:
|
||||
return int(rank)
|
||||
return 0
|
||||
|
||||
|
||||
def parse_version(ver):
|
||||
return version.parse(ver)
|
||||
@@ -12,8 +12,10 @@ from torch.utils.data import DataLoader
|
||||
import torchvision.transforms.functional as TF
|
||||
import numpy as np
|
||||
from safetensors.torch import load_file
|
||||
from einops import rearrange
|
||||
|
||||
from plyfile import PlyData
|
||||
from PIL import Image
|
||||
|
||||
from .mesh_processer.mesh import Mesh
|
||||
from .mesh_processer.mesh_utils import (
|
||||
@@ -30,17 +32,22 @@ from .algorithms.dmtet import DMTetMesh
|
||||
from .algorithms.triplane_gaussian_transformers import TGS
|
||||
from .algorithms.large_multiview_gaussian_model import LGM
|
||||
from .algorithms.nerf_marching_cubes_converter import GSConverterNeRFMarchingCubes
|
||||
from .tgs.utils.config import ExperimentConfig, load_config
|
||||
from .algorithms.NeuS_runner import NeuSParams, NeuSRunner
|
||||
from .tgs.utils.config import ExperimentConfig, load_config as load_config_tgs
|
||||
from .tgs.data import CustomImageOrbitDataset
|
||||
from .tgs.utils.misc import todevice, get_device
|
||||
from .lgm.core.options import config_defaults
|
||||
from .lgm.mvdream.pipeline_mvdream import MVDreamPipeline
|
||||
from .mvdiffusion.pipelines.pipeline_mvdiffusion_image import MVDiffusionImagePipeline
|
||||
from .mvdiffusion.data.single_image_dataset import SingleImageDataset
|
||||
from .mvdiffusion.utils.misc import load_config as load_config_wonder3d
|
||||
|
||||
from .shared_utils.image_utils import prepare_torch_img
|
||||
from .shared_utils.common_utils import cstr, parse_save_filename, get_list_filenames
|
||||
|
||||
DIFFUSERS_PIPE_DICT = OrderedDict([
|
||||
("MVDreamPipeline", MVDreamPipeline),
|
||||
("Wonder3DMVDiffusionPipeline", MVDiffusionImagePipeline),
|
||||
])
|
||||
|
||||
ROOT_PATH = os.path.join(comfy_paths.get_folder_paths("custom_nodes")[0], "ComfyUI-3D-Pack")
|
||||
@@ -74,6 +81,8 @@ ELEVATION_MAX = 90.0
|
||||
AZIMUTH_MIN = -180.0
|
||||
AZIMUTH_MAX = 180.0
|
||||
|
||||
WEIGHT_DTYPE = torch.float16
|
||||
|
||||
class Preview_3DGS:
|
||||
|
||||
@classmethod
|
||||
@@ -314,6 +323,7 @@ class Switch_Mesh_Axis:
|
||||
"axis_x_to": (["+x", "-x", "+y", "-y", "+z", "-z"],),
|
||||
"axis_y_to": (["+y", "-y", "+z", "-z", "+x", "-x"],),
|
||||
"axis_z_to": (["+z", "-z", "+x", "-x", "+y", "-y"],),
|
||||
"flip_normal": ("BOOLEAN", {"default": False},),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -326,13 +336,13 @@ class Switch_Mesh_Axis:
|
||||
FUNCTION = "switch_axis_and_scale"
|
||||
CATEGORY = "Comfy3D/Preprocessor"
|
||||
|
||||
def switch_axis_and_scale(self, mesh, axis_x_to, axis_y_to, axis_z_to):
|
||||
def switch_axis_and_scale(self, mesh, axis_x_to, axis_y_to, axis_z_to, flip_normal):
|
||||
|
||||
switched_mesh = None
|
||||
|
||||
if axis_x_to[1] != axis_y_to[1] and axis_x_to[1] != axis_z_to[1] and axis_y_to[1] != axis_z_to[1]:
|
||||
target_axis, target_scale, coordinate_invert_count = get_target_axis_and_scale([axis_x_to, axis_y_to, axis_z_to])
|
||||
switched_mesh = switch_mesh_axis_and_scale(mesh, target_axis, target_scale)
|
||||
switched_mesh = switch_mesh_axis_and_scale(mesh, target_axis, target_scale, flip_normal)
|
||||
else:
|
||||
cstr(f"[{self.__class__.__name__}] axis_x_to: {axis_x_to}, axis_y_to: {axis_y_to}, axis_z_to: {axis_z_to} have to be on separated axis").error.print()
|
||||
|
||||
@@ -813,30 +823,32 @@ class Gaussian_Splatting:
|
||||
|
||||
with torch.inference_mode(False):
|
||||
|
||||
gs_params = GSParams(training_iterations,
|
||||
batch_size,
|
||||
ms_ssim_loss_weight,
|
||||
alpha_loss_weight,
|
||||
offset_loss_weight,
|
||||
offset_opacity_loss_weight,
|
||||
invert_background_probability,
|
||||
feature_learning_rate,
|
||||
opacity_learning_rate,
|
||||
scaling_learning_rate,
|
||||
rotation_learning_rate,
|
||||
position_learning_rate_init,
|
||||
position_learning_rate_final,
|
||||
position_learning_rate_delay_mult,
|
||||
position_learning_rate_max_steps,
|
||||
initial_gaussians_num,
|
||||
K_nearest_neighbors,
|
||||
percent_dense,
|
||||
density_start_iterations,
|
||||
density_end_iterations,
|
||||
densification_interval,
|
||||
opacity_reset_interval,
|
||||
densify_grad_threshold,
|
||||
gaussian_sh_degree)
|
||||
gs_params = GSParams(
|
||||
training_iterations,
|
||||
batch_size,
|
||||
ms_ssim_loss_weight,
|
||||
alpha_loss_weight,
|
||||
offset_loss_weight,
|
||||
offset_opacity_loss_weight,
|
||||
invert_background_probability,
|
||||
feature_learning_rate,
|
||||
opacity_learning_rate,
|
||||
scaling_learning_rate,
|
||||
rotation_learning_rate,
|
||||
position_learning_rate_init,
|
||||
position_learning_rate_final,
|
||||
position_learning_rate_delay_mult,
|
||||
position_learning_rate_max_steps,
|
||||
initial_gaussians_num,
|
||||
K_nearest_neighbors,
|
||||
percent_dense,
|
||||
density_start_iterations,
|
||||
density_end_iterations,
|
||||
densification_interval,
|
||||
opacity_reset_interval,
|
||||
densify_grad_threshold,
|
||||
gaussian_sh_degree
|
||||
)
|
||||
|
||||
|
||||
if points_cloud_to_initialize_gaussian is not None:
|
||||
@@ -896,21 +908,23 @@ class Bake_Texture_To_Mesh:
|
||||
FUNCTION = "bake_texture"
|
||||
CATEGORY = "Comfy3D/Algorithm"
|
||||
|
||||
def bake_texture(self,
|
||||
reference_images,
|
||||
reference_masks,
|
||||
reference_orbit_camera_poses,
|
||||
reference_orbit_camera_fovy,
|
||||
mesh,
|
||||
mesh_albedo_width,
|
||||
mesh_albedo_height,
|
||||
training_iterations,
|
||||
batch_size,
|
||||
texture_learning_rate,
|
||||
train_mesh_geometry,
|
||||
geometry_learning_rate,
|
||||
ms_ssim_loss_weight,
|
||||
force_cuda_rasterize):
|
||||
def bake_texture(
|
||||
self,
|
||||
reference_images,
|
||||
reference_masks,
|
||||
reference_orbit_camera_poses,
|
||||
reference_orbit_camera_fovy,
|
||||
mesh,
|
||||
mesh_albedo_width,
|
||||
mesh_albedo_height,
|
||||
training_iterations,
|
||||
batch_size,
|
||||
texture_learning_rate,
|
||||
train_mesh_geometry,
|
||||
geometry_learning_rate,
|
||||
ms_ssim_loss_weight,
|
||||
force_cuda_rasterize
|
||||
):
|
||||
|
||||
mesh.set_new_albedo(mesh_albedo_width, mesh_albedo_height)
|
||||
|
||||
@@ -983,16 +997,18 @@ class Deep_Marching_Tetrahedrons:
|
||||
reference_points_cloud=None, reference_images=None, reference_masks=None):
|
||||
|
||||
with torch.inference_mode(False):
|
||||
dmtet_mesh = DMTetMesh(training_iterations,
|
||||
points_cloud_fitting_weight,
|
||||
mesh_smoothing_weight,
|
||||
chamfer_faces_sample_scale,
|
||||
mesh_scale,
|
||||
grid_resolution,
|
||||
geometry_learning_rate,
|
||||
positional_encoding_multires,
|
||||
mlp_internal_dims,
|
||||
mlp_hidden_layer_num)
|
||||
dmtet_mesh = DMTetMesh(
|
||||
training_iterations,
|
||||
points_cloud_fitting_weight,
|
||||
mesh_smoothing_weight,
|
||||
chamfer_faces_sample_scale,
|
||||
mesh_scale,
|
||||
grid_resolution,
|
||||
geometry_learning_rate,
|
||||
positional_encoding_multires,
|
||||
mlp_internal_dims,
|
||||
mlp_hidden_layer_num
|
||||
)
|
||||
|
||||
if reference_points_cloud is None and reference_images is None:
|
||||
cstr(f"[{self.__class__.__name__}] reference_points_cloud and reference_images cannot both be None, you need at least provide one of them!").error.print()
|
||||
@@ -1001,13 +1017,13 @@ class Deep_Marching_Tetrahedrons:
|
||||
dmtet_mesh.training(reference_points_cloud, reference_images, reference_masks)
|
||||
mesh = dmtet_mesh.get_mesh()
|
||||
return (mesh, )
|
||||
|
||||
TGS_CONFIG_PATH = "configs/tgs_config.yaml"
|
||||
|
||||
class Load_Triplane_Gaussian_Transformers:
|
||||
|
||||
checkpoints_dir = "checkpoints/tgs"
|
||||
default_ckpt_name = "model_lvis_rel.ckpt"
|
||||
default_repo_id = "VAST-AI/TriplaneGaussian"
|
||||
tgs_config_path = "configs/tgs_config.yaml"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -1035,16 +1051,16 @@ class Load_Triplane_Gaussian_Transformers:
|
||||
|
||||
device = get_device()
|
||||
|
||||
config_path = os.path.join(ROOT_PATH, TGS_CONFIG_PATH)
|
||||
cfg: ExperimentConfig = load_config(config_path)
|
||||
config_path = os.path.join(ROOT_PATH, self.tgs_config_path)
|
||||
cfg: ExperimentConfig = load_config_tgs(config_path)
|
||||
|
||||
# Download pre-trained model if it not exist locally
|
||||
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
|
||||
if not os.path.isfile(ckpt_path):
|
||||
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo VAST-AI/TriplaneGaussian instead").warning.print()
|
||||
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo {self.default_repo_id} instead").warning.print()
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
hf_hub_download(repo_id="VAST-AI/TriplaneGaussian", local_dir=self.checkpoints_dir_abs, filename=self.default_ckpt_name, repo_type="model")
|
||||
hf_hub_download(repo_id=self.default_repo_id, local_dir=self.checkpoints_dir_abs, filename=self.default_ckpt_name, repo_type="model")
|
||||
ckpt_path = os.path.join(self.checkpoints_dir_abs, self.default_ckpt_name)
|
||||
|
||||
cfg.system.weights=ckpt_path
|
||||
@@ -1059,6 +1075,8 @@ class Load_Triplane_Gaussian_Transformers:
|
||||
|
||||
class Triplane_Gaussian_Transformers:
|
||||
|
||||
tgs_config_path = "configs/tgs_config.yaml"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
@@ -1080,8 +1098,8 @@ class Triplane_Gaussian_Transformers:
|
||||
CATEGORY = "Comfy3D/Algorithm"
|
||||
|
||||
def run_TGS(self, reference_image, reference_mask, tgs_model, cam_dist):
|
||||
config_path = os.path.join(ROOT_PATH, TGS_CONFIG_PATH)
|
||||
cfg: ExperimentConfig = load_config(config_path)
|
||||
config_path = os.path.join(ROOT_PATH, self.tgs_config_path)
|
||||
cfg: ExperimentConfig = load_config_tgs(config_path)
|
||||
|
||||
#save_path = os.path.join(ROOT_PATH, "outputs")
|
||||
#tgs_model.set_save_dir(save_path)
|
||||
@@ -1089,11 +1107,12 @@ class Triplane_Gaussian_Transformers:
|
||||
cfg.data.cond_camera_distance = cam_dist
|
||||
cfg.data.eval_camera_distance = cam_dist
|
||||
dataset = CustomImageOrbitDataset(reference_image, reference_mask, cfg.data)
|
||||
dataloader = DataLoader(dataset,
|
||||
batch_size=cfg.data.eval_batch_size,
|
||||
shuffle=False,
|
||||
collate_fn=dataset.collate
|
||||
)
|
||||
dataloader = DataLoader(
|
||||
dataset,
|
||||
batch_size=cfg.data.eval_batch_size,
|
||||
shuffle=False,
|
||||
collate_fn=dataset.collate
|
||||
)
|
||||
|
||||
|
||||
gs_ply = []
|
||||
@@ -1112,7 +1131,7 @@ class Triplane_Gaussian_Transformers:
|
||||
)
|
||||
"""
|
||||
|
||||
return (gs_ply[0],)
|
||||
return (gs_ply[0], )
|
||||
|
||||
class Load_Diffusers_Pipeline:
|
||||
|
||||
@@ -1121,7 +1140,7 @@ class Load_Diffusers_Pipeline:
|
||||
return {
|
||||
"required": {
|
||||
"diffusers_pipeline_name": (list(DIFFUSERS_PIPE_DICT.keys()),),
|
||||
"pretrained_model_name": ("STRING", {"default": "ashawkey/imagedream-ipmv-diffusers", "multiline": False}),
|
||||
"model_name": ("STRING", {"default": "ashawkey/imagedream-ipmv-diffusers", "multiline": False}),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -1134,21 +1153,118 @@ class Load_Diffusers_Pipeline:
|
||||
FUNCTION = "load_diffusers_pipe"
|
||||
CATEGORY = "Comfy3D/Import|Export"
|
||||
|
||||
def load_diffusers_pipe(self, diffusers_pipeline_name, pretrained_model_name):
|
||||
def load_diffusers_pipe(self, diffusers_pipeline_name, model_name):
|
||||
|
||||
# resume pretrained checkpoint
|
||||
ckpt_path = os.path.join(ROOT_PATH, "checkpoints", model_name)
|
||||
if not os.path.exists(ckpt_path):
|
||||
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo {model_name}").warning.print()
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download(repo_id=model_name, local_dir=ckpt_path, repo_type="model")
|
||||
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
diffusers_pipeline_class = DIFFUSERS_PIPE_DICT[diffusers_pipeline_name]
|
||||
|
||||
# load image dream
|
||||
|
||||
pipe = diffusers_pipeline_class.from_pretrained(
|
||||
pretrained_model_name, # remote weights
|
||||
torch_dtype=torch.float16,
|
||||
ckpt_path,
|
||||
torch_dtype=WEIGHT_DTYPE,
|
||||
trust_remote_code=True,
|
||||
# local_files_only=True,
|
||||
).to(device)
|
||||
|
||||
return (pipe,)
|
||||
pipe.enable_xformers_memory_efficient_attention()
|
||||
|
||||
return (pipe, )
|
||||
|
||||
class Wonder3D_MVDiffusion_Model:
|
||||
|
||||
wonder3d_config_path = "configs/wonder3d_config.yaml"
|
||||
fix_cam_pose_dir = "mvdiffusion/data/fixed_poses/nine_views"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"mvdiffusion_pipe": ("DIFFUSERS_PIPE",),
|
||||
"reference_image": ("IMAGE",),
|
||||
"reference_mask": ("MASK",),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"mv_guidance_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.01}),
|
||||
"num_inference_steps": ("INT", {"default": 50, "min": 1}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"IMAGE",
|
||||
"IMAGE",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"multiview_images",
|
||||
"multiview_normals",
|
||||
)
|
||||
FUNCTION = "run_mvdiffusion"
|
||||
CATEGORY = "Comfy3D/Algorithm"
|
||||
|
||||
def run_mvdiffusion(
|
||||
self,
|
||||
mvdiffusion_pipe,
|
||||
reference_image, # [1, H, W, 3]
|
||||
reference_mask, # [1, H, W]
|
||||
seed,
|
||||
mv_guidance_scale,
|
||||
num_inference_steps,
|
||||
):
|
||||
|
||||
config_path = os.path.join(ROOT_PATH, self.wonder3d_config_path)
|
||||
cfg = load_config_wonder3d(config_path)
|
||||
|
||||
batch = self.prepare_data(reference_image[0], reference_mask[0].unsqueeze(2))
|
||||
|
||||
mvdiffusion_pipe.set_progress_bar_config(disable=True)
|
||||
seed = int(seed)
|
||||
generator = torch.Generator(device=mvdiffusion_pipe.unet.device).manual_seed(seed)
|
||||
|
||||
# repeat (2B, Nv, 3, H, W)
|
||||
imgs_in = torch.cat([batch['imgs_in']] * 2, dim=0).to(WEIGHT_DTYPE)
|
||||
|
||||
# (2B, Nv, Nce)
|
||||
camera_embeddings = torch.cat([batch['camera_embeddings']] * 2, dim=0).to(WEIGHT_DTYPE)
|
||||
|
||||
task_embeddings = torch.cat([batch['normal_task_embeddings'], batch['color_task_embeddings']], dim=0).to(WEIGHT_DTYPE)
|
||||
|
||||
camera_embeddings = torch.cat([camera_embeddings, task_embeddings], dim=-1).to(WEIGHT_DTYPE)
|
||||
|
||||
# (B*Nv, 3, H, W)
|
||||
imgs_in = rearrange(imgs_in, "Nv C H W -> (Nv) C H W")
|
||||
# (B*Nv, Nce)
|
||||
# camera_embeddings = rearrange(camera_embeddings, "B Nv Nce -> (B Nv) Nce")
|
||||
|
||||
out = mvdiffusion_pipe(
|
||||
imgs_in,
|
||||
# camera_embeddings,
|
||||
generator=generator,
|
||||
guidance_scale=mv_guidance_scale,
|
||||
num_inference_steps=num_inference_steps,
|
||||
output_type='pt',
|
||||
num_images_per_prompt=1,
|
||||
**cfg.pipe_validation_kwargs,
|
||||
).images
|
||||
|
||||
num_views = out.shape[0] // 2
|
||||
mv_images = out[num_views:].permute(0, 2, 3, 1) # [N, 3, H, W] -> [N, H, W, 3]
|
||||
mv_normals = out[:num_views].permute(0, 2, 3, 1) # [N, 3, H, W] -> [N, H, W, 3]
|
||||
|
||||
return (mv_images, mv_normals, )
|
||||
|
||||
def prepare_data(self, ref_image, ref_mask):
|
||||
single_image = torch.cat((ref_image, ref_mask), dim=2).detach().cpu().numpy()
|
||||
single_image = Image.fromarray((single_image * 255).astype(np.uint8), mode="RGBA")
|
||||
abs_fix_cam_pose_dir = os.path.join(ROOT_PATH, self.fix_cam_pose_dir)
|
||||
dataset = SingleImageDataset(fix_cam_pose_dir=abs_fix_cam_pose_dir, num_views=6, img_wh=[256, 256], bg_color='white', single_image=single_image)
|
||||
return dataset[0]
|
||||
|
||||
class MVDream_Model:
|
||||
|
||||
@@ -1210,12 +1326,13 @@ class MVDream_Model:
|
||||
mv_images = mvdream_pipe(prompt, reference_image, generator=generator, negative_prompt=prompt_neg, guidance_scale=mv_guidance_scale, num_inference_steps=num_inference_steps, elevation=elevation)
|
||||
mv_images = torch.from_numpy(np.stack([mv_images[1], mv_images[2], mv_images[3], mv_images[0]], axis=0)).float() # [4, H, W, 3], float32
|
||||
|
||||
return (mv_images,)
|
||||
return (mv_images, )
|
||||
|
||||
class Load_Large_Multiview_Gaussian_Model:
|
||||
|
||||
checkpoints_dir = "checkpoints/lgm"
|
||||
default_ckpt_name = "model_fp16.safetensors"
|
||||
default_repo_id = "ashawkey/LGM"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -1247,10 +1364,10 @@ class Load_Large_Multiview_Gaussian_Model:
|
||||
# resume pretrained checkpoint
|
||||
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
|
||||
if not os.path.isfile(ckpt_path):
|
||||
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo ashawkey/LGM instead").warning.print()
|
||||
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo {self.default_repo_id} instead").warning.print()
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
hf_hub_download(repo_id="ashawkey/LGM", local_dir=self.checkpoints_dir_abs, filename=self.default_ckpt_name, repo_type="model")
|
||||
hf_hub_download(repo_id=self.default_repo_id, local_dir=self.checkpoints_dir_abs, filename=self.default_ckpt_name, repo_type="model")
|
||||
ckpt_path = os.path.join(self.checkpoints_dir_abs, self.default_ckpt_name)
|
||||
|
||||
if ckpt_path.endswith('safetensors'):
|
||||
@@ -1266,7 +1383,7 @@ class Load_Large_Multiview_Gaussian_Model:
|
||||
lgm_model = lgm_model.half().to(device)
|
||||
lgm_model.eval()
|
||||
|
||||
return (lgm_model,)
|
||||
return (lgm_model, )
|
||||
|
||||
IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
|
||||
IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
|
||||
@@ -1299,7 +1416,7 @@ class Large_Multiview_Gaussian_Model:
|
||||
rays_embeddings = lgm_model.prepare_default_rays(device)
|
||||
ref_image_torch = torch.cat([ref_image_torch, rays_embeddings], dim=1).unsqueeze(0) # [1, 4, 9, 256, 256]
|
||||
|
||||
with torch.autocast(device_type=device, dtype=torch.float16):
|
||||
with torch.autocast(device_type=device, dtype=WEIGHT_DTYPE):
|
||||
# generate gaussians
|
||||
gaussians = lgm_model.forward_gaussians(ref_image_torch)
|
||||
|
||||
@@ -1338,4 +1455,104 @@ class Convert_3DGS_to_Mesh_with_NeRF_and_Marching_Cubes:
|
||||
converter.fit_mesh()
|
||||
converter.fit_mesh_uv()
|
||||
|
||||
return(converter.get_mesh(), )
|
||||
return(converter.get_mesh(), )
|
||||
|
||||
class NeuS:
|
||||
NeuS_config_path = "configs/NeuS.conf"
|
||||
fix_cam_pose_dir = "NeuS/models/fixed_poses"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"reference_image": ("IMAGE",),
|
||||
"reference_mask": ("MASK",),
|
||||
"reference_normals": ("IMAGE",),
|
||||
"training_iterations": ("INT", {"default": 1000, "min": 1, "max": 100000}), # longer time, better result. 1w will be ok for most cases
|
||||
"batch_size": ("INT", {"default": 512, "min": 1, "max": 0xffffffffffffffff}),
|
||||
"learning_rate": ("FLOAT", {"default": 5e-4, "min": 0.00001, "step": 0.00001}),
|
||||
"learning_rate_alpha": ("FLOAT", {"default": 5e-4, "min": 0.00001, "step": 0.00001}),
|
||||
"color_loss_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.01}),
|
||||
"mesh_smoothing_weight": ("FLOAT", {"default": 0.1, "min": 0.0, "step": 0.01}),
|
||||
"mask_loss_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.01}),
|
||||
"normal_loss_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.01}),
|
||||
"sparse_loss_weight": ("FLOAT", {"default": 0.1, "min": 0.0, "step": 0.01}),
|
||||
"warm_up_end": ("INT", {"default": 500, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"anneal_end": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"use_white_background": ("BOOLEAN", {"default": True},),
|
||||
"geometry_extract_resolution": ("INT", {"default": 512, "min": 1, "max": 0xffffffffffffffff}),
|
||||
"marching_cude_threshold": ("FLOAT", {"default": 0.0, "min": 0.0, "step": 0.01}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"MESH",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"mesh",
|
||||
)
|
||||
FUNCTION = "run_NeuS"
|
||||
CATEGORY = "Comfy3D/Algorithm"
|
||||
|
||||
def run_NeuS(
|
||||
self,
|
||||
reference_image, # [num_views, H, W, 3]
|
||||
reference_mask, # [num_views, H, W]
|
||||
reference_normals, # [num_views, H, W, 3]
|
||||
training_iterations,
|
||||
batch_size,
|
||||
learning_rate,
|
||||
learning_rate_alpha,
|
||||
color_loss_weight,
|
||||
mesh_smoothing_weight,
|
||||
mask_loss_weight,
|
||||
normal_loss_weight,
|
||||
sparse_loss_weight,
|
||||
warm_up_end,
|
||||
anneal_end,
|
||||
use_white_background,
|
||||
geometry_extract_resolution,
|
||||
marching_cude_threshold
|
||||
):
|
||||
mesh = None
|
||||
num_views = reference_image.shape[0]
|
||||
|
||||
if num_views in [6, 5, 4]:
|
||||
config_path = os.path.join(ROOT_PATH, self.NeuS_config_path)
|
||||
abs_fix_cam_pose_dir = os.path.join(ROOT_PATH, self.fix_cam_pose_dir)
|
||||
|
||||
with torch.inference_mode(False):
|
||||
NeuS_params = NeuSParams(
|
||||
training_iterations,
|
||||
batch_size,
|
||||
learning_rate,
|
||||
learning_rate_alpha,
|
||||
color_loss_weight,
|
||||
mesh_smoothing_weight,
|
||||
mask_loss_weight,
|
||||
normal_loss_weight,
|
||||
sparse_loss_weight,
|
||||
warm_up_end,
|
||||
anneal_end,
|
||||
use_white_background
|
||||
)
|
||||
|
||||
runner = NeuSRunner(
|
||||
reference_image,
|
||||
reference_mask,
|
||||
reference_normals,
|
||||
config_path,
|
||||
abs_fix_cam_pose_dir,
|
||||
num_views,
|
||||
NeuS_params
|
||||
)
|
||||
|
||||
cstr(f"[{self.__class__.__name__}] Training NeuS...").msg.print()
|
||||
runner.train()
|
||||
|
||||
cstr(f"[{self.__class__.__name__}] Extracting Mesh...").msg.print()
|
||||
mesh = runner.extract_mesh(resolution=geometry_extract_resolution, threshold=marching_cude_threshold)
|
||||
else:
|
||||
cstr(f"[{self.__class__.__name__}] Number of views must be one of the following: 6, 5, 4. But got {num_views}").error.print()
|
||||
|
||||
return(mesh, )
|
||||
@@ -8,9 +8,11 @@ einops
|
||||
scipy
|
||||
kornia
|
||||
opencv-python
|
||||
pillow
|
||||
roma
|
||||
nerfacc>=0.5.3
|
||||
PyMCubes
|
||||
scikit-learn
|
||||
|
||||
# for use ML models
|
||||
diffusers>=0.26.1
|
||||
@@ -38,3 +40,4 @@ tqdm
|
||||
jaxtyping
|
||||
packaging
|
||||
OmegaConf
|
||||
pyhocon
|
||||
|
||||
@@ -2,8 +2,10 @@ import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
def prepare_torch_img(img, size_H, size_W, device="cuda"):
|
||||
def prepare_torch_img(img, size_H, size_W, device="cuda", keep_shape=False):
|
||||
# [N, H, W, C] -> [N, C, H, W]
|
||||
img_new = img.permute(0, 3, 1, 2).to(device)
|
||||
img_new = F.interpolate(img_new, (size_H, size_W), mode="bilinear", align_corners=False).contiguous()
|
||||
if keep_shape:
|
||||
img_new = img_new.permute(0, 2, 3, 1)
|
||||
return img_new
|
||||
+1
-3
@@ -6,9 +6,7 @@ from omegaconf import OmegaConf
|
||||
from .typing import *
|
||||
|
||||
# ============ Register OmegaConf Recolvers ============= #
|
||||
OmegaConf.register_new_resolver(
|
||||
"calc_exp_lr_decay_rate", lambda factor, n: factor ** (1.0 / n), replace=True
|
||||
)
|
||||
OmegaConf.register_new_resolver("calc_exp_lr_decay_rate", lambda factor, n: factor ** (1.0 / n), replace=True)
|
||||
OmegaConf.register_new_resolver("add", lambda a, b: a + b, replace=True)
|
||||
OmegaConf.register_new_resolver("sub", lambda a, b: a - b, replace=True)
|
||||
OmegaConf.register_new_resolver("mul", lambda a, b: a * b, replace=True)
|
||||
|
||||
@@ -107,6 +107,11 @@ async function main(filepath="") {
|
||||
loader.load( url, function ( obj ) {
|
||||
obj.scale.setScalar( 5 );
|
||||
scene.add( obj );
|
||||
obj.traverse(node => {
|
||||
if (node.material && node.material.map == null) {
|
||||
node.material.vertexColors = true;
|
||||
}
|
||||
});
|
||||
|
||||
}, onProgress, onError );
|
||||
|
||||
@@ -121,7 +126,7 @@ async function main(filepath="") {
|
||||
const model = gltf.scene;
|
||||
//model.position.set( 1, 1, 0 );
|
||||
//model.scale.set( 0.01, 0.01, 0.01 );
|
||||
model.scale.set( 3, 3, 3 );
|
||||
model.scale.set( 5, 5, 5 );
|
||||
|
||||
scene.add( model );
|
||||
|
||||
|
||||
Reference in New Issue
Block a user