from dataclasses import dataclass import math import copy import torch import torch.nn as nn from einops import repeat, rearrange from einops.layers.torch import Rearrange from easydict import EasyDict as edict import diffusionGS from diffusionGS.models.transformers.utils_transformer import DiTBlock,_init_weights from diffusionGS.models.gsrenderer.renderer import Renderer,SceneRenderer from diffusionGS.utils.checkpoint import checkpoint, load_denoiser_checkpoint from diffusionGS.utils.base import BaseModule from diffusionGS.utils.typing import * from diffusionGS.utils.ops import generate_dense_grid_points from tqdm import tqdm from torchvision.utils import save_image import numpy as np # import cupy as np def modulate(x, shift, scale): return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) ################################################################################# # Embedding Layers for Timesteps and Class Labels # ################################################################################# class TimestepEmbedder(nn.Module): """ Embeds scalar timesteps into vector representations. """ def __init__(self, hidden_size, frequency_embedding_size=256): super().__init__() self.mlp = nn.Sequential( nn.Linear(frequency_embedding_size, hidden_size, bias=True), nn.SiLU(), nn.Linear(hidden_size, hidden_size, bias=True), ) self.frequency_embedding_size = frequency_embedding_size ''' staticmethod 装饰器可以声明静态方法 无需实例化 TimestepEmbedder 类, 就可以直接调用 timestep_embedding 方法 ''' @staticmethod def timestep_embedding(t, dim, max_period=10000): """ Create sinusoidal timestep embeddings. :param t: a 1-D Tensor of N indices, one per batch element. These may be fractional. :param dim: the dimension of the output. :param max_period: controls the minimum frequency of the embeddings. :return: an (N, D) Tensor of positional embeddings. """ # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py half = dim // 2 freqs = torch.exp( -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half ).to(device=t.device) args = t[:, None].float() * freqs[None] embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) if dim % 2: embedding = torch.cat( [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 ) return embedding def forward(self, t): t_freq = self.timestep_embedding(t, self.frequency_embedding_size) t_emb = self.mlp(t_freq) return t_emb class GaussiansUpsampler(nn.Module): def __init__(self, config): super().__init__() self.config = config """ xyz : torch.tensor of shape (n_gaussians, 3) features : torch.tensor of shape (n_gaussians, (sh_degree + 1) ** 2, 3) scaling : torch.tensor of shape (n_gaussians, 3) rotation : torch.tensor of shape (n_gaussians, 4) opacity : torch.tensor of shape (n_gaussians, 1) """ self.layernorm = nn.LayerNorm(config.width, bias=False) # d - dimension - channel u = 1 #config.model.gaussians.upsampler.upsample_factor if u > 1: raise NotImplementedError("GaussiansUpsampler only supports u=1") else: self.linear = nn.Linear( config.width, 3 + (config.gaussians_sh_degree + 1) ** 2 * 3 + 3 + 4 + 1, # 直接转成 per-pixel Gaussian map bias=False, ) self.adaLN_modulation = nn.Sequential( nn.SiLU(), nn.Linear(config.width, 2 * config.width, bias=True) ) def to_gs(self, gaussians): # 把 dimenssion 分掉，分出对应的维度 """ gaussians: [b, n_gaussians, d] n_gaussians - 高斯的数量 d - 每个高斯的 attribute 的数量 """ xyz, features, scaling, rotation, opacity = gaussians.split( [3, (self.config.gaussians_sh_degree + 1) ** 2 * 3, 3, 4, 1], dim=2 ) features = features.reshape( features.size(0), features.size(1), (self.config.gaussians_sh_degree + 1) ** 2, 3, ) #breakpoint() # torch.exp(torch.tensor(-2.3)) = 0.1 # torch.exp(torch.tensor(-1.2)) = 0.3 scaling = (scaling - 2.3).clamp(max=-1.20) # torch.sigmoid(-torch.tensor(2.)) = 0.1192 opacity = opacity - 2.0 return xyz, features, scaling, rotation, opacity def forward(self, gaussians, t_embedding): """ gaussians: [b, n_gaussians, d] t_embedding: [b, d] output: [b, n_gaussians, dd] """ u = 1 #self.config.model.gaussians.upsampler.upsample_factor if u > 1: raise NotImplementedError("GaussiansUpsampler only supports u=1") shift, scale = self.adaLN_modulation(t_embedding).chunk(2, dim=1) gaussians = modulate(self.layernorm(gaussians), shift, scale) gaussians = self.linear(gaussians) return gaussians class ImageTokenDecoder(nn.Module): def __init__(self, config): super().__init__() self.config = config self.layernorm = nn.LayerNorm(config.width, bias=False) self.linear = nn.Linear( config.width, (config.patch_size**2) * (3 + (config.gaussians_sh_degree + 1) ** 2 * 3 + 3 + 4 + 1), bias=False, ) self.adaLN_modulation = nn.Sequential( nn.SiLU(), nn.Linear(config.width, 2 * config.width, bias=True) ) def forward(self, img_tokens, t_embedding): """ img_tokens: [b, n_patches, d] t_embedding: [b, d] output: [b, n_patches, dd] """ shift, scale = self.adaLN_modulation(t_embedding).chunk(2, dim=1) img_tokens = modulate(self.layernorm(img_tokens), shift, scale) img_tokens = self.linear(img_tokens) return img_tokens ###################### AutoEncoder @diffusionGS.register("diffusion-gs-model-scene") class DGSDenoiser(BaseModule): r""" An image to gaussian diffusion models. """ @dataclass class Config(BaseModule.Config): pretrained_model_name_or_path: str = "" use_downsample: bool = False num_latents: int = 256 width: int = 1024 in_channels: int = 3 patch_size: int = 16 n_gaussians: int = 2 dim_heads: int = 64 num_layers: int = 24 ray_pe_type: str = "relative_plk" hard_pixelalign: bool = True clip_xyz: bool = True ##gaussian relative gaussians_sh_degree: int = 0 use_gssplat: bool = False ##diffusion relative prior_distribution: str = "gaussian" #### use_flash: bool = False use_checkpoint: bool = True grad_checkpoint_every: int = 1 range_setting_type: str = "linear_depth" range_setting_near: int = 0 range_setting_far: int = 500 cfg: Config def configure(self) -> None: super().configure() ## ##time embedder self.t_embedder = TimestepEmbedder(self.cfg.width) # Initialize timestep embedding MLP: nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) ##image tokenlizer self.image_tokenizer = nn.Sequential( Rearrange( "b v c (hh ph) (ww pw) -> (b v) (hh ww) (ph pw c)", ph=self.cfg.patch_size, pw=self.cfg.patch_size, ), nn.Linear( self.cfg.in_channels * (self.cfg.patch_size**2), self.cfg.width, bias=False, ), ) self.image_tokenizer.apply(_init_weights) self.gaussians_pos_embedding = nn.Parameter( torch.randn( 1, self.cfg.n_gaussians, # 3d gaussian 的个数？n_gaussians = 2? self.cfg.width, # d = 1024 ) ) nn.init.trunc_normal_(self.gaussians_pos_embedding, std=0.02) self.transformer_input_layernorm = nn.LayerNorm( self.cfg.width, bias=False ) self.transformer = nn.ModuleList( [ DiTBlock( self.cfg.width, self.cfg.width // self.cfg.dim_heads ) for _ in range(self.cfg.num_layers) ] ) self.transformer.apply(_init_weights) self.upsampler = GaussiansUpsampler(self.cfg) self.upsampler.apply(_init_weights) self.image_token_decoder = ImageTokenDecoder(self.cfg) self.image_token_decoder.apply(_init_weights) # encoder self.gs_renderer = Renderer(self.cfg) #SceneRenderer # range func self.range_func = lambda t : torch.sigmoid(t) * (self.cfg.range_setting_far - self.cfg.range_setting_near) + self.cfg.range_setting_near #get_point_range_func(self.cfg) # renderer if self.cfg.pretrained_model_name_or_path != "": print(f"Loading pretrained shape model from {self.cfg.pretrained_model_name_or_path}") self.load_state_dict( load_denoiser_checkpoint(self.cfg.pretrained_model_name_or_path), strict=True, ) def forward(self, input_batch, timesteps): guassians_parameters,_ = self.image_to_gaussians(input_batch['image'], input_batch['ray_o'], input_batch['ray_d'], timesteps) rendered_images = self.render_gaussians(guassians_parameters, input_batch['c2w'], input_batch['fxfycxcy'], input_batch['image'].shape[3], input_batch['image'].shape[4]) return rendered_images, self.prepare_to_save(guassians_parameters) def prepare_to_save(self,gaussians_parameters): gaussians = [] for b in range(gaussians_parameters.xyz.size(0)): self.gs_renderer.gaussians_model.empty() gaussians_model = copy.deepcopy(self.gs_renderer.gaussians_model) gaussians.append( gaussians_model.set_data( gaussians_parameters.xyz[b].detach().float(), gaussians_parameters.features[b].detach().float(), gaussians_parameters.scaling[b].detach().float(), gaussians_parameters.rotation[b].detach().float(), gaussians_parameters.opacity[b].detach().float(), ) ) return gaussians def image_to_gaussians(self, images:torch.FloatTensor, ray_o:torch.FloatTensor, ray_d:torch.FloatTensor, t:torch.LongTensor, training: bool = False): # breakpoint() #do plucker pe if self.cfg.ray_pe_type == "relative_plk": # ray_o 和 ray_d 实际存储在一个 per-pixel ray direction map 上边, 为什么要将两者点乘后相加呢 o_dot_d = torch.sum(-ray_o * ray_d, dim=2, keepdim=True) # 这里假设 ray 是被归一化后的了 # nearest_pts 表示离坐标系原点最近的点，平方求导 nearest_pts = ray_o + o_dot_d * ray_d posed_images = torch.cat( [ images[:, :, :3, :, :] * 2.0 - 1.0, ray_d, nearest_pts, ], dim=2, ) else: o_cross_d = torch.cross(ray_o, ray_d, dim=2) posed_images = torch.cat( [ images[:, :, :3, :, :] * 2.0 - 1.0, o_cross_d, ray_d, ], dim=2, ) b, v, c, h, w = posed_images.size() img_tokens = self.image_tokenizer(posed_images) #img_tokens = torch.cat([self.cond_image_tokenizer(posed_images[:,:1]),self.image_tokenizer(posed_images[:,1:])],dim=0) # breakpoint() t = self.t_embedder(t) # [b, d] _, n_patches, d = img_tokens.size() # [b*v, n_patches, d] img_tokens = img_tokens.reshape(b, v * n_patches, d) # [b, v*n_patches, d] # [b, n_gaussians, d] gaussians_tokens = self.gaussians_pos_embedding.expand(b, -1, -1) # 复制 b 份 # 表示每隔多少层去检查一次梯度，清空释放掉一些梯度 checkpoint_every = self.cfg.grad_checkpoint_every concat_nerf_img_tokens = torch.cat((gaussians_tokens, img_tokens), dim=1) concat_nerf_img_tokens = self.transformer_input_layernorm( concat_nerf_img_tokens ) for i in range(0, len(self.transformer), checkpoint_every): concat_nerf_img_tokens = torch.utils.checkpoint.checkpoint( self.run_layers(i, i + checkpoint_every), concat_nerf_img_tokens, t, use_reentrant=False, ) gaussians_tokens, img_tokens = concat_nerf_img_tokens.split( [self.cfg.n_gaussians, v * n_patches], dim=1 ) gaussians = self.upsampler(gaussians_tokens, t) # [b, v*n_patches, p*p*gs] img_aligned_gaussians = self.image_token_decoder(img_tokens, t) img_aligned_gaussians = img_aligned_gaussians.reshape( b, -1, (3 + (self.cfg.gaussians_sh_degree + 1) ** 2 * 3 + 3 + 4 + 1), ) # [b, v*pixels, gs] n_img_aligned_gaussians = img_aligned_gaussians.size(1) # torch.cuda.synchronize() # toc = time.time() # print(f"input.image.size(): {input.image.size()}, time: {toc - tic:.3f} sec") all_gaussians = torch.cat((gaussians, img_aligned_gaussians), dim=1) xyz, features, scaling, rotation, opacity = self.upsampler.to_gs(all_gaussians) ###gs post process # 选择最后 n_img_aligned_gaussians 个元素 # n_img_aligned_gaussians = v * n_patches * patchsize * patchsize img_aligned_xyz = xyz[:, -n_img_aligned_gaussians:, :] # 把 image 对应的 Gaussians 模型 xyz 取出来 img_aligned_xyz = rearrange( img_aligned_xyz, "b (v h w ph pw) c -> b v c (h ph) (w pw)", v=v, h=h // self.cfg.patch_size, w=w // self.cfg.patch_size, ph=self.cfg.patch_size, pw=self.cfg.patch_size, ) # 对 image_aligned_xyz 进行矫正 if self.cfg.hard_pixelalign: #这是 true, 要执行的 img_aligned_xyz = img_aligned_xyz.mean(dim=2, keepdim=True) img_aligned_xyz = self.range_func(img_aligned_xyz) img_aligned_xyz = ray_o + img_aligned_xyz * ray_d img_aligned_xyz_reshape = rearrange( img_aligned_xyz, "b v c (h ph) (w pw) -> b (v h w ph pw) c", ph=self.cfg.patch_size, pw=self.cfg.patch_size, ) xyz = torch.cat( (xyz[:, :-n_img_aligned_gaussians, :], img_aligned_xyz_reshape), dim=1 ) result_softpa = edict( xyz=xyz, features=features, scaling=scaling, rotation=rotation, opacity=opacity, ) return result_softpa, img_aligned_xyz def render_gaussians(self, gaussian_params, c2w, fxfycxcy, height, width): render_input = self.gs_renderer( gaussian_params.xyz, gaussian_params.features, gaussian_params.scaling, gaussian_params.rotation, gaussian_params.opacity, height, width, C2W=c2w, # fxfycxcy=fxfycxcy, # ) return render_input @property def dtype(self): # 返回模型参数的数据类型 return next(self.parameters()).dtype def run_layers(self, start, end): def custom_forward(concat_nerf_img_tokens, t): for i in range(start, min(end, len(self.transformer))): concat_nerf_img_tokens = self.transformer[i](concat_nerf_img_tokens, t) return concat_nerf_img_tokens return custom_forward