""" https://github.com/CompVis/stable-diffusion/blob/21f890f9da3cfbeaba8e2ac3c425ee9e998d5229/ldm/models/diffusion/ddpm.py#L30 """ import random import pytorch3d import copy import point_cloud_utils as pcu import cv2 import matplotlib.pyplot as plt import torch import gc import functools import json import os from pathlib import Path from pdb import set_trace as st from typing import Any from click import prompt import einops import blobfile as bf import imageio import numpy as np import torch as th import torch.distributed as dist import torchvision from nsr.camera_utils import generate_input_camera, uni_mesh_path from PIL import Image from torch.nn.parallel.distributed import DistributedDataParallel as DDP from torch.optim import AdamW from torch.utils.tensorboard.writer import SummaryWriter from tqdm import tqdm import trimesh from nsr.camera_utils import generate_input_camera from guided_diffusion import dist_util, logger from guided_diffusion.fp16_util import MixedPrecisionTrainer from guided_diffusion.nn import update_ema from guided_diffusion.resample import LossAwareSampler, UniformSampler # from .train_util import TrainLoop3DRec from guided_diffusion.train_util import (TrainLoop, calc_average_loss, find_ema_checkpoint, find_resume_checkpoint, get_blob_logdir, log_loss_dict, log_rec3d_loss_dict, parse_resume_step_from_filename) from guided_diffusion.gaussian_diffusion import ModelMeanType from ldm.modules.encoders.modules import FrozenClipImageEmbedder, TextEmbedder, FrozenCLIPTextEmbedder, FrozenOpenCLIPImagePredictionEmbedder, FrozenOpenCLIPImageEmbedder import dnnlib from dnnlib.util import requires_grad from dnnlib.util import calculate_adaptive_weight from ..train_util_diffusion import TrainLoop3DDiffusion from ..cvD.nvsD_canoD import TrainLoop3DcvD_nvsD_canoD from guided_diffusion.continuous_diffusion_utils import get_mixed_prediction, different_p_q_objectives, kl_per_group_vada, kl_balancer # from .train_util_diffusion_lsgm_noD_joint import TrainLoop3DDiffusionLSGMJointnoD # joint diffusion and rec class # from .controlLDM import TrainLoop3DDiffusionLSGM_Control # joint diffusion and rec class from .train_util_diffusion_lsgm_noD_joint import TrainLoop3DDiffusionLSGMJointnoD # joint diffusion and rec class # ! add new schedulers from https://github.com/Stability-AI/generative-models from .crossattn_cldm import TrainLoop3DDiffusionLSGM_crossattn # Function to generate a rotation matrix for an arbitrary theta along the x-axis def rotation_matrix_x(theta_degrees): theta = np.radians(theta_degrees) # Convert degrees to radians cos_theta = np.cos(theta) sin_theta = np.sin(theta) rotation_matrix = np.array([[1, 0, 0], [0, cos_theta, -sin_theta], [0, sin_theta, cos_theta]]) return rotation_matrix def rotation_matrix_z(theta): """ Returns a 3x3 rotation matrix that rotates a point around the z-axis by theta radians. Parameters: theta (float): The angle of rotation in radians. Returns: numpy.ndarray: A 3x3 rotation matrix. """ return np.array([ [np.cos(theta), -np.sin(theta), 0], [np.sin(theta), np.cos(theta), 0], [0, 0, 1] ]) def rotation_matrix_y(theta): """ Returns a 3x3 rotation matrix that rotates a point around the y-axis by theta radians. Parameters: theta (float): The angle of rotation in radians. Returns: numpy.ndarray: A 3x3 rotation matrix. """ return np.array([ [np.cos(theta), 0, np.sin(theta)], [0, 1, 0 ], [-np.sin(theta), 0, np.cos(theta)] ]) # import SD stuffs from typing import Any, Dict, List, Optional, Tuple, Union from contextlib import contextmanager from omegaconf import ListConfig, OmegaConf from sgm.modules import UNCONDITIONAL_CONFIG from sgm.util import (default, disabled_train, get_obj_from_str, instantiate_from_config, log_txt_as_img) from transport import create_transport, Sampler import math # for gs rendering from utils.gs_utils.graphics_utils import getWorld2View2, getProjectionMatrix, getView2World from utils.general_utils import matrix_to_quaternion from utils.mesh_util import post_process_mesh, to_cam_open3d_compat from datasets.g_buffer_objaverse import focal2fov, fov2focal import open3d as o3d # from sgm.sampling_utils.demo.streamlit_helpers import init_sampling def sample_uniform_cameras_on_sphere(num_samples=1): # Step 1: Sample azimuth angles uniformly from [0, 2*pi) theta = np.random.uniform(0, 2 * np.pi, num_samples) # Step 2: Sample cos(phi) uniformly from [-1, 1] cos_phi = np.random.uniform(-1, 1, num_samples) # Step 3: Calculate the elevation angle (phi) from cos(phi) phi = np.arccos(cos_phi) # phi will be in [0, pi] # Step 4: Convert spherical coordinates to Cartesian coordinates (x, y, z) # x = np.sin(phi) * np.cos(theta) # y = np.sin(phi) * np.sin(theta) # z = np.cos(phi) # Combine the x, y, z coordinates into a single array # cameras = np.vstack((x, y, z)).T # Shape: (num_samples, 3) # return cameras return theta, phi class FlowMatchingEngine(TrainLoop3DDiffusionLSGM_crossattn): def __init__( self, *, rec_model, denoise_model, diffusion, sde_diffusion, control_model, control_key, only_mid_control, loss_class, data, eval_data, batch_size, microbatch, lr, ema_rate, log_interval, eval_interval, save_interval, resume_checkpoint, resume_cldm_checkpoint=None, use_fp16=False, fp16_scale_growth=0.001, schedule_sampler=None, weight_decay=0, lr_anneal_steps=0, iterations=10001, ignore_resume_opt=False, freeze_ae=False, denoised_ae=True, triplane_scaling_divider=10, use_amp=False, diffusion_input_size=224, normalize_clip_encoding=False, scale_clip_encoding=1, cfg_dropout_prob=0, cond_key='img_sr', use_eos_feature=False, compile=False, snr_type='lognorm', # denoiser_config, # conditioner_config: Union[None, Dict, ListConfig, # OmegaConf] = None, # sampler_config: Union[None, Dict, ListConfig, OmegaConf] = None, # loss_fn_config: Union[None, Dict, ListConfig, OmegaConf] = None, **kwargs): super().__init__(rec_model=rec_model, denoise_model=denoise_model, diffusion=diffusion, sde_diffusion=sde_diffusion, control_model=control_model, control_key=control_key, only_mid_control=only_mid_control, loss_class=loss_class, data=data, eval_data=eval_data, batch_size=batch_size, microbatch=microbatch, lr=lr, ema_rate=ema_rate, log_interval=log_interval, eval_interval=eval_interval, save_interval=save_interval, resume_checkpoint=resume_checkpoint, resume_cldm_checkpoint=resume_cldm_checkpoint, use_fp16=use_fp16, fp16_scale_growth=fp16_scale_growth, schedule_sampler=schedule_sampler, weight_decay=weight_decay, lr_anneal_steps=lr_anneal_steps, iterations=iterations, ignore_resume_opt=ignore_resume_opt, freeze_ae=freeze_ae, denoised_ae=denoised_ae, triplane_scaling_divider=triplane_scaling_divider, use_amp=use_amp, diffusion_input_size=diffusion_input_size, normalize_clip_encoding=normalize_clip_encoding, scale_clip_encoding=scale_clip_encoding, cfg_dropout_prob=cfg_dropout_prob, cond_key=cond_key, use_eos_feature=use_eos_feature, compile=compile, **kwargs) # ! sgm diffusion pipeline # ! reuse the conditioner self.snr_type = snr_type self.latent_key = 'latent' if self.cond_key == 'caption': # ! text pretrain if snr_type == 'stage2-t23d': ldm_configs = OmegaConf.load( 'sgm/configs/stage2-t23d.yaml')['ldm_configs'] elif snr_type == 'stage1-t23d': ldm_configs = OmegaConf.load( 'sgm/configs/stage1-t23d.yaml')['ldm_configs'] self.latent_key = 'normalized-fps-xyz' # learn xyz diff else: # just simple t23d, no xyz condition ldm_configs = OmegaConf.load( 'sgm/configs/t23d-clipl-compat-fm.yaml')['ldm_configs'] else: # # assert 'lognorm' in snr_type if snr_type == 'lognorm': # by default ldm_configs = OmegaConf.load( 'sgm/configs/img23d-clipl-compat-fm-lognorm.yaml')['ldm_configs'] # st() # if snr_type == 'lognorm-highres': # by default elif snr_type == 'img-uniform-gvp': # by default ldm_configs = OmegaConf.load( 'sgm/configs/img23d-clipl-compat-fm-lognorm-336-uniform.yaml')['ldm_configs'] # self.latent_key = 'fps-xyz' # xyz diffusion self.latent_key = 'normalized-fps-xyz' # to std elif snr_type == 'img-uniform-gvp-dino': # by default ldm_configs = OmegaConf.load( 'sgm/configs/img23d-clipl-compat-fm-lognorm-480-uniform-clay-dinoonly.yaml')['ldm_configs'] self.latent_key = 'normalized-fps-xyz' # to std # elif snr_type == 'img-uniform-gvp-dino-xl': # by default # ldm_configs = OmegaConf.load( # 'sgm/configs/img23d-clipl-compat-fm-lognorm-480-uniform-clay-dinoonly.yaml')['ldm_configs'] # self.latent_key = 'normalized-fps-xyz' # to std elif snr_type == 'img-uniform-gvp-dino-stage2': # by default ldm_configs = OmegaConf.load( 'sgm/configs/stage2-i23d.yaml')['ldm_configs'] # self.latent_key = 'normalized-fps-xyz' # to std elif snr_type == 'img-uniform-gvp-clay': # contains both text and image condition ldm_configs = OmegaConf.load( 'sgm/configs/img23d-clipl-compat-fm-lognorm-480-uniform-clay.yaml')['ldm_configs'] # self.latent_key = 'fps-xyz' # xyz diffusion self.latent_key = 'normalized-fps-xyz' # to std elif snr_type == 'pcd-cond-tex': ldm_configs = OmegaConf.load( 'sgm/configs/img23d-clipl-compat-fm-lognorm-336-uniform-pcdcond.yaml')['ldm_configs'] # 'sgm/configs/img23d-clipl-compat-fm-lognorm-336.yaml')['ldm_configs'] # ! stage-2 text-xyz conditioned elif snr_type == 'stage2-t23d': ldm_configs = OmegaConf.load( 'sgm/configs/stage2-t23d.yaml')['ldm_configs'] elif snr_type == 'lognorm-mv': ldm_configs = OmegaConf.load( 'sgm/configs/mv23d-clipl-compat-fm-lognorm.yaml')['ldm_configs'] # ! mv version elif snr_type == 'lognorm-mv-plucker': ldm_configs = OmegaConf.load( 'sgm/configs/mv23d-plucker-clipl-compat-fm-lognorm-noclip.yaml')['ldm_configs'] # 'sgm/configs/mv23d-plucker-clipl-compat-fm-lognorm.yaml')['ldm_configs'] elif snr_type == 'stage1-mv-t23dpt': self.latent_key = 'normalized-fps-xyz' # learn xyz diff ldm_configs = OmegaConf.load( 'sgm/configs/stage1-mv23d-t23dpt.yaml')['ldm_configs'] elif snr_type == 'stage1-mv-i23dpt': self.latent_key = 'normalized-fps-xyz' # learn xyz diff ldm_configs = OmegaConf.load( 'sgm/configs/stage1-mv23d-i23dpt.yaml')['ldm_configs'] elif snr_type == 'stage1-mv-i23dpt-noi23d': self.latent_key = 'normalized-fps-xyz' # learn xyz diff ldm_configs = OmegaConf.load( 'sgm/configs/stage1-mv23d-i23dpt-noi23d.yaml')['ldm_configs'] elif snr_type == 'stage2-mv-i23dpt': # self.latent_key = 'normalized-fps-xyz' # learn xyz diff ldm_configs = OmegaConf.load( 'sgm/configs/stage2-mv23d-i23dpt.yaml')['ldm_configs'] else: ldm_configs = OmegaConf.load( 'sgm/configs/img23d-clipl-compat-fm.yaml')['ldm_configs'] self.loss_fn = ( instantiate_from_config(ldm_configs.loss_fn_config) # if loss_fn_config is not None # else None ) # self.denoiser = instantiate_from_config( # ldm_configs.denoiser_config).to(dist_util.dev()) self.transport_sampler = Sampler(self.loss_fn.transport, guider_config=ldm_configs.guider_config) self.conditioner = instantiate_from_config( default(ldm_configs.conditioner_config, UNCONDITIONAL_CONFIG)).to(dist_util.dev()) # ! setup optimizer (with cond embedder params here) self._set_grad_flag() self._setup_opt2() self._load_model2() def _set_grad_flag(self): requires_grad(self.ddpm_model, True) # do not change this flag during training. def _setup_opt(self): pass # see below def _setup_opt2(self): # ! add trainable conditioner parameters # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/models/diffusion.py#L219 # params = list(self.ddpm_model.parameters()) # https://discuss.pytorch.org/t/how-the-pytorch-freeze-network-in-some-layers-only-the-rest-of-the-training/7088/7 self.opt = AdamW([{ 'name': 'ddpm', # 'params': self.ddpm_model.parameters(), 'params': filter(lambda p: p.requires_grad, self.ddpm_model.parameters()), # if you want to freeze some layers }, ], lr=self.lr, weight_decay=self.weight_decay) embedder_params = [] for embedder in self.conditioner.embedders: if embedder.is_trainable: embedder_params = embedder_params + list(embedder.parameters()) if len(embedder_params) != 0: self.opt.add_param_group( { 'name': 'embedder', 'params': embedder_params, 'lr': self.lr*0.5, # smaller lr to finetune dino/clip } ) print(self.opt) def save(self, mp_trainer=None, model_name='ddpm'): # save embedder params also super().save(mp_trainer, model_name) # save embedder ckpt if dist_util.get_rank() == 0: for embedder in self.conditioner.embedders: if embedder.is_trainable: # embedder_params = embedder_params + list(embedder.parameters()) model_name = embedder.__class__.__name__ filename = f"embedder_{model_name}{(self.step+self.resume_step):07d}.pt" with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f: th.save(embedder.state_dict(), f) dist_util.synchronize() def _load_model2(self): # ! load embedder for embedder in self.conditioner.embedders: if embedder.is_trainable: # embedder_params = embedder_params + list(embedder.parameters()) model_name = embedder.__class__.__name__ filename = f"embedder_{model_name}{(self.step+self.resume_step):07d}.pt" # embedder_FrozenDinov2ImageEmbedderMV2115000.pt # with bf.BlobFile(bf.join(get_blob_logdir(), filename), # "wb") as f: # th.save(embedder.state_dict(), f) split = self.resume_checkpoint.split("model") resume_checkpoint = str( Path(split[0]) / filename) if os.path.exists(resume_checkpoint): if dist.get_rank() == 0: logger.log( f"loading cond embedder from checkpoint: {resume_checkpoint}...") # if model is None: # model = self.model embedder.load_state_dict( dist_util.load_state_dict( resume_checkpoint, map_location=dist_util.dev(), )) else: logger.log(f'{resume_checkpoint} not found.') if dist_util.get_world_size() > 1: dist_util.sync_params(embedder.parameters()) def instantiate_cond_stage(self, normalize_clip_encoding, scale_clip_encoding, cfg_dropout_prob, use_eos_feature): # https://github.com/CompVis/stable-diffusion/blob/21f890f9da3cfbeaba8e2ac3c425ee9e998d5229/ldm/models/diffusion/ddpm.py#L509C1-L509C46 # self.cond_stage_model.train = disabled_train # type: ignore # if self.cond_key == 'caption': # self.cond_txt_model = TextEmbedder(dropout_prob=cfg_dropout_prob, # use_eos_feature=use_eos_feature) # elif self.cond_key == 'img': # self.cond_img_model = FrozenOpenCLIPImagePredictionEmbedder( # 1, 1, # FrozenOpenCLIPImageEmbedder(freeze=True, # device=dist_util.dev(), # init_device=dist_util.dev())) # else: # zero-shot Text to 3D using normalized clip latent # self.cond_stage_model = FrozenClipImageEmbedder( # 'ViT-L/14', # dropout_prob=cfg_dropout_prob, # normalize_encoding=normalize_clip_encoding, # scale_clip_encoding=scale_clip_encoding) # self.cond_stage_model.freeze() # self.cond_txt_model = FrozenCLIPTextEmbedder( # dropout_prob=cfg_dropout_prob, # scale_clip_encoding=scale_clip_encoding) # self.cond_txt_model.freeze() pass # initialized in the self.__init__() using SD api # ! already merged def prepare_ddpm(self, eps, mode='p'): raise NotImplementedError('already implemented in self.denoiser') # merged from noD.py # use sota denoiser, loss_fn etc. def ldm_train_step(self, batch, behaviour='cano', *args, **kwargs): # ! enable the gradient of both models # requires_grad(self.ddpm_model, True) self._set_grad_flag() # more flexible self.mp_trainer.zero_grad() # !!!! if 'img' in batch: batch_size = batch['img'].shape[0] else: batch_size = len(batch['caption']) for i in range(0, batch_size, self.microbatch): micro = { k: v[i:i + self.microbatch].to(dist_util.dev()) if isinstance( v, th.Tensor) else v[i:i+self.microbatch] for k, v in batch.items() } # move condition to self.dtype # =================================== ae part =================================== # with th.cuda.amp.autocast(dtype=th.bfloat16, with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): loss = th.tensor(0.).to(dist_util.dev()) assert 'latent' in micro # st() # torchvision.utils.save_image(micro['img'], 'tmp/img.png', normalize=True, value_range=(-1,1)) # vae_out = {self.latent_name: micro['latent']} # else: # vae_out = self.ddp_rec_model( # img=micro['img_to_encoder'], # c=micro['c'], # behaviour='encoder_vae', # ) # pred: (B, 3, 64, 64) # eps = vae_out[self.latent_name] / self.triplane_scaling_divider # ! if training xyz only # eps = vae_out[self.latent_name][..., -3:] / self.triplane_scaling_divider # ! if training texture only eps = micro[self.latent_key] / self.triplane_scaling_divider if self.cond_key == 'img-c': micro['img-c'] = { # 'img': micro['img'].to(self.dtype), 'img': micro['mv_img'].to(self.dtype), # for compat issue 'c': micro['c'].to(self.dtype), } # log_rec3d_loss_dict({ # f"mv-alpha/{i}": self.ddpm_model.blocks[i].mv_alpha[0] for i in range(len(self.ddpm_model.blocks)) # }) loss, loss_other_info = self.loss_fn(self.ddp_ddpm_model, # self.denoiser, self.conditioner, eps.to(self.dtype), micro) # type: ignore loss = loss.mean() # log_rec3d_loss_dict({}) log_rec3d_loss_dict({ # 'eps_mean': # eps.mean(), # 'eps_std': # eps.std([1, 2, 3]).mean(0), # 'pred_x0_std': # loss_other_info['model_output'].std([1, 2, 3]).mean(0), "p_loss": loss, }) self.mp_trainer.backward(loss) # joint gradient descent # update ddpm accordingly self.mp_trainer.optimize(self.opt) # ! directly eval_cldm() for sampling. # if dist_util.get_rank() == 0 and self.step % 500 == 0: # self.log_control_images(vae_out, micro, loss_other_info) @th.inference_mode() def log_control_images(self, vae_out, micro, ddpm_ret): if 'posterior' in vae_out: vae_out.pop('posterior') # for calculating kl loss vae_out_for_pred = {self.latent_name: vae_out[self.latent_name][0:1].to(self.dtype)} with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): pred = self.ddp_rec_model(latent=vae_out_for_pred, c=micro['c'][0:1], behaviour=self.render_latent_behaviour) assert isinstance(pred, dict) pred_img = pred['image_raw'] if 'img' in micro: gt_img = micro['img'] else: gt_img = th.zeros_like(pred['image_raw']) if 'depth' in micro: gt_depth = micro['depth'] if gt_depth.ndim == 3: gt_depth = gt_depth.unsqueeze(1) gt_depth = (gt_depth - gt_depth.min()) / (gt_depth.max() - gt_depth.min()) else: gt_depth = th.zeros_like(gt_img[:, 0:1, ...]) if 'image_depth' in pred: pred_depth = pred['image_depth'] pred_depth = (pred_depth - pred_depth.min()) / (pred_depth.max() - pred_depth.min()) else: pred_depth = th.zeros_like(gt_depth) gt_img = self.pool_128(gt_img) gt_depth = self.pool_128(gt_depth) # cond = self.get_c_input(micro) # hint = th.cat(cond['c_concat'], 1) gt_vis = th.cat( [ gt_img, gt_img, gt_img, # self.pool_128(hint), # gt_img, gt_depth.repeat_interleave(3, dim=1) ], dim=-1)[0:1] # TODO, fail to load depth. range [0, 1] # eps_t_p_3D = eps_t_p.reshape(batch_size, eps_t_p.shape[1]//3, 3, -1) # B C 3 L # self.sampler noised_latent, sigmas, x_start = [ ddpm_ret[k] for k in ['noised_input', 'sigmas', 'model_output'] ] noised_latent = { 'latent_normalized_2Ddiffusion': noised_latent[0:1].to(self.dtype) * self.triplane_scaling_divider, } denoised_latent = { 'latent_normalized_2Ddiffusion': x_start[0:1].to(self.dtype) * self.triplane_scaling_divider, } with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): noised_ae_pred = self.ddp_rec_model( img=None, c=micro['c'][0:1], latent=noised_latent, behaviour=self.render_latent_behaviour) # pred_x0 = self.sde_diffusion._predict_x0_from_eps( # eps_t_p, pred_eps_p, logsnr_p) # for VAE loss, denosied latent # pred_xstart_3D denoised_ae_pred = self.ddp_rec_model( img=None, c=micro['c'][0:1], latent=denoised_latent, # latent=pred_x0[0:1] * self. # triplane_scaling_divider, # TODO, how to define the scale automatically? behaviour=self.render_latent_behaviour) pred_vis = th.cat( [ self.pool_128(img) for img in ( pred_img[0:1], noised_ae_pred['image_raw'][0:1], denoised_ae_pred['image_raw'][0:1], # controlnet result pred_depth[0:1].repeat_interleave(3, dim=1)) ], dim=-1) # B, 3, H, W if 'img' in micro: vis = th.cat([gt_vis, pred_vis], dim=-2)[0].permute(1, 2, 0).cpu() # ! pred in range[-1, 1] else: vis = pred_vis[0].permute(1, 2, 0).cpu() # vis_grid = torchvision.utils.make_grid(vis) # HWC vis = vis.numpy() * 127.5 + 127.5 vis = vis.clip(0, 255).astype(np.uint8) img_save_path = f'{logger.get_dir()}/{self.step+self.resume_step}denoised_{sigmas[0].item():3}.jpg' Image.fromarray(vis).save(img_save_path) # if self.cond_key == 'caption': # with open(f'{logger.get_dir()}/{self.step+self.resume_step}caption_{t_p[0].item():3}.txt', 'w') as f: # f.write(micro['caption'][0]) print('log denoised vis to: ', img_save_path) th.cuda.empty_cache() @th.no_grad() def sample( self, cond: Dict, uc: Union[Dict, None] = None, batch_size: int = 16, shape: Union[None, Tuple, List] = None, use_cfg=True, # cfg_scale=4, # default value in SiT # cfg_scale=1.5, # default value in SiT cfg_scale=4.0, # default value in SiT seed=42, **kwargs, ): # self.sampler sample_fn = self.transport_sampler.sample_ode(num_steps=250, cfg=True) # default ode sampling setting. logger.log(f'cfg_scale: {cfg_scale}, seed: {seed}') th.manual_seed(seed) # to reproduce result zs = th.randn(batch_size, *shape).to(dist_util.dev()).to(self.dtype) # st() assert use_cfg # sample_model_kwargs = {'uc': uc, 'cond': cond} model_fn = self.ddpm_model.forward_with_cfg # default # ! prepare_inputs in VanillaCFG, for compat issue c_out = {} for k in cond: # if k in ["vector", "crossattn", "concat", 'fps-xyz']: c_out[k] = th.cat((cond[k], uc[k]), 0) # else: # assert cond[k] == uc[k] # c_out[k] = cond[k] sample_model_kwargs = {'context': c_out, 'cfg_scale': cfg_scale} zs = th.cat([zs, zs], 0) with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): samples = sample_fn(zs, model_fn, **sample_model_kwargs)[-1] samples, _ = samples.chunk(2, dim=0) # Remove null class samples # return samples return samples * self.triplane_scaling_divider @th.inference_mode() def eval_cldm( self, prompt="", # use_ddim=False, # unconditional_guidance_scale=1.0, unconditional_guidance_scale=4.0, seed=42, save_img=False, use_train_trajectory=False, camera=None, num_samples=1, num_instances=1, export_mesh=False, ): # ! slightly modified for new API. combined with # /cpfs01/shared/V2V/V2V_hdd/yslan/Repo/generative-models/sgm/models/diffusion.py:249 log_images() # TODO, support batch_size > 1 self.ddpm_model.eval() args = dnnlib.EasyDict( dict( batch_size=1, image_size=self.diffusion_input_size, denoise_in_channels=self.rec_model.decoder.triplane_decoder. out_chans, # type: ignore clip_denoised=False, class_cond=False)) model_kwargs = {} uc = None log = dict() ucg_keys = [self.cond_key] # i23d # if self.cond_key == 'caption': if self.cond_key in ['caption', 'img-xyz']: # batch_c = {self.cond_key: prompt} # batch_c = {self.cond_key: prompt} batch_c = next(self.data) # ! use training set to evaluate t23d for now. elif self.cond_key == 'img-caption': batch_c = {'caption': prompt, 'img': batch['img'].to(dist_util.dev()).to(self.dtype)} else: batch = next(self.data) # random cond here if self.cond_key == 'img-c': batch_c = { self.cond_key: { # 'img': batch['img'].to(self.dtype).to(dist_util.dev()), 'img': batch['mv_img'].to(self.dtype).to(dist_util.dev()), 'c': batch['c'].to(self.dtype).to(dist_util.dev()), }, 'img': batch['img'].to(self.dtype).to(dist_util.dev()) # required by clip } else: batch_c = {self.cond_key: batch[self.cond_key].to(dist_util.dev()).to(self.dtype)} with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): c, uc = self.conditioner.get_unconditional_conditioning( batch_c, force_uc_zero_embeddings=ucg_keys if len(self.conditioner.embedders) > 0 else [], ) sampling_kwargs = {'seed': seed, 'cfg_scale': unconditional_guidance_scale} N = 3 # hard coded, to update z_shape = ( N, self.ddpm_model.in_channels if not self.ddpm_model.roll_out else 3 * self.ddpm_model.in_channels, # type: ignore self.diffusion_input_size, self.diffusion_input_size) for k in c: if isinstance(c[k], th.Tensor): c[k], uc[k] = map(lambda y: y[k][:N].to(dist_util.dev()), (c, uc)) samples = self.sample(c, shape=z_shape[1:], uc=uc, batch_size=N, **sampling_kwargs) # st() # do rendering first # ! get c if 'img' in self.cond_key: img_save_path = f'{logger.get_dir()}/{self.step+self.resume_step}_imgcond.jpg' if 'c' in self.cond_key: torchvision.utils.save_image(batch_c['img'][0], img_save_path, value_range=(-1,1), normalize=True, padding=0) # torch.Size([24, 6, 3, 256, 256]) th.save(batch_c['img-c']['c'][0], f'{logger.get_dir()}/{self.step+self.resume_step}_c.pt') else: torchvision.utils.save_image(batch_c['img'][0:1], img_save_path, value_range=(-1,1), normalize=True, padding=0) assert camera is not None batch = {'c': camera.clone()} # else: # if use_train_trajectory: # batch = next(iter(self.data)) # else: # try: # batch = next(self.eval_data) # except Exception as e: # self.eval_data = iter(self.eval_data) # batch = next(self.eval_data) # if camera is not None: # batch['c'] = camera.clone() # rendering for i in range(samples.shape[0]): th.cuda.empty_cache() # ! render sampled latent name_prefix = f'{self.step + self.resume_step}_{i}' if self.cond_key == 'caption': name_prefix = f'{name_prefix}_{prompt}' with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): self.render_video_given_triplane( samples[i:i+1].to(self.dtype), # default version self.rec_model, # compatible with join_model name_prefix=name_prefix, save_img=save_img, render_reference=batch, export_mesh=False) self.ddpm_model.train() class FlowMatchingEngine_gs(FlowMatchingEngine): def __init__( self, *, rec_model, denoise_model, diffusion, sde_diffusion, control_model, control_key, only_mid_control, loss_class, data, eval_data, batch_size, microbatch, lr, ema_rate, log_interval, eval_interval, save_interval, resume_checkpoint, resume_cldm_checkpoint=None, use_fp16=False, fp16_scale_growth=0.001, schedule_sampler=None, weight_decay=0, lr_anneal_steps=0, iterations=10001, ignore_resume_opt=False, freeze_ae=False, denoised_ae=True, triplane_scaling_divider=10, use_amp=False, diffusion_input_size=224, normalize_clip_encoding=False, scale_clip_encoding=1, cfg_dropout_prob=0, cond_key='img_sr', use_eos_feature=False, compile=False, snr_type='lognorm', **kwargs): super().__init__(rec_model=rec_model, denoise_model=denoise_model, diffusion=diffusion, sde_diffusion=sde_diffusion, control_model=control_model, control_key=control_key, only_mid_control=only_mid_control, loss_class=loss_class, data=data, eval_data=eval_data, batch_size=batch_size, microbatch=microbatch, lr=lr, ema_rate=ema_rate, log_interval=log_interval, eval_interval=eval_interval, save_interval=save_interval, resume_checkpoint=resume_checkpoint, resume_cldm_checkpoint=resume_cldm_checkpoint, use_fp16=use_fp16, fp16_scale_growth=fp16_scale_growth, schedule_sampler=schedule_sampler, weight_decay=weight_decay, lr_anneal_steps=lr_anneal_steps, iterations=iterations, ignore_resume_opt=ignore_resume_opt, freeze_ae=freeze_ae, denoised_ae=denoised_ae, triplane_scaling_divider=triplane_scaling_divider, use_amp=use_amp, diffusion_input_size=diffusion_input_size, normalize_clip_encoding=normalize_clip_encoding, scale_clip_encoding=scale_clip_encoding, cfg_dropout_prob=cfg_dropout_prob, cond_key=cond_key, use_eos_feature=use_eos_feature, compile=compile, snr_type=snr_type, **kwargs) self.gs_bg_color=th.tensor([1,1,1], dtype=th.float32, device=dist_util.dev()) self.latent_name = 'latent_normalized' # normalized triplane latent # self.pcd_unnormalize_fn = lambda x: x.clip(-1,1) * 0.45 # [-1,1] -> [-0.45, 0.45] as in g-buffer dataset. # self.pcd_unnormalize_fn = lambda x: (x * 0.1862).clip(-0.45, 0.45) # [-1,1] -> [-0.45, 0.45] as in g-buffer dataset. # /cpfs01/user/lanyushi.p/logs/nips24/LSGM/t23d/FM/9cls/gs/i23d/dit-b/gpu4-batch32-lr1e-4-gs_surf_latent_224-drop0.33-same # self.pcd_unnormalize_fn = lambda x: (x * 0.158).clip(-0.45, 0.45) # [-1,1] -> [-0.45, 0.45] as in g-buffer dataset. # self.feat_scale_factor = th.Tensor([0.99227685, 1.014337 , 0.20842505, 0.98727155, 0.3305389 , # 0.38729668, 1.0155401 , 0.9728264 , 1.0009694 , 0.97328585, # 0.2881106 , 0.1652732 , 0.3482468 , 0.9971449 , 0.99895126, # 0.18491288]).float().reshape(1,1,-1) # stat for normalization # self.xyz_mean = torch.Tensor([-0.00053714, 0.08095618, -0.01914407] ).reshape(1, 3).float() # self.xyz_std = th.Tensor([0.14593576, 0.15753542, 0.18873914] ).reshape(1,3).float().to(dist_util.dev()) self.xyz_std = 0.164 # ! for debug self.kl_mean = th.Tensor([ 0.0184, 0.0024, 0.0926, 0.0517, 0.1781, 0.7137, -0.0355, 0.0267, 0.0183, 0.0164, -0.5090, 0.2406, 0.2733, -0.0256, -0.0285, 0.0761]).reshape(1,16).float().to(dist_util.dev()) self.kl_std = th.Tensor([1.0018, 1.0309, 1.3001, 1.0160, 0.8182, 0.8023, 1.0591, 0.9789, 0.9966, 0.9448, 0.8908, 1.4595, 0.7957, 0.9871, 1.0236, 1.2923]).reshape(1,16).float().to(dist_util.dev()) # ! for surfel-gs rendering self.zfar = 100.0 self.znear = 0.01 def unnormalize_pcd_act(self, x): return x * self.xyz_std def unnormalize_kl_feat(self, latent): # return latent / self.feat_scale_factor # return (latent-self.kl_mean) / self.kl_std return (latent * self.kl_std) + self.kl_mean # def unnormalize_kl_feat(self, latent): # return latent * self.feat_scale_factor @th.inference_mode() def eval_cldm( self, prompt="Yellow rubber duck", # use_ddim=False, # unconditional_guidance_scale=1.0, save_img=False, use_train_trajectory=False, camera=None, num_samples=1, num_instances=1, export_mesh=False, ): self.ddpm_model.eval() args = dnnlib.EasyDict( dict( batch_size=1, image_size=self.diffusion_input_size, denoise_in_channels=self.rec_model.decoder.triplane_decoder. out_chans, # type: ignore clip_denoised=False, class_cond=False)) model_kwargs = {} uc = None log = dict() ucg_keys = [self.cond_key] # i23d if self.cond_key == 'caption': if prompt == '': batch = next(self.data) # random cond here batch_c = {self.cond_key: prompt, 'fps-xyz': batch['fps-xyz'].to(self.dtype).to(dist_util.dev()), } else: # ! TODO, update the cascaded generation fps-xyz loading. Manual load for now. batch_c = { self.cond_key: prompt } if self.latent_key == 'latent': # stage 2 # hard-coded path for now # fps_xyz_output_prefix = '/nas/shared/public/yslan/logs/nips24/LSGM/t23d/FM/9cls/gs-disentangle/cascade_check/clay/stage1/eval/clip_text/50w-iter/' # stage1_num_steps = '500000' batch = next(self.data) # random cond here batch_c[self.cond_key] = batch[self.cond_key][0:1] # colorizing GT xyz gt_xyz = batch['fps-xyz'][0:1] gt_kl_latent = batch['latent'][0:1] # cascaded = False # st() # if self.step % 1e4 == 0: # cascaded = True # else: cascaded = False prompt = batch[self.cond_key][0:1] # replace with on-the-fly GT point clouds batch_c[self.cond_key] = prompt # ! for logging two-stage cascaded result. change the path to your stage-1 output pcd logdir. if cascaded: # ! use stage-1 as output # fps_xyz_output_prefix = '/nas/shared/public/yslan/logs/nips24/LSGM/t23d/FM/9cls/gs-disentangle/cascade_check/clay/stage1/eval/clip_text/60w-iter/' fps_xyz_output_prefix = '' stage1_num_steps = '600000' fps_xyz = torch.from_numpy(pcu.load_mesh_v(f'{fps_xyz_output_prefix}/{stage1_num_steps}_0_{prompt}.ply') ).clip(-0.45,0.45).unsqueeze(0) batch_c.update({ 'fps-xyz': fps_xyz.to(self.dtype).to(dist_util.dev()) }) else: # use gt as condition # batch_c = {k: v[0:1].to(self.dtype).to(dist_util.dev()) for k, v in batch_c.items() if k in [self.cond_key, 'fps-xyz']} for k in ['fps-xyz']: batch_c[k] = batch[k][0:1].to(self.dtype).to(dist_util.dev()) batch_c[self.cond_key] = prompt else: batch = next(self.data) # random cond here #! debugging, get GT xyz and KL latent for disentangled debugging if self.cond_key == 'img-c': prompt = batch['caption'][0:1] batch_c = { self.cond_key: { 'img': batch['mv_img'][0:1].to(self.dtype).to(dist_util.dev()), 'c': batch['c'][0:1].to(self.dtype).to(dist_util.dev()), }, 'img': batch['img'][0:1].to(self.dtype).to(dist_util.dev()), 'caption': prompt, 'fps-xyz': batch['fps-xyz'][0:1].to(self.dtype).to(dist_util.dev()) } elif self.cond_key == 'img-caption': batch_c = {'caption': prompt, 'img': batch['img'].to(dist_util.dev()).to(self.dtype)} elif self.cond_key == 'img-xyz': # load local xyz here # fps_xyz = torch.from_numpy(pcu.load_mesh_v('/cpfs01/user/lanyushi.p/Repo/diffusion-3d/tmp/sampled-0.ply') ).clip(-0.45,0.45).unsqueeze(0) # fps_xyz = torch.from_numpy(pcu.load_mesh_v('/cpfs01/user/lanyushi.p/Repo/diffusion-3d/tmp/sampled-2.ply') ).clip(-0.45,0.45).unsqueeze(0) # fps_xyz = torch.from_numpy(pcu.load_mesh_v('/cpfs01/user/lanyushi.p/Repo/diffusion-3d/tmp/sampled-1.ply') ).clip(-0.45,0.45).unsqueeze(0) # fps_xyz = torch.from_numpy(pcu.load_mesh_v('/cpfs01/user/lanyushi.p/Repo/diffusion-3d/tmp/sampled-3.ply') ).clip(-0.45,0.45).unsqueeze(0) # fps_xyz = torch.from_numpy(pcu.load_mesh_v('/nas/shared/V2V/yslan/logs/nips24/LSGM/t23d/FM/9cls/gs-disentangle/cascade_check/xyz_output_fullset_stillclip_but448_eval/1725000_0_0.ply') ).clip(-0.45,0.45).unsqueeze(0) # fps_xyz = torch.from_numpy(pcu.load_mesh_v('/nas/shared/V2V/yslan/logs/nips24/LSGM/t23d/FM/9cls/gs-disentangle/cascade_check/xyz_output_fullset_stillclip_but448_eval/1725000_0_0.ply') ).clip(-0.45,0.45).unsqueeze(0) # fps_xyz = torch.from_numpy(pcu.load_mesh_v('/nas/shared/public/yslan/logs/nips24/LSGM/t23d/FM/9cls/gs-disentangle/cascade_check/clay/stage1/eval/dino_img/debug/1875000_0.ply') ).clip(-0.45,0.45).unsqueeze(0) # fps_xyz = torch.from_numpy(pcu.load_mesh_v('/nas/shared/public/yslan/logs/nips24/LSGM/t23d/FM/9cls/gs-disentangle/cascade_check/clay/stage1/eval/dino_img/debug/1875000_0_0.ply') ).clip(-0.45,0.45).unsqueeze(0) # ! edit # st() # fps_xyz[..., 2:3] *= 4 # fps_xyz[..., 2:3] *= 3 # fps_xyz = torch.from_numpy(pcu.load_mesh_v('/nas/shared/V2V/yslan/logs/nips24/LSGM/t23d/FM/9cls/gs-disentangle/cascade_check/xyz_output_fullset_stillclip_but448_eval/1725000_0_1.ply') ).clip(-0.45,0.45).unsqueeze(0) batch_c = { # 'img': batch['img'][[1,0]].to(self.dtype).to(dist_util.dev()), 'img': batch['img'][0:1].to(self.dtype).to(dist_util.dev()), 'fps-xyz': batch['fps-xyz'][0:1].to(self.dtype).to(dist_util.dev()), # 'caption': batch['caption'] # 'fps-xyz': fps_xyz.repeat(batch['img'].shape[0],1,1).to(self.dtype).to(dist_util.dev()), } else: # gt_xyz = batch['fps-xyz'][0:1] # gt_kl_latent = batch['latent'][0:1] batch_c = {self.cond_key: batch[self.cond_key][0:1].to(dist_util.dev()).to(self.dtype), } # swap for more results, hard-coded here. # if 'img' in batch_c: # batch_c['img'] = batch_c['img'][[1,0]] # batch_c['fps-xyz'] = batch_c['fps-xyz'][[1,0]] with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): c, uc = self.conditioner.get_unconditional_conditioning( batch_c, force_uc_zero_embeddings=ucg_keys if len(self.conditioner.embedders) > 0 else [], ) sampling_kwargs = {} N = num_samples # hard coded, to update z_shape = (N, 768, self.ddpm_model.in_channels) for k in c: if isinstance(c[k], th.Tensor): # c[k], uc[k] = map(lambda y: y[k][:N].to(dist_util.dev()), # (c, uc)) assert c[k].shape[0] == 1 # ! support batch inference c[k], uc[k] = map(lambda y: y[k].repeat_interleave(N, 0).to(dist_util.dev()), (c, uc)) # support bs>1 sampling given a condition samples = self.sample(c, shape=z_shape[1:], uc=uc, batch_size=N, **sampling_kwargs) # ! get c if 'img' in self.cond_key: img_save_path = f'{logger.get_dir()}/{self.step+self.resume_step}_imgcond.jpg' if 'c' in self.cond_key: mv_img_save_path = f'{logger.get_dir()}/{self.step+self.resume_step}_mv-imgcond.jpg' torchvision.utils.save_image(batch_c['img-c']['img'][0], mv_img_save_path, value_range=(-1,1), normalize=True, padding=0) # torch.Size([24, 6, 3, 256, 256]) torchvision.utils.save_image(batch_c['img'][0], img_save_path, value_range=(-1,1), normalize=True, padding=0) # torch.Size([24, 6, 3, 256, 256]) else: torchvision.utils.save_image(batch_c['img'], img_save_path, value_range=(-1,1), normalize=True, padding=0) assert camera is not None batch = {'c': camera.clone()} # rendering for i in range(samples.shape[0]): th.cuda.empty_cache() # ! render sampled latent name_prefix = f'{self.step + self.resume_step}_{i}' # if self.cond_key in ['caption', 'img-c']: if self.cond_key in ['caption']: if isinstance(prompt, list): name_prefix = f'{name_prefix}_{"-".join(prompt[0].split())}' else: name_prefix = f'{name_prefix}_{"-".join(prompt.split())}' with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): # # ! todo, transform to gs camera if self.latent_key != 'latent': # normalized-xyz pcu.save_mesh_v( f'{logger.get_dir()}/{name_prefix}.ply', self.unnormalize_pcd_act(samples[i]).detach().cpu().float().numpy()) logger.log(f'point cloud saved to {logger.get_dir()}/{name_prefix}.ply') else: # ! editing debug self.render_gs_video_given_latent( # samples[i:i+1].to(self.dtype), # default version # th.cat([gt_kl_latent.to(samples), gt_xyz.to(samples)], dim=-1), # ! xyz-cond kl feature gen: th.cat([samples[i:i+1], batch_c['fps-xyz'][i:i+1]], dim=-1), # ! debugging xyz diffusion # ! xyz debugging # th.cat([gt_kl_latent.to(samples), samples[i:i+1]], dim=-1), # ! debugging xyz diffusion # th.cat([samples[i:i+1], gt_xyz.to(samples), ], dim=-1) # ! debugging kl feature diffusion self.rec_model, # compatible with join_model name_prefix=name_prefix, save_img=save_img, render_reference=batch, export_mesh=False) # st() pass # for noise_scale in np.linspace(0,0.1, 10): # per_scale_name_prefix = f'{name_prefix}_{noise_scale}' # self.render_gs_video_given_latent( th.cat([gt_kl_latent.to(samples), (gt_xyz+noise_scale*th.randn_like(gt_xyz)).to(samples)], dim=-1), self.rec_model, name_prefix=per_scale_name_prefix, save_img=save_img, render_reference=batch, export_mesh=False) # self.render_gs_video_given_latent( th.cat([gt_kl_latent.to(samples), (batch_c['fps-xyz'][0:1]).to(samples)], dim=-1), self.rec_model, name_prefix=per_scale_name_prefix, save_img=save_img, render_reference=batch, export_mesh=False) # pcu.save_mesh_v( f'{logger.get_dir()}/sampled-4.ply', self.unnormalize_pcd_act(samples[0]).detach().cpu().float().numpy()) # st() # pcu.save_mesh_v( f'tmp/sampled-3.ply', self.unnormalize_pcd_act(samples[0]).detach().cpu().float().numpy()) gc.collect() self.ddpm_model.train() @torch.no_grad() def export_mesh_from_2dgs(self, all_rgbs, all_depths, all_alphas, cam_pathes, idx, i): # https://github.com/autonomousvision/LaRa/blob/main/evaluation.py n_thread = 1 # avoid TSDF cpu hanging bug. os.environ["MKL_NUM_THREADS"] = f"{n_thread}" os.environ["NUMEXPR_NUM_THREADS"] = f"{n_thread}" os.environ["OMP_NUM_THREADS"] = f"4" os.environ["VECLIB_MAXIMUM_THREADS"] = f"{n_thread}" os.environ["OPENBLAS_NUM_THREADS"] = f"{n_thread}" # copied from: https://github.com/hbb1/2d-gaussian-splatting/blob/19eb5f1e091a582e911b4282fe2832bac4c89f0f/render.py#L23 logger.log("exporting mesh ...") # os.makedirs(train_dir, exist_ok=True) train_dir = logger.get_dir() # for g-objv # aabb = [-0.5,-0.5,-0.5,0.5,0.5,0.5] # aabb = None aabb = [-0.45,-0.45,-0.45,0.45,0.45,0.45] self.aabb = np.array(aabb).reshape(2,3)*1.1 # center = self.aabb.mean(0) # radius = np.linalg.norm(self.aabb[1] - self.aabb[0]) * 0.5 # voxel_size = radius / 256 # sdf_trunc = voxel_size * 2 # print("using aabb") # set the active_sh to 0 to export only diffuse texture # gaussExtractor.gaussians.active_sh_degree = 0 # gaussExtractor.reconstruction(scene.getTrainCameras()) # extract the mesh and save # if args.unbounded: # name = 'fuse_unbounded.ply' # mesh = gaussExtractor.extract_mesh_unbounded(resolution=args.mesh_res) # else: # name = f'{idx}-{i}-fuse.ply' # name = f'mesh.obj' name = f'{idx}/{i}-mesh_raw.obj' # st() # depth_trunc = (radius * 2.0) if depth_trunc < 0 else depth_trunc # voxel_size = (depth_trunc / mesh_res) if voxel_size < 0 else voxel_size # sdf_trunc = 5.0 * voxel_size if sdf_trunc < 0 else sdf_trunc # mesh = self.extract_mesh_bounded(all_rgbs, all_depths, all_alphas, cam_pathes, voxel_size=voxel_size, sdf_trunc=sdf_trunc, depth_trunc=depth_trunc, mask_backgrond=False) mesh = self.extract_mesh_bounded(all_rgbs, all_depths, all_alphas, cam_pathes) o3d.io.write_triangle_mesh(os.path.join(train_dir, name), mesh) logger.log("mesh saved at {}".format(os.path.join(train_dir, name))) # post-process the mesh and save, saving the largest N clusters # mesh_post = post_process_mesh(mesh, cluster_to_keep=num_cluster) mesh_post = post_process_mesh(mesh) mesh_vertices = np.asarray(mesh_post.vertices) # Convert vertices to a numpy array rotated_vertices = mesh_vertices @ rotation_matrix_x(-90).T # rotated_vertices = rotated_vertices @ rotation_matrix_z(np.pi).T rotated_vertices = rotated_vertices @ rotation_matrix_y(np.pi).T mesh_post.vertices = o3d.utility.Vector3dVector(rotated_vertices) # Update vertices post_mesh_path = os.path.join(train_dir, name.replace('_raw.obj', '.obj')) o3d.io.write_triangle_mesh(post_mesh_path, mesh_post) logger.log("mesh post processed saved at {}".format(post_mesh_path)) return post_mesh_path @torch.no_grad() def extract_mesh_bounded(self, rgbmaps, depthmaps, alpha_maps, cam_pathes, voxel_size=0.004, sdf_trunc=0.02, depth_trunc=3, alpha_thres=0.08, mask_backgrond=False): """ Perform TSDF fusion given a fixed depth range, used in the paper. voxel_size: the voxel size of the volume sdf_trunc: truncation value depth_trunc: maximum depth range, should depended on the scene's scales mask_backgrond: whether to mask backgroud, only works when the dataset have masks return o3d.mesh """ # if self.aabb is not None: # as in lara. # center = self.aabb.mean(0) # radius = np.linalg.norm(self.aabb[1] - self.aabb[0]) * 0.5 # voxel_size = radius / 256 # sdf_trunc = voxel_size * 2 # print("using aabb") assert self.aabb is not None # as in lara. center = self.aabb.mean(0) radius = np.linalg.norm(self.aabb[1] - self.aabb[0]) * 0.5 voxel_size = radius / 160 # less holes sdf_trunc = voxel_size * 12 print("using aabb") volume = o3d.pipelines.integration.ScalableTSDFVolume( voxel_length= voxel_size, sdf_trunc=sdf_trunc, color_type=o3d.pipelines.integration.TSDFVolumeColorType.RGB8 ) print("Running tsdf volume integration ...") print(f'voxel_size: {voxel_size}') print(f'sdf_trunc: {sdf_trunc}') print(f'depth_truc: {depth_trunc}') # render_reference = th.load('eval_pose.pt', map_location='cpu').numpy() # ! use uni_mesh_path, from Lara, Chen et al, ECCV 24' # ''' # for i, cam_o3d in tqdm(enumerate(to_cam_open3d(self.viewpoint_stack)), desc="TSDF integration progress"): for i, cam in tqdm(enumerate(cam_pathes), desc="TSDF integration progress"): # rgb = self.rgbmaps[i] # depth = self.depthmaps[i] cam = self.c_to_3dgs_format(cam) cam_o3d = to_cam_open3d_compat(cam) rgb = rgbmaps[i][0] depth = depthmaps[i][0] alpha = alpha_maps[i][0] # if we have mask provided, use it # if mask_backgrond and (self.viewpoint_stack[i].gt_alpha_mask is not None): # depth[(self.viewpoint_stack[i].gt_alpha_mask < 0.5)] = 0 depth[(alpha < alpha_thres)] = 0 if self.aabb is not None: campos = cam['cam_pos'].cpu().numpy() depth_trunc = np.linalg.norm(campos - center, axis=-1) + radius # make open3d rgbd rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth( o3d.geometry.Image(np.asarray(np.clip(rgb.permute(1,2,0).cpu().numpy(), 0.0, 1.0) * 255, order="C", dtype=np.uint8)), o3d.geometry.Image(np.asarray(depth.permute(1,2,0).cpu().numpy(), order="C")), depth_trunc = depth_trunc, convert_rgb_to_intensity=False, depth_scale = 1.0 ) volume.integrate(rgbd, intrinsic=cam_o3d.intrinsic, extrinsic=cam_o3d.extrinsic) mesh = volume.extract_triangle_mesh() return mesh @th.inference_mode() def render_gs_video_given_latent(self, planes, rec_model, name_prefix='0', save_img=False, render_reference=None, export_mesh=False, output_dir=None, for_fid=False,): if output_dir is None: output_dir = logger.get_dir() batch_size, L, C = planes.shape # ddpm_latent = { self.latent_name: planes[..., :-3] * self.feat_scale_factor.to(planes), # kl-reg latent # 'query_pcd_xyz': self.pcd_unnormalize_fn(planes[..., -3:]) } # ddpm_latent = { self.latent_name: self.unnormalize_kl_feat(planes[..., :-3]), # kl-reg latent # ddpm_latent = { self.latent_name: planes[..., :-3], # kl-reg latent # 'query_pcd_xyz': self.unnormalize_pcd_act(planes[..., -3:]) } ddpm_latent = { self.latent_name: planes[..., :-3], # kl-reg latent 'query_pcd_xyz': planes[..., -3:]} ddpm_latent.update(rec_model(latent=ddpm_latent, behaviour='decode_gs_after_vae_no_render')) # ! editing debug, raw scaling # for beacon # edited_fps_xyz[..., 2] *= 1.5 # edited_fps_xyz[..., :2] *= 0.75 # z_mask = edited_fps_xyz[..., 2] > 0 # edited_fps_xyz[..., 2] *= 1.25 # only apply to upper points # z_dim_coord = edited_fps_xyz[..., 2] # edited_fps_xyz[..., 2] = th.where(z_dim_coord>0, z_dim_coord*1.25, z_dim_coord) # edited_fps_xyz[..., :2] *= 0.6 fine_scale = 'gaussians_upsampled_3' # ddpm_latent[fine_scale][..., :2] *= 1.5 # ddpm_latent[fine_scale][..., 2:3] *= 0.75 # ddpm_latent[fine_scale][..., :2] *= 3 # ddpm_latent[fine_scale][..., 2:3] *= 0.75 # z_dim_coord = ddpm_latent[fine_scale][..., 2] # ddpm_latent[fine_scale][..., 2] = th.where(z_dim_coord>0.24, z_dim_coord+0.1, z_dim_coord) # pcu.save_mesh_v(f'{output_dir}/gaussian.ply', ddpm_latent['gaussians_upsampled'][0, ..., :3].cpu().numpy()) # fps-downsampling? # pred_gaussians_xyz = ddpm_latent['gaussians_upsampled_3'][..., :3] fine_gs = ddpm_latent[fine_scale] fine_gs_numpy = fine_gs.cpu().numpy() vtx = np.transpose(rotation_matrix_x(-90) @ np.transpose(fine_gs_numpy[0, :, :3])) # for gradio visualization # vtx = vtx @ rotation_matrix_z(np.pi).T vtx = vtx @ rotation_matrix_y(np.pi).T cloud = trimesh.PointCloud(vtx, colors=fine_gs_numpy[0, :, 10:13]) # Save the point cloud to an OBJ file rgb_xyz_path_forgradio = f'{output_dir}/{name_prefix}-gaussian-pcd.glb' # gradio only accepts glb for visualization _ = cloud.export(rgb_xyz_path_forgradio) rgb_xyz_path_formeshlab = f'{output_dir}/{name_prefix}-gaussian-pcd.ply' # for meshlab visualization _ = cloud.export(rgb_xyz_path_formeshlab) # K=4096 # query_pcd_xyz, fps_idx = pytorch3d.ops.sample_farthest_points( # pred_gaussians_xyz, K=K, # # random_start_point=False) # B self.latent_num # random_start_point=True) # B self.latent_num # pcu.save_mesh_v(f'{output_dir}/{name_prefix}-gaussian-{K}.ply', query_pcd_xyz[0].cpu().numpy()) np.save(f'{output_dir}/{name_prefix}-gaussian.npy', fine_gs_numpy) video_path = f'{output_dir}/{name_prefix}-gs.mp4' # return None, None try: # video_out = imageio.get_writer( # f'{output_dir}/gs_{name_prefix}.mp4', # mode='I', # fps=15, # codec='libx264') video_out = imageio.get_writer( video_path, mode='I', fps=15, codec='libx264') except Exception as e: logger.log(e) # return # some caption are too tired and cannot be parsed as file name # !for FID ''' # if for uniform FID rendering. Will not adopt this later. azimuths = [] elevations = [] frame_number = 10 for i in range(frame_number): # 1030 * 5 * 10, for FID 50K azi, elevation = sample_uniform_cameras_on_sphere() # azi, elevation = azi[0] / np.pi * 180, elevation[0] / np.pi * 180 # azi, elevation = azi[0] / np.pi * 180, 0 azi, elevation = azi[0] / np.pi * 180, (elevation[0]-np.pi*0.5) / np.pi * 180 # [-0.5 pi, 0.5 pi] azimuths.append(azi) elevations.append(elevation) azimuths = np.array(azimuths) elevations = np.array(elevations) # azimuths = np.array(list(range(0,360,30))).astype(float) # frame_number = azimuths.shape[0] # elevations = np.array([10]*azimuths.shape[0]).astype(float) zero123pp_pose, _ = generate_input_camera(1.8, [[elevations[i], azimuths[i]] for i in range(frame_number)], fov=30) K = th.Tensor([1.3889, 0.0000, 0.5000, 0.0000, 1.3889, 0.5000, 0.0000, 0.0000, 0.0039]).to(zero123pp_pose) # keeps the same render_reference = th.cat([zero123pp_pose.reshape(frame_number,-1), K.unsqueeze(0).repeat(frame_number,1)], dim=-1).cpu().numpy() ''' assert render_reference is not None # render_reference = th.load('eval_pose.pt', map_location='cpu').numpy()[:24] # rand_start_idx = random.randint(0,2) # render_reference = render_reference[rand_start_idx::3] # randomly render 8 views, maintain fixed azimuths # assert len(render_reference)==8 # assert render_reference is None # render_reference = self.eval_data # compat # else: # use train_traj # for key in ['ins', 'bbox', 'caption']: # if key in render_reference: # render_reference.pop(key) # render_reference = [ { k:v[idx:idx+1] for k, v in render_reference.items() } for idx in range(40) ] all_rgbs, all_depths, all_alphas = [], [], [] # for i, batch in enumerate(tqdm(self.eval_data)): for i, micro_c in enumerate(tqdm(render_reference)): # micro = { # k: v.to(dist_util.dev()) if isinstance(v, th.Tensor) else v # for k, v in batch.items() # } # c = self.eval_data.post_process.c_to_3dgs_format(micro_c) c = self.c_to_3dgs_format(micro_c) for k in c.keys(): # to cuda if isinstance(c[k], th.Tensor) and k != 'tanfov': c[k] = c[k].unsqueeze(0).unsqueeze(0).to(dist_util.dev()) # actually, could render 40 views together. c['tanfov'] = th.tensor(c['tanfov']).to(dist_util.dev()) pred = rec_model( img=None, c=c, # TODO, to dict latent=ddpm_latent, # render gs behaviour='triplane_dec', bg_color=self.gs_bg_color, render_all_scale=True, # for better visualization ) # ! if visualizing a single scale fine_scale_key = list(pred.keys())[-1] # pred = pred[fine_scale_key] # for k in pred.keys(): # pred[k] = einops.rearrange(pred[k], 'B V ... -> (B V) ...') # merge # pred_vis = self._make_vis_img(pred) # vis = pred_vis.permute(0, 2, 3, 1).cpu().numpy() # vis = vis * 127.5 + 127.5 # vis = vis.clip(0, 255).astype(np.uint8) # # if not save_img: # for j in range(vis.shape[0] # ): # ! currently only export one plane at a time # video_out.append_data(vis[j]) # save multi-scale rendering all_rgbs.append(einops.rearrange(pred[fine_scale_key]['image'], 'B V ... -> (B V) ...')) all_depths.append(einops.rearrange(pred[fine_scale_key]['depth'], 'B V ... -> (B V) ...')) all_alphas.append(einops.rearrange(pred[fine_scale_key]['alpha'], 'B V ... -> (B V) ...')) all_pred_vis = {} # for key in pred.keys(): for key in ['gaussians_base', fine_scale_key]: # only show two LoDs pred_scale = pred[key] # only show finest result here for k in pred_scale.keys(): pred_scale[k] = einops.rearrange(pred_scale[k], 'B V ... -> (B V) ...') # merge pred_vis = self._make_vis_img(pred_scale, ignore_depth=True) vis = pred_vis.permute(0, 2, 3, 1).cpu().numpy() vis = vis * 127.5 + 127.5 vis = vis.clip(0, 255).astype(np.uint8) all_pred_vis[key] = vis # all_pred_vis_concat = np.concatenate([cv2.resize(all_pred_vis[k][0], (384*3, 384)) for k in ['gaussians_base', 'gaussians_upsampled', 'gaussians_upsampled_2']], axis=0) # all_pred_vis_concat = np.concatenate([cv2.resize(all_pred_vis[k][0], (256*3, 256)) for k in ['gaussians_base', 'gaussians_upsampled', 'gaussians_upsampled_2']], axis=0) # all_pred_vis_concat = np.concatenate([cv2.resize(all_pred_vis[k][0], (384*len(all_pred_vis.keys()), 384)) for k in all_pred_vis.keys()], axis=0) # all_pred_vis_concat = np.concatenate([cv2.resize(all_pred_vis[k][0], (384*3, 384)) for k in all_pred_vis.keys()], axis=0) # all_pred_vis_concat = np.concatenate([cv2.resize(all_pred_vis[k][0], (512*3, 512)) for k in all_pred_vis.keys()], axis=0) all_pred_vis_concat = np.concatenate([cv2.resize(all_pred_vis[k][0], (512*2, 512)) for k in all_pred_vis.keys()], axis=0) video_out.append_data(all_pred_vis_concat) if save_img: # for fid for idx in range(len(all_rgbs)): sampled_img = Image.fromarray( (all_rgbs[idx][0].permute(1, 2, 0).cpu().numpy() * 255).clip(0, 255).astype(np.uint8)) sampled_img.save(os.path.join(output_dir,f'{name_prefix}-{idx}.jpg')) # if not save_img: video_out.close() print('logged video to: ', f'{output_dir}/{name_prefix}.mp4') del video_out, pred, pred_vis, vis # return all_rgbs, all_depths, all_alphas return all_rgbs, all_depths, all_alphas, video_path, rgb_xyz_path_forgradio @th.no_grad() def _make_vis_img(self, pred, ignore_depth=False): gen_img = pred['image_raw'] rend_normal = pred['rend_normal'] # if True: if not ignore_depth: pred_depth = pred['image_depth'] pred_depth = (pred_depth - pred_depth.min()) / (pred_depth.max() - pred_depth.min()) pred_depth = pred_depth.cpu()[0].permute(1, 2, 0).numpy() pred_depth = (plt.cm.viridis(pred_depth[..., 0])[..., :3]) * 2 - 1 pred_depth = th.from_numpy(pred_depth).to( pred['image_raw'].device).permute(2, 0, 1).unsqueeze(0) pred_vis = th.cat( [ gen_img, rend_normal, pred_depth, ], dim=-1) # B, 3, H, W else: pred_vis = th.cat( [ gen_img, rend_normal, ], dim=-1) # B, 3, H, W return pred_vis def _set_grad_flag(self): requires_grad(self.ddpm_model, True) # @th.inference_mode() def sample_and_save(self, batch_c, ucg_keys, num_samples, camera, save_img, idx=0, save_dir='', export_mesh=False, stage1_idx=0, cfg_scale=4.0, seed=42): with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): c, uc = self.conditioner.get_unconditional_conditioning( batch_c, force_uc_zero_embeddings=ucg_keys if len(self.conditioner.embedders) > 0 else [], ) sampling_kwargs = { 'cfg_scale': cfg_scale, # default value in SiT 'seed': seed, } N = num_samples # hard coded, to update z_shape = (N, 768, self.ddpm_model.in_channels) for k in c: if isinstance(c[k], th.Tensor): # c[k], uc[k] = map(lambda y: y[k][:N].to(dist_util.dev()), # (c, uc)) assert c[k].shape[0] == 1 # ! support batch inference c[k], uc[k] = map(lambda y: y[k].repeat_interleave(N, 0).to(dist_util.dev()), (c, uc)) # support bs>1 sampling given a condition samples = self.sample(c, shape=z_shape[1:], uc=uc, batch_size=N, **sampling_kwargs) # ! get c if save_dir == '': save_dir = logger.get_dir() if 'img' in self.cond_key: # img_save_path = f'{save_dir}/{idx}_imgcond.jpg' img_save_path = f'{save_dir}/{idx}/imgcond.jpg' os.makedirs(f'{save_dir}/{idx}', exist_ok=True) if 'c' in self.cond_key: torchvision.utils.save_image(batch_c['img-c']['img'][0], img_save_path, value_range=(-1,1), normalize=True, padding=0) # torch.Size([24, 6, 3, 256, 256]) else: torchvision.utils.save_image(batch_c['img'], img_save_path, value_range=(-1,1), normalize=True, padding=0) assert camera is not None # batch = {'c': camera.clone()} # rendering for i in range(samples.shape[0]): th.cuda.empty_cache() if self.cond_key in ['caption']: name_prefix = f'{batch_c["caption"]}_sample-{stage1_idx}-{i}' else: # ! render sampled latent # name_prefix = f'{idx}_sample-{i}' name_prefix = f'{idx}/sample-{stage1_idx}-{i}' # if self.cond_key in ['caption', 'img-c']: cam_pathes = uni_mesh_path(10) with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): # # ! todo, transform to gs camera if self.latent_key != 'latent': # normalized-xyz pcd_export_dir = f'{save_dir}/{name_prefix}.glb' # pcu fails on py=3.9 vtx = self.unnormalize_pcd_act(samples[i]).detach().cpu().float().numpy() cloud = trimesh.PointCloud(vtx @ rotation_matrix_x(-90).T, colors=np.ones_like(vtx)*0.1) # since white background _ = cloud.export(pcd_export_dir) # for gradio display logger.log(f'stage-1 glb point cloud saved to {pcd_export_dir}') pcd_export_dir_forstage1 = f'{save_dir}/{name_prefix}.ply' pcu.save_mesh_v(pcd_export_dir_forstage1, self.unnormalize_pcd_act(samples[i]).detach().cpu().float().numpy()) logger.log(f'point cloud saved to {pcd_export_dir}') return pcd_export_dir else: # ! editing debug all_rgbs, all_depths, all_alphas, video_path, rgb_xyz_path = self.render_gs_video_given_latent( th.cat([samples[i:i+1], batch_c['fps-xyz'][0:1]], dim=-1), # ! debugging xyz diffusion self.rec_model, # compatible with join_model name_prefix=name_prefix, save_img=save_img, render_reference=cam_pathes, export_mesh=False,) # for_fid=False) if export_mesh: post_mesh_path=self.export_mesh_from_2dgs(all_rgbs, all_depths, all_alphas, cam_pathes, idx, i) else: post_mesh_path = '' return video_path, rgb_xyz_path, post_mesh_path # mesh = self.extract_mesh_bounded(all_rgbs, all_depths, all_alphas, cam_pathes, voxel_size=voxel_size, sdf_trunc=sdf_trunc, depth_trunc=depth_trunc, mask_backgrond=False) @th.inference_mode() def eval_and_export( self, prompt="Yellow rubber duck", # use_ddim=False, # unconditional_guidance_scale=1.0, save_img=False, use_train_trajectory=False, camera=None, num_samples=1, stage_1_output_dir='', num_instances=1, export_mesh=False, ): self.ddpm_model.eval() args = dnnlib.EasyDict( dict( batch_size=1, image_size=self.diffusion_input_size, denoise_in_channels=self.rec_model.decoder.triplane_decoder. out_chans, # type: ignore clip_denoised=False, class_cond=False)) model_kwargs = {} uc = None log = dict() ucg_keys = [self.cond_key] # i23d def sample_and_save(batch_c, idx=0): with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): c, uc = self.conditioner.get_unconditional_conditioning( batch_c, force_uc_zero_embeddings=ucg_keys if len(self.conditioner.embedders) > 0 else [], ) sampling_kwargs = {} N = num_samples # hard coded, to update z_shape = (N, 768, self.ddpm_model.in_channels) for k in c: if isinstance(c[k], th.Tensor): # c[k], uc[k] = map(lambda y: y[k][:N].to(dist_util.dev()), # (c, uc)) assert c[k].shape[0] == 1 # ! support batch inference c[k], uc[k] = map(lambda y: y[k].repeat_interleave(N, 0).to(dist_util.dev()), (c, uc)) # support bs>1 sampling given a condition samples = self.sample(c, shape=z_shape[1:], uc=uc, batch_size=N, **sampling_kwargs) # ! get c if 'img' in self.cond_key: img_save_path = f'{logger.get_dir()}/{idx}_imgcond.jpg' if 'c' in self.cond_key: torchvision.utils.save_image(batch_c['img-c']['img'][0], img_save_path, value_range=(-1,1), normalize=True, padding=0) # torch.Size([24, 6, 3, 256, 256]) else: torchvision.utils.save_image(batch_c['img'], img_save_path, value_range=(-1,1), normalize=True, padding=0) assert camera is not None batch = {'c': camera.clone()} # rendering for i in range(samples.shape[0]): th.cuda.empty_cache() if self.cond_key in ['caption']: name_prefix = f'{batch_c["caption"]}_sample-{idx}-{i}' else: # ! render sampled latent name_prefix = f'{idx}_sample-{i}' # if self.cond_key in ['caption', 'img-c']: with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): # # ! todo, transform to gs camera if self.latent_key != 'latent': # normalized-xyz pcu.save_mesh_v( f'{logger.get_dir()}/{name_prefix}.ply', self.unnormalize_pcd_act(samples[i]).detach().cpu().float().numpy()) logger.log(f'point cloud saved to {logger.get_dir()}/{name_prefix}.ply') else: # ! editing debug all_rgbs, all_depths, all_alphas = self.render_gs_video_given_latent( # samples[i:i+1].to(self.dtype), # default version # th.cat([gt_kl_latent.to(samples), gt_xyz.to(samples)], dim=-1), # ! xyz-cond kl feature gen: # th.cat([samples[i:i+1], batch_c['fps-xyz'][i:i+1]], dim=-1), # ! debugging xyz diffusion th.cat([samples[i:i+1], batch_c['fps-xyz'][0:1]], dim=-1), # ! debugging xyz diffusion # ! xyz debugging # th.cat([gt_kl_latent.to(samples), samples[i:i+1]], dim=-1), # ! debugging xyz diffusion # th.cat([samples[i:i+1], gt_xyz.to(samples), ], dim=-1) # ! debugging kl feature diffusion self.rec_model, # compatible with join_model name_prefix=name_prefix, save_img=save_img, render_reference=batch, export_mesh=False) if export_mesh: self.export_mesh_from_2dgs(all_rgbs, all_depths, idx, i) if self.cond_key == 'caption': assert prompt != '' batch_c = {self.cond_key: prompt} if self.latent_key == 'latent': # t23d, stage-2 for i in range(2): # 8 * num_samples here fps_xyz = torch.from_numpy(pcu.load_mesh_v(f'{stage_1_output_dir}/{prompt}_sample-0-{i}.ply') ).clip(-0.45,0.45).unsqueeze(0) # ! if editing, change the latent points accordingly. # edited_fps_xyz = fps_xyz.clone() # B N 3 # z_dim_coord = edited_fps_xyz[..., 2] # edited_fps_xyz[..., 2] = th.where(z_dim_coord>0.24, z_dim_coord+0.075, z_dim_coord) batch_c.update({ 'fps-xyz': fps_xyz.to(self.dtype).to(dist_util.dev()) # 'fps-xyz': edited_fps_xyz.to(self.dtype).to(dist_util.dev()) }) sample_and_save(batch_c, idx=i) else: sample_and_save(batch_c) @th.inference_mode() def eval_t23d_and_export( self, prompt="Yellow rubber duck", # use_ddim=False, # unconditional_guidance_scale=1.0, save_img=False, use_train_trajectory=False, camera=None, num_samples=1, stage_1_output_dir='', num_instances=1, export_mesh=False, ): self.ddpm_model.eval() args = dnnlib.EasyDict( dict( batch_size=1, image_size=self.diffusion_input_size, denoise_in_channels=self.rec_model.decoder.triplane_decoder. out_chans, # type: ignore clip_denoised=False, class_cond=False)) model_kwargs = {} uc = None log = dict() ucg_keys = [self.cond_key] # i23d assert self.cond_key == 'caption' and prompt != '' batch_c = {self.cond_key: prompt} if self.latent_key == 'latent': # t23d, stage-2 fps_xyz = torch.from_numpy(pcu.load_mesh_v(f'{stage_1_output_dir}/{prompt}_sample-0.ply') ).clip(-0.45,0.45).unsqueeze(0) batch_c.update({ 'fps-xyz': fps_xyz.to(self.dtype).to(dist_util.dev()) }) self.sample_and_save(batch_c, ucg_keys, num_samples, camera,) @th.inference_mode() def eval_i23d_and_export_gradio( self, inp_img, seed=42, cfg_scale=4.0, # default value in neural ode save_img=False, **kwargs, ): # logger.log( # unconditional_guidance_scale, # seed, # mesh_size, # mesh_thres, # ) sampling_kwargs = { 'cfg_scale': cfg_scale, # default value in SiT 'seed': seed, } camera = th.load('assets/objv_eval_pose.pt', map_location=dist_util.dev())[:24] inp_img = th.from_numpy(inp_img).permute(2,0,1).unsqueeze(0) / 127.5 - 1 # to [-1,1] num_samples=1 export_mesh=True # self.ddpm_model.eval() # args = dnnlib.EasyDict( # dict( # batch_size=1, # image_size=self.diffusion_input_size, # denoise_in_channels=self.rec_model.decoder.triplane_decoder. # out_chans, # type: ignore # clip_denoised=False, # class_cond=False)) # model_kwargs = {} ucg_keys = [self.cond_key] # i23d ins_name = 'house2-input' # for debug here if self.cond_key == 'img-xyz': # stage-2 i = 0 # for gradio only # for i in range(1): stage_1_output_dir="./logs/i23d/stage-1/dino_img/" stage1_pcd_output_path = f'{stage_1_output_dir}/{ins_name}/sample-0-{i}.ply' fps_xyz = trimesh.load(stage1_pcd_output_path).vertices # pcu may fail on py=3.9 fps_xyz = torch.from_numpy(fps_xyz).clip(-0.45,0.45).unsqueeze(0) logger.log('loading stage-1 point cloud from: ', stage1_pcd_output_path) # fps_xyz = None # ! TODO, load from local directory # batch_c = { # 'img': batch['img'][0:1].to(self.dtype).to(dist_util.dev()), # 'fps-xyz': fps_xyz[0:1].to(self.dtype).to(dist_util.dev()), # } batch_c = {'img': inp_img.to(dist_util.dev()).to(self.dtype), 'fps-xyz': fps_xyz[0:1].to(self.dtype).to(dist_util.dev())} # no need to return here? video_path, rgb_xyz_path, post_mesh_path = self.sample_and_save(batch_c, ucg_keys, num_samples, camera, save_img, idx=ins_name, export_mesh=export_mesh, stage1_idx=i, **sampling_kwargs) # type: ignore # video_path = './logs/i23d/stage-2/dino_img/house2-input/sample-0-0-gs.mp4' # rgb_xyz_path = './logs/i23d/stage-2/dino_img/low-poly-model-of-a-green-pine-tree,-also-resembling-a-Christmas-tree.-vc.ply' assert post_mesh_path != '' return video_path, rgb_xyz_path, post_mesh_path else: # stage-1 data # batch_c = {self.cond_key: batch[self.cond_key][0:1].to(dist_util.dev()).to(self.dtype), } # raise NotImplementedError('stage-2 only') batch_c = {'img': inp_img.to(dist_util.dev()).to(self.dtype)} pcd_export_dir = self.sample_and_save(batch_c, ucg_keys, num_samples, camera, save_img, idx=ins_name, export_mesh=export_mesh, **sampling_kwargs) # type: ignore return pcd_export_dir @th.inference_mode() def eval_i23d_and_export( self, prompt="Yellow rubber duck", # use_ddim=False, unconditional_guidance_scale=4.0, save_img=False, seed=42, # cfg_scale=4.0, # default value in neural ode camera=None, num_samples=1, stage_1_output_dir='', # num_instances=1, export_mesh=False, ): self.ddpm_model.eval() # args = dnnlib.EasyDict( # dict( # batch_size=1, # image_size=self.diffusion_input_size, # denoise_in_channels=self.rec_model.decoder.triplane_decoder. # out_chans, # type: ignore # clip_denoised=False, # class_cond=False)) # model_kwargs = {} sampling_kwargs = { 'cfg_scale': unconditional_guidance_scale, # default value in SiT 'seed': seed, } uc = None log = dict() ucg_keys = [self.cond_key] # i23d for idx, batch in enumerate(tqdm(self.data)): ins = batch['ins'][0] # obj_folder, _, frame = ins.split('/') ins = ins.split('/') # obj_folder, frame = ins[0], ins[-1] # for gso if len(ins) >2: if ins[1] == 'render_mvs_25': # gso obj_folder, frame = ins[0], int(ins[-1].split('.')[0]) ins_name = f'{obj_folder}/{str(frame)}' else: obj_folder, frame = os.path.join(ins[1], ins[2]), ins[-1] # for objv frame = int(frame.split('.')[0]) ins_name = f'{obj_folder}/{str(frame)}' else: # folder of images, e.g., instantmesh ins_name = ins[0].split('.')[0] # pcd_export_dir = f'{logger.get_dir()}/{ins_name}/sample-0.ply' # if os.path.exists(pcd_export_dir): # continue #! debugging, get GT xyz and KL latent for disentangled debugging if self.cond_key == 'img-c': # mv23d prompt = batch['caption'][0:1] batch_c = { self.cond_key: { 'img': batch['mv_img'][0:1].to(self.dtype).to(dist_util.dev()), 'c': batch['c'][0:1].to(self.dtype).to(dist_util.dev()), }, 'img': batch['img'][0:1].to(self.dtype).to(dist_util.dev()), 'caption': prompt, } if self.latent_key == 'latent': # stage-2 fps_xyz = torch.from_numpy(pcu.load_mesh_v(f'{stage_1_output_dir}/{idx}_sample-0.ply') ).clip(-0.45,0.45).unsqueeze(0) # fps_xyz = torch.from_numpy(pcu.load_mesh_v(f'{stage_1_output_dir}/sample-0.ply') ).clip(-0.45,0.45).unsqueeze(0) batch_c.update({ 'fps-xyz': fps_xyz[0:1].to(self.dtype).to(dist_util.dev()), }) # elif self.cond_key == 'img-caption': # batch_c = {'caption': prompt, 'img': batch['img'].to(dist_util.dev()).to(self.dtype)} elif self.cond_key == 'img-xyz': # stage-2 for i in range(2): stage1_pcd_output_path = f'{stage_1_output_dir}/{ins_name}/sample-0-{i}.ply' fps_xyz = trimesh.load(stage1_pcd_output_path).vertices # pcu may fail on py=3.9 fps_xyz = torch.from_numpy(fps_xyz).clip(-0.45,0.45).unsqueeze(0) # fps_xyz = None # ! TODO, load from local directory batch_c = { 'img': batch['img'][0:1].to(self.dtype).to(dist_util.dev()), 'fps-xyz': fps_xyz[0:1].to(self.dtype).to(dist_util.dev()), } self.sample_and_save(batch_c, ucg_keys, num_samples, camera, save_img, idx=ins_name, export_mesh=export_mesh, stage1_idx=i,**sampling_kwargs) # type: ignore else: # stage-1 data batch_c = {self.cond_key: batch[self.cond_key][0:1].to(dist_util.dev()).to(self.dtype), } if self.cond_key == 'caption' and self.latent_key == 'latent': # t23d, stage-2 fps_xyz = torch.from_numpy(pcu.load_mesh_v(f'{stage_1_output_dir}/{idx}_sample-0.ply') ).clip(-0.45,0.45).unsqueeze(0) batch_c.update({ 'fps-xyz': fps_xyz.to(self.dtype).to(dist_util.dev()) }) # save_dir = f'{logger.get_dir()}/{ins}' # os.mkdir(save_dir, exists_ok=True, parents=True) # self.sample_and_save(batch_c, ucg_keys, num_samples, camera, save_img, idx=f'{idx}-{ins}', export_mesh=export_mesh) self.sample_and_save(batch_c, ucg_keys, num_samples, camera, save_img, idx=ins_name, export_mesh=export_mesh,**sampling_kwargs) # type: ignore gc.collect() def get_source_cw2wT(self, source_cameras_view_to_world): return matrix_to_quaternion( source_cameras_view_to_world[:3, :3].transpose(0, 1)) def c_to_3dgs_format(self, pose): # TODO, switch to torch version (batched later) c2w = pose[:16].reshape(4, 4) # 3x4 # ! load cam w2c = np.linalg.inv(c2w) R = np.transpose( w2c[:3, :3]) # R is stored transposed due to 'glm' in CUDA code T = w2c[:3, 3] fx = pose[16] FovX = focal2fov(fx, 1) FovY = focal2fov(fx, 1) tanfovx = math.tan(FovX * 0.5) tanfovy = math.tan(FovY * 0.5) assert tanfovx == tanfovy trans = np.array([0.0, 0.0, 0.0]) scale = 1.0 world_view_transform = torch.tensor(getWorld2View2(R, T, trans, scale)).transpose( 0, 1) projection_matrix = getProjectionMatrix(znear=self.znear, zfar=self.zfar, fovX=FovX, fovY=FovY).transpose(0, 1) full_proj_transform = (world_view_transform.unsqueeze(0).bmm( projection_matrix.unsqueeze(0))).squeeze(0) camera_center = world_view_transform.inverse()[3, :3] view_world_transform = torch.tensor(getView2World(R, T, trans, scale)).transpose( 0, 1) # item.update(viewpoint_cam=[viewpoint_cam]) c = {} c["source_cv2wT_quat"] = self.get_source_cw2wT(view_world_transform) c.update( projection_matrix=projection_matrix, # K cam_view=world_view_transform, # world_view_transform cam_view_proj=full_proj_transform, # full_proj_transform cam_pos=camera_center, tanfov=tanfovx, # TODO, fix in the renderer # orig_c2w=c2w, # orig_w2c=w2c, orig_pose=torch.from_numpy(pose), orig_c2w=torch.from_numpy(c2w), orig_w2c=torch.from_numpy(w2c), # tanfovy=tanfovy, ) return c # dict for gs rendering class FlowMatchingEngine_gs_clay(FlowMatchingEngine_gs): def __init__( self, *, rec_model, denoise_model, diffusion, sde_diffusion, control_model, control_key, only_mid_control, loss_class, data, eval_data, batch_size, microbatch, lr, ema_rate, log_interval, eval_interval, save_interval, resume_checkpoint, resume_cldm_checkpoint=None, use_fp16=False, fp16_scale_growth=0.001, schedule_sampler=None, weight_decay=0, lr_anneal_steps=0, iterations=10001, ignore_resume_opt=False, freeze_ae=False, denoised_ae=True, triplane_scaling_divider=10, use_amp=False, diffusion_input_size=224, normalize_clip_encoding=False, scale_clip_encoding=1, cfg_dropout_prob=0, cond_key='img_sr', use_eos_feature=False, compile=False, snr_type='lognorm', **kwargs): super().__init__(rec_model=rec_model, denoise_model=denoise_model, diffusion=diffusion, sde_diffusion=sde_diffusion, control_model=control_model, control_key=control_key, only_mid_control=only_mid_control, loss_class=loss_class, data=data, eval_data=eval_data, batch_size=batch_size, microbatch=microbatch, lr=lr, ema_rate=ema_rate, log_interval=log_interval, eval_interval=eval_interval, save_interval=save_interval, resume_checkpoint=resume_checkpoint, resume_cldm_checkpoint=resume_cldm_checkpoint, use_fp16=use_fp16, fp16_scale_growth=fp16_scale_growth, schedule_sampler=schedule_sampler, weight_decay=weight_decay, lr_anneal_steps=lr_anneal_steps, iterations=iterations, ignore_resume_opt=ignore_resume_opt, freeze_ae=freeze_ae, denoised_ae=denoised_ae, triplane_scaling_divider=triplane_scaling_divider, use_amp=use_amp, diffusion_input_size=diffusion_input_size, normalize_clip_encoding=normalize_clip_encoding, scale_clip_encoding=scale_clip_encoding, cfg_dropout_prob=cfg_dropout_prob, cond_key=cond_key, use_eos_feature=use_eos_feature, compile=compile, snr_type=snr_type, **kwargs) # self._init_new_ca_weight() # after ckpt loading def _set_grad_flag(self): # unfree CA only requires_grad(self.ddpm_model, True) # # for k, v in self.ddpm_model.named_parameters(): # # if 'cross_attn_dino' in k: # if 'mv' in k: # for mv dino # v.requires_grad_(True) # if self.step == 0: # logger.log(k) # else: # v.requires_grad_(False) def _init_new_ca_weight(self): blks_to_copy = ['cross_attn_dino', 'prenorm_ca_dino'] for blk in self.ddpm_model.blocks: for param_name in blks_to_copy: try: getattr(blk, param_name.replace('dino', 'dino_mv')).load_state_dict(getattr(blk, param_name).state_dict()) except Exception as e: logger.log(e) # some key misalignment