Spaces:

CaiYuanhao
/

DiffusionGS-Object

Running on Zero

DiffusionGS-Object / diffusionGS /systems /diffusion_gs_system.py

Deploy DiffusionGS online generation demos

3718945 verified 2 days ago

10.7 kB

from dataclasses import dataclass, field

import numpy as np
import json
import copy
import torch
import torch.nn.functional as F
from skimage import measure
from einops import repeat
from tqdm import tqdm
from PIL import Image
import os
import cv2

import diffusionGS
from easydict import EasyDict as edict
from einops import repeat, rearrange
from diffusionGS.utils.misc import get_rank
from diffusionGS.systems.base import BaseSystem
from diffusionGS.models.diffusion import create_diffusion
from diffusionGS.utils.typing import *
from diffusionGS.utils.losses import LossComputer
from diffusionGS.systems.utils import *
#from diffusionGS.models.denoiser.denoiser_utils import TransformInput
import random
# DEBUG = True
@diffusionGS.register("diffusion-gs-system")
class PointDiffusionSystem(BaseSystem):
    @dataclass
    class Config(BaseSystem.Config):
        num_inference_steps: int = 50
        # shape vae model
        shape_model_type: str = None
        shape_model: dict = field(default_factory=dict)
        # noise scheduler
        noise_scheduler_type: str = None
        noise_scheduler: dict = field(default_factory=dict)

    cfg: Config

    def configure(self):
        super().configure()

        self.shape_model = diffusionGS.find(self.cfg.shape_model_type)(self.cfg.shape_model)
        self.noise_scheduler = diffusionGS.find(self.cfg.noise_scheduler_type)(**self.cfg.noise_scheduler)
        self.diffusion_training = create_diffusion(str(1000), predict_xstart=True)
        self.diffusion_inference = create_diffusion(str(self.cfg.num_inference_steps), predict_xstart=True)
        self.loss_computer = LossComputer()
    
    def get_example_data(self):
        example_data_path = "examp_data/example/debug_objaverse_dataset/data_examples/batch_objaverse_example.json"
        with open(example_data_path, 'r') as file:
            batch = json.load(file)
        batch['c2w'] = torch.tensor(np.stack(batch['c2w']))[:4].float().unsqueeze(0)
        batch['fxfycxcy'] = torch.tensor(np.stack(batch['fxfycxcy']))[:4].float().unsqueeze(0)
        batch['image'] = torch.tensor(np.stack(batch['image']))[:4].float().unsqueeze(0)
        batch['index'] = torch.tensor(np.stack(batch['index']))[:4].float().unsqueeze(0)
        batch['depth'] = torch.tensor(np.stack(batch['depth']))[:4].float().unsqueeze(0)
        batch['normal'] = torch.tensor(np.stack(batch['normal']))[:4].float().unsqueeze(0)
        batch = {k: v.cuda() for k, v in batch.items() if k!="camera_json_path" and k!="view_choices"}
        return batch #torch.tensor(images,dtype=torch.float32).unsqueeze(0).to('cuda'), torch.tensor(c2ws,dtype=torch.float32).unsqueeze(0).to('cuda'), torch.tensor(fxfycxcys,dtype=torch.float32).unsqueeze(0).to('cuda')


    def select_train_views(self,images,c2ws,fxfycxcys):
        sel_index = np.random.choice(images.shape[1], 4,replace=False)
        sel_c2ws = c2ws[:,sel_index]
        sel_images = images[:,sel_index]
        sel_fxfycxcys = fxfycxcys[:,sel_index]
        return sel_images, sel_c2ws, sel_fxfycxcys

    def forward(self, batch: Dict[str, Any], skip_noise=False) -> Dict[str, Any]:
        #get training input
        sel_images = batch['rgbs_input'] #(batch['rgbs_input']*2.) - 1. ##归一化
        ray_o, ray_d = TransformInput(sel_images, batch['c2ws_input'], batch['fxfycxcys_input'])
        bs, v, c, h, w = sel_images.shape
        # 3. sample noise that we"ll add to the latents
        noise = torch.randn_like(sel_images)#.to(sel_images)
        # 4. Sample a random timestep for each motion
        timesteps = torch.randint(
            0,
            self.cfg.noise_scheduler.num_train_timesteps,
            (bs,),
            device=sel_images.device,
        )
        timesteps = timesteps.long()
        # breakpoint()
        sel_images[:,1:] = self.diffusion_training.q_sample(sel_images[:,1:], timesteps, noise=noise[:,1:])
        gt_img_aligned_xyz = ray_o + ray_d * batch['depths_input']
        #breakpoint()
        guassians_parameters, img_aligned_xyz = self.shape_model.image_to_gaussians(sel_images, ray_o, ray_d, timesteps)
        ### 5. render image from gen gaussians
        rendered_images = self.shape_model.render_gaussians(guassians_parameters, batch['c2ws'], batch['fxfycxcys'], sel_images.shape[3],sel_images.shape[4])
        #rendered_images,rendered_depths = rendered_images[:,:,:3],rendered_images[:,:,3:]
        l2_loss, lpips_loss, ssim_loss, pointsdist_loss, l2_loss_xyz = self.loss_computer(
                                                                        rendered_images,
                                                                        batch['rgbs'],
                                                                        batch['masks'],
                                                                        batch['masks_input'],
                                                                        ray_o,
                                                                        img_aligned_xyz=img_aligned_xyz,
                                                                        gt_img_aligned_xyz=gt_img_aligned_xyz,
                                                                        )
        # breakpoint()
        return {
            "loss_diffusion": l2_loss.mean(),
            "loss_lpips": lpips_loss.mean(),
            "loss_ssim": ssim_loss.mean(),
            "loss_xyz": l2_loss_xyz.mean(),
            "loss_pointsdist": pointsdist_loss.mean(),
            "latents": sel_images,
            "x_t": sel_images,
            "noise": noise,
            "noise_pred": rendered_images,
            "timesteps": timesteps,
            }


    def training_step(self, batch, batch_idx):
        out = self(batch)
        loss = 0.
        for name, value in out.items():
            if name.startswith("loss_"):
                self.log(f"train/{name}", value)
                loss += value * self.C(self.cfg.loss[name.replace("loss_", "lambda_")])

        for name, value in self.cfg.loss.items():
            if name.startswith("lambda_"):
                self.log(f"train_params/{name}", self.C(value))
        return {"loss": loss}

    @torch.no_grad()
    def validation_step(self, batch, batch_idx):
        self.eval()
        # TODO: write a save and denoise process
        # out = self(batch)      
            # if get_rank() == 0:
            ######################## for debug ########################
        # input_batch_test = self.get_example_data()
        # ray_o, ray_d = TransformInput(input_batch_test['image'], input_batch_test['c2w'], input_batch_test['fxfycxcy'])
        # input_batch = edict(
        #     image = input_batch_test['image'][:,:1],
        #     c2w = input_batch_test['c2w'],
        #     fxfycxcy = input_batch_test['fxfycxcy'],
        #     ray_o=ray_o,
        #     ray_d=ray_d,
        #     )
        # input_images = input_batch_test['image'] #batch['rgbs_input']
        # image_condition = input_batch_test['image'][0,:1]
            ######################## for debug ########################
        ######################## for real run ########################     
        ray_o, ray_d = TransformInput(batch['rgbs_input'], batch['c2ws_input'], batch['fxfycxcys_input'])
        input_batch = edict(
            image = batch['rgbs_input'][:,:1],
            c2w = batch['c2ws_input'],
            fxfycxcy = batch['fxfycxcys_input'],
            ray_o=ray_o,
            ray_d=ray_d,
            )
        input_images = batch['rgbs_input']  #[:,:1]
        b, v, c, h, w = input_images.shape      # b = n
        sample_noise = torch.randn(b, v-1, c, h, w, device=self.device)  # 在 timestep 为 T 的时候采的 noise
        input_batch["image_noisy"] = sample_noise
        ####################### for real run ########################
        image_condition = input_images[:, 0, ...].unsqueeze(1)
        traj_timesteps = list(self.diffusion_inference.timestep_map)[::-1]         # 提前设计好的 timestep 路线
        traj_len = len(traj_timesteps)

        traj_samples = []
        traj_pred_xstart = []
        final_out = None

        for out in self.diffusion_inference.p_sample_loop_progressive(
            self.shape_model,
            sample_noise.shape,
            input_batch,
            clip_denoised=False,
            progress=True,
            device=self.device,
        ):
            samples = out["sample"]
            pred_xstart = out["pred_xstart"]
            # 把 image condition 和 samples concate 起来
            samples_present = torch.cat((image_condition, samples), dim=1)
            pred_xstart_present = torch.cat((image_condition, pred_xstart), dim=1)
            traj_samples.append(samples_present)
            traj_pred_xstart.append(pred_xstart_present)
            final_out = out

        traj_samples = torch.stack(traj_samples, 0)  # [T, B, V, C, H, W]                   [250, 8, 4, 3, 256, 256]
        traj_pred_xstart = torch.stack(traj_pred_xstart, 0)  # [T, B, V, C, H, W]           [250, 8, 4, 3, 256, 256]

        traj_samples = rearrange(
            traj_samples, "t n v c h w -> n t h (v w) c"
        ).contiguous().cpu().numpy()

        traj_pred_xstart = rearrange(
            traj_pred_xstart, "t n v c h w -> n t h (v w) c"
        ).contiguous().cpu().numpy()

        # save multiple videos
        for i in range(traj_samples.shape[0]):
            vis = traj_samples[i]
            vis = display_timestep_on_video(vis, traj_timesteps)
            self.save_videos(f"it{self.true_global_step}/{batch['uid'][0]}_{batch['sel_idx'][0] if 'sel_idx' in batch.keys() else 0}_{i:04d}_traj_xt.mp4", vis, fps=24, quality=8)

        # 将输出的图片保存为视频
        image_grid = []
        for i in range(traj_pred_xstart.shape[0]):
            vis = traj_pred_xstart[i]
            vis = display_timestep_on_video(vis, traj_timesteps)
            self.save_videos(f"it{self.true_global_step}/{batch['uid'][0]}_{batch['sel_idx'][0] if 'sel_idx' in batch.keys() else 0}_{i:04d}_traj_xstart.mp4", vis, fps=24, quality=8)
        self.save_guassians_ply(f"it{self.true_global_step}/{batch['uid'][0]}_{batch['sel_idx'][0] if 'sel_idx' in batch.keys() else 0}.ply", final_out['denoiser_output_dict']['pred_gaussians'][0],render_video=True)
        self.save_torch_images(f"it{self.true_global_step}/{batch['uid'][0]}_{batch['sel_idx'][0] if 'sel_idx' in batch.keys() else 0}.png", torch.cat([image_condition[0],final_out['denoiser_output_dict']['render_images'][0]], dim=0))
        return {"val/loss": 123}

    def on_validation_epoch_end(self):
        pass

	from dataclasses import dataclass, field

	import numpy as np
	import json
	import copy
	import torch
	import torch.nn.functional as F
	from skimage import measure
	from einops import repeat
	from tqdm import tqdm
	from PIL import Image
	import os
	import cv2

	import diffusionGS
	from easydict import EasyDict as edict
	from einops import repeat, rearrange
	from diffusionGS.utils.misc import get_rank
	from diffusionGS.systems.base import BaseSystem
	from diffusionGS.models.diffusion import create_diffusion
	from diffusionGS.utils.typing import *
	from diffusionGS.utils.losses import LossComputer
	from diffusionGS.systems.utils import *
	#from diffusionGS.models.denoiser.denoiser_utils import TransformInput
	import random
	# DEBUG = True
	@diffusionGS.register("diffusion-gs-system")
	class PointDiffusionSystem(BaseSystem):
	@dataclass
	class Config(BaseSystem.Config):
	num_inference_steps: int = 50
	# shape vae model
	shape_model_type: str = None
	shape_model: dict = field(default_factory=dict)
	# noise scheduler
	noise_scheduler_type: str = None
	noise_scheduler: dict = field(default_factory=dict)

	cfg: Config

	def configure(self):
	super().configure()

	self.shape_model = diffusionGS.find(self.cfg.shape_model_type)(self.cfg.shape_model)
	self.noise_scheduler = diffusionGS.find(self.cfg.noise_scheduler_type)(**self.cfg.noise_scheduler)
	self.diffusion_training = create_diffusion(str(1000), predict_xstart=True)
	self.diffusion_inference = create_diffusion(str(self.cfg.num_inference_steps), predict_xstart=True)
	self.loss_computer = LossComputer()

	def get_example_data(self):
	example_data_path = "examp_data/example/debug_objaverse_dataset/data_examples/batch_objaverse_example.json"
	with open(example_data_path, 'r') as file:
	batch = json.load(file)
	batch['c2w'] = torch.tensor(np.stack(batch['c2w']))[:4].float().unsqueeze(0)
	batch['fxfycxcy'] = torch.tensor(np.stack(batch['fxfycxcy']))[:4].float().unsqueeze(0)
	batch['image'] = torch.tensor(np.stack(batch['image']))[:4].float().unsqueeze(0)
	batch['index'] = torch.tensor(np.stack(batch['index']))[:4].float().unsqueeze(0)
	batch['depth'] = torch.tensor(np.stack(batch['depth']))[:4].float().unsqueeze(0)
	batch['normal'] = torch.tensor(np.stack(batch['normal']))[:4].float().unsqueeze(0)
	batch = {k: v.cuda() for k, v in batch.items() if k!="camera_json_path" and k!="view_choices"}
	return batch #torch.tensor(images,dtype=torch.float32).unsqueeze(0).to('cuda'), torch.tensor(c2ws,dtype=torch.float32).unsqueeze(0).to('cuda'), torch.tensor(fxfycxcys,dtype=torch.float32).unsqueeze(0).to('cuda')


	def select_train_views(self,images,c2ws,fxfycxcys):
	sel_index = np.random.choice(images.shape[1], 4,replace=False)
	sel_c2ws = c2ws[:,sel_index]
	sel_images = images[:,sel_index]
	sel_fxfycxcys = fxfycxcys[:,sel_index]
	return sel_images, sel_c2ws, sel_fxfycxcys

	def forward(self, batch: Dict[str, Any], skip_noise=False) -> Dict[str, Any]:
	#get training input
	sel_images = batch['rgbs_input'] #(batch['rgbs_input']*2.) - 1. ##归一化
	ray_o, ray_d = TransformInput(sel_images, batch['c2ws_input'], batch['fxfycxcys_input'])
	bs, v, c, h, w = sel_images.shape
	# 3. sample noise that we"ll add to the latents
	noise = torch.randn_like(sel_images)#.to(sel_images)
	# 4. Sample a random timestep for each motion
	timesteps = torch.randint(
	0,
	self.cfg.noise_scheduler.num_train_timesteps,
	(bs,),
	device=sel_images.device,
	)
	timesteps = timesteps.long()
	# breakpoint()
	sel_images[:,1:] = self.diffusion_training.q_sample(sel_images[:,1:], timesteps, noise=noise[:,1:])
	gt_img_aligned_xyz = ray_o + ray_d * batch['depths_input']
	#breakpoint()
	guassians_parameters, img_aligned_xyz = self.shape_model.image_to_gaussians(sel_images, ray_o, ray_d, timesteps)
	### 5. render image from gen gaussians
	rendered_images = self.shape_model.render_gaussians(guassians_parameters, batch['c2ws'], batch['fxfycxcys'], sel_images.shape[3],sel_images.shape[4])
	#rendered_images,rendered_depths = rendered_images[:,:,:3],rendered_images[:,:,3:]
	l2_loss, lpips_loss, ssim_loss, pointsdist_loss, l2_loss_xyz = self.loss_computer(
	rendered_images,
	batch['rgbs'],
	batch['masks'],
	batch['masks_input'],
	ray_o,
	img_aligned_xyz=img_aligned_xyz,
	gt_img_aligned_xyz=gt_img_aligned_xyz,
	)
	# breakpoint()
	return {
	"loss_diffusion": l2_loss.mean(),
	"loss_lpips": lpips_loss.mean(),
	"loss_ssim": ssim_loss.mean(),
	"loss_xyz": l2_loss_xyz.mean(),
	"loss_pointsdist": pointsdist_loss.mean(),
	"latents": sel_images,
	"x_t": sel_images,
	"noise": noise,
	"noise_pred": rendered_images,
	"timesteps": timesteps,
	}


	def training_step(self, batch, batch_idx):
	out = self(batch)
	loss = 0.
	for name, value in out.items():
	if name.startswith("loss_"):
	self.log(f"train/{name}", value)
	loss += value * self.C(self.cfg.loss[name.replace("loss_", "lambda_")])

	for name, value in self.cfg.loss.items():
	if name.startswith("lambda_"):
	self.log(f"train_params/{name}", self.C(value))
	return {"loss": loss}

	@torch.no_grad()
	def validation_step(self, batch, batch_idx):
	self.eval()
	# TODO: write a save and denoise process
	# out = self(batch)
	# if get_rank() == 0:
	######################## for debug ########################
	# input_batch_test = self.get_example_data()
	# ray_o, ray_d = TransformInput(input_batch_test['image'], input_batch_test['c2w'], input_batch_test['fxfycxcy'])
	# input_batch = edict(
	# image = input_batch_test['image'][:,:1],
	# c2w = input_batch_test['c2w'],
	# fxfycxcy = input_batch_test['fxfycxcy'],
	# ray_o=ray_o,
	# ray_d=ray_d,
	# )
	# input_images = input_batch_test['image'] #batch['rgbs_input']
	# image_condition = input_batch_test['image'][0,:1]
	######################## for debug ########################
	######################## for real run ########################
	ray_o, ray_d = TransformInput(batch['rgbs_input'], batch['c2ws_input'], batch['fxfycxcys_input'])
	input_batch = edict(
	image = batch['rgbs_input'][:,:1],
	c2w = batch['c2ws_input'],
	fxfycxcy = batch['fxfycxcys_input'],
	ray_o=ray_o,
	ray_d=ray_d,
	)
	input_images = batch['rgbs_input'] #[:,:1]
	b, v, c, h, w = input_images.shape # b = n
	sample_noise = torch.randn(b, v-1, c, h, w, device=self.device) # 在 timestep 为 T 的时候采的 noise
	input_batch["image_noisy"] = sample_noise
	####################### for real run ########################
	image_condition = input_images[:, 0, ...].unsqueeze(1)
	traj_timesteps = list(self.diffusion_inference.timestep_map)[::-1] # 提前设计好的 timestep 路线
	traj_len = len(traj_timesteps)

	traj_samples = []
	traj_pred_xstart = []
	final_out = None

	for out in self.diffusion_inference.p_sample_loop_progressive(
	self.shape_model,
	sample_noise.shape,
	input_batch,
	clip_denoised=False,
	progress=True,
	device=self.device,
	):
	samples = out["sample"]
	pred_xstart = out["pred_xstart"]
	# 把 image condition 和 samples concate 起来
	samples_present = torch.cat((image_condition, samples), dim=1)
	pred_xstart_present = torch.cat((image_condition, pred_xstart), dim=1)
	traj_samples.append(samples_present)
	traj_pred_xstart.append(pred_xstart_present)
	final_out = out

	traj_samples = torch.stack(traj_samples, 0) # [T, B, V, C, H, W] [250, 8, 4, 3, 256, 256]
	traj_pred_xstart = torch.stack(traj_pred_xstart, 0) # [T, B, V, C, H, W] [250, 8, 4, 3, 256, 256]

	traj_samples = rearrange(
	traj_samples, "t n v c h w -> n t h (v w) c"
	).contiguous().cpu().numpy()

	traj_pred_xstart = rearrange(
	traj_pred_xstart, "t n v c h w -> n t h (v w) c"
	).contiguous().cpu().numpy()

	# save multiple videos
	for i in range(traj_samples.shape[0]):
	vis = traj_samples[i]
	vis = display_timestep_on_video(vis, traj_timesteps)
	self.save_videos(f"it{self.true_global_step}/{batch['uid'][0]}_{batch['sel_idx'][0] if 'sel_idx' in batch.keys() else 0}_{i:04d}_traj_xt.mp4", vis, fps=24, quality=8)

	# 将输出的图片保存为视频
	image_grid = []
	for i in range(traj_pred_xstart.shape[0]):
	vis = traj_pred_xstart[i]
	vis = display_timestep_on_video(vis, traj_timesteps)
	self.save_videos(f"it{self.true_global_step}/{batch['uid'][0]}_{batch['sel_idx'][0] if 'sel_idx' in batch.keys() else 0}_{i:04d}_traj_xstart.mp4", vis, fps=24, quality=8)
	self.save_guassians_ply(f"it{self.true_global_step}/{batch['uid'][0]}_{batch['sel_idx'][0] if 'sel_idx' in batch.keys() else 0}.ply", final_out['denoiser_output_dict']['pred_gaussians'][0],render_video=True)
	self.save_torch_images(f"it{self.true_global_step}/{batch['uid'][0]}_{batch['sel_idx'][0] if 'sel_idx' in batch.keys() else 0}.png", torch.cat([image_condition[0],final_out['denoiser_output_dict']['render_images'][0]], dim=0))
	return {"val/loss": 123}

	def on_validation_epoch_end(self):
	pass