Spaces:

CaiYuanhao
/

DiffusionGS-Object

Running on Zero

DiffusionGS-Object / diffusionGS /models /denoiser /denoiser_scene.py

Deploy DiffusionGS online generation demos

98e7453 verified 3 days ago

16.8 kB

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

	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, vn_patches, pp*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