Blame for diffusionGS/models/denoiser/denoiser.py

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