import os
import numpy as np
import torch
from einops import rearrange
from imageio import imwrite
from pydantic import validator
import imageio
import tempfile
import gradio as gr
from PIL import Image
from my.utils import (
tqdm, EventStorage, HeartBeat, EarlyLoopBreak,
get_event_storage, get_heartbeat, read_stats
)
from my.config import BaseConf, dispatch, optional_load_config
from my.utils.seed import seed_everything
from adapt import ScoreAdapter
from run_img_sampling import SD
from misc import torch_samps_to_imgs
from pose import PoseConfig
from run_nerf import VoxConfig
from voxnerf.utils import every
from voxnerf.render import (
as_torch_tsrs, rays_from_img, ray_box_intersect, render_ray_bundle
)
from voxnerf.vis import stitch_vis, bad_vis as nerf_vis
from pytorch3d.renderer import PointsRasterizationSettings
from semantic_coding import semantic_coding, semantic_karlo, semantic_sd
from pc_project import point_e, render_depth_from_cloud
device_glb = torch.device("cuda")
def tsr_stats(tsr):
return {
"mean": tsr.mean().item(),
"std": tsr.std().item(),
"max": tsr.max().item(),
}
class SJC_3DFuse(BaseConf):
family: str = "sd"
sd: SD = SD(
variant="v1",
prompt="a comfortable bed",
scale=100.0,
dir="./results",
alpha=0.3
)
lr: float = 0.05
n_steps: int = 10000
vox: VoxConfig = VoxConfig(
model_type="V_SD", grid_size=100, density_shift=-1.0, c=3,
blend_bg_texture=False , bg_texture_hw=4,
bbox_len=1.0
)
pose: PoseConfig = PoseConfig(rend_hw=64, FoV=60.0, R=1.5)
emptiness_scale: int = 10
emptiness_weight: int = 1e4
emptiness_step: float = 0.5
emptiness_multiplier: float = 20.0
depth_weight: int = 0
var_red: bool = True
exp_dir: str = "./results"
ti_step: int = 800
pt_step: int = 800
initial: str = ""
random_seed: int = 0
semantic_model: str = "Karlo"
bg_preprocess: bool = True
num_initial_image: int = 4
@validator("vox")
def check_vox(cls, vox_cfg, values):
family = values['family']
if family == "sd":
vox_cfg.c = 4
return vox_cfg
def run(self):
raise Exception("This version is for huggingface demo, which doesn't support CLI. Please visit https://github.com/KU-CVLAB/3DFuse")
def run_gradio(self, points, images):
cfgs = self.dict()
initial = cfgs.pop('initial')
exp_dir=os.path.join(cfgs.pop('exp_dir'),initial)
# Optimization and pivotal tuning for LoRA
yield gr.update(value=None), "Tuning for the LoRA layer is starting now. It will take approximately ~10 mins.", gr.update(value=None)
state=semantic_coding(images, cfgs,self.sd,initial)
self.sd.dir=state
# Load SD with Consistency Injection Module
family = cfgs.pop("family")
model = getattr(self, family).make()
print(model.prompt)
cfgs.pop("vox")
vox = self.vox.make()
cfgs.pop("pose")
poser = self.pose.make()
# Score distillation
yield from fuse_3d(**cfgs, poser=poser,model=model,vox=vox,exp_dir=exp_dir, points=points, is_gradio=True)
def fuse_3d(
poser, vox, model: ScoreAdapter,
lr, n_steps, emptiness_scale, emptiness_weight, emptiness_step, emptiness_multiplier,
depth_weight, var_red, exp_dir, points, is_gradio, **kwargs
):
del kwargs
if is_gradio:
yield gr.update(visible=True), "LoRA layers tuning has just finished. \nScore distillation has started.", gr.update(visible=True)
assert model.samps_centered()
_, target_H, target_W = model.data_shape()
bs = 1
aabb = vox.aabb.T.cpu().numpy()
vox = vox.to(device_glb)
opt = torch.optim.Adamax(vox.opt_params(), lr=lr)
H, W = poser.H, poser.W
Ks_, poses_, prompt_prefixes_, angles_list = poser.sample_train(n_steps,device_glb)
ts = model.us[30:-10]
fuse = EarlyLoopBreak(5)
raster_settings = PointsRasterizationSettings(
image_size= 800,
radius = 0.02,
points_per_pixel = 10
)
ts = model.us[30:-10]
calibration_value=0.0
with tqdm(total=n_steps) as pbar:
# HeartBeat(pbar) as hbeat, \
# EventStorage(output_dir=os.path.join(exp_dir,'3d')) as metric:
for i in range(len(poses_)):
if fuse.on_break():
break
depth_map = render_depth_from_cloud(points, angles_list[i], raster_settings, device_glb,calibration_value)
y, depth, ws = render_one_view(vox, aabb, H, W, Ks_[i], poses_[i], return_w=True)
p = f"{prompt_prefixes_[i]} {model.prompt}"
score_conds = model.prompts_emb([p])
score_conds['c']=score_conds['c'].repeat(bs,1,1)
score_conds['uc']=score_conds['uc'].repeat(bs,1,1)
opt.zero_grad()
with torch.no_grad():
chosen_σs = np.random.choice(ts, bs, replace=False)
chosen_σs = chosen_σs.reshape(-1, 1, 1, 1)
chosen_σs = torch.as_tensor(chosen_σs, device=model.device, dtype=torch.float32)
noise = torch.randn(bs, *y.shape[1:], device=model.device)
zs = y + chosen_σs * noise
Ds = model.denoise(zs, chosen_σs,depth_map.unsqueeze(dim=0),**score_conds)
if var_red:
grad = (Ds - y) / chosen_σs
else:
grad = (Ds - zs) / chosen_σs
grad = grad.mean(0, keepdim=True)
y.backward(-grad, retain_graph=True)
if depth_weight > 0:
center_depth = depth[7:-7, 7:-7]
border_depth_mean = (depth.sum() - center_depth.sum()) / (64*64-50*50)
center_depth_mean = center_depth.mean()
depth_diff = center_depth_mean - border_depth_mean
depth_loss = - torch.log(depth_diff + 1e-12)
depth_loss = depth_weight * depth_loss
depth_loss.backward(retain_graph=True)
emptiness_loss = torch.log(1 + emptiness_scale * ws).mean()
emptiness_loss = emptiness_weight * emptiness_loss
if emptiness_step * n_steps <= i:
emptiness_loss *= emptiness_multiplier
emptiness_loss.backward()
opt.step()
# metric.put_scalars(**tsr_stats(y))
if every(pbar, percent=2):
with torch.no_grad():
y = model.decode(y)
# vis_routine(metric, y, depth,p,depth_map[0])
if is_gradio :
yield torch_samps_to_imgs(y)[0], f"Progress: {pbar.n}/{pbar.total} \nAfter the generation is complete, the video results will be displayed below.", gr.update(value=None)
# metric.step()
pbar.update()
pbar.set_description(p)
# hbeat.beat()
# metric.put_artifact(
# "ckpt", ".pt","", lambda fn: torch.save(vox.state_dict(), fn)
# )
# with EventStorage("result"):
out=evaluate(model, vox, poser)
if is_gradio:
yield gr.update(visible=False), f"Generation complete. Please check the video below.", gr.update(value=out)
else :
yield None
# metric.step()
# hbeat.done()
@torch.no_grad()
def evaluate(score_model, vox, poser):
H, W = poser.H, poser.W
vox.eval()
K, poses = poser.sample_test(100)
fuse = EarlyLoopBreak(5)
# metric = get_event_storage()
# hbeat = get_heartbeat()
aabb = vox.aabb.T.cpu().numpy()
vox = vox.to(device_glb)
num_imgs = len(poses)
frames=[]
for i in (pbar := tqdm(range(num_imgs))):
if fuse.on_break():
break
pose = poses[i]
y, depth = render_one_view(vox, aabb, H, W, K, pose)
y = score_model.decode(y)
# vis_routine(metric, y, depth,"",None)
y=torch_samps_to_imgs(y)[0]
frames.append(y)
# metric.step()
# hbeat.beat()
# metric.flush_history()
# metric.put_artifact(
# "video", ".mp4","",
# lambda fn: stitch_vis(fn, read_stats(metric.output_dir, "img")[1])
# )
out_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
writer = imageio.get_writer(out_file.name, fps=10)
for img in frames:
writer.append_data(img)
writer.close()
# metric.step()
return out_file.name
def render_one_view(vox, aabb, H, W, K, pose, return_w=False):
N = H * W
ro, rd = rays_from_img(H, W, K, pose)
ro, rd, t_min, t_max = scene_box_filter_(ro, rd, aabb)
assert len(ro) == N, "for now all pixels must be in"
ro, rd, t_min, t_max = as_torch_tsrs(vox.device, ro, rd, t_min, t_max)
rgbs, depth, weights = render_ray_bundle(vox, ro, rd, t_min, t_max)
rgbs = rearrange(rgbs, "(h w) c -> 1 c h w", h=H, w=W)
depth = rearrange(depth, "(h w) 1 -> h w", h=H, w=W)
if return_w:
return rgbs, depth, weights
else:
return rgbs, depth
def scene_box_filter_(ro, rd, aabb):
_, t_min, t_max = ray_box_intersect(ro, rd, aabb)
# do not render what's behind the ray origin
t_min, t_max = np.maximum(t_min, 0), np.maximum(t_max, 0)
return ro, rd, t_min, t_max
def vis_routine(metric, y, depth,prompt,depth_map):
pane = nerf_vis(y, depth, final_H=256)
im = torch_samps_to_imgs(y)[0]
depth = depth.cpu().numpy()
metric.put_artifact("view", ".png","",lambda fn: imwrite(fn, pane))
metric.put_artifact("img", ".png",prompt, lambda fn: imwrite(fn, im))
if depth_map != None:
metric.put_artifact("PC_depth", ".png",prompt, lambda fn: imwrite(fn, depth_map.cpu().squeeze()))
metric.put_artifact("depth", ".npy","",lambda fn: np.save(fn, depth))
if __name__ == "__main__":
dispatch(SJC_3DFuse)