Spaces:

CaiYuanhao
/

DiffusionGS-Object

Running on Zero

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

Deploy DiffusionGS online generation demos

98e7453 verified 1 day ago

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


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,
        )
        scaling = (scaling - 2.3).clamp(max=-1.20)
        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")
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

    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
                # 这是 patchify and linear layer
        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)
        ##gaussian pos embedding
        self.gaussians_pos_embedding = nn.Parameter(
            torch.randn(
                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)

        # 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):
        if self.cfg.ray_pe_type == "relative_plk":
            o_dot_d = torch.sum(-ray_o * ray_d, dim=2, keepdim=True)
            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)
        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)
        all_gaussians = torch.cat((gaussians, img_aligned_gaussians), dim=1)
        xyz, features, scaling, rotation, opacity = self.upsampler.to_gs(all_gaussians)
        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, 要执行的
            depth_preact_bias = 0. 
            img_aligned_xyz = torch.sigmoid(
                img_aligned_xyz.mean(dim=2, keepdim=True) + depth_preact_bias
            )
            # stx()
            #breakpoint()
            if self.cfg.ray_pe_type == 'relative_plk':
                img_aligned_xyz = (2.0 * img_aligned_xyz - 1.0) * 1.8 + o_dot_d
                # print(f"Using augmented plucker coordinates to compute xyz")
            img_aligned_xyz = ray_o + img_aligned_xyz * ray_d
            # breakpoint()
            # 将坐标限制在 -1 到 1 之间
            if self.cfg.clip_xyz and training:
                img_aligned_xyz = img_aligned_xyz.clamp(-1.0, 1.0)

            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):
        # breakpoint()
        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,       # 
                )
        #self.gs_renderer()
        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
	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


	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,
	)
	scaling = (scaling - 2.3).clamp(max=-1.20)
	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")
	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

	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
	# 这是 patchify and linear layer
	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)
	##gaussian pos embedding
	self.gaussians_pos_embedding = nn.Parameter(
	torch.randn(
	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)

	# 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):
	if self.cfg.ray_pe_type == "relative_plk":
	o_dot_d = torch.sum(-ray_o * ray_d, dim=2, keepdim=True)
	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)
	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)
	all_gaussians = torch.cat((gaussians, img_aligned_gaussians), dim=1)
	xyz, features, scaling, rotation, opacity = self.upsampler.to_gs(all_gaussians)
	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, 要执行的
	depth_preact_bias = 0.
	img_aligned_xyz = torch.sigmoid(
	img_aligned_xyz.mean(dim=2, keepdim=True) + depth_preact_bias
	)
	# stx()
	#breakpoint()
	if self.cfg.ray_pe_type == 'relative_plk':
	img_aligned_xyz = (2.0 * img_aligned_xyz - 1.0) * 1.8 + o_dot_d
	# print(f"Using augmented plucker coordinates to compute xyz")
	img_aligned_xyz = ray_o + img_aligned_xyz * ray_d
	# breakpoint()
	# 将坐标限制在 -1 到 1 之间
	if self.cfg.clip_xyz and training:
	img_aligned_xyz = img_aligned_xyz.clamp(-1.0, 1.0)

	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):
	# breakpoint()
	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, #
	)
	#self.gs_renderer()
	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