init

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,35 @@
+__pycache__/
+build/
+*.egg-info/
+*.so
+venv_*/
+.vs/
+.vscode/
+.idea/
+
+tmp_*
+data?
+data??
+scripts2
+
+model_cache
+
+logs
+videos
+images
+*.mp4
+
+vis_data*/
+logs*/
+data*/
+eval_data*/
+
+
+*.sh
+*.out
+batchscript*
+
+pretrained/
+diff-gaussian-rasterization/
+
+tmp_data/

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "diff-gaussian-rasterization"]
+	path = diff-gaussian-rasterization
+	url = https://github.com/ashawkey/diff-gaussian-rasterization

app.py

@@ -0,0 +1,135 @@
+import gradio as gr
+import os
+from PIL import Image
+import subprocess
+from gradio_model4dgs import Model4DGS
+import numpy
+import hashlib
+
+os.system('pip install -e ./simple-knn')
+os.system('pip install -e ./diff-gaussian-rasterization')
+
+from huggingface_hub import hf_hub_download
+ckpt_path = hf_hub_download(repo_id="ashawkey/LGM", filename="model_fp16_fixrot.safetensors")
+
+js_func = """
+function refresh() {
+    const url = new URL(window.location);
+
+    if (url.searchParams.get('__theme') !== 'light') {
+        url.searchParams.set('__theme', 'light');
+        window.location.href = url.href;
+    }
+}
+"""
+
+# check if there is a picture uploaded or selected
+def check_img_input(control_image):
+    if control_image is None:
+        raise gr.Error("Please select or upload an input image")
+
+# check if there is a picture uploaded or selected
+def check_video_input(image_block: Image.Image):
+    img_hash = hashlib.sha256(image_block.tobytes()).hexdigest()
+    if not os.path.exists(os.path.join('tmp_data', f'{img_hash}_rgba_generated.mp4')):
+        raise gr.Error("Please generate a video first")
+
+
+def optimize_stage_1(image_block: Image.Image, preprocess_chk: bool, seed_slider: int):
+    if not os.path.exists('tmp_data'):
+        os.makedirs('tmp_data')
+    img_hash = hashlib.sha256(image_block.tobytes()).hexdigest()
+    if preprocess_chk:
+        # save image to a designated path
+        image_block.save(os.path.join('tmp_data', f'{img_hash}.png'))
+
+        # preprocess image
+        print(f'python scripts/process.py {os.path.join("tmp_data", f"{img_hash}.png")}')
+        subprocess.run(f'python scripts/process.py {os.path.join("tmp_data", f"{img_hash}.png")}', shell=True)
+    else:
+        image_block.save(os.path.join('tmp_data', f'{img_hash}_rgba.png'))
+
+    # stage 1
+    subprocess.run(f'export MKL_THREADING_LAYER=GNU;export MKL_SERVICE_FORCE_INTEL=1;python scripts/gen_vid.py --path tmp_data/{img_hash}_rgba.png --seed {seed_slider} --bg white', shell=True)
+
+    # return [os.path.join('logs', 'tmp_rgba_model.ply')]
+    return os.path.join('tmp_data', f'{img_hash}_rgba_generated.mp4')
+
+
+def optimize_stage_2(image_block: Image.Image, seed_slider: int):
+    img_hash = hashlib.sha256(image_block.tobytes()).hexdigest()
+    subprocess.run(f'python lgm/infer.py big --resume {ckpt_path} --test_path tmp_data/{img_hash}_rgba.png', shell=True)
+    # stage 2
+    subprocess.run(f'python main_4d.py --config {os.path.join("configs", "4d_demo.yaml")} input={os.path.join("tmp_data", f"{img_hash}_rgba.png")}', shell=True)
+    # os.rename(os.path.join('logs', f'{img_hash}_rgba_frames'), os.path.join('logs', f'{img_hash}_{seed_slider:03d}_rgba_frames'))
+    image_dir = os.path.join('logs', f'{img_hash}_rgba_frames')
+    # return 'vis_data/tmp_rgba.mp4', [os.path.join(image_dir, file) for file in os.listdir(image_dir) if file.endswith('.ply')]
+    return [image_dir+f'/{t:03d}.ply' for t in range(28)]
+
+
+if __name__ == "__main__":
+    _TITLE = '''DreamGaussian4D: Generative 4D Gaussian Splatting'''
+
+    _DESCRIPTION = '''
+    <div>
+    <a style="display:inline-block" href="https://jiawei-ren.github.io/projects/dreamgaussian4d/"><img src='https://img.shields.io/badge/public_website-8A2BE2'></a>
+    <a style="display:inline-block; margin-left: .5em" href="https://arxiv.org/abs/2312.17142"><img src="https://img.shields.io/badge/2312.17142-f9f7f7?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAADcAAABMCAYAAADJPi9EAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAAuIwAALiMBeKU/dgAAABl0RVh0U29mdHdhcmUAd3d3Lmlua3NjYXBlLm9yZ5vuPBoAAAa2SURBVHja3Zt7bBRFGMAXUCDGF4rY7m7bAwuhlggKStFgLBgFEkCIIRJEEoOBYHwRFYKilUgEReVNJEGCJJpehHI3M9vZvd3bUP1DjNhEIRQQsQgSHiJgQZ5dv7krWEvvdmZ7d7vHJN+ft/f99pv5XvOtJMFCqvoCUpTdIEeRLC+L9Ox5i3Q9LACaCeK0kXoSChVcD3C/tQPHpAEsquQ73IkUcEz2kcLCknyGW5MGjkljRFVL8xJOKyi4CwCOuQAeAkfTP1+tNxLkogvgEbDgffkJqKqvuMA5ifOpqg/5qWecRstNg7xoUTI1Fovdxg8oy2s5AP8CGeYHmGngeZaOL4I4LXLcpHg4149/GDz4xqgsb+UAbMKKUpkrqHA43MUyyJpWUK0EHeG2YKRXr7tB+QMcgGewLD+ebTDbtrtbBt7UPlhS4rV4IvcDI7J8P1OeA/AcAI7LHljN7aB8XTowJmZt9EFRD/o0SDMH4HlwMhMyDWZZSAHFf3YDs3RS49WDLuaAY3IJq+qzmQKLxXAZKN7oDoYbdV3v5elPqiSpMyiOuAEVZVqHXb1OhloUH+MA+ztO0cAO/RkrfyBE7OAEbAZvO8vzVtTRWFD6DAfY5biBM3PWiaL0a4lvXICwnV8WjmE6ntYmhqX2jjp5LbMZjCw/wbYeN6CizOa2GMVzQOlmHjB4Ceuyk6LJ8huccEmR5Xddg7OOV/NAtchW+E3XbOag60QA4Qwuarca0bRuEJyr+cFQwzcY98huxhAKdQelt4kAQpj4qJ3gvFXAYn+aJumXk1yPlpQUgtIHhbYoFMUstNRRWgjnpl4A7IKlayNymqFHFaWCpV9CFry3LGxR1CgA5kB5M8OX2goApwpaz6mdOMGxtAgXWJySxb4WuQD4qTDgU+N5AAnzpr7ChSWpCyisiQJqY0Y7FtmSKpbV23b45kC0KHBxcQ9QeI8w4KgnHRPVtIU7rOtbioLVg5Hl/qDwSVFAMqLSMSObroCdZYlzIJtMRFVHCaRo/wFWPgaAXzdbBpkc2A4aKzCNd97+URQuESYGDDhIVfWOQIKZJu4D2+oXlgDTV1865gUQZDts756BArMNMoR1oa46BYqbyPixZz1ZUFV3sgwoGBajuBKATl3btIn8QYYMuezRgrsiRUWyr2BxA40EkPMpA/Hm6gbUu7fjEXA3azP6AsbKD9bxdUuhjM9W7fII52BF+daRpE4+WA3P501+jbfmHvQKyFqMuXf7Ot4mkN2fr50y+bRH61X7AXdUpHSxaPQ4GVbR5AGw3g+434XgQGKfr72I+vQRhfsu92dOx7WicInzt3CBg1RVpMm0NveWo2SqFzgmdNZMbriILD+S+zoueWf2vSdAipzacWN5nMl6XxNlUHa/J8DoJodUDE0HR8Ll5V0lPxcrLEHZPV4AzS83OLis7FowVa3RSku7BSNxJqQAlN3hBTC2apmDSkpaw22wJemGQFUG7J4MlP3JC6A+f96V7vRyX9It3nzT/GrjIU8edM7rMSnIi10f476lzbE1K7yEiEuWro0OJBguLCwDuFOJc1Na6sRWL/cCeMIwUN9ggSVbe3v/5/EgzTKWLvEAiBrYRUkgwNI2ZaFQNT75UDxEUEx97zYnzpmiLEmbaYCbNxYtFAb0/Z4AztgUrhyxuNgxPnhfHFDHz/vTgFWUQZxTRkkJhQ6YNdVUEPAfO6ZV5BRss6LcCVb7VaAma9giy0XJZBt9IQh42NY0NSdgbLIPlLUF6rEdrdt0CUCK1wsCbkcI3ZSLc7ZSwGLbmJXbPsNxnE5xilYKAobZ77LpGZ8TAIun+/iCKQoF71IxQDI3K2CCd+ARNvXg9sykBcnHAoCZG4u66hlDoQLe6QV4CRtFSxZQ+D0BwNO2jgdkzoGoah1nj3FVlSR19taTSYxI8QLut23U8dsgzqHulJNCQpcqBnpTALCuQ6NSYLHpmR5i42gZzuIdcrMMvMJbQlxe3jXxyZnLACl7ARm/FjPIDOY8ODtpM71sxwfcZpvBeUzKWmfNINM5AS+wO0Khh7dMqKccu4+qatarZjYAwDlgetzStHtEt+XedsBOQtU9XMrRgjg4KTnc5nr+dmqadit/4C4uLm8DuA9koJTj1TL7fI5nDL+qqoo/FLGAzL7dYT17PzvAcQONYSUQRxW/QMrHZVIyik0ZuQA2mzp+Ji8BW4YM3Mbzm9inaHkJCGfrUZZjujiYailfFwA8DHIy3acwUj4v9vUVa+SmgNsl5fuyDTKovW9/IAmfLV0Pi2UncA515kjYdrwC9i9rpuHiq3JwtAAAAABJRU5ErkJggg=="></a>
+    <a style="display:inline-block; margin-left: .5em" href='https://github.com/jiawei-ren/dreamgaussian4d'><img src='https://img.shields.io/github/stars/jiawei-ren/dreamgaussian4d?style=social'/></a>
+    </div>
+    We present DreamGausssion4D, an efficient 4D generation framework that builds on Gaussian Splatting. 
+    '''
+    _IMG_USER_GUIDE = "Please upload an image in the block above (or choose an example above), select a random seed, and click **Generate Video**. After having the video generated, please click **Generate 4D**."
+
+    # load images in 'data' folder as examples
+    example_folder = os.path.join(os.path.dirname(__file__), 'data')
+    example_fns = os.listdir(example_folder)
+    example_fns.sort()
+    examples_full = [os.path.join(example_folder, x) for x in example_fns if x.endswith('.png')]
+
+    # Compose demo layout & data flow
+    with gr.Blocks(title=_TITLE, theme=gr.themes.Soft(), js=js_func) as demo:
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown('# ' + _TITLE)
+        gr.Markdown(_DESCRIPTION)
+
+        # Image-to-3D
+        with gr.Row(variant='panel'):
+            with gr.Column(scale=4):
+                image_block = gr.Image(type='pil', image_mode='RGBA', height=290, label='Input image')
+
+                # elevation_slider = gr.Slider(-90, 90, value=0, step=1, label='Estimated elevation angle')
+                seed_slider = gr.Slider(0, 100000, value=0, step=1, label='Random Seed')
+                gr.Markdown(
+                    "random seed for video generation.")
+
+                preprocess_chk = gr.Checkbox(True,
+                                             label='Preprocess image automatically (remove background and recenter object)')
+
+                gr.Examples(
+                    examples=examples_full,  # NOTE: elements must match inputs list!
+                    inputs=[image_block],
+                    outputs=[image_block],
+                    cache_examples=False,
+                    label='Examples (click one of the images below to start)',
+                    examples_per_page=40
+                )
+                img_run_btn = gr.Button("Generate Video")
+                fourd_run_btn = gr.Button("Generate 4D")
+                img_guide_text = gr.Markdown(_IMG_USER_GUIDE, visible=True)
+
+            with gr.Column(scale=5):
+                obj3d = gr.Video(label="video",height=290)
+                obj4d = Model4DGS(label="4D Model", height=500, fps=14)
+
+            img_run_btn.click(check_img_input, inputs=[image_block], queue=False).success(optimize_stage_1,
+                                                                                          inputs=[image_block,
+                                                                                                  preprocess_chk,
+                                                                                                  seed_slider],
+                                                                                          outputs=[
+                                                                                              obj3d])
+            fourd_run_btn.click(check_video_input, inputs=[image_block], queue=False).success(optimize_stage_2, inputs=[image_block, seed_slider], outputs=[obj4d])
+
+    # demo.queue().launch(share=True)
+    demo.queue(max_size=10)  # <-- Sets up a queue with default parameters
+    demo.launch(share=True)

cam_utils.py

@@ -0,0 +1,146 @@
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+
+import torch
+
+def dot(x, y):
+    if isinstance(x, np.ndarray):
+        return np.sum(x * y, -1, keepdims=True)
+    else:
+        return torch.sum(x * y, -1, keepdim=True)
+
+
+def length(x, eps=1e-20):
+    if isinstance(x, np.ndarray):
+        return np.sqrt(np.maximum(np.sum(x * x, axis=-1, keepdims=True), eps))
+    else:
+        return torch.sqrt(torch.clamp(dot(x, x), min=eps))
+
+
+def safe_normalize(x, eps=1e-20):
+    return x / length(x, eps)
+
+
+def look_at(campos, target, opengl=True):
+    # campos: [N, 3], camera/eye position
+    # target: [N, 3], object to look at
+    # return: [N, 3, 3], rotation matrix
+    if not opengl:
+        # camera forward aligns with -z
+        forward_vector = safe_normalize(target - campos)
+        up_vector = np.array([0, 1, 0], dtype=np.float32)
+        right_vector = safe_normalize(np.cross(forward_vector, up_vector))
+        up_vector = safe_normalize(np.cross(right_vector, forward_vector))
+    else:
+        # camera forward aligns with +z
+        forward_vector = safe_normalize(campos - target)
+        up_vector = np.array([0, 1, 0], dtype=np.float32)
+        right_vector = safe_normalize(np.cross(up_vector, forward_vector))
+        up_vector = safe_normalize(np.cross(forward_vector, right_vector))
+    R = np.stack([right_vector, up_vector, forward_vector], axis=1)
+    return R
+
+
+# elevation & azimuth to pose (cam2world) matrix
+def orbit_camera(elevation, azimuth, radius=1, is_degree=True, target=None, opengl=True):
+    # radius: scalar
+    # elevation: scalar, in (-90, 90), from +y to -y is (-90, 90)
+    # azimuth: scalar, in (-180, 180), from +z to +x is (0, 90)
+    # return: [4, 4], camera pose matrix
+    if is_degree:
+        elevation = np.deg2rad(elevation)
+        azimuth = np.deg2rad(azimuth)
+    x = radius * np.cos(elevation) * np.sin(azimuth)
+    y = - radius * np.sin(elevation)
+    z = radius * np.cos(elevation) * np.cos(azimuth)
+    if target is None:
+        target = np.zeros([3], dtype=np.float32)
+    campos = np.array([x, y, z]) + target  # [3]
+    T = np.eye(4, dtype=np.float32)
+    T[:3, :3] = look_at(campos, target, opengl)
+    T[:3, 3] = campos
+    return T
+
+
+class OrbitCamera:
+    def __init__(self, W, H, r=2, fovy=60, near=0.01, far=100):
+        self.W = W
+        self.H = H
+        self.radius = r  # camera distance from center
+        self.fovy = np.deg2rad(fovy)  # deg 2 rad
+        self.near = near
+        self.far = far
+        self.center = np.array([0, 0, 0], dtype=np.float32)  # look at this point
+        self.rot = R.from_matrix(np.eye(3))
+        self.up = np.array([0, 1, 0], dtype=np.float32)  # need to be normalized!
+
+    @property
+    def fovx(self):
+        return 2 * np.arctan(np.tan(self.fovy / 2) * self.W / self.H)
+
+    @property
+    def campos(self):
+        return self.pose[:3, 3]
+
+    # pose (c2w)
+    @property
+    def pose(self):
+        # first move camera to radius
+        res = np.eye(4, dtype=np.float32)
+        res[2, 3] = self.radius  # opengl convention...
+        # rotate
+        rot = np.eye(4, dtype=np.float32)
+        rot[:3, :3] = self.rot.as_matrix()
+        res = rot @ res
+        # translate
+        res[:3, 3] -= self.center
+        return res
+
+    # view (w2c)
+    @property
+    def view(self):
+        return np.linalg.inv(self.pose)
+
+    # projection (perspective)
+    @property
+    def perspective(self):
+        y = np.tan(self.fovy / 2)
+        aspect = self.W / self.H
+        return np.array(
+            [
+                [1 / (y * aspect), 0, 0, 0],
+                [0, -1 / y, 0, 0],
+                [
+                    0,
+                    0,
+                    -(self.far + self.near) / (self.far - self.near),
+                    -(2 * self.far * self.near) / (self.far - self.near),
+                ],
+                [0, 0, -1, 0],
+            ],
+            dtype=np.float32,
+        )
+
+    # intrinsics
+    @property
+    def intrinsics(self):
+        focal = self.H / (2 * np.tan(self.fovy / 2))
+        return np.array([focal, focal, self.W // 2, self.H // 2], dtype=np.float32)
+
+    @property
+    def mvp(self):
+        return self.perspective @ np.linalg.inv(self.pose)  # [4, 4]
+
+    def orbit(self, dx, dy):
+        # rotate along camera up/side axis!
+        side = self.rot.as_matrix()[:3, 0]
+        rotvec_x = self.up * np.radians(-0.05 * dx)
+        rotvec_y = side * np.radians(-0.05 * dy)
+        self.rot = R.from_rotvec(rotvec_x) * R.from_rotvec(rotvec_y) * self.rot
+
+    def scale(self, delta):
+        self.radius *= 1.1 ** (-delta)
+
+    def pan(self, dx, dy, dz=0):
+        # pan in camera coordinate system (careful on the sensitivity!)
+        self.center += 0.0005 * self.rot.as_matrix()[:3, :3] @ np.array([-dx, -dy, dz])

configs/4d.yaml

@@ -0,0 +1,121 @@
+### Input
+# input rgba image path (default to None, can be load in GUI too)
+input: 
+# input text prompt (default to None, can be input in GUI too)
+prompt:
+# input mesh for stage 2 (auto-search from stage 1 output path if None)
+mesh:
+# estimated elevation angle for input image 
+elevation: 0
+# reference image resolution
+ref_size: 256
+# density thresh for mesh extraction
+density_thresh: 0.5
+
+### Output
+outdir: logs
+mesh_format: frames
+save_path: ''
+save_model: False
+
+### Training
+# guidance loss weights (0 to disable)
+mvdream: False
+imagedream: False
+lambda_sd: 0
+lambda_zero123: 1
+# use stable-zero123 instead of zero123-xl
+stable_zero123: True 
+lambda_svd: 0
+# training batch size per iter
+batch_size: 14
+# training iterations for stage 1
+iters: 500
+# training iterations for stage 2
+iters_refine: 50
+# training camera radius
+radius: 1.5
+# training camera fovy
+fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
+# training camera min elevation
+min_ver: -30
+# training camera max elevation
+max_ver: 30
+# checkpoint to load for stage 1 (should be a ply file)
+load:
+# whether allow geom training in stage 2
+train_geo: False
+# prob to invert background color during training (0 = always black, 1 = always white)
+invert_bg_prob: 0.
+n_views: 4
+t_max: 0.5
+
+
+### GUI
+gui: False
+force_cuda_rast: False
+# GUI resolution
+H: 800
+W: 800
+
+### Gaussian splatting
+optimize_gaussians: True
+position_lr_init: 0.001
+position_lr_final: 0.00002
+position_lr_delay_mult: 0.02
+position_lr_max_steps: 500
+feature_lr: 0.01
+opacity_lr: 0.05
+scaling_lr: 0.005
+rotation_lr: 0.005
+
+num_pts: 5000
+sh_degree: 0
+percent_dense: 0.1
+density_start_iter: 3000
+density_end_iter: 3000
+densification_interval: 100
+opacity_reset_interval: 700
+densify_grad_threshold: 0.05
+
+# deformation field
+deformation_lr_init: 0.00064
+deformation_lr_final: 0.00064
+deformation_lr_delay_mult: 0.01
+grid_lr_init: 0.0064
+grid_lr_final: 0.0064
+
+### Textured Mesh
+geom_lr: 0.0001
+texture_lr: 0.2
+
+deformation:
+    net_width: 64
+    timebase_pe: 4
+    defor_depth: 1
+    posebase_pe: 10
+    scale_rotation_pe: 2
+    opacity_pe: 2
+    timenet_width: 64
+    timenet_output: 32
+    bounds: 1.6
+    plane_tv_weight: 0.0001
+    time_smoothness_weight: 0.01
+    l1_time_planes: 0.0001
+    kplanes_config:
+        grid_dimensions: 2
+        input_coordinate_dim: 4
+        output_coordinate_dim: 32
+        resolution: [32, 32, 32, 12]
+    multires: [1]
+    no_grid: False
+    no_mlp: False
+    no_ds: False
+    no_dr: False
+    no_do: True
+    use_res: True
+
+data_mode: svd
+downsample_rate: 1
+# data_mode: c4d
+# downsample_rate: 2

configs/4d_c4d.yaml

@@ -0,0 +1,119 @@
+### Input
+# input rgba image path (default to None, can be load in GUI too)
+input: 
+# input text prompt (default to None, can be input in GUI too)
+prompt:
+# input mesh for stage 2 (auto-search from stage 1 output path if None)
+mesh:
+# estimated elevation angle for input image 
+elevation: 0
+# reference image resolution
+ref_size: 256
+# density thresh for mesh extraction
+density_thresh: 0.5
+
+### Output
+outdir: logs
+mesh_format: frames
+save_path: ''
+save_model: False
+
+### Training
+# guidance loss weights (0 to disable)
+mvdream: False
+imagedream: False
+lambda_sd: 0
+lambda_zero123: 1
+# use stable-zero123 instead of zero123-xl
+stable_zero123: True 
+lambda_svd: 0
+# training batch size per iter
+batch_size: 32
+# training iterations for stage 1
+iters: 500
+# training iterations for stage 2
+iters_refine: 50
+# training camera radius
+radius: 1.5
+# training camera fovy
+fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
+# training camera min elevation
+min_ver: -30
+# training camera max elevation
+max_ver: 30
+# checkpoint to load for stage 1 (should be a ply file)
+load:
+# whether allow geom training in stage 2
+train_geo: False
+# prob to invert background color during training (0 = always black, 1 = always white)
+invert_bg_prob: 0.
+n_views: 4
+t_max: 0.5
+
+
+### GUI
+gui: False
+force_cuda_rast: False
+# GUI resolution
+H: 800
+W: 800
+
+### Gaussian splatting
+optimize_gaussians: True
+position_lr_init: 0.001
+position_lr_final: 0.00002
+position_lr_delay_mult: 0.02
+position_lr_max_steps: 500
+feature_lr: 0.01
+opacity_lr: 0.05
+scaling_lr: 0.005
+rotation_lr: 0.005
+
+num_pts: 5000
+sh_degree: 0
+percent_dense: 0.1
+density_start_iter: 3000
+density_end_iter: 3000
+densification_interval: 100
+opacity_reset_interval: 700
+densify_grad_threshold: 0.05
+
+# deformation field
+deformation_lr_init: 0.00064
+deformation_lr_final: 0.00064
+deformation_lr_delay_mult: 0.01
+grid_lr_init: 0.0064
+grid_lr_final: 0.0064
+
+### Textured Mesh
+geom_lr: 0.0001
+texture_lr: 0.2
+
+deformation:
+    net_width: 64
+    timebase_pe: 4
+    defor_depth: 1
+    posebase_pe: 10
+    scale_rotation_pe: 2
+    opacity_pe: 2
+    timenet_width: 64
+    timenet_output: 32
+    bounds: 1.6
+    plane_tv_weight: 0.0001
+    time_smoothness_weight: 0.01
+    l1_time_planes: 0.0001
+    kplanes_config:
+        grid_dimensions: 2
+        input_coordinate_dim: 4
+        output_coordinate_dim: 32
+        resolution: [32, 32, 32, 32]
+    multires: [1]
+    no_grid: False
+    no_mlp: False
+    no_ds: False
+    no_dr: False
+    no_do: True
+    use_res: True
+
+data_mode: c4d
+downsample_rate: 1

configs/4d_c4d_low.yaml

@@ -0,0 +1,119 @@
+### Input
+# input rgba image path (default to None, can be load in GUI too)
+input: 
+# input text prompt (default to None, can be input in GUI too)
+prompt:
+# input mesh for stage 2 (auto-search from stage 1 output path if None)
+mesh:
+# estimated elevation angle for input image 
+elevation: 0
+# reference image resolution
+ref_size: 256
+# density thresh for mesh extraction
+density_thresh: 0.5
+
+### Output
+outdir: logs
+mesh_format: frames
+save_path: ''
+save_model: False
+
+### Training
+# guidance loss weights (0 to disable)
+mvdream: False
+imagedream: False
+lambda_sd: 0
+lambda_zero123: 1
+# use stable-zero123 instead of zero123-xl
+stable_zero123: True 
+lambda_svd: 0
+# training batch size per iter
+batch_size: 8
+# training iterations for stage 1
+iters: 500
+# training iterations for stage 2
+iters_refine: 50
+# training camera radius
+radius: 1.5
+# training camera fovy
+fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
+# training camera min elevation
+min_ver: -30
+# training camera max elevation
+max_ver: 30
+# checkpoint to load for stage 1 (should be a ply file)
+load:
+# whether allow geom training in stage 2
+train_geo: False
+# prob to invert background color during training (0 = always black, 1 = always white)
+invert_bg_prob: 0.
+n_views: 1
+t_max: 0.5
+
+
+### GUI
+gui: False
+force_cuda_rast: False
+# GUI resolution
+H: 800
+W: 800
+
+### Gaussian splatting
+optimize_gaussians: True
+position_lr_init: 0.001
+position_lr_final: 0.00002
+position_lr_delay_mult: 0.02
+position_lr_max_steps: 500
+feature_lr: 0.01
+opacity_lr: 0.05
+scaling_lr: 0.005
+rotation_lr: 0.005
+
+num_pts: 5000
+sh_degree: 0
+percent_dense: 0.1
+density_start_iter: 3000
+density_end_iter: 3000
+densification_interval: 100
+opacity_reset_interval: 700
+densify_grad_threshold: 0.05
+
+# deformation field
+deformation_lr_init: 0.00064
+deformation_lr_final: 0.00064
+deformation_lr_delay_mult: 0.01
+grid_lr_init: 0.0064
+grid_lr_final: 0.0064
+
+### Textured Mesh
+geom_lr: 0.0001
+texture_lr: 0.2
+
+deformation:
+    net_width: 64
+    timebase_pe: 4
+    defor_depth: 1
+    posebase_pe: 10
+    scale_rotation_pe: 2
+    opacity_pe: 2
+    timenet_width: 64
+    timenet_output: 32
+    bounds: 1.6
+    plane_tv_weight: 0.0001
+    time_smoothness_weight: 0.01
+    l1_time_planes: 0.0001
+    kplanes_config:
+        grid_dimensions: 2
+        input_coordinate_dim: 4
+        output_coordinate_dim: 32
+        resolution: [32, 32, 32, 12]
+    multires: [1]
+    no_grid: False
+    no_mlp: False
+    no_ds: False
+    no_dr: False
+    no_do: True
+    use_res: True
+
+data_mode: c4d
+downsample_rate: 4

configs/4d_demo.yaml

@@ -0,0 +1,121 @@
+### Input
+# input rgba image path (default to None, can be load in GUI too)
+input: 
+# input text prompt (default to None, can be input in GUI too)
+prompt:
+# input mesh for stage 2 (auto-search from stage 1 output path if None)
+mesh:
+# estimated elevation angle for input image 
+elevation: 0
+# reference image resolution
+ref_size: 256
+# density thresh for mesh extraction
+density_thresh: 0.5
+
+### Output
+outdir: logs
+mesh_format: frames
+save_path: ''
+save_model: False
+
+### Training
+# guidance loss weights (0 to disable)
+mvdream: False
+imagedream: False
+lambda_sd: 0
+lambda_zero123: 1
+# use stable-zero123 instead of zero123-xl
+stable_zero123: True 
+lambda_svd: 0
+# training batch size per iter
+batch_size: 7
+# training iterations for stage 1
+iters: 500
+# training iterations for stage 2
+iters_refine: 50
+# training camera radius
+radius: 1.5
+# training camera fovy
+fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
+# training camera min elevation
+min_ver: -30
+# training camera max elevation
+max_ver: 30
+# checkpoint to load for stage 1 (should be a ply file)
+load:
+# whether allow geom training in stage 2
+train_geo: False
+# prob to invert background color during training (0 = always black, 1 = always white)
+invert_bg_prob: 0.
+n_views: 1
+t_max: 0.5
+
+
+### GUI
+gui: False
+force_cuda_rast: False
+# GUI resolution
+H: 800
+W: 800
+
+### Gaussian splatting
+optimize_gaussians: True
+position_lr_init: 0.001
+position_lr_final: 0.00002
+position_lr_delay_mult: 0.02
+position_lr_max_steps: 500
+feature_lr: 0.01
+opacity_lr: 0.05
+scaling_lr: 0.005
+rotation_lr: 0.005
+
+num_pts: 5000
+sh_degree: 0
+percent_dense: 0.1
+density_start_iter: 3000
+density_end_iter: 3000
+densification_interval: 100
+opacity_reset_interval: 700
+densify_grad_threshold: 0.05
+
+# deformation field
+deformation_lr_init: 0.00064
+deformation_lr_final: 0.00064
+deformation_lr_delay_mult: 0.01
+grid_lr_init: 0.0064
+grid_lr_final: 0.0064
+
+### Textured Mesh
+geom_lr: 0.0001
+texture_lr: 0.2
+
+deformation:
+    net_width: 64
+    timebase_pe: 4
+    defor_depth: 1
+    posebase_pe: 10
+    scale_rotation_pe: 2
+    opacity_pe: 2
+    timenet_width: 64
+    timenet_output: 32
+    bounds: 1.6
+    plane_tv_weight: 0.0001
+    time_smoothness_weight: 0.01
+    l1_time_planes: 0.0001
+    kplanes_config:
+        grid_dimensions: 2
+        input_coordinate_dim: 4
+        output_coordinate_dim: 32
+        resolution: [32, 32, 32, 12]
+    multires: [1]
+    no_grid: False
+    no_mlp: False
+    no_ds: False
+    no_dr: False
+    no_do: True
+    use_res: True
+
+data_mode: svd
+downsample_rate: 2
+# data_mode: c4d
+# downsample_rate: 2

configs/4d_low.yaml

@@ -0,0 +1,121 @@
+### Input
+# input rgba image path (default to None, can be load in GUI too)
+input: 
+# input text prompt (default to None, can be input in GUI too)
+prompt:
+# input mesh for stage 2 (auto-search from stage 1 output path if None)
+mesh:
+# estimated elevation angle for input image 
+elevation: 0
+# reference image resolution
+ref_size: 256
+# density thresh for mesh extraction
+density_thresh: 0.5
+
+### Output
+outdir: logs
+mesh_format: frames
+save_path: ''
+save_model: False
+
+### Training
+# guidance loss weights (0 to disable)
+mvdream: False
+imagedream: False
+lambda_sd: 0
+lambda_zero123: 1
+# use stable-zero123 instead of zero123-xl
+stable_zero123: True 
+lambda_svd: 0
+# training batch size per iter
+batch_size: 14
+# training iterations for stage 1
+iters: 500
+# training iterations for stage 2
+iters_refine: 50
+# training camera radius
+radius: 1.5
+# training camera fovy
+fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
+# training camera min elevation
+min_ver: -30
+# training camera max elevation
+max_ver: 30
+# checkpoint to load for stage 1 (should be a ply file)
+load:
+# whether allow geom training in stage 2
+train_geo: False
+# prob to invert background color during training (0 = always black, 1 = always white)
+invert_bg_prob: 0.
+n_views: 1
+t_max: 0.5
+
+
+### GUI
+gui: False
+force_cuda_rast: False
+# GUI resolution
+H: 800
+W: 800
+
+### Gaussian splatting
+optimize_gaussians: True
+position_lr_init: 0.001
+position_lr_final: 0.00002
+position_lr_delay_mult: 0.02
+position_lr_max_steps: 500
+feature_lr: 0.01
+opacity_lr: 0.05
+scaling_lr: 0.005
+rotation_lr: 0.005
+
+num_pts: 5000
+sh_degree: 0
+percent_dense: 0.1
+density_start_iter: 3000
+density_end_iter: 3000
+densification_interval: 100
+opacity_reset_interval: 700
+densify_grad_threshold: 0.05
+
+# deformation field
+deformation_lr_init: 0.00064
+deformation_lr_final: 0.00064
+deformation_lr_delay_mult: 0.01
+grid_lr_init: 0.0064
+grid_lr_final: 0.0064
+
+### Textured Mesh
+geom_lr: 0.0001
+texture_lr: 0.2
+
+deformation:
+    net_width: 64
+    timebase_pe: 4
+    defor_depth: 1
+    posebase_pe: 10
+    scale_rotation_pe: 2
+    opacity_pe: 2
+    timenet_width: 64
+    timenet_output: 32
+    bounds: 1.6
+    plane_tv_weight: 0.0001
+    time_smoothness_weight: 0.01
+    l1_time_planes: 0.0001
+    kplanes_config:
+        grid_dimensions: 2
+        input_coordinate_dim: 4
+        output_coordinate_dim: 32
+        resolution: [32, 32, 32, 22]
+    multires: [1]
+    no_grid: False
+    no_mlp: False
+    no_ds: False
+    no_dr: False
+    no_do: True
+    use_res: True
+
+data_mode: svd
+downsample_rate: 1
+# data_mode: c4d
+# downsample_rate: 2

configs/dg.yaml

@@ -0,0 +1,85 @@
+### Input
+# input rgba image path (default to None, can be load in GUI too)
+input: 
+# input text prompt (default to None, can be input in GUI too)
+prompt:
+negative_prompt:
+# input mesh for stage 2 (auto-search from stage 1 output path if None)
+mesh:
+# estimated elevation angle for input image 
+elevation: 0
+# reference image resolution
+ref_size: 256
+# density thresh for mesh extraction
+density_thresh: 1
+
+### Output
+outdir: logs
+mesh_format: obj
+save_path: ''
+
+### Training
+# use mvdream instead of sd 2.1
+mvdream: False
+# use imagedream
+imagedream: False
+# use stable-zero123 instead of zero123-xl
+stable_zero123: False 
+# guidance loss weights (0 to disable)
+lambda_sd: 0
+lambda_zero123: 1
+# warmup rgb supervision for image-to-3d
+warmup_rgb_loss: True
+# training batch size per iter
+batch_size: 1
+# training iterations for stage 1
+iters: 500
+# whether to linearly anneal timestep
+anneal_timestep: True
+# training iterations for stage 2
+iters_refine: 50
+# training camera radius
+radius: 2
+# training camera fovy
+fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
+# training camera min elevation
+min_ver: -30
+# training camera max elevation
+max_ver: 30
+# checkpoint to load for stage 1 (should be a ply file)
+load:
+# whether allow geom training in stage 2
+train_geo: False
+# prob to invert background color during training (0 = always black, 1 = always white)
+invert_bg_prob: 0.5
+
+
+### GUI
+gui: False
+force_cuda_rast: False
+# GUI resolution
+H: 800
+W: 800
+
+### Gaussian splatting
+num_pts: 5000
+sh_degree: 0
+position_lr_init: 0.001
+position_lr_final: 0.00002
+position_lr_delay_mult: 0.02
+position_lr_max_steps: 500
+feature_lr: 0.01
+opacity_lr: 0.05
+scaling_lr: 0.005
+rotation_lr: 0.005
+percent_dense: 0.1
+density_start_iter: 100
+density_end_iter: 3000
+densification_interval: 100
+opacity_reset_interval: 700
+densify_grad_threshold: 0.5
+
+
+### Textured Mesh
+geom_lr: 0.0001
+texture_lr: 0.2

configs/dghd.yaml

@@ -0,0 +1,72 @@
+### Input
+# input rgba image path (default to None, can be load in GUI too)
+input: 
+# input text prompt (default to None, can be input in GUI too)
+prompt:
+# input mesh for stage 2 (auto-search from stage 1 output path if None)
+mesh:
+# estimated elevation angle for input image 
+elevation: 0
+# reference image resolution
+ref_size: 256
+# density thresh for mesh extraction
+density_thresh: 1
+
+### Output
+outdir: logs
+mesh_format: obj
+save_path: ''
+
+### Training
+# guidance loss weights (0 to disable)
+lambda_sd: 0
+mvdream: False
+lambda_zero123: 1
+# use stable-zero123 instead of zero123-xl
+stable_zero123: False 
+# training batch size per iter
+batch_size: 16
+# training iterations for stage 1
+iters: 500
+# training iterations for stage 2
+iters_refine: 50
+# training camera radius
+radius: 2
+# training camera fovy
+fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
+# checkpoint to load for stage 1 (should be a ply file)
+load:
+# whether allow geom training in stage 2
+train_geo: False
+# prob to invert background color during training (0 = always black, 1 = always white)
+invert_bg_prob: 0.
+
+
+### GUI
+gui: False
+force_cuda_rast: False
+# GUI resolution
+H: 800
+W: 800
+
+### Gaussian splatting
+num_pts: 5000
+sh_degree: 0
+position_lr_init: 0.001
+position_lr_final: 0.00002
+position_lr_delay_mult: 0.02
+position_lr_max_steps: 500
+feature_lr: 0.01
+opacity_lr: 0.05
+scaling_lr: 0.005
+rotation_lr: 0.005
+percent_dense: 0.1
+density_start_iter: 0
+density_end_iter: 3000
+densification_interval: 100
+opacity_reset_interval: 700
+densify_grad_threshold: 0.05
+
+### Textured Mesh
+geom_lr: 0.0001
+texture_lr: 0.2

configs/refine.yaml

@@ -0,0 +1,79 @@
+### Input
+# input rgba image path (default to None, can be load in GUI too)
+input: 
+# input text prompt (default to None, can be input in GUI too)
+prompt:
+# input mesh for stage 2 (auto-search from stage 1 output path if None)
+mesh:
+# estimated elevation angle for input image 
+elevation: 0
+# reference image resolution
+ref_size: 256
+# density thresh for mesh extraction
+density_thresh: 1
+
+### Output
+outdir: logs
+mesh_format: obj
+save_path: ''
+
+### Training
+# guidance loss weights (0 to disable)
+lambda_sd: 0
+mvdream: False
+lambda_zero123: 0
+# use stable-zero123 instead of zero123-xl
+stable_zero123: False 
+lambda_svd: 1
+# training batch size per iter
+batch_size: 1
+# training iterations for stage 1
+iters: 500
+# training iterations for stage 2
+iters_refine: 50
+# training camera radius
+radius: 1.5
+# training camera fovy
+fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
+# checkpoint to load for stage 1 (should be a ply file)
+load:
+# whether allow geom training in stage 2
+train_geo: False
+# prob to invert background color during training (0 = always black, 1 = always white)
+invert_bg_prob: 0.5
+
+
+### GUI
+gui: False
+force_cuda_rast: False
+# GUI resolution
+H: 800
+W: 800
+
+### Gaussian splatting
+num_pts: 5000
+sh_degree: 0
+position_lr_init: 0.001
+position_lr_final: 0.00002
+position_lr_delay_mult: 0.02
+position_lr_max_steps: 500
+feature_lr: 0.01
+opacity_lr: 0.05
+scaling_lr: 0.005
+rotation_lr: 0.005
+percent_dense: 0.1
+density_start_iter: 100
+density_end_iter: 3000
+densification_interval: 100
+opacity_reset_interval: 700
+densify_grad_threshold: 0.5
+
+### Textured Mesh
+geom_lr: 0.0001
+texture_lr: 0.2
+
+static_model: lgm
+data_mode: svd
+downsample_rate: 2
+
+oom_hack: False

gaussian_model_4d.py

@@ -0,0 +1,773 @@
+#
+# Copyright (C) 2023, Inria
+# GRAPHDECO research group, https://team.inria.fr/graphdeco
+# All rights reserved.
+#
+# This software is free for non-commercial, research and evaluation use 
+# under the terms of the LICENSE.md file.
+#
+# For inquiries contact  [email protected]
+#
+
+import torch
+import numpy as np
+from utils.general_utils import inverse_sigmoid, get_expon_lr_func, build_rotation
+from torch import nn
+import os
+from utils.system_utils import mkdir_p
+from plyfile import PlyData, PlyElement
+from random import randint
+from utils.sh_utils import RGB2SH
+from simple_knn._C import distCUDA2
+from utils.graphics_utils import BasicPointCloud
+from utils.general_utils import strip_symmetric, build_scaling_rotation
+from scene.deformation import deform_network
+from scene.regulation import compute_plane_smoothness
+
+
+def gaussian_3d_coeff(xyzs, covs):
+    # xyzs: [N, 3]
+    # covs: [N, 6]
+    x, y, z = xyzs[:, 0], xyzs[:, 1], xyzs[:, 2]
+    a, b, c, d, e, f = covs[:, 0], covs[:, 1], covs[:, 2], covs[:, 3], covs[:, 4], covs[:, 5]
+
+    # eps must be small enough !!!
+    inv_det = 1 / (a * d * f + 2 * e * c * b - e**2 * a - c**2 * d - b**2 * f + 1e-24)
+    inv_a = (d * f - e**2) * inv_det
+    inv_b = (e * c - b * f) * inv_det
+    inv_c = (e * b - c * d) * inv_det
+    inv_d = (a * f - c**2) * inv_det
+    inv_e = (b * c - e * a) * inv_det
+    inv_f = (a * d - b**2) * inv_det
+
+    power = -0.5 * (x**2 * inv_a + y**2 * inv_d + z**2 * inv_f) - x * y * inv_b - x * z * inv_c - y * z * inv_e
+
+    power[power > 0] = -1e10 # abnormal values... make weights 0
+        
+    return torch.exp(power)
+
+class GaussianModel:
+
+    def setup_functions(self):
+        def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation):
+            L = build_scaling_rotation(scaling_modifier * scaling, rotation)
+            actual_covariance = L @ L.transpose(1, 2)
+            symm = strip_symmetric(actual_covariance)
+            return symm
+        
+        self.scaling_activation = torch.exp
+        self.scaling_inverse_activation = torch.log
+
+        self.covariance_activation = build_covariance_from_scaling_rotation
+
+        self.opacity_activation = torch.sigmoid
+        self.inverse_opacity_activation = inverse_sigmoid
+
+        self.rotation_activation = torch.nn.functional.normalize
+
+
+    def __init__(self, sh_degree : int, args):
+        self.active_sh_degree = 0
+        self.max_sh_degree = sh_degree  
+        self._xyz = torch.empty(0)
+        # self._deformation =  torch.empty(0)
+        self._deformation = deform_network(args)
+        # self.grid = TriPlaneGrid()
+        self._features_dc = torch.empty(0)
+        self._features_rest = torch.empty(0)
+        self._scaling = torch.empty(0)
+        self._rotation = torch.empty(0)
+        self._opacity = torch.empty(0)
+        self.max_radii2D = torch.empty(0)
+        self.xyz_gradient_accum = torch.empty(0)
+        self.denom = torch.empty(0)
+        self.optimizer = None
+        self.percent_dense = 0
+        self.spatial_lr_scale = 0
+        self._deformation_table = torch.empty(0)
+        self.setup_functions()
+
+    def capture(self):
+        return (
+            self.active_sh_degree,
+            self._xyz,
+            self._deformation.state_dict(),
+            self._deformation_table,
+            # self.grid,
+            self._features_dc,
+            self._features_rest,
+            self._scaling,
+            self._rotation,
+            self._opacity,
+            self.max_radii2D,
+            self.xyz_gradient_accum,
+            self.denom,
+            self.optimizer.state_dict(),
+            self.spatial_lr_scale,
+        )
+    
+    def restore(self, model_args, training_args):
+        (self.active_sh_degree, 
+        self._xyz, 
+        self._deformation_table,
+        self._deformation,
+        # self.grid,
+        self._features_dc, 
+        self._features_rest,
+        self._scaling, 
+        self._rotation, 
+        self._opacity,
+        self.max_radii2D, 
+        xyz_gradient_accum, 
+        denom,
+        opt_dict, 
+        self.spatial_lr_scale) = model_args
+        self.training_setup(training_args)
+        self.xyz_gradient_accum = xyz_gradient_accum
+        self.denom = denom
+        self.optimizer.load_state_dict(opt_dict)
+
+    @property
+    def get_scaling(self):
+        return self.scaling_activation(self._scaling)
+    
+    @property
+    def get_rotation(self):
+        return self.rotation_activation(self._rotation)
+    
+    @property
+    def get_xyz(self):
+        return self._xyz
+
+    @property
+    def get_features(self):
+        features_dc = self._features_dc
+        features_rest = self._features_rest
+        return torch.cat((features_dc, features_rest), dim=1)
+    
+    @property
+    def get_opacity(self):
+        return self.opacity_activation(self._opacity)
+
+
+    def get_deformed_everything(self, time):
+        means3D = self.get_xyz
+        time = torch.tensor(time).to(means3D.device).repeat(means3D.shape[0],1)
+        time = ((time.float() / self.T) - 0.5) * 2
+
+        opacity = self._opacity
+        scales = self._scaling
+        rotations = self._rotation
+
+        deformation_point = self._deformation_table
+        means3D_deform, scales_deform, rotations_deform, opacity_deform = self._deformation(means3D[deformation_point], scales[deformation_point], 
+                                                    rotations[deformation_point], opacity[deformation_point],
+                                                    time[deformation_point])
+
+        means3D_final =  means3D + means3D_deform
+        rotations_final =  rotations + rotations_deform
+        scales_final =  scales + scales_deform
+        opacity_final = opacity
+
+        return means3D_final, rotations_final, scales_final, opacity_final
+
+
+    
+    @torch.no_grad()
+    def extract_fields_t(self, resolution=128, num_blocks=16, relax_ratio=1.5, t=0):
+        # resolution: resolution of field
+        
+        block_size = 2 / num_blocks
+
+        assert resolution % block_size == 0
+        split_size = resolution // num_blocks
+
+        xyzs, rotation, scale, opacities = self.get_deformed_everything(t)
+
+        scale = self.scaling_activation(scale)
+        opacities = self.opacity_activation(opacities)
+
+        # pre-filter low opacity gaussians to save computation
+        mask = (opacities > 0.005).squeeze(1)
+
+        opacities = opacities[mask]
+        xyzs = xyzs[mask]
+        stds = scale[mask]
+        
+        # normalize to ~ [-1, 1]
+        mn, mx = xyzs.amin(0), xyzs.amax(0)
+        self.center = (mn + mx) / 2
+        self.scale = 1.8 / (mx - mn).amax().item()
+
+        xyzs = (xyzs - self.center) * self.scale
+        stds = stds * self.scale
+
+        covs = self.covariance_activation(stds, 1, rotation[mask])
+
+        # tile
+        device = opacities.device
+        occ = torch.zeros([resolution] * 3, dtype=torch.float32, device=device)
+
+        X = torch.linspace(-1, 1, resolution).split(split_size)
+        Y = torch.linspace(-1, 1, resolution).split(split_size)
+        Z = torch.linspace(-1, 1, resolution).split(split_size)
+
+
+        # loop blocks (assume max size of gaussian is small than relax_ratio * block_size !!!)
+        for xi, xs in enumerate(X):
+            for yi, ys in enumerate(Y):
+                for zi, zs in enumerate(Z):
+                    xx, yy, zz = torch.meshgrid(xs, ys, zs)
+                    # sample points [M, 3]
+                    pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1).to(device)
+                    # in-tile gaussians mask
+                    vmin, vmax = pts.amin(0), pts.amax(0)
+                    vmin -= block_size * relax_ratio
+                    vmax += block_size * relax_ratio
+                    mask = (xyzs < vmax).all(-1) & (xyzs > vmin).all(-1)
+                    # if hit no gaussian, continue to next block
+                    if not mask.any():
+                        continue
+                    mask_xyzs = xyzs[mask] # [L, 3]
+                    mask_covs = covs[mask] # [L, 6]
+                    mask_opas = opacities[mask].view(1, -1) # [L, 1] --> [1, L]
+
+                    # query per point-gaussian pair.
+                    g_pts = pts.unsqueeze(1).repeat(1, mask_covs.shape[0], 1) - mask_xyzs.unsqueeze(0) # [M, L, 3]
+                    g_covs = mask_covs.unsqueeze(0).repeat(pts.shape[0], 1, 1) # [M, L, 6]
+
+                    # batch on gaussian to avoid OOM
+                    batch_g = 1024
+                    val = 0
+                    for start in range(0, g_covs.shape[1], batch_g):
+                        end = min(start + batch_g, g_covs.shape[1])
+                        w = gaussian_3d_coeff(g_pts[:, start:end].reshape(-1, 3), g_covs[:, start:end].reshape(-1, 6)).reshape(pts.shape[0], -1) # [M, l]
+                        val += (mask_opas[:, start:end] * w).sum(-1)
+                    
+                    # kiui.lo(val, mask_opas, w)
+                
+                    occ[xi * split_size: xi * split_size + len(xs), 
+                        yi * split_size: yi * split_size + len(ys), 
+                        zi * split_size: zi * split_size + len(zs)] = val.reshape(len(xs), len(ys), len(zs)) 
+        return occ
+
+    def extract_mesh_t(self, path, density_thresh=1, t=0, resolution=128, decimate_target=1e5):
+        from mesh import Mesh
+        from mesh_utils import decimate_mesh, clean_mesh
+
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+
+        occ = self.extract_fields_t(resolution, t=t).detach().cpu().numpy()
+
+        import mcubes
+        vertices, triangles = mcubes.marching_cubes(occ, density_thresh)
+        vertices = vertices / (resolution - 1.0) * 2 - 1
+
+        # transform back to the original space
+        vertices = vertices / self.scale + self.center.detach().cpu().numpy()
+
+        vertices, triangles = clean_mesh(vertices, triangles, remesh=True, remesh_size=0.015)
+        if decimate_target > 0 and triangles.shape[0] > decimate_target:
+            vertices, triangles = decimate_mesh(vertices, triangles, decimate_target)
+
+        v = torch.from_numpy(vertices.astype(np.float32)).contiguous().cuda()
+        f = torch.from_numpy(triangles.astype(np.int32)).contiguous().cuda()
+
+        print(
+            f"[INFO] marching cubes result: {v.shape} ({v.min().item()}-{v.max().item()}), {f.shape}"
+        )
+
+        mesh = Mesh(v=v, f=f, device='cuda')
+
+        return mesh
+    
+    def get_covariance(self, scaling_modifier = 1):
+        return self.covariance_activation(self.get_scaling, scaling_modifier, self._rotation)
+
+    def oneupSHdegree(self):
+        if self.active_sh_degree < self.max_sh_degree:
+            self.active_sh_degree += 1
+
+    def create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : float, time_line: int):
+        self.spatial_lr_scale = spatial_lr_scale
+        fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda()
+        fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda())
+        features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda()
+        features[:, :3, 0 ] = fused_color
+        features[:, 3:, 1:] = 0.0
+
+        print("Number of points at initialisation : ", fused_point_cloud.shape[0])
+
+        dist2 = torch.clamp_min(distCUDA2(torch.from_numpy(np.asarray(pcd.points)).float().cuda()), 0.0000001)
+        scales = torch.log(torch.sqrt(dist2))[...,None].repeat(1, 3)
+        rots = torch.zeros((fused_point_cloud.shape[0], 4), device="cuda")
+        rots[:, 0] = 1
+
+        opacities = inverse_sigmoid(0.1 * torch.ones((fused_point_cloud.shape[0], 1), dtype=torch.float, device="cuda"))
+
+        self._xyz = nn.Parameter(fused_point_cloud.requires_grad_(True))
+        self._deformation = self._deformation.to("cuda") 
+        # self.grid = self.grid.to("cuda")
+        self._features_dc = nn.Parameter(features[:,:,0:1].transpose(1, 2).contiguous().requires_grad_(True))
+        self._features_rest = nn.Parameter(features[:,:,1:].transpose(1, 2).contiguous().requires_grad_(True))
+        self._scaling = nn.Parameter(scales.requires_grad_(True))
+        self._rotation = nn.Parameter(rots.requires_grad_(True))
+        self._opacity = nn.Parameter(opacities.requires_grad_(True))
+        self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
+        self._deformation_table = torch.gt(torch.ones((self.get_xyz.shape[0]),device="cuda"),0)
+
+    def training_setup(self, training_args):
+        self.percent_dense = training_args.percent_dense
+        self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+        self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+        self._deformation_accum = torch.zeros((self.get_xyz.shape[0],3),device="cuda")
+        self.T = training_args.batch_size
+        
+        if training_args.optimize_gaussians:
+            l = [
+                {'params': [self._xyz], 'lr': training_args.position_lr_init * self.spatial_lr_scale, "name": "xyz"},
+                {'params': list(self._deformation.get_mlp_parameters()), 'lr': training_args.deformation_lr_init * self.spatial_lr_scale, "name": "deformation"},
+                {'params': list(self._deformation.get_grid_parameters()), 'lr': training_args.grid_lr_init * self.spatial_lr_scale, "name": "grid"},
+                {'params': [self._features_dc], 'lr': training_args.feature_lr, "name": "f_dc"},
+                {'params': [self._features_rest], 'lr': training_args.feature_lr / 20.0, "name": "f_rest"},
+                {'params': [self._opacity], 'lr': training_args.opacity_lr, "name": "opacity"},
+                {'params': [self._scaling], 'lr': training_args.scaling_lr, "name": "scaling"},
+                {'params': [self._rotation], 'lr': training_args.rotation_lr, "name": "rotation"}
+            ]
+        else:
+            l = [
+                {'params': list(self._deformation.get_mlp_parameters()), 'lr': training_args.deformation_lr_init * self.spatial_lr_scale, "name": "deformation"},
+                {'params': list(self._deformation.get_grid_parameters()), 'lr': training_args.grid_lr_init * self.spatial_lr_scale, "name": "grid"},
+            ]
+
+        self.optimizer = torch.optim.Adam(l, lr=0.0, eps=1e-15)
+        self.xyz_scheduler_args = get_expon_lr_func(lr_init=training_args.position_lr_init*self.spatial_lr_scale,
+                                                    lr_final=training_args.position_lr_final*self.spatial_lr_scale,
+                                                    lr_delay_mult=training_args.position_lr_delay_mult,
+                                                    max_steps=training_args.position_lr_max_steps)
+        self.deformation_scheduler_args = get_expon_lr_func(lr_init=training_args.deformation_lr_init*self.spatial_lr_scale,
+                                                    lr_final=training_args.deformation_lr_final*self.spatial_lr_scale,
+                                                    lr_delay_mult=training_args.deformation_lr_delay_mult,
+                                                    max_steps=training_args.position_lr_max_steps)    
+        self.grid_scheduler_args = get_expon_lr_func(lr_init=training_args.grid_lr_init*self.spatial_lr_scale,
+                                                    lr_final=training_args.grid_lr_final*self.spatial_lr_scale,
+                                                    lr_delay_mult=training_args.deformation_lr_delay_mult,
+                                                    max_steps=training_args.position_lr_max_steps)    
+
+    def update_learning_rate(self, iteration):
+        ''' Learning rate scheduling per step '''
+        for param_group in self.optimizer.param_groups:
+            if param_group["name"] == "xyz":
+                lr = self.xyz_scheduler_args(iteration)
+                param_group['lr'] = lr
+                # return lr
+            if  "grid" in param_group["name"]:
+                lr = self.grid_scheduler_args(iteration)
+                param_group['lr'] = lr
+                # return lr
+            elif param_group["name"] == "deformation":
+                lr = self.deformation_scheduler_args(iteration)
+                param_group['lr'] = lr
+                # return lr
+
+    def construct_list_of_attributes(self):
+        l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
+        # All channels except the 3 DC
+        for i in range(self._features_dc.shape[1]*self._features_dc.shape[2]):
+            l.append('f_dc_{}'.format(i))
+        for i in range(self._features_rest.shape[1]*self._features_rest.shape[2]):
+            l.append('f_rest_{}'.format(i))
+        l.append('opacity')
+        for i in range(self._scaling.shape[1]):
+            l.append('scale_{}'.format(i))
+        for i in range(self._rotation.shape[1]):
+            l.append('rot_{}'.format(i))
+        return l
+    def compute_deformation(self,time):
+        
+        deform = self._deformation[:,:,:time].sum(dim=-1)
+        xyz = self._xyz + deform
+        return xyz
+
+    def load_model(self, path, name):
+        print("loading model from exists{}".format(path))
+        weight_dict = torch.load(os.path.join(path, name+"_deformation.pth"),map_location="cuda")
+        self._deformation.load_state_dict(weight_dict)
+        self._deformation = self._deformation.to("cuda")
+        self._deformation_table = torch.gt(torch.ones((self.get_xyz.shape[0]),device="cuda"),0)
+        self._deformation_accum = torch.zeros((self.get_xyz.shape[0],3),device="cuda")
+        if os.path.exists(os.path.join(path, name+"_deformation_table.pth")):
+            self._deformation_table = torch.load(os.path.join(path, name+"_deformation_table.pth"),map_location="cuda")
+        if os.path.exists(os.path.join(path,name+"_deformation_accum.pth")):
+            self._deformation_accum = torch.load(os.path.join(path, name+"_deformation_accum.pth"),map_location="cuda")
+        self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
+
+    def save_deformation(self, path, name):
+        torch.save(self._deformation.state_dict(),os.path.join(path, name+"_deformation.pth"))
+        torch.save(self._deformation_table,os.path.join(path, name+"_deformation_table.pth"))
+        torch.save(self._deformation_accum,os.path.join(path, name+"_deformation_accum.pth"))
+
+    def save_ply(self, path):
+        mkdir_p(os.path.dirname(path))
+
+        xyz = self._xyz.detach().cpu().numpy()
+        normals = np.zeros_like(xyz)
+        f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+        f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+        opacities = self._opacity.detach().cpu().numpy()
+        scale = self._scaling.detach().cpu().numpy()
+        rotation = self._rotation.detach().cpu().numpy()
+        
+        dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()]
+
+        elements = np.empty(xyz.shape[0], dtype=dtype_full)
+        attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1)
+        elements[:] = list(map(tuple, attributes))
+        el = PlyElement.describe(elements, 'vertex')
+        PlyData([el]).write(path)
+
+    def save_frame_ply(self, path, t):
+        mkdir_p(os.path.dirname(path))
+
+        xyzs, rotation, scale, opacities = self.get_deformed_everything(t)
+
+        xyz = xyzs.detach().cpu().numpy()
+        normals = np.zeros_like(xyz)
+        f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+        f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+        opacities = opacities.detach().cpu().numpy()
+        scale = scale.detach().cpu().numpy()
+        rotation = rotation.detach().cpu().numpy()
+        
+        dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()]
+
+        elements = np.empty(xyz.shape[0], dtype=dtype_full)
+        attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1)
+        elements[:] = list(map(tuple, attributes))
+        el = PlyElement.describe(elements, 'vertex')
+        PlyData([el]).write(path)
+    # def save_frame_ply(self, path, t):
+    #     mkdir_p(os.path.dirname(path))
+
+    #     xyz = self._xyz.detach().cpu().numpy()
+    #     normals = np.zeros_like(xyz)
+    #     f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+    #     f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+    #     opacities = self._opacity.detach().cpu().numpy()
+    #     scale = self._scaling.detach().cpu().numpy()
+    #     rotation = self._rotation.detach().cpu().numpy()
+        
+    #     dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()]
+
+    #     elements = np.empty(xyz.shape[0], dtype=dtype_full)
+    #     attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1)
+    #     elements[:] = list(map(tuple, attributes))
+    #     el = PlyElement.describe(elements, 'vertex')
+    #     PlyData([el]).write(path)
+        
+    def reset_opacity(self):
+        opacities_new = inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01))
+        optimizable_tensors = self.replace_tensor_to_optimizer(opacities_new, "opacity")
+        self._opacity = optimizable_tensors["opacity"]
+
+    def load_ply(self, path):
+        self.spatial_lr_scale = 1
+        plydata = PlyData.read(path)
+
+        xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
+                        np.asarray(plydata.elements[0]["y"]),
+                        np.asarray(plydata.elements[0]["z"])),  axis=1)
+        opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]
+
+        features_dc = np.zeros((xyz.shape[0], 3, 1))
+        features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
+        features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"])
+        features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"])
+
+        extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")]
+        extra_f_names = sorted(extra_f_names, key = lambda x: int(x.split('_')[-1]))
+        assert len(extra_f_names)==3*(self.max_sh_degree + 1) ** 2 - 3
+        features_extra = np.zeros((xyz.shape[0], len(extra_f_names)))
+        for idx, attr_name in enumerate(extra_f_names):
+            features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name])
+        # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC)
+        features_extra = features_extra.reshape((features_extra.shape[0], 3, (self.max_sh_degree + 1) ** 2 - 1))
+
+        scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
+        scale_names = sorted(scale_names, key = lambda x: int(x.split('_')[-1]))
+        scales = np.zeros((xyz.shape[0], len(scale_names)))
+        for idx, attr_name in enumerate(scale_names):
+            scales[:, idx] = np.asarray(plydata.elements[0][attr_name])
+
+        rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")]
+        rot_names = sorted(rot_names, key = lambda x: int(x.split('_')[-1]))
+        rots = np.zeros((xyz.shape[0], len(rot_names)))
+        for idx, attr_name in enumerate(rot_names):
+            rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
+
+        self._xyz = nn.Parameter(torch.tensor(xyz, dtype=torch.float, device="cuda").requires_grad_(True))
+        self._features_dc = nn.Parameter(torch.tensor(features_dc, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True))
+        self._features_rest = nn.Parameter(torch.tensor(features_extra, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True))
+        self._opacity = nn.Parameter(torch.tensor(opacities, dtype=torch.float, device="cuda").requires_grad_(True))
+        self._scaling = nn.Parameter(torch.tensor(scales, dtype=torch.float, device="cuda").requires_grad_(True))
+        self._rotation = nn.Parameter(torch.tensor(rots, dtype=torch.float, device="cuda").requires_grad_(True))
+        self.active_sh_degree = self.max_sh_degree
+        self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
+        self._deformation = self._deformation.to("cuda")
+        self._deformation_table = torch.gt(torch.ones((self.get_xyz.shape[0]),device="cuda"),0) # everything deformed
+
+        print(self._xyz.shape)
+
+
+    def replace_tensor_to_optimizer(self, tensor, name):
+        optimizable_tensors = {}
+        for group in self.optimizer.param_groups:
+            if group["name"] == name:
+                stored_state = self.optimizer.state.get(group['params'][0], None)
+                stored_state["exp_avg"] = torch.zeros_like(tensor)
+                stored_state["exp_avg_sq"] = torch.zeros_like(tensor)
+
+                del self.optimizer.state[group['params'][0]]
+                group["params"][0] = nn.Parameter(tensor.requires_grad_(True))
+                self.optimizer.state[group['params'][0]] = stored_state
+
+                optimizable_tensors[group["name"]] = group["params"][0]
+        return optimizable_tensors
+
+    def _prune_optimizer(self, mask):
+        optimizable_tensors = {}
+        for group in self.optimizer.param_groups:
+            if len(group["params"]) > 1:
+                continue
+            stored_state = self.optimizer.state.get(group['params'][0], None)
+            if stored_state is not None:
+                stored_state["exp_avg"] = stored_state["exp_avg"][mask]
+                stored_state["exp_avg_sq"] = stored_state["exp_avg_sq"][mask]
+
+                del self.optimizer.state[group['params'][0]]
+                group["params"][0] = nn.Parameter((group["params"][0][mask].requires_grad_(True)))
+                self.optimizer.state[group['params'][0]] = stored_state
+
+                optimizable_tensors[group["name"]] = group["params"][0]
+            else:
+                group["params"][0] = nn.Parameter(group["params"][0][mask].requires_grad_(True))
+                optimizable_tensors[group["name"]] = group["params"][0]
+        return optimizable_tensors
+
+    def prune_points(self, mask):
+        valid_points_mask = ~mask
+        optimizable_tensors = self._prune_optimizer(valid_points_mask)
+
+        self._xyz = optimizable_tensors["xyz"]
+        self._features_dc = optimizable_tensors["f_dc"]
+        self._features_rest = optimizable_tensors["f_rest"]
+        self._opacity = optimizable_tensors["opacity"]
+        self._scaling = optimizable_tensors["scaling"]
+        self._rotation = optimizable_tensors["rotation"]
+        self._deformation_accum = self._deformation_accum[valid_points_mask]
+        self.xyz_gradient_accum = self.xyz_gradient_accum[valid_points_mask]
+        self._deformation_table = self._deformation_table[valid_points_mask]
+        self.denom = self.denom[valid_points_mask]
+        self.max_radii2D = self.max_radii2D[valid_points_mask]
+
+    def cat_tensors_to_optimizer(self, tensors_dict):
+        optimizable_tensors = {}
+        for group in self.optimizer.param_groups:
+            if len(group["params"])>1:continue
+            assert len(group["params"]) == 1
+            extension_tensor = tensors_dict[group["name"]]
+            stored_state = self.optimizer.state.get(group['params'][0], None)
+            if stored_state is not None:
+
+                stored_state["exp_avg"] = torch.cat((stored_state["exp_avg"], torch.zeros_like(extension_tensor)), dim=0)
+                stored_state["exp_avg_sq"] = torch.cat((stored_state["exp_avg_sq"], torch.zeros_like(extension_tensor)), dim=0)
+
+                del self.optimizer.state[group['params'][0]]
+                group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))
+                self.optimizer.state[group['params'][0]] = stored_state
+
+                optimizable_tensors[group["name"]] = group["params"][0]
+            else:
+                group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))
+                optimizable_tensors[group["name"]] = group["params"][0]
+
+        return optimizable_tensors
+
+    def densification_postfix(self, new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation, new_deformation_table):
+        d = {"xyz": new_xyz,
+        "f_dc": new_features_dc,
+        "f_rest": new_features_rest,
+        "opacity": new_opacities,
+        "scaling" : new_scaling,
+        "rotation" : new_rotation,
+        # "deformation": new_deformation
+       }
+
+        optimizable_tensors = self.cat_tensors_to_optimizer(d)
+        self._xyz = optimizable_tensors["xyz"]
+        self._features_dc = optimizable_tensors["f_dc"]
+        self._features_rest = optimizable_tensors["f_rest"]
+        self._opacity = optimizable_tensors["opacity"]
+        self._scaling = optimizable_tensors["scaling"]
+        self._rotation = optimizable_tensors["rotation"]
+        # self._deformation = optimizable_tensors["deformation"]
+        
+        self._deformation_table = torch.cat([self._deformation_table,new_deformation_table],-1)
+        self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+        self._deformation_accum = torch.zeros((self.get_xyz.shape[0], 3), device="cuda")
+        self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+        self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
+
+    def densify_and_split(self, grads, grad_threshold, scene_extent, N=2):
+        n_init_points = self.get_xyz.shape[0]
+        # Extract points that satisfy the gradient condition
+        padded_grad = torch.zeros((n_init_points), device="cuda")
+        padded_grad[:grads.shape[0]] = grads.squeeze()
+        selected_pts_mask = torch.where(padded_grad >= grad_threshold, True, False)
+        selected_pts_mask = torch.logical_and(selected_pts_mask,
+                                              torch.max(self.get_scaling, dim=1).values > self.percent_dense*scene_extent)
+        if not selected_pts_mask.any():
+            return
+        stds = self.get_scaling[selected_pts_mask].repeat(N,1)
+        means =torch.zeros((stds.size(0), 3),device="cuda")
+        samples = torch.normal(mean=means, std=stds)
+        rots = build_rotation(self._rotation[selected_pts_mask]).repeat(N,1,1)
+        new_xyz = torch.bmm(rots, samples.unsqueeze(-1)).squeeze(-1) + self.get_xyz[selected_pts_mask].repeat(N, 1)
+        new_scaling = self.scaling_inverse_activation(self.get_scaling[selected_pts_mask].repeat(N,1) / (0.8*N))
+        new_rotation = self._rotation[selected_pts_mask].repeat(N,1)
+        new_features_dc = self._features_dc[selected_pts_mask].repeat(N,1,1)
+        new_features_rest = self._features_rest[selected_pts_mask].repeat(N,1,1)
+        new_opacity = self._opacity[selected_pts_mask].repeat(N,1)
+        new_deformation_table = self._deformation_table[selected_pts_mask].repeat(N)
+        self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacity, new_scaling, new_rotation, new_deformation_table)
+
+        prune_filter = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device="cuda", dtype=bool)))
+        self.prune_points(prune_filter)
+
+    def densify_and_clone(self, grads, grad_threshold, scene_extent):
+        # Extract points that satisfy the gradient condition
+        selected_pts_mask = torch.where(torch.norm(grads, dim=-1) >= grad_threshold, True, False)
+        selected_pts_mask = torch.logical_and(selected_pts_mask,
+                                              torch.max(self.get_scaling, dim=1).values <= self.percent_dense*scene_extent)
+        
+        new_xyz = self._xyz[selected_pts_mask] 
+        # - 0.001 * self._xyz.grad[selected_pts_mask]
+        new_features_dc = self._features_dc[selected_pts_mask]
+        new_features_rest = self._features_rest[selected_pts_mask]
+        new_opacities = self._opacity[selected_pts_mask]
+        new_scaling = self._scaling[selected_pts_mask]
+        new_rotation = self._rotation[selected_pts_mask]
+        new_deformation_table = self._deformation_table[selected_pts_mask]
+
+        self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation, new_deformation_table)
+        
+    def prune(self, min_opacity, extent, max_screen_size):
+        prune_mask = (self.get_opacity < min_opacity).squeeze()
+        # prune_mask_2 = torch.logical_and(self.get_opacity <= inverse_sigmoid(0.101 , dtype=torch.float, device="cuda"), self.get_opacity >= inverse_sigmoid(0.999 , dtype=torch.float, device="cuda"))
+        # prune_mask = torch.logical_or(prune_mask, prune_mask_2)
+        # deformation_sum = abs(self._deformation).sum(dim=-1).mean(dim=-1) 
+        # deformation_mask = (deformation_sum < torch.quantile(deformation_sum, torch.tensor([0.5]).to("cuda")))
+        # prune_mask = prune_mask & deformation_mask
+        if max_screen_size:
+            big_points_vs = self.max_radii2D > max_screen_size
+            big_points_ws = self.get_scaling.max(dim=1).values > 0.1 * extent
+            prune_mask = torch.logical_or(prune_mask, big_points_vs)
+
+            prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws)
+        self.prune_points(prune_mask)
+
+        torch.cuda.empty_cache()
+    def densify(self, max_grad, min_opacity, extent, max_screen_size):
+        grads = self.xyz_gradient_accum / self.denom
+        grads[grads.isnan()] = 0.0
+
+        self.densify_and_clone(grads, max_grad, extent)
+        self.densify_and_split(grads, max_grad, extent)
+    def standard_constaint(self):
+        
+        means3D = self._xyz.detach()
+        scales = self._scaling.detach()
+        rotations = self._rotation.detach()
+        opacity = self._opacity.detach()
+        time =  torch.tensor(0).to("cuda").repeat(means3D.shape[0],1)
+        means3D_deform, scales_deform, rotations_deform, _ = self._deformation(means3D, scales, rotations, opacity, time)
+        position_error = (means3D_deform - means3D)**2
+        rotation_error = (rotations_deform - rotations)**2 
+        scaling_erorr = (scales_deform - scales)**2
+        return position_error.mean() + rotation_error.mean() + scaling_erorr.mean()
+
+
+    def add_densification_stats(self, viewspace_point_tensor, update_filter):
+        self.xyz_gradient_accum[update_filter] += torch.norm(viewspace_point_tensor[update_filter,:2], dim=-1, keepdim=True)
+        self.denom[update_filter] += 1
+    @torch.no_grad()
+    def update_deformation_table(self,threshold):
+        # print("origin deformation point nums:",self._deformation_table.sum())
+        self._deformation_table = torch.gt(self._deformation_accum.max(dim=-1).values/100,threshold)
+    def print_deformation_weight_grad(self):
+        for name, weight in self._deformation.named_parameters():
+            if weight.requires_grad:
+                if weight.grad is None:
+                    
+                    print(name," :",weight.grad)
+                else:
+                    if weight.grad.mean() != 0:
+                        print(name," :",weight.grad.mean(), weight.grad.min(), weight.grad.max())
+        print("-"*50)
+    def _plane_regulation(self):
+        multi_res_grids = self._deformation.deformation_net.grid.grids
+        total = 0
+        # model.grids is 6 x [1, rank * F_dim, reso, reso]
+        for grids in multi_res_grids:
+            if len(grids) == 3:
+                time_grids = []
+            else:
+                time_grids =  [0,1,3]
+            for grid_id in time_grids:
+                total += compute_plane_smoothness(grids[grid_id])
+        return total
+    def _time_regulation(self):
+        multi_res_grids = self._deformation.deformation_net.grid.grids
+        total = 0
+        # model.grids is 6 x [1, rank * F_dim, reso, reso]
+        for grids in multi_res_grids:
+            if len(grids) == 3:
+                time_grids = []
+            else:
+                time_grids =[2, 4, 5]
+            for grid_id in time_grids:
+                total += compute_plane_smoothness(grids[grid_id])
+        return total
+    def _l1_regulation(self):
+                # model.grids is 6 x [1, rank * F_dim, reso, reso]
+        multi_res_grids = self._deformation.deformation_net.grid.grids
+
+        total = 0.0
+        for grids in multi_res_grids:
+            if len(grids) == 3:
+                continue
+            else:
+                # These are the spatiotemporal grids
+                spatiotemporal_grids = [2, 4, 5]
+            for grid_id in spatiotemporal_grids:
+                total += torch.abs(1 - grids[grid_id]).mean()
+        return total
+    def compute_regulation(self, time_smoothness_weight, l1_time_planes_weight, plane_tv_weight):
+        return plane_tv_weight * self._plane_regulation() + time_smoothness_weight * self._time_regulation() + l1_time_planes_weight * self._l1_regulation()
+    
+
+    def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_size):
+        grads = self.xyz_gradient_accum / self.denom
+        grads[grads.isnan()] = 0.0
+
+        self.densify_and_clone(grads, max_grad, extent)
+        self.densify_and_split(grads, max_grad, extent)
+
+        prune_mask = (self.get_opacity < min_opacity).squeeze()
+        if max_screen_size:
+            big_points_vs = self.max_radii2D > max_screen_size
+            big_points_ws = self.get_scaling.max(dim=1).values > 0.1 * extent
+            prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws)
+        self.prune_points(prune_mask)
+
+        torch.cuda.empty_cache()

grid_put.py

@@ -0,0 +1,300 @@
+import torch
+import torch.nn.functional as F
+
+def stride_from_shape(shape):
+    stride = [1]
+    for x in reversed(shape[1:]):
+        stride.append(stride[-1] * x) 
+    return list(reversed(stride))
+
+
+def scatter_add_nd(input, indices, values):
+    # input: [..., C], D dimension + C channel
+    # indices: [N, D], long
+    # values: [N, C]
+
+    D = indices.shape[-1]
+    C = input.shape[-1]
+    size = input.shape[:-1]
+    stride = stride_from_shape(size)
+
+    assert len(size) == D
+
+    input = input.view(-1, C)  # [HW, C]
+    flatten_indices = (indices * torch.tensor(stride, dtype=torch.long, device=indices.device)).sum(-1)  # [N]
+
+    input.scatter_add_(0, flatten_indices.unsqueeze(1).repeat(1, C), values)
+
+    return input.view(*size, C)
+
+
+def scatter_add_nd_with_count(input, count, indices, values, weights=None):
+    # input: [..., C], D dimension + C channel
+    # count: [..., 1], D dimension
+    # indices: [N, D], long
+    # values: [N, C]
+
+    D = indices.shape[-1]
+    C = input.shape[-1]
+    size = input.shape[:-1]
+    stride = stride_from_shape(size)
+
+    assert len(size) == D
+
+    input = input.view(-1, C)  # [HW, C]
+    count = count.view(-1, 1)
+
+    flatten_indices = (indices * torch.tensor(stride, dtype=torch.long, device=indices.device)).sum(-1)  # [N]
+
+    if weights is None:
+        weights = torch.ones_like(values[..., :1]) 
+
+    input.scatter_add_(0, flatten_indices.unsqueeze(1).repeat(1, C), values)
+    count.scatter_add_(0, flatten_indices.unsqueeze(1), weights)
+
+    return input.view(*size, C), count.view(*size, 1)
+
+def nearest_grid_put_2d(H, W, coords, values, return_count=False):
+    # coords: [N, 2], float in [-1, 1]
+    # values: [N, C]
+
+    C = values.shape[-1]
+
+    indices = (coords * 0.5 + 0.5) * torch.tensor(
+        [H - 1, W - 1], dtype=torch.float32, device=coords.device
+    )
+    indices = indices.round().long()  # [N, 2]
+
+    result = torch.zeros(H, W, C, device=values.device, dtype=values.dtype)  # [H, W, C]
+    count = torch.zeros(H, W, 1, device=values.device, dtype=values.dtype)  # [H, W, 1]
+    weights = torch.ones_like(values[..., :1])  # [N, 1]
+    
+    result, count = scatter_add_nd_with_count(result, count, indices, values, weights)
+
+    if return_count:
+        return result, count
+
+    mask = (count.squeeze(-1) > 0)
+    result[mask] = result[mask] / count[mask].repeat(1, C)
+
+    return result
+
+
+def linear_grid_put_2d(H, W, coords, values, return_count=False):
+    # coords: [N, 2], float in [-1, 1]
+    # values: [N, C]
+
+    C = values.shape[-1]
+
+    indices = (coords * 0.5 + 0.5) * torch.tensor(
+        [H - 1, W - 1], dtype=torch.float32, device=coords.device
+    )
+    indices_00 = indices.floor().long()  # [N, 2]
+    indices_00[:, 0].clamp_(0, H - 2)
+    indices_00[:, 1].clamp_(0, W - 2)
+    indices_01 = indices_00 + torch.tensor(
+        [0, 1], dtype=torch.long, device=indices.device
+    )
+    indices_10 = indices_00 + torch.tensor(
+        [1, 0], dtype=torch.long, device=indices.device
+    )
+    indices_11 = indices_00 + torch.tensor(
+        [1, 1], dtype=torch.long, device=indices.device
+    )
+
+    h = indices[..., 0] - indices_00[..., 0].float()
+    w = indices[..., 1] - indices_00[..., 1].float()
+    w_00 = (1 - h) * (1 - w)
+    w_01 = (1 - h) * w
+    w_10 = h * (1 - w)
+    w_11 = h * w
+
+    result = torch.zeros(H, W, C, device=values.device, dtype=values.dtype)  # [H, W, C]
+    count = torch.zeros(H, W, 1, device=values.device, dtype=values.dtype)  # [H, W, 1]
+    weights = torch.ones_like(values[..., :1])  # [N, 1]
+    
+    result, count = scatter_add_nd_with_count(result, count, indices_00, values * w_00.unsqueeze(1), weights* w_00.unsqueeze(1))
+    result, count = scatter_add_nd_with_count(result, count, indices_01, values * w_01.unsqueeze(1), weights* w_01.unsqueeze(1))
+    result, count = scatter_add_nd_with_count(result, count, indices_10, values * w_10.unsqueeze(1), weights* w_10.unsqueeze(1))
+    result, count = scatter_add_nd_with_count(result, count, indices_11, values * w_11.unsqueeze(1), weights* w_11.unsqueeze(1))
+
+    if return_count:
+        return result, count
+
+    mask = (count.squeeze(-1) > 0)
+    result[mask] = result[mask] / count[mask].repeat(1, C)
+
+    return result
+
+def mipmap_linear_grid_put_2d(H, W, coords, values, min_resolution=32, return_count=False):
+    # coords: [N, 2], float in [-1, 1]
+    # values: [N, C]
+
+    C = values.shape[-1]
+
+    result = torch.zeros(H, W, C, device=values.device, dtype=values.dtype)  # [H, W, C]
+    count = torch.zeros(H, W, 1, device=values.device, dtype=values.dtype)  # [H, W, 1]
+
+    cur_H, cur_W = H, W
+    
+    while min(cur_H, cur_W) > min_resolution:
+
+        # try to fill the holes
+        mask = (count.squeeze(-1) == 0)
+        if not mask.any():
+            break
+
+        cur_result, cur_count = linear_grid_put_2d(cur_H, cur_W, coords, values, return_count=True)
+        result[mask] = result[mask] + F.interpolate(cur_result.permute(2,0,1).unsqueeze(0).contiguous(), (H, W), mode='bilinear', align_corners=False).squeeze(0).permute(1,2,0).contiguous()[mask]
+        count[mask] = count[mask] + F.interpolate(cur_count.view(1, 1, cur_H, cur_W), (H, W), mode='bilinear', align_corners=False).view(H, W, 1)[mask]
+        cur_H //= 2
+        cur_W //= 2
+    
+    if return_count:
+        return result, count
+
+    mask = (count.squeeze(-1) > 0)
+    result[mask] = result[mask] / count[mask].repeat(1, C)
+
+    return result
+
+def nearest_grid_put_3d(H, W, D, coords, values, return_count=False):
+    # coords: [N, 3], float in [-1, 1]
+    # values: [N, C]
+
+    C = values.shape[-1]
+
+    indices = (coords * 0.5 + 0.5) * torch.tensor(
+        [H - 1, W - 1, D - 1], dtype=torch.float32, device=coords.device
+    )
+    indices = indices.round().long()  # [N, 2]
+
+    result = torch.zeros(H, W, D, C, device=values.device, dtype=values.dtype)  # [H, W, C]
+    count = torch.zeros(H, W, D, 1, device=values.device, dtype=values.dtype)  # [H, W, 1]
+    weights = torch.ones_like(values[..., :1])  # [N, 1]
+
+    result, count = scatter_add_nd_with_count(result, count, indices, values, weights)
+    
+    if return_count:
+        return result, count
+
+    mask = (count.squeeze(-1) > 0)
+    result[mask] = result[mask] / count[mask].repeat(1, C)
+
+    return result
+
+
+def linear_grid_put_3d(H, W, D, coords, values, return_count=False):
+    # coords: [N, 3], float in [-1, 1]
+    # values: [N, C]
+
+    C = values.shape[-1]
+
+    indices = (coords * 0.5 + 0.5) * torch.tensor(
+        [H - 1, W - 1, D - 1], dtype=torch.float32, device=coords.device
+    )
+    indices_000 = indices.floor().long()  # [N, 3]
+    indices_000[:, 0].clamp_(0, H - 2)
+    indices_000[:, 1].clamp_(0, W - 2)
+    indices_000[:, 2].clamp_(0, D - 2)
+
+    indices_001 = indices_000 + torch.tensor([0, 0, 1], dtype=torch.long, device=indices.device)
+    indices_010 = indices_000 + torch.tensor([0, 1, 0], dtype=torch.long, device=indices.device)
+    indices_011 = indices_000 + torch.tensor([0, 1, 1], dtype=torch.long, device=indices.device)
+    indices_100 = indices_000 + torch.tensor([1, 0, 0], dtype=torch.long, device=indices.device)
+    indices_101 = indices_000 + torch.tensor([1, 0, 1], dtype=torch.long, device=indices.device)
+    indices_110 = indices_000 + torch.tensor([1, 1, 0], dtype=torch.long, device=indices.device)
+    indices_111 = indices_000 + torch.tensor([1, 1, 1], dtype=torch.long, device=indices.device)
+
+    h = indices[..., 0] - indices_000[..., 0].float()
+    w = indices[..., 1] - indices_000[..., 1].float()
+    d = indices[..., 2] - indices_000[..., 2].float()
+    
+    w_000 = (1 - h) * (1 - w) * (1 - d)
+    w_001 = (1 - h) * w * (1 - d)
+    w_010 = h * (1 - w) * (1 - d)
+    w_011 = h * w * (1 - d)
+    w_100 = (1 - h) * (1 - w) * d
+    w_101 = (1 - h) * w * d
+    w_110 = h * (1 - w) * d
+    w_111 = h * w * d
+
+    result = torch.zeros(H, W, D, C, device=values.device, dtype=values.dtype)  # [H, W, D, C]
+    count = torch.zeros(H, W, D, 1, device=values.device, dtype=values.dtype)  # [H, W, D, 1]
+    weights = torch.ones_like(values[..., :1])  # [N, 1]
+    
+    result, count = scatter_add_nd_with_count(result, count, indices_000, values * w_000.unsqueeze(1), weights * w_000.unsqueeze(1))
+    result, count = scatter_add_nd_with_count(result, count, indices_001, values * w_001.unsqueeze(1), weights * w_001.unsqueeze(1))
+    result, count = scatter_add_nd_with_count(result, count, indices_010, values * w_010.unsqueeze(1), weights * w_010.unsqueeze(1))
+    result, count = scatter_add_nd_with_count(result, count, indices_011, values * w_011.unsqueeze(1), weights * w_011.unsqueeze(1))
+    result, count = scatter_add_nd_with_count(result, count, indices_100, values * w_100.unsqueeze(1), weights * w_100.unsqueeze(1))
+    result, count = scatter_add_nd_with_count(result, count, indices_101, values * w_101.unsqueeze(1), weights * w_101.unsqueeze(1))
+    result, count = scatter_add_nd_with_count(result, count, indices_110, values * w_110.unsqueeze(1), weights * w_110.unsqueeze(1))
+    result, count = scatter_add_nd_with_count(result, count, indices_111, values * w_111.unsqueeze(1), weights * w_111.unsqueeze(1))
+
+    if return_count:
+        return result, count
+
+    mask = (count.squeeze(-1) > 0)
+    result[mask] = result[mask] / count[mask].repeat(1, C)
+
+    return result
+
+def mipmap_linear_grid_put_3d(H, W, D, coords, values, min_resolution=32, return_count=False):
+    # coords: [N, 3], float in [-1, 1]
+    # values: [N, C]
+
+    C = values.shape[-1]
+
+    result = torch.zeros(H, W, D, C, device=values.device, dtype=values.dtype)  # [H, W, D, C]
+    count = torch.zeros(H, W, D, 1, device=values.device, dtype=values.dtype)  # [H, W, D, 1]
+    cur_H, cur_W, cur_D = H, W, D
+    
+    while min(min(cur_H, cur_W), cur_D) > min_resolution:
+
+        # try to fill the holes
+        mask = (count.squeeze(-1) == 0)
+        if not mask.any():
+            break
+
+        cur_result, cur_count = linear_grid_put_3d(cur_H, cur_W, cur_D, coords, values, return_count=True)
+        result[mask] = result[mask] + F.interpolate(cur_result.permute(3,0,1,2).unsqueeze(0).contiguous(), (H, W, D), mode='trilinear', align_corners=False).squeeze(0).permute(1,2,3,0).contiguous()[mask]
+        count[mask] = count[mask] + F.interpolate(cur_count.view(1, 1, cur_H, cur_W, cur_D), (H, W, D), mode='trilinear', align_corners=False).view(H, W, D, 1)[mask]
+        cur_H //= 2
+        cur_W //= 2
+        cur_D //= 2
+    
+    if return_count:
+        return result, count
+
+    mask = (count.squeeze(-1) > 0)
+    result[mask] = result[mask] / count[mask].repeat(1, C)
+
+    return result
+
+
+def grid_put(shape, coords, values, mode='linear-mipmap', min_resolution=32, return_raw=False):
+    # shape: [D], list/tuple
+    # coords: [N, D], float in [-1, 1]
+    # values: [N, C]
+
+    D = len(shape)
+    assert D in [2, 3], f'only support D == 2 or 3, but got D == {D}'
+
+    if mode == 'nearest':
+        if D == 2:
+            return nearest_grid_put_2d(*shape, coords, values, return_raw)
+        else:
+            return nearest_grid_put_3d(*shape, coords, values, return_raw)
+    elif mode == 'linear':
+        if D == 2:
+            return linear_grid_put_2d(*shape, coords, values, return_raw)
+        else:
+            return linear_grid_put_3d(*shape, coords, values, return_raw)
+    elif mode == 'linear-mipmap':
+        if D == 2:
+            return mipmap_linear_grid_put_2d(*shape, coords, values, min_resolution, return_raw)
+        else:
+            return mipmap_linear_grid_put_3d(*shape, coords, values, min_resolution, return_raw)
+    else:
+        raise NotImplementedError(f"got mode {mode}")    

gs_renderer_4d.py

@@ -0,0 +1,277 @@
+import math
+import numpy as np
+
+import torch
+
+from diff_gaussian_rasterization import (
+    GaussianRasterizationSettings,
+    GaussianRasterizer,
+)
+
+from sh_utils import eval_sh, SH2RGB, RGB2SH
+
+from gaussian_model_4d import GaussianModel, BasicPointCloud
+
+def getProjectionMatrix(znear, zfar, fovX, fovY):
+    tanHalfFovY = math.tan((fovY / 2))
+    tanHalfFovX = math.tan((fovX / 2))
+
+    P = torch.zeros(4, 4)
+
+    z_sign = 1.0
+
+    P[0, 0] = 1 / tanHalfFovX
+    P[1, 1] = 1 / tanHalfFovY
+    P[3, 2] = z_sign
+    P[2, 2] = z_sign * zfar / (zfar - znear)
+    P[2, 3] = -(zfar * znear) / (zfar - znear)
+    return P
+
+
+class MiniCam:
+    def __init__(self, c2w, width, height, fovy, fovx, znear, zfar, time=0, gs_convention=True):
+        # c2w (pose) should be in NeRF convention.
+
+        self.image_width = width
+        self.image_height = height
+        self.FoVy = fovy
+        self.FoVx = fovx
+        self.znear = znear
+        self.zfar = zfar
+
+        w2c = np.linalg.inv(c2w)
+
+        if gs_convention:
+            # rectify...
+            w2c[1:3, :3] *= -1
+            w2c[:3, 3] *= -1
+
+        self.world_view_transform = torch.tensor(w2c).transpose(0, 1).cuda()
+        self.projection_matrix = (
+            getProjectionMatrix(
+                znear=self.znear, zfar=self.zfar, fovX=self.FoVx, fovY=self.FoVy
+            )
+            .transpose(0, 1)
+            .cuda()
+        )
+        self.full_proj_transform = self.world_view_transform @ self.projection_matrix
+        self.camera_center = -torch.tensor(c2w[:3, 3]).cuda()
+
+        self.time = time
+
+
+class Renderer:
+    def __init__(self, opt, sh_degree=3, white_background=True, radius=1):
+        
+        self.sh_degree = sh_degree
+        self.white_background = white_background
+        self.radius = radius
+        self.opt = opt
+        self.T = self.opt.batch_size
+
+        self.gaussians = GaussianModel(sh_degree, opt.deformation)
+
+        self.bg_color = torch.tensor(
+            [1, 1, 1] if white_background else [0, 0, 0],
+            dtype=torch.float32,
+            device="cuda",
+        )
+        self.means3D_deform_T = None
+        self.opacity_deform_T = None
+        self.scales_deform_T = None
+        self.rotations_deform_T = None
+
+
+    
+    def initialize(self, input=None, num_pts=5000, radius=0.5):
+        # load checkpoint
+        if input is None:
+            # init from random point cloud
+            
+            phis = np.random.random((num_pts,)) * 2 * np.pi
+            costheta = np.random.random((num_pts,)) * 2 - 1
+            thetas = np.arccos(costheta)
+            mu = np.random.random((num_pts,))
+            radius = radius * np.cbrt(mu)
+            x = radius * np.sin(thetas) * np.cos(phis)
+            y = radius * np.sin(thetas) * np.sin(phis)
+            z = radius * np.cos(thetas)
+            xyz = np.stack((x, y, z), axis=1)
+            # xyz = np.random.random((num_pts, 3)) * 2.6 - 1.3
+
+            shs = np.random.random((num_pts, 3)) / 255.0
+            pcd = BasicPointCloud(
+                points=xyz, colors=SH2RGB(shs), normals=np.zeros((num_pts, 3))
+            )
+            # self.gaussians.create_from_pcd(pcd, 10)
+            self.gaussians.create_from_pcd(pcd, 10, 1)
+        elif isinstance(input, BasicPointCloud):
+            # load from a provided pcd
+            self.gaussians.create_from_pcd(input, 1)
+        else:
+            # load from saved ply
+            self.gaussians.load_ply(input)
+
+    def prepare_render(
+        self,
+    ):
+        means3D = self.gaussians.get_xyz
+        opacity = self.gaussians._opacity
+        scales = self.gaussians._scaling
+        rotations = self.gaussians._rotation
+
+        means3D_T = []
+        opacity_T = []
+        scales_T = []
+        rotations_T = []
+        time_T = []
+
+        for t in range(self.T):
+            time = torch.tensor(t).to(means3D.device).repeat(means3D.shape[0],1)
+            time = ((time.float() / self.T) - 0.5) * 2
+
+            means3D_T.append(means3D)
+            opacity_T.append(opacity)
+            scales_T.append(scales)
+            rotations_T.append(rotations)
+            time_T.append(time)
+
+        means3D_T = torch.cat(means3D_T)
+        opacity_T = torch.cat(opacity_T)
+        scales_T = torch.cat(scales_T)
+        rotations_T = torch.cat(rotations_T)
+        time_T = torch.cat(time_T)
+
+
+        means3D_deform_T, scales_deform_T, rotations_deform_T, opacity_deform_T = self.gaussians._deformation(means3D_T, scales_T, 
+                                                            rotations_T, opacity_T,
+                                                            time_T) #  time is not none
+        self.means3D_deform_T = means3D_deform_T.reshape([self.T, means3D_deform_T.shape[0]//self.T, -1])
+        self.opacity_deform_T = opacity_deform_T.reshape([self.T, means3D_deform_T.shape[0]//self.T, -1])
+        self.scales_deform_T = scales_deform_T.reshape([self.T, means3D_deform_T.shape[0]//self.T, -1])
+        self.rotations_deform_T = rotations_deform_T.reshape([self.T, means3D_deform_T.shape[0]//self.T, -1])
+        
+
+    def render(
+        self,
+        viewpoint_camera,
+        scaling_modifier=1.0,
+        bg_color=None,
+        override_color=None,
+        compute_cov3D_python=False,
+        convert_SHs_python=False,
+    ):
+        # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
+        screenspace_points = (
+            torch.zeros_like(
+                self.gaussians.get_xyz,
+                dtype=self.gaussians.get_xyz.dtype,
+                requires_grad=True,
+                device="cuda",
+            )
+            + 0
+        )
+        try:
+            screenspace_points.retain_grad()
+        except:
+            pass
+
+        # Set up rasterization configuration
+        tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
+        tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)
+
+        raster_settings = GaussianRasterizationSettings(
+            image_height=int(viewpoint_camera.image_height),
+            image_width=int(viewpoint_camera.image_width),
+            tanfovx=tanfovx,
+            tanfovy=tanfovy,
+            bg=self.bg_color if bg_color is None else bg_color,
+            scale_modifier=scaling_modifier,
+            viewmatrix=viewpoint_camera.world_view_transform,
+            projmatrix=viewpoint_camera.full_proj_transform,
+            sh_degree=self.gaussians.active_sh_degree,
+            campos=viewpoint_camera.camera_center,
+            prefiltered=False,
+            debug=False,
+        )
+
+        rasterizer = GaussianRasterizer(raster_settings=raster_settings)
+
+        means3D = self.gaussians.get_xyz
+        time = torch.tensor(viewpoint_camera.time).to(means3D.device).repeat(means3D.shape[0],1)
+        time = ((time.float() / self.T) - 0.5) * 2
+
+        means2D = screenspace_points
+        opacity = self.gaussians._opacity
+
+        # If precomputed 3d covariance is provided, use it. If not, then it will be computed from
+        # scaling / rotation by the rasterizer.
+        scales = None
+        rotations = None
+        cov3D_precomp = None
+        if compute_cov3D_python:
+            cov3D_precomp = self.gaussians.get_covariance(scaling_modifier)
+        else:
+            scales = self.gaussians._scaling
+            rotations = self.gaussians._rotation
+
+        means3D_deform, scales_deform, rotations_deform, opacity_deform = self.means3D_deform_T[viewpoint_camera.time], self.scales_deform_T[viewpoint_camera.time], self.rotations_deform_T[viewpoint_camera.time], self.opacity_deform_T[viewpoint_camera.time]
+
+
+        means3D_final =  means3D + means3D_deform
+        rotations_final =  rotations + rotations_deform
+        scales_final =  scales + scales_deform
+        opacity_final = opacity + opacity_deform
+
+
+
+        scales_final = self.gaussians.scaling_activation(scales_final)
+        rotations_final = self.gaussians.rotation_activation(rotations_final)
+        opacity = self.gaussians.opacity_activation(opacity)
+
+
+        # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
+        # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
+        shs = None
+        colors_precomp = None
+        if colors_precomp is None:
+            if convert_SHs_python:
+                shs_view = self.gaussians.get_features.transpose(1, 2).view(
+                    -1, 3, (self.gaussians.max_sh_degree + 1) ** 2
+                )
+                dir_pp = self.gaussians.get_xyz - viewpoint_camera.camera_center.repeat(
+                    self.gaussians.get_features.shape[0], 1
+                )
+                dir_pp_normalized = dir_pp / dir_pp.norm(dim=1, keepdim=True)
+                sh2rgb = eval_sh(
+                    self.gaussians.active_sh_degree, shs_view, dir_pp_normalized
+                )
+                colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)
+            else:
+                shs = self.gaussians.get_features
+        else:
+            colors_precomp = override_color
+
+        rendered_image, radii, rendered_depth, rendered_alpha = rasterizer(
+        means3D = means3D_final,
+        means2D = means2D,
+        shs = shs,
+        colors_precomp = colors_precomp,
+        opacities = opacity,
+        scales = scales_final,
+        rotations = rotations_final,
+        cov3D_precomp = cov3D_precomp)
+
+
+        rendered_image = rendered_image.clamp(0, 1)
+
+        # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
+        # They will be excluded from value updates used in the splitting criteria.
+        return {
+            "image": rendered_image,
+            "depth": rendered_depth,
+            "alpha": rendered_alpha,
+            "viewspace_points": screenspace_points,
+            "visibility_filter": radii > 0,
+            "radii": radii,
+        }

guidance/imagedream_utils.py

@@ -0,0 +1,326 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+
+from imagedream.camera_utils import get_camera, convert_opengl_to_blender, normalize_camera
+from imagedream.model_zoo import build_model
+from imagedream.ldm.models.diffusion.ddim import DDIMSampler
+
+from diffusers import DDIMScheduler
+
+class ImageDream(nn.Module):
+    def __init__(
+        self,
+        device,
+        model_name='sd-v2.1-base-4view-ipmv',
+        ckpt_path=None,
+        t_range=[0.02, 0.98],
+    ):
+        super().__init__()
+
+        self.device = device
+        self.model_name = model_name
+        self.ckpt_path = ckpt_path
+
+        self.model = build_model(self.model_name, ckpt_path=self.ckpt_path).eval().to(self.device)
+        self.model.device = device
+        for p in self.model.parameters():
+            p.requires_grad_(False)
+
+        self.dtype = torch.float32
+
+        self.num_train_timesteps = 1000
+        self.min_step = int(self.num_train_timesteps * t_range[0])
+        self.max_step = int(self.num_train_timesteps * t_range[1])
+
+        self.image_embeddings = {}
+        self.embeddings = {}
+
+        self.scheduler = DDIMScheduler.from_pretrained(
+            "stabilityai/stable-diffusion-2-1-base", subfolder="scheduler", torch_dtype=self.dtype
+        )
+
+    @torch.no_grad()
+    def get_image_text_embeds(self, image, prompts, negative_prompts):
+
+        image = F.interpolate(image, (256, 256), mode='bilinear', align_corners=False)
+        image_pil = TF.to_pil_image(image[0])
+        image_embeddings = self.model.get_learned_image_conditioning(image_pil).repeat(5,1,1) # [5, 257, 1280]
+        self.image_embeddings['pos'] = image_embeddings
+        self.image_embeddings['neg'] = torch.zeros_like(image_embeddings)
+
+        self.image_embeddings['ip_img'] = self.encode_imgs(image)
+        self.image_embeddings['neg_ip_img'] = torch.zeros_like(self.image_embeddings['ip_img'])
+
+        pos_embeds = self.encode_text(prompts).repeat(5,1,1)
+        neg_embeds = self.encode_text(negative_prompts).repeat(5,1,1)
+        self.embeddings['pos'] = pos_embeds
+        self.embeddings['neg'] = neg_embeds
+
+        return self.image_embeddings['pos'], self.image_embeddings['neg'], self.image_embeddings['ip_img'], self.image_embeddings['neg_ip_img'], self.embeddings['pos'], self.embeddings['neg']
+    
+    def encode_text(self, prompt):
+        # prompt: [str]
+        embeddings = self.model.get_learned_conditioning(prompt).to(self.device)
+        return embeddings
+    
+    @torch.no_grad()
+    def refine(self, pred_rgb, camera,
+               guidance_scale=5, steps=50, strength=0.8,
+        ):
+
+        batch_size = pred_rgb.shape[0]
+        real_batch_size = batch_size // 4
+        pred_rgb_256 = F.interpolate(pred_rgb, (256, 256), mode='bilinear', align_corners=False)
+        latents = self.encode_imgs(pred_rgb_256.to(self.dtype))
+
+        self.scheduler.set_timesteps(steps)
+        init_step = int(steps * strength)
+        latents = self.scheduler.add_noise(latents, torch.randn_like(latents), self.scheduler.timesteps[init_step])
+        
+        camera = camera[:, [0, 2, 1, 3]] # to blender convention (flip y & z axis)
+        camera[:, 1] *= -1
+        camera = normalize_camera(camera).view(batch_size, 16)
+
+        # extra view
+        camera = camera.view(real_batch_size, 4, 16)
+        camera = torch.cat([camera, torch.zeros_like(camera[:, :1])], dim=1) # [rB, 5, 16]
+        camera = camera.view(real_batch_size * 5, 16)
+
+        camera = camera.repeat(2, 1)
+        embeddings = torch.cat([self.embeddings['neg'].repeat(real_batch_size, 1, 1), self.embeddings['pos'].repeat(real_batch_size, 1, 1)], dim=0)
+        image_embeddings = torch.cat([self.image_embeddings['neg'].repeat(real_batch_size, 1, 1), self.image_embeddings['pos'].repeat(real_batch_size, 1, 1)], dim=0)
+        ip_img_embeddings= torch.cat([self.image_embeddings['neg_ip_img'].repeat(real_batch_size, 1, 1, 1), self.image_embeddings['ip_img'].repeat(real_batch_size, 1, 1, 1)], dim=0)
+        
+        context = {
+            "context": embeddings, 
+            "ip": image_embeddings, 
+            "ip_img": ip_img_embeddings,
+            "camera": camera, 
+            "num_frames": 4 + 1
+        }
+
+        for i, t in enumerate(self.scheduler.timesteps[init_step:]):
+
+            # extra view
+            
+            latents = latents.view(real_batch_size, 4, 4, 32, 32)
+            latents = torch.cat([latents, torch.zeros_like(latents[:, :1])], dim=1).view(-1, 4, 32, 32)
+            latent_model_input = torch.cat([latents] * 2)
+
+            tt = torch.cat([t.unsqueeze(0).repeat(real_batch_size * 5)] * 2).to(self.device)
+
+            noise_pred = self.model.apply_model(latent_model_input, tt, context)
+
+            noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+            
+            # remove extra view
+            noise_pred_uncond = noise_pred_uncond.reshape(real_batch_size, 5, 4, 32, 32)[:, :-1].reshape(-1, 4, 32, 32)
+            noise_pred_cond = noise_pred_cond.reshape(real_batch_size, 5, 4, 32, 32)[:, :-1].reshape(-1, 4, 32, 32)
+            latents = latents.reshape(real_batch_size, 5, 4, 32, 32)[:, :-1].reshape(-1, 4, 32, 32)
+
+            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+            
+            latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+
+        imgs = self.decode_latents(latents) # [1, 3, 512, 512]
+        return imgs
+
+    def train_step(
+        self,
+        pred_rgb, # [B, C, H, W]
+        camera, # [B, 4, 4]
+        step_ratio=None,
+        guidance_scale=5,
+        as_latent=False,
+    ):
+        
+        batch_size = pred_rgb.shape[0]
+        real_batch_size = batch_size // 4
+        pred_rgb = pred_rgb.to(self.dtype)
+
+        if as_latent:
+            latents = F.interpolate(pred_rgb, (32, 32), mode="bilinear", align_corners=False) * 2 - 1
+        else:
+            # interp to 256x256 to be fed into vae.
+            pred_rgb_256 = F.interpolate(pred_rgb, (256, 256), mode="bilinear", align_corners=False)
+            # encode image into latents with vae, requires grad!
+            latents = self.encode_imgs(pred_rgb_256)
+
+        if step_ratio is not None:
+            # dreamtime-like
+            # t = self.max_step - (self.max_step - self.min_step) * np.sqrt(step_ratio)
+            t = np.round((1 - step_ratio) * self.num_train_timesteps).clip(self.min_step, self.max_step)
+            t = torch.full((batch_size,), t, dtype=torch.long, device=self.device)
+        else:
+            t = torch.randint(self.min_step, self.max_step + 1, (real_batch_size,), dtype=torch.long, device=self.device).repeat(4)
+
+        camera = camera[:, [0, 2, 1, 3]] # to blender convention (flip y & z axis)
+        camera[:, 1] *= -1
+        camera = normalize_camera(camera).view(batch_size, 16)
+
+        # extra view
+        camera = camera.view(real_batch_size, 4, 16)
+        camera = torch.cat([camera, torch.zeros_like(camera[:, :1])], dim=1) # [rB, 5, 16]
+        camera = camera.view(real_batch_size * 5, 16)
+
+        camera = camera.repeat(2, 1)
+        embeddings = torch.cat([self.embeddings['neg'].repeat(real_batch_size, 1, 1), self.embeddings['pos'].repeat(real_batch_size, 1, 1)], dim=0)
+        image_embeddings = torch.cat([self.image_embeddings['neg'].repeat(real_batch_size, 1, 1), self.image_embeddings['pos'].repeat(real_batch_size, 1, 1)], dim=0)
+        ip_img_embeddings= torch.cat([self.image_embeddings['neg_ip_img'].repeat(real_batch_size, 1, 1, 1), self.image_embeddings['ip_img'].repeat(real_batch_size, 1, 1, 1)], dim=0)
+        
+        context = {
+            "context": embeddings, 
+            "ip": image_embeddings, 
+            "ip_img": ip_img_embeddings,
+            "camera": camera, 
+            "num_frames": 4 + 1
+        }
+
+        # predict the noise residual with unet, NO grad!
+        with torch.no_grad():
+            # add noise
+            noise = torch.randn_like(latents)
+            latents_noisy = self.model.q_sample(latents, t, noise) # [B=4, 4, 32, 32]
+            # extra view
+            t = t.view(real_batch_size, 4)
+            t = torch.cat([t, t[:, :1]], dim=1).view(-1)
+            latents_noisy = latents_noisy.view(real_batch_size, 4, 4, 32, 32)
+            latents_noisy = torch.cat([latents_noisy, torch.zeros_like(latents_noisy[:, :1])], dim=1).view(-1, 4, 32, 32)
+            # pred noise
+            latent_model_input = torch.cat([latents_noisy] * 2)
+            tt = torch.cat([t] * 2)
+
+            # import kiui
+            # kiui.lo(latent_model_input, t, context['context'], context['camera'])
+            
+            noise_pred = self.model.apply_model(latent_model_input, tt, context)
+
+            # perform guidance (high scale from paper!)
+            noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+
+            # remove extra view
+            noise_pred_uncond = noise_pred_uncond.reshape(real_batch_size, 5, 4, 32, 32)[:, :-1].reshape(-1, 4, 32, 32)
+            noise_pred_cond = noise_pred_cond.reshape(real_batch_size, 5, 4, 32, 32)[:, :-1].reshape(-1, 4, 32, 32)
+
+            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+
+        grad = (noise_pred - noise)
+        grad = torch.nan_to_num(grad)
+
+        target = (latents - grad).detach()
+        loss = 0.5 * F.mse_loss(latents.float(), target, reduction='sum') / latents.shape[0]
+
+        return loss
+
+    def decode_latents(self, latents):
+        imgs = self.model.decode_first_stage(latents)
+        imgs = ((imgs + 1) / 2).clamp(0, 1)
+        return imgs
+
+    def encode_imgs(self, imgs):
+        # imgs: [B, 3, 256, 256]
+        imgs = 2 * imgs - 1
+        latents = self.model.get_first_stage_encoding(self.model.encode_first_stage(imgs))
+        return latents # [B, 4, 32, 32]
+
+    @torch.no_grad()
+    def prompt_to_img(
+        self,
+        image,
+        prompts,
+        negative_prompts="",
+        height=256,
+        width=256,
+        num_inference_steps=50,
+        guidance_scale=5.0,
+        latents=None,
+        elevation=0,
+        azimuth_start=0,
+    ):
+        if isinstance(prompts, str):
+            prompts = [prompts]
+
+        if isinstance(negative_prompts, str):
+            negative_prompts = [negative_prompts]
+        
+        real_batch_size = len(prompts)
+        batch_size = len(prompts) * 5
+
+        # Text embeds -> img latents
+        sampler = DDIMSampler(self.model)
+        shape = [4, height // 8, width // 8]
+
+        c_ = {"context": self.encode_text(prompts).repeat(5,1,1)}
+        uc_ = {"context": self.encode_text(negative_prompts).repeat(5,1,1)}
+
+        # image embeddings
+        image = F.interpolate(image, (256, 256), mode='bilinear', align_corners=False)
+        image_pil = TF.to_pil_image(image[0])
+        image_embeddings = self.model.get_learned_image_conditioning(image_pil).repeat(5,1,1).to(self.device)
+        c_["ip"] = image_embeddings
+        uc_["ip"] = torch.zeros_like(image_embeddings)
+            
+        ip_img = self.encode_imgs(image)
+        c_["ip_img"] = ip_img
+        uc_["ip_img"] = torch.zeros_like(ip_img)
+
+        camera = get_camera(4, elevation=elevation, azimuth_start=azimuth_start, extra_view=True)
+        camera = camera.repeat(real_batch_size, 1).to(self.device)
+
+        c_["camera"] = uc_["camera"] = camera
+        c_["num_frames"] = uc_["num_frames"] = 5
+
+        kiui.lo(image_embeddings, ip_img, camera)
+
+        latents, _ = sampler.sample(S=num_inference_steps, conditioning=c_,
+                                        batch_size=batch_size, shape=shape,
+                                        verbose=False, 
+                                        unconditional_guidance_scale=guidance_scale,
+                                        unconditional_conditioning=uc_,
+                                        eta=0, x_T=None)
+
+        # Img latents -> imgs
+        imgs = self.decode_latents(latents)  # [4, 3, 256, 256]
+
+        kiui.lo(latents, imgs)
+        
+        # Img to Numpy
+        imgs = imgs.detach().cpu().permute(0, 2, 3, 1).numpy()
+        imgs = (imgs * 255).round().astype("uint8")
+
+        return imgs
+
+
+if __name__ == "__main__":
+    import argparse
+    import matplotlib.pyplot as plt
+    import kiui
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("image", type=str)
+    parser.add_argument("prompt", type=str)
+    parser.add_argument("--negative", default="", type=str)
+    parser.add_argument("--steps", type=int, default=30)
+    opt = parser.parse_args()
+
+    device = torch.device("cuda")
+
+    sd = ImageDream(device)
+
+    image = kiui.read_image(opt.image, mode='tensor')
+    image = image.permute(2, 0, 1).unsqueeze(0).to(device)
+
+    while True:
+        imgs = sd.prompt_to_img(image, opt.prompt, opt.negative, num_inference_steps=opt.steps)
+
+        grid = np.concatenate([
+            np.concatenate([imgs[0], imgs[1]], axis=1),
+            np.concatenate([imgs[2], imgs[3]], axis=1),
+        ], axis=0)
+
+        # visualize image
+        plt.imshow(grid)
+        plt.show()

guidance/mvdream_utils.py

@@ -0,0 +1,271 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from mvdream.camera_utils import get_camera, convert_opengl_to_blender, normalize_camera
+from mvdream.model_zoo import build_model
+from mvdream.ldm.models.diffusion.ddim import DDIMSampler
+
+from diffusers import DDIMScheduler
+
+class MVDream(nn.Module):
+    def __init__(
+        self,
+        device,
+        model_name='sd-v2.1-base-4view',
+        ckpt_path=None,
+        t_range=[0.02, 0.98],
+    ):
+        super().__init__()
+
+        self.device = device
+        self.model_name = model_name
+        self.ckpt_path = ckpt_path
+
+        self.model = build_model(self.model_name, ckpt_path=self.ckpt_path).eval().to(self.device)
+        self.model.device = device
+        for p in self.model.parameters():
+            p.requires_grad_(False)
+
+        self.dtype = torch.float32
+
+        self.num_train_timesteps = 1000
+        self.min_step = int(self.num_train_timesteps * t_range[0])
+        self.max_step = int(self.num_train_timesteps * t_range[1])
+
+        self.embeddings = None
+
+        self.scheduler = DDIMScheduler.from_pretrained(
+            "stabilityai/stable-diffusion-2-1-base", subfolder="scheduler", torch_dtype=self.dtype
+        )
+
+    @torch.no_grad()
+    def get_text_embeds(self, prompts, negative_prompts):
+        pos_embeds = self.encode_text(prompts).repeat(4,1,1)  # [1, 77, 768]
+        neg_embeds = self.encode_text(negative_prompts).repeat(4,1,1)
+        self.embeddings = torch.cat([neg_embeds, pos_embeds], dim=0)  # [2, 77, 768]
+    
+    def encode_text(self, prompt):
+        # prompt: [str]
+        embeddings = self.model.get_learned_conditioning(prompt).to(self.device)
+        return embeddings
+    
+    @torch.no_grad()
+    def refine(self, pred_rgb, camera,
+               guidance_scale=100, steps=50, strength=0.8,
+        ):
+
+        batch_size = pred_rgb.shape[0]
+        pred_rgb_256 = F.interpolate(pred_rgb, (256, 256), mode='bilinear', align_corners=False)
+        latents = self.encode_imgs(pred_rgb_256.to(self.dtype))
+        # latents = torch.randn((1, 4, 64, 64), device=self.device, dtype=self.dtype)
+
+        self.scheduler.set_timesteps(steps)
+        init_step = int(steps * strength)
+        latents = self.scheduler.add_noise(latents, torch.randn_like(latents), self.scheduler.timesteps[init_step])
+
+        camera = camera[:, [0, 2, 1, 3]] # to blender convention (flip y & z axis)
+        camera[:, 1] *= -1
+        camera = normalize_camera(camera).view(batch_size, 16)
+        camera = camera.repeat(2, 1)
+        context = {"context": self.embeddings, "camera": camera, "num_frames": 4}
+
+        for i, t in enumerate(self.scheduler.timesteps[init_step:]):
+    
+            latent_model_input = torch.cat([latents] * 2)
+            
+            tt = torch.cat([t.unsqueeze(0).repeat(batch_size)] * 2).to(self.device)
+
+            noise_pred = self.model.apply_model(latent_model_input, tt, context)
+
+            noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+            
+            latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+
+        imgs = self.decode_latents(latents) # [1, 3, 512, 512]
+        return imgs
+
+    def train_step(
+        self,
+        pred_rgb, # [B, C, H, W], B is multiples of 4
+        camera, # [B, 4, 4]
+        step_ratio=None,
+        guidance_scale=50,
+        as_latent=False,
+    ):
+        
+        batch_size = pred_rgb.shape[0]
+        pred_rgb = pred_rgb.to(self.dtype)
+
+        if as_latent:
+            latents = F.interpolate(pred_rgb, (32, 32), mode="bilinear", align_corners=False) * 2 - 1
+        else:
+            # interp to 256x256 to be fed into vae.
+            pred_rgb_256 = F.interpolate(pred_rgb, (256, 256), mode="bilinear", align_corners=False)
+            # encode image into latents with vae, requires grad!
+            latents = self.encode_imgs(pred_rgb_256)
+
+        if step_ratio is not None:
+            # dreamtime-like
+            # t = self.max_step - (self.max_step - self.min_step) * np.sqrt(step_ratio)
+            t = np.round((1 - step_ratio) * self.num_train_timesteps).clip(self.min_step, self.max_step)
+            t = torch.full((batch_size,), t, dtype=torch.long, device=self.device)
+        else:
+            t = torch.randint(self.min_step, self.max_step + 1, (batch_size,), dtype=torch.long, device=self.device)
+
+        # camera = convert_opengl_to_blender(camera)
+        # flip_yz = torch.tensor([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]).unsqueeze(0)
+        # camera = torch.matmul(flip_yz.to(camera), camera)
+        camera = camera[:, [0, 2, 1, 3]] # to blender convention (flip y & z axis)
+        camera[:, 1] *= -1
+        camera = normalize_camera(camera).view(batch_size, 16)
+
+        ###############
+        # sampler = DDIMSampler(self.model)
+        # shape = [4, 32, 32]
+        # c_ = {"context": self.embeddings[4:]}
+        # uc_ = {"context": self.embeddings[:4]}
+
+        # # print(camera)
+
+        # # camera = get_camera(4, elevation=0, azimuth_start=0)
+        # # camera = camera.repeat(batch_size // 4, 1).to(self.device)
+
+        # # print(camera)
+
+        # c_["camera"] = uc_["camera"] = camera
+        # c_["num_frames"] = uc_["num_frames"] = 4
+
+        # latents_, _ = sampler.sample(S=30, conditioning=c_,
+        #                                 batch_size=batch_size, shape=shape,
+        #                                 verbose=False, 
+        #                                 unconditional_guidance_scale=guidance_scale,
+        #                                 unconditional_conditioning=uc_,
+        #                                 eta=0, x_T=None)
+
+        # # Img latents -> imgs
+        # imgs = self.decode_latents(latents_)  # [4, 3, 256, 256]
+        # import kiui
+        # kiui.vis.plot_image(imgs)
+        ###############
+
+        camera = camera.repeat(2, 1)
+        context = {"context": self.embeddings, "camera": camera, "num_frames": 4}
+
+        # predict the noise residual with unet, NO grad!
+        with torch.no_grad():
+            # add noise
+            noise = torch.randn_like(latents)
+            latents_noisy = self.model.q_sample(latents, t, noise)
+            # pred noise
+            latent_model_input = torch.cat([latents_noisy] * 2)
+            tt = torch.cat([t] * 2)
+
+            # import kiui
+            # kiui.lo(latent_model_input, t, context['context'], context['camera'])
+            
+            noise_pred = self.model.apply_model(latent_model_input, tt, context)
+
+            # perform guidance (high scale from paper!)
+            noise_pred_uncond, noise_pred_pos = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_pos - noise_pred_uncond)
+
+        grad = (noise_pred - noise)
+        grad = torch.nan_to_num(grad)
+
+        # seems important to avoid NaN...
+        # grad = grad.clamp(-1, 1)
+
+        target = (latents - grad).detach()
+        loss = 0.5 * F.mse_loss(latents.float(), target, reduction='sum') / latents.shape[0]
+
+        return loss
+
+    def decode_latents(self, latents):
+        imgs = self.model.decode_first_stage(latents)
+        imgs = ((imgs + 1) / 2).clamp(0, 1)
+        return imgs
+
+    def encode_imgs(self, imgs):
+        # imgs: [B, 3, 256, 256]
+        imgs = 2 * imgs - 1
+        latents = self.model.get_first_stage_encoding(self.model.encode_first_stage(imgs))
+        return latents # [B, 4, 32, 32]
+
+    @torch.no_grad()
+    def prompt_to_img(
+        self,
+        prompts,
+        negative_prompts="",
+        height=256,
+        width=256,
+        num_inference_steps=50,
+        guidance_scale=7.5,
+        latents=None,
+        elevation=0,
+        azimuth_start=0,
+    ):
+        if isinstance(prompts, str):
+            prompts = [prompts]
+
+        if isinstance(negative_prompts, str):
+            negative_prompts = [negative_prompts]
+        
+        batch_size = len(prompts) * 4
+
+        # Text embeds -> img latents
+        sampler = DDIMSampler(self.model)
+        shape = [4, height // 8, width // 8]
+        c_ = {"context": self.encode_text(prompts).repeat(4,1,1)}
+        uc_ = {"context": self.encode_text(negative_prompts).repeat(4,1,1)}
+
+        camera = get_camera(4, elevation=elevation, azimuth_start=azimuth_start)
+        camera = camera.repeat(batch_size // 4, 1).to(self.device)
+
+        c_["camera"] = uc_["camera"] = camera
+        c_["num_frames"] = uc_["num_frames"] = 4
+
+        latents, _ = sampler.sample(S=num_inference_steps, conditioning=c_,
+                                        batch_size=batch_size, shape=shape,
+                                        verbose=False, 
+                                        unconditional_guidance_scale=guidance_scale,
+                                        unconditional_conditioning=uc_,
+                                        eta=0, x_T=None)
+
+        # Img latents -> imgs
+        imgs = self.decode_latents(latents)  # [4, 3, 256, 256]
+        
+        # Img to Numpy
+        imgs = imgs.detach().cpu().permute(0, 2, 3, 1).numpy()
+        imgs = (imgs * 255).round().astype("uint8")
+
+        return imgs
+
+
+if __name__ == "__main__":
+    import argparse
+    import matplotlib.pyplot as plt
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("prompt", type=str)
+    parser.add_argument("--negative", default="", type=str)
+    parser.add_argument("--steps", type=int, default=30)
+    opt = parser.parse_args()
+
+    device = torch.device("cuda")
+
+    sd = MVDream(device)
+
+    while True:
+        imgs = sd.prompt_to_img(opt.prompt, opt.negative, num_inference_steps=opt.steps)
+
+        grid = np.concatenate([
+            np.concatenate([imgs[0], imgs[1]], axis=1),
+            np.concatenate([imgs[2], imgs[3]], axis=1),
+        ], axis=0)
+
+        # visualize image
+        plt.imshow(grid)
+        plt.show()

guidance/sd_utils.py

@@ -0,0 +1,334 @@
+from transformers import CLIPTextModel, CLIPTokenizer, logging
+from diffusers import (
+    AutoencoderKL,
+    UNet2DConditionModel,
+    PNDMScheduler,
+    DDIMScheduler,
+    StableDiffusionPipeline,
+)
+from diffusers.utils.import_utils import is_xformers_available
+
+# suppress partial model loading warning
+logging.set_verbosity_error()
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def seed_everything(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    # torch.backends.cudnn.deterministic = True
+    # torch.backends.cudnn.benchmark = True
+
+
+class StableDiffusion(nn.Module):
+    def __init__(
+        self,
+        device,
+        fp16=True,
+        vram_O=False,
+        sd_version="2.1",
+        hf_key=None,
+        t_range=[0.02, 0.98],
+    ):
+        super().__init__()
+
+        self.device = device
+        self.sd_version = sd_version
+
+        if hf_key is not None:
+            print(f"[INFO] using hugging face custom model key: {hf_key}")
+            model_key = hf_key
+        elif self.sd_version == "2.1":
+            model_key = "stabilityai/stable-diffusion-2-1-base"
+        elif self.sd_version == "2.0":
+            model_key = "stabilityai/stable-diffusion-2-base"
+        elif self.sd_version == "1.5":
+            model_key = "runwayml/stable-diffusion-v1-5"
+        else:
+            raise ValueError(
+                f"Stable-diffusion version {self.sd_version} not supported."
+            )
+
+        self.dtype = torch.float16 if fp16 else torch.float32
+
+        # Create model
+        pipe = StableDiffusionPipeline.from_pretrained(
+            model_key, torch_dtype=self.dtype
+        )
+
+        if vram_O:
+            pipe.enable_sequential_cpu_offload()
+            pipe.enable_vae_slicing()
+            pipe.unet.to(memory_format=torch.channels_last)
+            pipe.enable_attention_slicing(1)
+            # pipe.enable_model_cpu_offload()
+        else:
+            pipe.to(device)
+
+        self.vae = pipe.vae
+        self.tokenizer = pipe.tokenizer
+        self.text_encoder = pipe.text_encoder
+        self.unet = pipe.unet
+
+        self.scheduler = DDIMScheduler.from_pretrained(
+            model_key, subfolder="scheduler", torch_dtype=self.dtype
+        )
+
+        del pipe
+
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+        self.min_step = int(self.num_train_timesteps * t_range[0])
+        self.max_step = int(self.num_train_timesteps * t_range[1])
+        self.alphas = self.scheduler.alphas_cumprod.to(self.device)  # for convenience
+
+        self.embeddings = None
+
+    @torch.no_grad()
+    def get_text_embeds(self, prompts, negative_prompts):
+        pos_embeds = self.encode_text(prompts)  # [1, 77, 768]
+        neg_embeds = self.encode_text(negative_prompts)
+        self.embeddings = torch.cat([neg_embeds, pos_embeds], dim=0)  # [2, 77, 768]
+    
+    def encode_text(self, prompt):
+        # prompt: [str]
+        inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            return_tensors="pt",
+        )
+        embeddings = self.text_encoder(inputs.input_ids.to(self.device))[0]
+        return embeddings
+
+    @torch.no_grad()
+    def refine(self, pred_rgb,
+               guidance_scale=100, steps=50, strength=0.8,
+        ):
+
+        batch_size = pred_rgb.shape[0]
+        pred_rgb_512 = F.interpolate(pred_rgb, (512, 512), mode='bilinear', align_corners=False)
+        latents = self.encode_imgs(pred_rgb_512.to(self.dtype))
+        # latents = torch.randn((1, 4, 64, 64), device=self.device, dtype=self.dtype)
+
+        self.scheduler.set_timesteps(steps)
+        init_step = int(steps * strength)
+        latents = self.scheduler.add_noise(latents, torch.randn_like(latents), self.scheduler.timesteps[init_step])
+
+        for i, t in enumerate(self.scheduler.timesteps[init_step:]):
+    
+            latent_model_input = torch.cat([latents] * 2)
+
+            noise_pred = self.unet(
+                latent_model_input, t, encoder_hidden_states=self.embeddings,
+            ).sample
+
+            noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+            
+            latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+
+        imgs = self.decode_latents(latents) # [1, 3, 512, 512]
+        return imgs
+
+    def train_step(
+        self,
+        pred_rgb,
+        step_ratio=None,
+        guidance_scale=100,
+        as_latent=False,
+    ):
+        
+        batch_size = pred_rgb.shape[0]
+        pred_rgb = pred_rgb.to(self.dtype)
+
+        if as_latent:
+            latents = F.interpolate(pred_rgb, (64, 64), mode="bilinear", align_corners=False) * 2 - 1
+        else:
+            # interp to 512x512 to be fed into vae.
+            pred_rgb_512 = F.interpolate(pred_rgb, (512, 512), mode="bilinear", align_corners=False)
+            # encode image into latents with vae, requires grad!
+            latents = self.encode_imgs(pred_rgb_512)
+
+        if step_ratio is not None:
+            # dreamtime-like
+            # t = self.max_step - (self.max_step - self.min_step) * np.sqrt(step_ratio)
+            t = np.round((1 - step_ratio) * self.num_train_timesteps).clip(self.min_step, self.max_step)
+            t = torch.full((batch_size,), t, dtype=torch.long, device=self.device)
+        else:
+            t = torch.randint(self.min_step, self.max_step + 1, (batch_size,), dtype=torch.long, device=self.device)
+
+        # w(t), sigma_t^2
+        w = (1 - self.alphas[t]).view(batch_size, 1, 1, 1)
+
+        # predict the noise residual with unet, NO grad!
+        with torch.no_grad():
+            # add noise
+            noise = torch.randn_like(latents)
+            latents_noisy = self.scheduler.add_noise(latents, noise, t)
+            # pred noise
+            latent_model_input = torch.cat([latents_noisy] * 2)
+            tt = torch.cat([t] * 2)
+
+            noise_pred = self.unet(
+                latent_model_input, tt, encoder_hidden_states=self.embeddings.repeat(batch_size, 1, 1)
+            ).sample
+
+            # perform guidance (high scale from paper!)
+            noise_pred_uncond, noise_pred_pos = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (
+                noise_pred_pos - noise_pred_uncond
+            )
+
+        grad = w * (noise_pred - noise)
+        grad = torch.nan_to_num(grad)
+
+        # seems important to avoid NaN...
+        # grad = grad.clamp(-1, 1)
+
+        target = (latents - grad).detach()
+        loss = 0.5 * F.mse_loss(latents.float(), target, reduction='sum') / latents.shape[0]
+
+        return loss
+
+    @torch.no_grad()
+    def produce_latents(
+        self,
+        height=512,
+        width=512,
+        num_inference_steps=50,
+        guidance_scale=7.5,
+        latents=None,
+    ):
+        if latents is None:
+            latents = torch.randn(
+                (
+                    self.embeddings.shape[0] // 2,
+                    self.unet.in_channels,
+                    height // 8,
+                    width // 8,
+                ),
+                device=self.device,
+            )
+
+        self.scheduler.set_timesteps(num_inference_steps)
+
+        for i, t in enumerate(self.scheduler.timesteps):
+            # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
+            latent_model_input = torch.cat([latents] * 2)
+            # predict the noise residual
+            noise_pred = self.unet(
+                latent_model_input, t, encoder_hidden_states=self.embeddings
+            ).sample
+
+            # perform guidance
+            noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (
+                noise_pred_cond - noise_pred_uncond
+            )
+
+            # compute the previous noisy sample x_t -> x_t-1
+            latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+
+        return latents
+
+    def decode_latents(self, latents):
+        latents = 1 / self.vae.config.scaling_factor * latents
+
+        imgs = self.vae.decode(latents).sample
+        imgs = (imgs / 2 + 0.5).clamp(0, 1)
+
+        return imgs
+
+    def encode_imgs(self, imgs):
+        # imgs: [B, 3, H, W]
+
+        imgs = 2 * imgs - 1
+
+        posterior = self.vae.encode(imgs).latent_dist
+        latents = posterior.sample() * self.vae.config.scaling_factor
+
+        return latents
+
+    def prompt_to_img(
+        self,
+        prompts,
+        negative_prompts="",
+        height=512,
+        width=512,
+        num_inference_steps=50,
+        guidance_scale=7.5,
+        latents=None,
+    ):
+        if isinstance(prompts, str):
+            prompts = [prompts]
+
+        if isinstance(negative_prompts, str):
+            negative_prompts = [negative_prompts]
+
+        # Prompts -> text embeds
+        self.get_text_embeds(prompts, negative_prompts)
+        
+        # Text embeds -> img latents
+        latents = self.produce_latents(
+            height=height,
+            width=width,
+            latents=latents,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+        )  # [1, 4, 64, 64]
+
+        # Img latents -> imgs
+        imgs = self.decode_latents(latents)  # [1, 3, 512, 512]
+
+        # Img to Numpy
+        imgs = imgs.detach().cpu().permute(0, 2, 3, 1).numpy()
+        imgs = (imgs * 255).round().astype("uint8")
+
+        return imgs
+
+
+if __name__ == "__main__":
+    import argparse
+    import matplotlib.pyplot as plt
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("prompt", type=str)
+    parser.add_argument("--negative", default="", type=str)
+    parser.add_argument(
+        "--sd_version",
+        type=str,
+        default="2.1",
+        choices=["1.5", "2.0", "2.1"],
+        help="stable diffusion version",
+    )
+    parser.add_argument(
+        "--hf_key",
+        type=str,
+        default=None,
+        help="hugging face Stable diffusion model key",
+    )
+    parser.add_argument("--fp16", action="store_true", help="use float16 for training")
+    parser.add_argument(
+        "--vram_O", action="store_true", help="optimization for low VRAM usage"
+    )
+    parser.add_argument("-H", type=int, default=512)
+    parser.add_argument("-W", type=int, default=512)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--steps", type=int, default=50)
+    opt = parser.parse_args()
+
+    seed_everything(opt.seed)
+
+    device = torch.device("cuda")
+
+    sd = StableDiffusion(device, opt.fp16, opt.vram_O, opt.sd_version, opt.hf_key)
+
+    imgs = sd.prompt_to_img(opt.prompt, opt.negative, opt.H, opt.W, opt.steps)
+
+    # visualize image
+    plt.imshow(imgs[0])
+    plt.show()

guidance/svd_utils.py

@@ -0,0 +1,221 @@
+import torch
+
+from svd import StableVideoDiffusionPipeline
+from diffusers import DDIMScheduler
+
+from PIL import Image
+import numpy as np
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class StableVideoDiffusion:
+    def __init__(
+        self,
+        device,
+        fp16=True,
+        t_range=[0.02, 0.98],
+    ):
+        super().__init__()
+
+        self.guidance_type = [
+            'sds',
+            'pixel reconstruction',
+            'latent reconstruction'
+        ][1]
+
+        self.device = device
+        self.dtype = torch.float16 if fp16 else torch.float32
+
+        # Create model
+        pipe = StableVideoDiffusionPipeline.from_pretrained(
+            "stabilityai/stable-video-diffusion-img2vid", torch_dtype=torch.float16, variant="fp16"
+        )
+        pipe.to(device)
+
+        self.pipe = pipe
+
+        self.num_train_timesteps = self.pipe.scheduler.config.num_train_timesteps if self.guidance_type == 'sds' else 25
+        self.pipe.scheduler.set_timesteps(self.num_train_timesteps, device=device)  # set sigma for euler discrete scheduling
+
+        self.min_step = int(self.num_train_timesteps * t_range[0])
+        self.max_step = int(self.num_train_timesteps * t_range[1])
+        self.alphas = self.pipe.scheduler.alphas_cumprod.to(self.device)  # for convenience
+
+        self.embeddings = None
+        self.image = None
+        self.target_cache = None
+
+    @torch.no_grad()
+    def get_img_embeds(self, image):
+        self.image = Image.fromarray(np.uint8(image*255))
+
+    def encode_image(self, image):
+        image = image * 2 -1
+        latents = self.pipe._encode_vae_image(image, self.device, num_videos_per_prompt=1, do_classifier_free_guidance=False)
+        latents = self.pipe.vae.config.scaling_factor * latents
+        return latents
+    
+    def refine(self,
+        pred_rgb,
+        steps=25, strength=0.8,
+        min_guidance_scale: float = 1.0,
+        max_guidance_scale: float = 3.0,
+    ):
+        # strength = 0.8
+        batch_size = pred_rgb.shape[0]
+        pred_rgb = pred_rgb.to(self.dtype)
+
+        # interp to 512x512 to be fed into vae.
+        pred_rgb_512 = F.interpolate(pred_rgb, (512, 512), mode="bilinear", align_corners=False)
+        # encode image into latents with vae, requires grad!
+
+        # latents = []
+        # for i in range(batch_size):
+        #     latent =  self.encode_image(pred_rgb_512[i:i+1])
+        #     latents.append(latent)
+        # latents = torch.cat(latents, 0)
+
+        latents = self.encode_image(pred_rgb_512)
+        latents = latents.unsqueeze(0)
+
+        if strength == 0:
+            init_step = 0
+            latents = torch.randn_like(latents)
+        else:
+            init_step = int(steps * strength)
+            latents = self.pipe.scheduler.add_noise(latents, torch.randn_like(latents), self.pipe.scheduler.timesteps[init_step:init_step+1])
+
+        target = self.pipe(
+            image=self.image,
+            height=512,
+            width=512,
+            latents=latents,
+            denoise_beg=init_step,
+            denoise_end=steps,
+            output_type='frame', 
+            num_frames=batch_size,
+            min_guidance_scale=min_guidance_scale,
+            max_guidance_scale=max_guidance_scale,
+            num_inference_steps=steps,
+            decode_chunk_size=1
+        ).frames[0]
+        target = (target + 1) * 0.5
+        target  = target.permute(1,0,2,3)
+        return target
+            
+        # frames = self.pipe(
+        #     image=self.image,
+        #     height=512,
+        #     width=512,
+        #     latents=latents,
+        #     denoise_beg=init_step,
+        #     denoise_end=steps,
+        #     num_frames=batch_size,
+        #     min_guidance_scale=min_guidance_scale,
+        #     max_guidance_scale=max_guidance_scale,
+        #     num_inference_steps=steps,
+        #     decode_chunk_size=1
+        # ).frames[0]
+        # export_to_gif(frames, f"tmp.gif")
+        # raise
+ 
+    def train_step(
+        self,
+        pred_rgb,
+        step_ratio=None,
+        min_guidance_scale: float = 1.0,
+        max_guidance_scale: float = 3.0,
+    ):
+        
+        batch_size = pred_rgb.shape[0]
+        pred_rgb = pred_rgb.to(self.dtype)
+
+        # interp to 512x512 to be fed into vae.
+        pred_rgb_512 = F.interpolate(pred_rgb, (512, 512), mode="bilinear", align_corners=False)
+        # encode image into latents with vae, requires grad!
+        # latents = self.pipe._encode_image(pred_rgb_512, self.device, num_videos_per_prompt=1, do_classifier_free_guidance=True)
+        latents = self.encode_image(pred_rgb_512)
+        latents = latents.unsqueeze(0)
+
+        if step_ratio is not None:
+            # dreamtime-like
+            # t = self.max_step - (self.max_step - self.min_step) * np.sqrt(step_ratio)
+            t = np.round((1 - step_ratio) * self.num_train_timesteps).clip(self.min_step, self.max_step)
+            t = torch.full((1,), t, dtype=torch.long, device=self.device)
+        else:
+            t = torch.randint(self.min_step, self.max_step + 1, (1,), dtype=torch.long, device=self.device)
+        # print(t)
+
+        w = (1 - self.alphas[t]).view(1, 1, 1, 1)
+
+
+        if self.guidance_type == 'sds':
+            # predict the noise residual with unet, NO grad!
+            with torch.no_grad():
+                t = self.num_train_timesteps - t.item()
+                # add noise
+                noise = torch.randn_like(latents)
+                latents_noisy = self.pipe.scheduler.add_noise(latents, noise, self.pipe.scheduler.timesteps[t:t+1]) # t=0 noise;t=999 clean
+                noise_pred = self.pipe(
+                    image=self.image,
+                    # image_embeddings=self.embeddings, 
+                    height=512,
+                    width=512,
+                    latents=latents_noisy,
+                    output_type='noise', 
+                    denoise_beg=t,
+                    denoise_end=t + 1,
+                    min_guidance_scale=min_guidance_scale,
+                    max_guidance_scale=max_guidance_scale,
+                    num_frames=batch_size,
+                    num_inference_steps=self.num_train_timesteps
+                ).frames[0]
+            
+            grad = w * (noise_pred - noise)
+            grad = torch.nan_to_num(grad)
+
+            target = (latents - grad).detach()
+            loss = 0.5 * F.mse_loss(latents.float(), target, reduction='sum') / latents.shape[1]
+            print(loss.item())
+            return loss
+        
+        elif self.guidance_type == 'pixel reconstruction':
+            # pixel space reconstruction
+            if self.target_cache is None:
+                with torch.no_grad():
+                    self.target_cache = self.pipe(
+                        image=self.image,
+                        height=512,
+                        width=512,
+                        output_type='frame', 
+                        num_frames=batch_size,
+                        num_inference_steps=self.num_train_timesteps,
+                        decode_chunk_size=1
+                    ).frames[0]
+                    self.target_cache = (self.target_cache + 1) * 0.5
+                    self.target_cache  = self.target_cache.permute(1,0,2,3)
+
+            loss = 0.5 * F.mse_loss(pred_rgb_512.float(), self.target_cache.detach().float(), reduction='sum') / latents.shape[1]
+            print(loss.item())
+
+            return loss
+
+        elif self.guidance_type == 'latent reconstruction':
+            # latent space reconstruction
+            if self.target_cache is None:
+                with torch.no_grad():
+                    self.target_cache = self.pipe(
+                        image=self.image,
+                        height=512,
+                        width=512,
+                        output_type='latent', 
+                        num_frames=batch_size,
+                        num_inference_steps=self.num_train_timesteps,
+                    ).frames[0]
+
+            loss = 0.5 * F.mse_loss(latents.float(), self.target_cache.detach().float(), reduction='sum') / latents.shape[1]
+            print(loss.item())
+
+            return loss

guidance/zero123_utils.py

@@ -0,0 +1,237 @@
+from diffusers import DDIMScheduler
+import torchvision.transforms.functional as TF
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import sys
+sys.path.append('./')
+
+from zero123 import Zero123Pipeline
+
+
+class Zero123(nn.Module):
+    def __init__(self, device, fp16=True, t_range=[0.02, 0.98], model_key="ashawkey/zero123-xl-diffusers"):
+        super().__init__()
+
+        self.device = device
+        self.fp16 = fp16
+        self.dtype = torch.float16 if fp16 else torch.float32
+
+        assert self.fp16, 'Only zero123 fp16 is supported for now.'
+
+        # model_key = "ashawkey/zero123-xl-diffusers"
+        # model_key = './model_cache/stable_zero123_diffusers'
+
+        self.pipe = Zero123Pipeline.from_pretrained(
+            model_key,
+            torch_dtype=self.dtype,
+            trust_remote_code=True,
+        ).to(self.device)
+
+        # stable-zero123 has a different camera embedding
+        self.use_stable_zero123 = 'stable' in model_key
+
+        self.pipe.image_encoder.eval()
+        self.pipe.vae.eval()
+        self.pipe.unet.eval()
+        self.pipe.clip_camera_projection.eval()
+
+        self.vae = self.pipe.vae
+        self.unet = self.pipe.unet
+
+        self.pipe.set_progress_bar_config(disable=True)
+
+        self.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config)
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+
+        self.min_step = int(self.num_train_timesteps * t_range[0])
+        self.max_step = int(self.num_train_timesteps * t_range[1])
+        self.alphas = self.scheduler.alphas_cumprod.to(self.device) # for convenience
+
+        self.embeddings = None
+
+    @torch.no_grad()
+    def get_img_embeds(self, x):
+        # x: image tensor in [0, 1]
+        x = F.interpolate(x, (256, 256), mode='bilinear', align_corners=False)
+        x_pil = [TF.to_pil_image(image) for image in x]
+        x_clip = self.pipe.feature_extractor(images=x_pil, return_tensors="pt").pixel_values.to(device=self.device, dtype=self.dtype)
+        c = self.pipe.image_encoder(x_clip).image_embeds
+        v = self.encode_imgs(x.to(self.dtype)) / self.vae.config.scaling_factor
+        self.embeddings = [c, v]
+        return c, v
+    
+    def get_cam_embeddings(self, elevation, azimuth, radius, default_elevation=0):
+        if self.use_stable_zero123:
+            T = np.stack([np.deg2rad(elevation), np.sin(np.deg2rad(azimuth)), np.cos(np.deg2rad(azimuth)), np.deg2rad([90 + default_elevation] * len(elevation))], axis=-1)
+        else:
+            # original zero123 camera embedding
+            T = np.stack([np.deg2rad(elevation), np.sin(np.deg2rad(azimuth)), np.cos(np.deg2rad(azimuth)), radius], axis=-1)
+        T = torch.from_numpy(T).unsqueeze(1).to(dtype=self.dtype, device=self.device) # [8, 1, 4]
+        return T
+
+    @torch.no_grad()
+    def refine(self, pred_rgb, elevation, azimuth, radius, 
+               guidance_scale=5, steps=50, strength=0.8, default_elevation=0,
+        ):
+
+        batch_size = pred_rgb.shape[0]
+
+        self.scheduler.set_timesteps(steps)
+
+        if strength == 0:
+            init_step = 0
+            latents = torch.randn((1, 4, 32, 32), device=self.device, dtype=self.dtype)
+        else:
+            init_step = int(steps * strength)
+            pred_rgb_256 = F.interpolate(pred_rgb, (256, 256), mode='bilinear', align_corners=False)
+            latents = self.encode_imgs(pred_rgb_256.to(self.dtype))
+            latents = self.scheduler.add_noise(latents, torch.randn_like(latents), self.scheduler.timesteps[init_step])
+
+        T = self.get_cam_embeddings(elevation, azimuth, radius, default_elevation)
+        cc_emb = torch.cat([self.embeddings[0].repeat(batch_size, 1, 1), T], dim=-1)
+        cc_emb = self.pipe.clip_camera_projection(cc_emb)
+        cc_emb = torch.cat([cc_emb, torch.zeros_like(cc_emb)], dim=0)
+
+        vae_emb = self.embeddings[1].repeat(batch_size, 1, 1, 1)
+        vae_emb = torch.cat([vae_emb, torch.zeros_like(vae_emb)], dim=0)
+
+        for i, t in enumerate(self.scheduler.timesteps[init_step:]):
+            
+            x_in = torch.cat([latents] * 2)
+            t_in = t.view(1).to(self.device)
+
+            noise_pred = self.unet(
+                torch.cat([x_in, vae_emb], dim=1),
+                t_in.to(self.unet.dtype),
+                encoder_hidden_states=cc_emb,
+            ).sample
+
+            noise_pred_cond, noise_pred_uncond = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+            
+            latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+
+        imgs = self.decode_latents(latents) # [1, 3, 256, 256]
+        return imgs
+    
+    def train_step(self, pred_rgb, elevation, azimuth, radius, step_ratio=None, guidance_scale=5, as_latent=False, default_elevation=0):
+        # pred_rgb: tensor [1, 3, H, W] in [0, 1]
+
+        batch_size = pred_rgb.shape[0]
+
+        if as_latent:
+            latents = F.interpolate(pred_rgb, (32, 32), mode='bilinear', align_corners=False) * 2 - 1
+        else:
+            pred_rgb_256 = F.interpolate(pred_rgb, (256, 256), mode='bilinear', align_corners=False)
+            latents = self.encode_imgs(pred_rgb_256.to(self.dtype))
+
+        if step_ratio is not None:
+            # dreamtime-like
+            # t = self.max_step - (self.max_step - self.min_step) * np.sqrt(step_ratio)
+            # t = self.max_step - (self.max_step - self.min_step) * (step_ratio ** 2)
+            t = np.round((1 - step_ratio) * self.num_train_timesteps).clip(self.min_step, self.max_step)
+            t = torch.full((batch_size,), t, dtype=torch.long, device=self.device)
+        else:
+            t = torch.randint(self.min_step, self.max_step + 1, (batch_size,), dtype=torch.long, device=self.device)
+
+        w = (1 - self.alphas[t]).view(batch_size, 1, 1, 1)
+
+        with torch.no_grad():
+            noise = torch.randn_like(latents)
+            latents_noisy = self.scheduler.add_noise(latents, noise, t)
+
+            x_in = torch.cat([latents_noisy] * 2)
+            t_in = torch.cat([t] * 2)
+
+            T = self.get_cam_embeddings(elevation, azimuth, radius, default_elevation)
+            cc_emb = torch.cat([self.embeddings[0].unsqueeze(1), T], dim=-1)
+            cc_emb = self.pipe.clip_camera_projection(cc_emb)
+            cc_emb = torch.cat([cc_emb, torch.zeros_like(cc_emb)], dim=0)
+
+            vae_emb = self.embeddings[1]
+            vae_emb = torch.cat([vae_emb, torch.zeros_like(vae_emb)], dim=0)
+
+            noise_pred = self.unet(
+                torch.cat([x_in, vae_emb], dim=1),
+                t_in.to(self.unet.dtype),
+                encoder_hidden_states=cc_emb,
+            ).sample
+
+        noise_pred_cond, noise_pred_uncond = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+
+        grad = w * (noise_pred - noise)
+        grad = torch.nan_to_num(grad)
+
+        target = (latents - grad).detach()
+        loss = 0.5 * F.mse_loss(latents.float(), target, reduction='sum')
+
+        return loss
+
+    def decode_latents(self, latents):
+        latents = 1 / self.vae.config.scaling_factor * latents
+
+        imgs = self.vae.decode(latents).sample
+        imgs = (imgs / 2 + 0.5).clamp(0, 1)
+
+        return imgs
+
+    def encode_imgs(self, imgs, mode=False):
+        # imgs: [B, 3, H, W]
+
+        imgs = 2 * imgs - 1
+
+        posterior = self.vae.encode(imgs).latent_dist
+        if mode:
+            latents = posterior.mode()
+        else:
+            latents = posterior.sample() 
+        latents = latents * self.vae.config.scaling_factor
+
+        return latents
+    
+    
+if __name__ == '__main__':
+    import cv2
+    import argparse
+    import numpy as np
+    import matplotlib.pyplot as plt
+    import kiui
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('input', type=str)
+    parser.add_argument('--elevation', type=float, default=0, help='delta elevation angle in [-90, 90]')
+    parser.add_argument('--azimuth', type=float, default=0, help='delta azimuth angle in [-180, 180]')
+    parser.add_argument('--radius', type=float, default=0, help='delta camera radius multiplier in [-0.5, 0.5]')
+    parser.add_argument('--stable', action='store_true')
+
+    opt = parser.parse_args()
+
+    device = torch.device('cuda')
+
+    print(f'[INFO] loading image from {opt.input} ...')
+    image = kiui.read_image(opt.input, mode='tensor')
+    image = image.permute(2, 0, 1).unsqueeze(0).contiguous().to(device)
+    image = F.interpolate(image, (256, 256), mode='bilinear', align_corners=False)
+
+    print(f'[INFO] loading model ...')
+    
+    if opt.stable:
+        zero123 = Zero123(device, model_key='ashawkey/stable-zero123-diffusers')
+    else:
+        zero123 = Zero123(device, model_key='ashawkey/zero123-xl-diffusers')
+
+    print(f'[INFO] running model ...')
+    zero123.get_img_embeds(image)
+
+    azimuth = opt.azimuth
+    while True:
+        outputs = zero123.refine(image, elevation=[opt.elevation], azimuth=[azimuth], radius=[opt.radius], strength=0)
+        plt.imshow(outputs.float().cpu().numpy().transpose(0, 2, 3, 1)[0])
+        plt.show()
+        azimuth = (azimuth + 10) % 360

lgm/core/attention.py

@@ -0,0 +1,156 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+import os
+import warnings
+
+from torch import Tensor
+from torch import nn
+
+XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
+try:
+    if XFORMERS_ENABLED:
+        from xformers.ops import memory_efficient_attention, unbind
+
+        XFORMERS_AVAILABLE = True
+        warnings.warn("xFormers is available (Attention)")
+    else:
+        warnings.warn("xFormers is disabled (Attention)")
+        raise ImportError
+except ImportError:
+    XFORMERS_AVAILABLE = False
+    warnings.warn("xFormers is not available (Attention)")
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+
+        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+        attn = q @ k.transpose(-2, -1)
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class MemEffAttention(Attention):
+    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
+        if not XFORMERS_AVAILABLE:
+            if attn_bias is not None:
+                raise AssertionError("xFormers is required for using nested tensors")
+            return super().forward(x)
+
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+
+        q, k, v = unbind(qkv, 2)
+
+        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+        x = x.reshape([B, N, C])
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class CrossAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        dim_q: int,
+        dim_k: int,
+        dim_v: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.to_q = nn.Linear(dim_q, dim, bias=qkv_bias)
+        self.to_k = nn.Linear(dim_k, dim, bias=qkv_bias)
+        self.to_v = nn.Linear(dim_v, dim, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
+        # q: [B, N, Cq]
+        # k: [B, M, Ck]
+        # v: [B, M, Cv]
+        # return: [B, N, C]
+
+        B, N, _ = q.shape
+        M = k.shape[1]
+        
+        q = self.scale * self.to_q(q).reshape(B, N, self.num_heads, self.dim // self.num_heads).permute(0, 2, 1, 3) # [B, nh, N, C/nh]
+        k = self.to_k(k).reshape(B, M, self.num_heads, self.dim // self.num_heads).permute(0, 2, 1, 3) # [B, nh, M, C/nh]
+        v = self.to_v(v).reshape(B, M, self.num_heads, self.dim // self.num_heads).permute(0, 2, 1, 3) # [B, nh, M, C/nh]
+
+        attn = q @ k.transpose(-2, -1) # [B, nh, N, M]
+
+        attn = attn.softmax(dim=-1) # [B, nh, N, M]
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, -1) # [B, nh, N, M] @ [B, nh, M, C/nh] --> [B, nh, N, C/nh] --> [B, N, nh, C/nh] --> [B, N, C]
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class MemEffCrossAttention(CrossAttention):
+    def forward(self, q: Tensor, k: Tensor, v: Tensor, attn_bias=None) -> Tensor:
+        if not XFORMERS_AVAILABLE:
+            if attn_bias is not None:
+                raise AssertionError("xFormers is required for using nested tensors")
+            return super().forward(x)
+
+        B, N, _ = q.shape
+        M = k.shape[1]
+
+        q = self.scale * self.to_q(q).reshape(B, N, self.num_heads, self.dim // self.num_heads) # [B, N, nh, C/nh]
+        k = self.to_k(k).reshape(B, M, self.num_heads, self.dim // self.num_heads) # [B, M, nh, C/nh]
+        v = self.to_v(v).reshape(B, M, self.num_heads, self.dim // self.num_heads) # [B, M, nh, C/nh]
+
+        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+        x = x.reshape(B, N, -1)
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x

lgm/core/gs.py

@@ -0,0 +1,190 @@
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diff_gaussian_rasterization import (
+    GaussianRasterizationSettings,
+    GaussianRasterizer,
+)
+
+from core.options import Options
+
+import kiui
+
+class GaussianRenderer:
+    def __init__(self, opt: Options):
+        
+        self.opt = opt
+        self.bg_color = torch.tensor([1, 1, 1], dtype=torch.float32, device="cuda")
+        
+        # intrinsics
+        self.tan_half_fov = np.tan(0.5 * np.deg2rad(self.opt.fovy))
+        self.proj_matrix = torch.zeros(4, 4, dtype=torch.float32)
+        self.proj_matrix[0, 0] = 1 / self.tan_half_fov
+        self.proj_matrix[1, 1] = 1 / self.tan_half_fov
+        self.proj_matrix[2, 2] = (opt.zfar + opt.znear) / (opt.zfar - opt.znear)
+        self.proj_matrix[3, 2] = - (opt.zfar * opt.znear) / (opt.zfar - opt.znear)
+        self.proj_matrix[2, 3] = 1
+        
+    def render(self, gaussians, cam_view, cam_view_proj, cam_pos, bg_color=None, scale_modifier=1):
+        # gaussians: [B, N, 14]
+        # cam_view, cam_view_proj: [B, V, 4, 4]
+        # cam_pos: [B, V, 3]
+
+        device = gaussians.device
+        B, V = cam_view.shape[:2]
+
+        # loop of loop...
+        images = []
+        alphas = []
+        for b in range(B):
+
+            # pos, opacity, scale, rotation, shs
+            means3D = gaussians[b, :, 0:3].contiguous().float()
+            opacity = gaussians[b, :, 3:4].contiguous().float()
+            scales = gaussians[b, :, 4:7].contiguous().float()
+            rotations = gaussians[b, :, 7:11].contiguous().float()
+            rgbs = gaussians[b, :, 11:].contiguous().float() # [N, 3]
+
+            for v in range(V):
+                
+                # render novel views
+                view_matrix = cam_view[b, v].float()
+                view_proj_matrix = cam_view_proj[b, v].float()
+                campos = cam_pos[b, v].float()
+
+                raster_settings = GaussianRasterizationSettings(
+                    image_height=self.opt.output_size,
+                    image_width=self.opt.output_size,
+                    tanfovx=self.tan_half_fov,
+                    tanfovy=self.tan_half_fov,
+                    bg=self.bg_color if bg_color is None else bg_color,
+                    scale_modifier=scale_modifier,
+                    viewmatrix=view_matrix,
+                    projmatrix=view_proj_matrix,
+                    sh_degree=0,
+                    campos=campos,
+                    prefiltered=False,
+                    debug=False,
+                )
+
+                rasterizer = GaussianRasterizer(raster_settings=raster_settings)
+
+                # Rasterize visible Gaussians to image, obtain their radii (on screen).
+                rendered_image, radii, rendered_depth, rendered_alpha = rasterizer(
+                    means3D=means3D,
+                    means2D=torch.zeros_like(means3D, dtype=torch.float32, device=device),
+                    shs=None,
+                    colors_precomp=rgbs,
+                    opacities=opacity,
+                    scales=scales,
+                    rotations=rotations,
+                    cov3D_precomp=None,
+                )
+
+                rendered_image = rendered_image.clamp(0, 1)
+
+                images.append(rendered_image)
+                alphas.append(rendered_alpha)
+
+        images = torch.stack(images, dim=0).view(B, V, 3, self.opt.output_size, self.opt.output_size)
+        alphas = torch.stack(alphas, dim=0).view(B, V, 1, self.opt.output_size, self.opt.output_size)
+
+        return {
+            "image": images, # [B, V, 3, H, W]
+            "alpha": alphas, # [B, V, 1, H, W]
+        }
+
+
+    def save_ply(self, gaussians, path, compatible=True):
+        # gaussians: [B, N, 14]
+        # compatible: save pre-activated gaussians as in the original paper
+
+        assert gaussians.shape[0] == 1, 'only support batch size 1'
+
+        from plyfile import PlyData, PlyElement
+     
+        means3D = gaussians[0, :, 0:3].contiguous().float()
+        opacity = gaussians[0, :, 3:4].contiguous().float()
+        scales = gaussians[0, :, 4:7].contiguous().float()
+        rotations = gaussians[0, :, 7:11].contiguous().float()
+        shs = gaussians[0, :, 11:].unsqueeze(1).contiguous().float() # [N, 1, 3]
+
+        # prune by opacity
+        mask = opacity.squeeze(-1) >= 0.005
+        means3D = means3D[mask]
+        opacity = opacity[mask]
+        scales = scales[mask]
+        rotations = rotations[mask]
+        shs = shs[mask]
+
+        # invert activation to make it compatible with the original ply format
+        if compatible:
+            opacity = kiui.op.inverse_sigmoid(opacity)
+            scales = torch.log(scales + 1e-8)
+            shs = (shs - 0.5) / 0.28209479177387814
+
+        xyzs = means3D.detach().cpu().numpy()
+        f_dc = shs.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+        opacities = opacity.detach().cpu().numpy()
+        scales = scales.detach().cpu().numpy()
+        rotations = rotations.detach().cpu().numpy()
+
+        l = ['x', 'y', 'z']
+        # All channels except the 3 DC
+        for i in range(f_dc.shape[1]):
+            l.append('f_dc_{}'.format(i))
+        l.append('opacity')
+        for i in range(scales.shape[1]):
+            l.append('scale_{}'.format(i))
+        for i in range(rotations.shape[1]):
+            l.append('rot_{}'.format(i))
+
+        dtype_full = [(attribute, 'f4') for attribute in l]
+
+        elements = np.empty(xyzs.shape[0], dtype=dtype_full)
+        attributes = np.concatenate((xyzs, f_dc, opacities, scales, rotations), axis=1)
+        elements[:] = list(map(tuple, attributes))
+        el = PlyElement.describe(elements, 'vertex')
+
+        PlyData([el]).write(path)
+    
+    def load_ply(self, path, compatible=True):
+
+        from plyfile import PlyData, PlyElement
+
+        plydata = PlyData.read(path)
+
+        xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
+                        np.asarray(plydata.elements[0]["y"]),
+                        np.asarray(plydata.elements[0]["z"])),  axis=1)
+        print("Number of points at loading : ", xyz.shape[0])
+
+        opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]
+
+        shs = np.zeros((xyz.shape[0], 3))
+        shs[:, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
+        shs[:, 1] = np.asarray(plydata.elements[0]["f_dc_1"])
+        shs[:, 2] = np.asarray(plydata.elements[0]["f_dc_2"])
+
+        scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
+        scales = np.zeros((xyz.shape[0], len(scale_names)))
+        for idx, attr_name in enumerate(scale_names):
+            scales[:, idx] = np.asarray(plydata.elements[0][attr_name])
+
+        rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot_")]
+        rots = np.zeros((xyz.shape[0], len(rot_names)))
+        for idx, attr_name in enumerate(rot_names):
+            rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
+          
+        gaussians = np.concatenate([xyz, opacities, scales, rots, shs], axis=1)
+        gaussians = torch.from_numpy(gaussians).float() # cpu
+
+        if compatible:
+            gaussians[..., 3:4] = torch.sigmoid(gaussians[..., 3:4])
+            gaussians[..., 4:7] = torch.exp(gaussians[..., 4:7])
+            gaussians[..., 11:] = 0.28209479177387814 * gaussians[..., 11:] + 0.5
+
+        return gaussians

lgm/core/models.py

@@ -0,0 +1,171 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+import kiui
+from kiui.lpips import LPIPS
+
+from core.unet import UNet
+from core.options import Options
+from core.gs import GaussianRenderer
+
+
+class LGM(nn.Module):
+    def __init__(
+        self,
+        opt: Options,
+    ):
+        super().__init__()
+
+        self.opt = opt
+
+        # unet
+        self.unet = UNet(
+            9, 14, 
+            down_channels=self.opt.down_channels,
+            down_attention=self.opt.down_attention,
+            mid_attention=self.opt.mid_attention,
+            up_channels=self.opt.up_channels,
+            up_attention=self.opt.up_attention,
+        )
+
+        # last conv
+        self.conv = nn.Conv2d(14, 14, kernel_size=1) # NOTE: maybe remove it if train again
+
+        # Gaussian Renderer
+        self.gs = GaussianRenderer(opt)
+
+        # activations...
+        self.pos_act = lambda x: x.clamp(-1, 1)
+        self.scale_act = lambda x: 0.1 * F.softplus(x)
+        self.opacity_act = lambda x: torch.sigmoid(x)
+        self.rot_act = lambda x: F.normalize(x, dim=-1)
+        self.rgb_act = lambda x: 0.5 * torch.tanh(x) + 0.5 # NOTE: may use sigmoid if train again
+
+        # LPIPS loss
+        if self.opt.lambda_lpips > 0:
+            self.lpips_loss = LPIPS(net='vgg')
+            self.lpips_loss.requires_grad_(False)
+
+
+    def state_dict(self, **kwargs):
+        # remove lpips_loss
+        state_dict = super().state_dict(**kwargs)
+        for k in list(state_dict.keys()):
+            if 'lpips_loss' in k:
+                del state_dict[k]
+        return state_dict
+
+
+    def prepare_default_rays(self, device, elevation=0):
+        
+        from kiui.cam import orbit_camera
+        from core.utils import get_rays
+
+        cam_poses = np.stack([
+            orbit_camera(elevation, 0, radius=self.opt.cam_radius),
+            orbit_camera(elevation, 90, radius=self.opt.cam_radius),
+            orbit_camera(elevation, 180, radius=self.opt.cam_radius),
+            orbit_camera(elevation, 270, radius=self.opt.cam_radius),
+        ], axis=0) # [4, 4, 4]
+        cam_poses = torch.from_numpy(cam_poses)
+
+        rays_embeddings = []
+        for i in range(cam_poses.shape[0]):
+            rays_o, rays_d = get_rays(cam_poses[i], self.opt.input_size, self.opt.input_size, self.opt.fovy) # [h, w, 3]
+            rays_plucker = torch.cat([torch.cross(rays_o, rays_d, dim=-1), rays_d], dim=-1) # [h, w, 6]
+            rays_embeddings.append(rays_plucker)
+
+            ## visualize rays for plotting figure
+            # kiui.vis.plot_image(rays_d * 0.5 + 0.5, save=True)
+
+        rays_embeddings = torch.stack(rays_embeddings, dim=0).permute(0, 3, 1, 2).contiguous().to(device) # [V, 6, h, w]
+        
+        return rays_embeddings
+        
+
+    def forward_gaussians(self, images):
+        # images: [B, 4, 9, H, W]
+        # return: Gaussians: [B, dim_t]
+
+        B, V, C, H, W = images.shape
+        images = images.view(B*V, C, H, W)
+
+        x = self.unet(images) # [B*4, 14, h, w]
+        x = self.conv(x) # [B*4, 14, h, w]
+
+        x = x.reshape(B, 4, 14, self.opt.splat_size, self.opt.splat_size)
+        
+        ## visualize multi-view gaussian features for plotting figure
+        # tmp_alpha = self.opacity_act(x[0, :, 3:4])
+        # tmp_img_rgb = self.rgb_act(x[0, :, 11:]) * tmp_alpha + (1 - tmp_alpha)
+        # tmp_img_pos = self.pos_act(x[0, :, 0:3]) * 0.5 + 0.5
+        # kiui.vis.plot_image(tmp_img_rgb, save=True)
+        # kiui.vis.plot_image(tmp_img_pos, save=True)
+
+        x = x.permute(0, 1, 3, 4, 2).reshape(B, -1, 14)
+        
+        pos = self.pos_act(x[..., 0:3]) # [B, N, 3]
+        opacity = self.opacity_act(x[..., 3:4])
+        scale = self.scale_act(x[..., 4:7])
+        rotation = self.rot_act(x[..., 7:11])
+        rgbs = self.rgb_act(x[..., 11:])
+
+        gaussians = torch.cat([pos, opacity, scale, rotation, rgbs], dim=-1) # [B, N, 14]
+        
+        return gaussians
+
+    
+    def forward(self, data, step_ratio=1):
+        # data: output of the dataloader
+        # return: loss
+
+        results = {}
+        loss = 0
+
+        images = data['input'] # [B, 4, 9, h, W], input features
+        
+        # use the first view to predict gaussians
+        gaussians = self.forward_gaussians(images) # [B, N, 14]
+
+        results['gaussians'] = gaussians
+
+        # always use white bg
+        bg_color = torch.ones(3, dtype=torch.float32, device=gaussians.device)
+        
+        # use the other views for rendering and supervision
+        results = self.gs.render(gaussians, data['cam_view'], data['cam_view_proj'], data['cam_pos'], bg_color=bg_color)
+        pred_images = results['image'] # [B, V, C, output_size, output_size]
+        pred_alphas = results['alpha'] # [B, V, 1, output_size, output_size]
+
+        results['images_pred'] = pred_images
+        results['alphas_pred'] = pred_alphas
+
+        gt_images = data['images_output'] # [B, V, 3, output_size, output_size], ground-truth novel views
+        gt_masks = data['masks_output'] # [B, V, 1, output_size, output_size], ground-truth masks
+
+        gt_images = gt_images * gt_masks + bg_color.view(1, 1, 3, 1, 1) * (1 - gt_masks)
+
+        loss_mse = F.mse_loss(pred_images, gt_images) + F.mse_loss(pred_alphas, gt_masks)
+        loss = loss + loss_mse
+
+        if self.opt.lambda_lpips > 0:
+            loss_lpips = self.lpips_loss(
+                # gt_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1,
+                # pred_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1,
+                # downsampled to at most 256 to reduce memory cost
+                F.interpolate(gt_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1, (256, 256), mode='bilinear', align_corners=False), 
+                F.interpolate(pred_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1, (256, 256), mode='bilinear', align_corners=False),
+            ).mean()
+            results['loss_lpips'] = loss_lpips
+            loss = loss + self.opt.lambda_lpips * loss_lpips
+            
+        results['loss'] = loss
+
+        # metric
+        with torch.no_grad():
+            psnr = -10 * torch.log10(torch.mean((pred_images.detach() - gt_images) ** 2))
+            results['psnr'] = psnr
+
+        return results

lgm/core/options.py

@@ -0,0 +1,120 @@
+import tyro
+from dataclasses import dataclass
+from typing import Tuple, Literal, Dict, Optional
+
+
+@dataclass
+class Options:
+    ### model
+    # Unet image input size
+    input_size: int = 256
+    # Unet definition
+    down_channels: Tuple[int, ...] = (64, 128, 256, 512, 1024, 1024)
+    down_attention: Tuple[bool, ...] = (False, False, False, True, True, True)
+    mid_attention: bool = True
+    up_channels: Tuple[int, ...] = (1024, 1024, 512, 256)
+    up_attention: Tuple[bool, ...] = (True, True, True, False)
+    # Unet output size, dependent on the input_size and U-Net structure!
+    splat_size: int = 64
+    # gaussian render size
+    output_size: int = 256
+
+    ### dataset
+    # data mode (only support s3 now)
+    data_mode: Literal['s3'] = 's3'
+    # fovy of the dataset
+    fovy: float = 49.1
+    # camera near plane
+    znear: float = 0.5
+    # camera far plane
+    zfar: float = 2.5
+    # number of all views (input + output)
+    num_views: int = 12
+    # number of views
+    num_input_views: int = 4
+    # camera radius
+    cam_radius: float = 1.5 # to better use [-1, 1]^3 space
+    # num workers
+    num_workers: int = 8
+
+    ### training
+    # workspace
+    workspace: str = './workspace'
+    # resume
+    resume: Optional[str] = 'pretrained/model_fp16_fixrot.safetensors'
+    # batch size (per-GPU)
+    batch_size: int = 8
+    # gradient accumulation
+    gradient_accumulation_steps: int = 1
+    # training epochs
+    num_epochs: int = 30
+    # lpips loss weight
+    lambda_lpips: float = 1.0
+    # gradient clip
+    gradient_clip: float = 1.0
+    # mixed precision
+    mixed_precision: str = 'bf16'
+    # learning rate
+    lr: float = 4e-4
+    # augmentation prob for grid distortion
+    prob_grid_distortion: float = 0.5
+    # augmentation prob for camera jitter
+    prob_cam_jitter: float = 0.5
+
+    ### testing
+    # test image path
+    test_path: Optional[str] = None
+
+    ### misc
+    # nvdiffrast backend setting
+    force_cuda_rast: bool = False
+    # render fancy video with gaussian scaling effect
+    fancy_video: bool = False
+    
+
+# all the default settings
+config_defaults: Dict[str, Options] = {}
+config_doc: Dict[str, str] = {}
+
+config_doc['lrm'] = 'the default settings for LGM'
+config_defaults['lrm'] = Options()
+
+config_doc['small'] = 'small model with lower resolution Gaussians'
+config_defaults['small'] = Options(
+    input_size=256,
+    splat_size=64,
+    output_size=256,
+    batch_size=8,
+    gradient_accumulation_steps=1,
+    mixed_precision='bf16',
+)
+
+config_doc['big'] = 'big model with higher resolution Gaussians'
+config_defaults['big'] = Options(
+    input_size=256,
+    up_channels=(1024, 1024, 512, 256, 128), # one more decoder
+    up_attention=(True, True, True, False, False),
+    splat_size=128,
+    output_size=512, # render & supervise Gaussians at a higher resolution.
+    batch_size=8,
+    num_views=8,
+    gradient_accumulation_steps=1,
+    mixed_precision='bf16',
+)
+
+config_doc['tiny'] = 'tiny model for ablation'
+config_defaults['tiny'] = Options(
+    input_size=256, 
+    down_channels=(32, 64, 128, 256, 512),
+    down_attention=(False, False, False, False, True),
+    up_channels=(512, 256, 128),
+    up_attention=(True, False, False, False),
+    splat_size=64,
+    output_size=256,
+    batch_size=16,
+    num_views=8,
+    gradient_accumulation_steps=1,
+    mixed_precision='bf16',
+)
+
+AllConfigs = tyro.extras.subcommand_type_from_defaults(config_defaults, config_doc)

lgm/core/unet.py

@@ -0,0 +1,319 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import numpy as np
+from typing import Tuple, Literal
+from functools import partial
+
+from core.attention import MemEffAttention
+
+class MVAttention(nn.Module):
+    def __init__(
+        self, 
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        groups: int = 32,
+        eps: float = 1e-5,
+        residual: bool = True,
+        skip_scale: float = 1,
+        num_frames: int = 4, # WARN: hardcoded!
+    ):
+        super().__init__()
+
+        self.residual = residual
+        self.skip_scale = skip_scale
+        self.num_frames = num_frames
+
+        self.norm = nn.GroupNorm(num_groups=groups, num_channels=dim, eps=eps, affine=True)
+        self.attn = MemEffAttention(dim, num_heads, qkv_bias, proj_bias, attn_drop, proj_drop)
+
+    def forward(self, x):
+        # x: [B*V, C, H, W]
+        BV, C, H, W = x.shape
+        B = BV // self.num_frames # assert BV % self.num_frames == 0
+
+        res = x
+        x = self.norm(x)
+
+        x = x.reshape(B, self.num_frames, C, H, W).permute(0, 1, 3, 4, 2).reshape(B, -1, C)
+        x = self.attn(x)
+        x = x.reshape(B, self.num_frames, H, W, C).permute(0, 1, 4, 2, 3).reshape(BV, C, H, W)
+
+        if self.residual:
+            x = (x + res) * self.skip_scale
+        return x
+
+class ResnetBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        resample: Literal['default', 'up', 'down'] = 'default',
+        groups: int = 32,
+        eps: float = 1e-5,
+        skip_scale: float = 1, # multiplied to output
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.skip_scale = skip_scale
+
+        self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+        self.norm2 = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+        self.act = F.silu
+
+        self.resample = None
+        if resample == 'up':
+            self.resample = partial(F.interpolate, scale_factor=2.0, mode="nearest")
+        elif resample == 'down':
+            self.resample = nn.AvgPool2d(kernel_size=2, stride=2)
+        
+        self.shortcut = nn.Identity()
+        if self.in_channels != self.out_channels:
+            self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=True)
+
+    
+    def forward(self, x):
+        res = x
+
+        x = self.norm1(x)
+        x = self.act(x)
+
+        if self.resample:
+            res = self.resample(res)
+            x = self.resample(x)
+        
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.act(x)
+        x = self.conv2(x)
+
+        x = (x + self.shortcut(res)) * self.skip_scale
+
+        return x
+
+class DownBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        num_layers: int = 1,
+        downsample: bool = True,
+        attention: bool = True,
+        attention_heads: int = 16,
+        skip_scale: float = 1,
+    ):
+        super().__init__()
+ 
+        nets = []
+        attns = []
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            nets.append(ResnetBlock(in_channels, out_channels, skip_scale=skip_scale))
+            if attention:
+                attns.append(MVAttention(out_channels, attention_heads, skip_scale=skip_scale))
+            else:
+                attns.append(None)
+        self.nets = nn.ModuleList(nets)
+        self.attns = nn.ModuleList(attns)
+
+        self.downsample = None
+        if downsample:
+            self.downsample = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1)
+
+    def forward(self, x):
+        xs = []
+
+        for attn, net in zip(self.attns, self.nets):
+            x = net(x)
+            if attn:
+                x = attn(x)
+            xs.append(x)
+
+        if self.downsample:
+            x = self.downsample(x)
+            xs.append(x)
+  
+        return x, xs
+
+
+class MidBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        num_layers: int = 1,
+        attention: bool = True,
+        attention_heads: int = 16,
+        skip_scale: float = 1,
+    ):
+        super().__init__()
+
+        nets = []
+        attns = []
+        # first layer
+        nets.append(ResnetBlock(in_channels, in_channels, skip_scale=skip_scale))
+        # more layers
+        for i in range(num_layers):
+            nets.append(ResnetBlock(in_channels, in_channels, skip_scale=skip_scale))
+            if attention:
+                attns.append(MVAttention(in_channels, attention_heads, skip_scale=skip_scale))
+            else:
+                attns.append(None)
+        self.nets = nn.ModuleList(nets)
+        self.attns = nn.ModuleList(attns)
+        
+    def forward(self, x):
+        x = self.nets[0](x)
+        for attn, net in zip(self.attns, self.nets[1:]):
+            if attn:
+                x = attn(x)
+            x = net(x)
+        return x
+
+
+class UpBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        prev_out_channels: int,
+        out_channels: int,
+        num_layers: int = 1,
+        upsample: bool = True,
+        attention: bool = True,
+        attention_heads: int = 16,
+        skip_scale: float = 1,
+    ):
+        super().__init__()
+
+        nets = []
+        attns = []
+        for i in range(num_layers):
+            cin = in_channels if i == 0 else out_channels
+            cskip = prev_out_channels if (i == num_layers - 1) else out_channels
+
+            nets.append(ResnetBlock(cin + cskip, out_channels, skip_scale=skip_scale))
+            if attention:
+                attns.append(MVAttention(out_channels, attention_heads, skip_scale=skip_scale))
+            else:
+                attns.append(None)
+        self.nets = nn.ModuleList(nets)
+        self.attns = nn.ModuleList(attns)
+
+        self.upsample = None
+        if upsample:
+            self.upsample = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x, xs):
+
+        for attn, net in zip(self.attns, self.nets):
+            res_x = xs[-1]
+            xs = xs[:-1]
+            x = torch.cat([x, res_x], dim=1)
+            x = net(x)
+            if attn:
+                x = attn(x)
+            
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2.0, mode='nearest')
+            x = self.upsample(x)
+
+        return x
+
+
+# it could be asymmetric!
+class UNet(nn.Module):
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        down_channels: Tuple[int, ...] = (64, 128, 256, 512, 1024),
+        down_attention: Tuple[bool, ...] = (False, False, False, True, True),
+        mid_attention: bool = True,
+        up_channels: Tuple[int, ...] = (1024, 512, 256),
+        up_attention: Tuple[bool, ...] = (True, True, False),
+        layers_per_block: int = 2,
+        skip_scale: float = np.sqrt(0.5),
+    ):
+        super().__init__()
+
+        # first
+        self.conv_in = nn.Conv2d(in_channels, down_channels[0], kernel_size=3, stride=1, padding=1)
+
+        # down
+        down_blocks = []
+        cout = down_channels[0]
+        for i in range(len(down_channels)):
+            cin = cout
+            cout = down_channels[i]
+
+            down_blocks.append(DownBlock(
+                cin, cout, 
+                num_layers=layers_per_block, 
+                downsample=(i != len(down_channels) - 1), # not final layer
+                attention=down_attention[i],
+                skip_scale=skip_scale,
+            ))
+        self.down_blocks = nn.ModuleList(down_blocks)
+
+        # mid
+        self.mid_block = MidBlock(down_channels[-1], attention=mid_attention, skip_scale=skip_scale)
+
+        # up
+        up_blocks = []
+        cout = up_channels[0]
+        for i in range(len(up_channels)):
+            cin = cout
+            cout = up_channels[i]
+            cskip = down_channels[max(-2 - i, -len(down_channels))] # for assymetric
+
+            up_blocks.append(UpBlock(
+                cin, cskip, cout, 
+                num_layers=layers_per_block + 1, # one more layer for up
+                upsample=(i != len(up_channels) - 1), # not final layer
+                attention=up_attention[i],
+                skip_scale=skip_scale,
+            ))
+        self.up_blocks = nn.ModuleList(up_blocks)
+
+        # last
+        self.norm_out = nn.GroupNorm(num_channels=up_channels[-1], num_groups=32, eps=1e-5)
+        self.conv_out = nn.Conv2d(up_channels[-1], out_channels, kernel_size=3, stride=1, padding=1)
+
+
+    def forward(self, x):
+        # x: [B, Cin, H, W]
+
+        # first
+        x = self.conv_in(x)
+        
+        # down
+        xss = [x]
+        for block in self.down_blocks:
+            x, xs = block(x)
+            xss.extend(xs)
+        
+        # mid
+        x = self.mid_block(x)
+
+        # up
+        for block in self.up_blocks:
+            xs = xss[-len(block.nets):]
+            xss = xss[:-len(block.nets)]
+            x = block(x, xs)
+
+        # last
+        x = self.norm_out(x)
+        x = F.silu(x)
+        x = self.conv_out(x) # [B, Cout, H', W']
+
+        return x

lgm/core/utils.py

@@ -0,0 +1,108 @@
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import roma
+from kiui.op import safe_normalize
+
+def get_rays(pose, h, w, fovy, opengl=True):
+
+    x, y = torch.meshgrid(
+        torch.arange(w, device=pose.device),
+        torch.arange(h, device=pose.device),
+        indexing="xy",
+    )
+    x = x.flatten()
+    y = y.flatten()
+
+    cx = w * 0.5
+    cy = h * 0.5
+
+    focal = h * 0.5 / np.tan(0.5 * np.deg2rad(fovy))
+
+    camera_dirs = F.pad(
+        torch.stack(
+            [
+                (x - cx + 0.5) / focal,
+                (y - cy + 0.5) / focal * (-1.0 if opengl else 1.0),
+            ],
+            dim=-1,
+        ),
+        (0, 1),
+        value=(-1.0 if opengl else 1.0),
+    )  # [hw, 3]
+
+    rays_d = camera_dirs @ pose[:3, :3].transpose(0, 1)  # [hw, 3]
+    rays_o = pose[:3, 3].unsqueeze(0).expand_as(rays_d) # [hw, 3]
+
+    rays_o = rays_o.view(h, w, 3)
+    rays_d = safe_normalize(rays_d).view(h, w, 3)
+
+    return rays_o, rays_d
+
+def orbit_camera_jitter(poses, strength=0.1):
+    # poses: [B, 4, 4], assume orbit camera in opengl format
+    # random orbital rotate
+
+    B = poses.shape[0]
+    rotvec_x = poses[:, :3, 1] * strength * np.pi * (torch.rand(B, 1, device=poses.device) * 2 - 1)
+    rotvec_y = poses[:, :3, 0] * strength * np.pi / 2 * (torch.rand(B, 1, device=poses.device) * 2 - 1)
+
+    rot = roma.rotvec_to_rotmat(rotvec_x) @ roma.rotvec_to_rotmat(rotvec_y)
+    R = rot @ poses[:, :3, :3]
+    T = rot @ poses[:, :3, 3:]
+
+    new_poses = poses.clone()
+    new_poses[:, :3, :3] = R
+    new_poses[:, :3, 3:] = T
+    
+    return new_poses
+
+def grid_distortion(images, strength=0.5):
+    # images: [B, C, H, W]
+    # num_steps: int, grid resolution for distortion
+    # strength: float in [0, 1], strength of distortion
+
+    B, C, H, W = images.shape
+
+    num_steps = np.random.randint(8, 17)
+    grid_steps = torch.linspace(-1, 1, num_steps)
+
+    # have to loop batch...
+    grids = []
+    for b in range(B):
+        # construct displacement
+        x_steps = torch.linspace(0, 1, num_steps) # [num_steps], inclusive
+        x_steps = (x_steps + strength * (torch.rand_like(x_steps) - 0.5) / (num_steps - 1)).clamp(0, 1) # perturb
+        x_steps = (x_steps * W).long() # [num_steps]
+        x_steps[0] = 0
+        x_steps[-1] = W
+        xs = []
+        for i in range(num_steps - 1):
+            xs.append(torch.linspace(grid_steps[i], grid_steps[i + 1], x_steps[i + 1] - x_steps[i]))
+        xs = torch.cat(xs, dim=0) # [W]
+
+        y_steps = torch.linspace(0, 1, num_steps) # [num_steps], inclusive
+        y_steps = (y_steps + strength * (torch.rand_like(y_steps) - 0.5) / (num_steps - 1)).clamp(0, 1) # perturb
+        y_steps = (y_steps * H).long() # [num_steps]
+        y_steps[0] = 0
+        y_steps[-1] = H
+        ys = []
+        for i in range(num_steps - 1):
+            ys.append(torch.linspace(grid_steps[i], grid_steps[i + 1], y_steps[i + 1] - y_steps[i]))
+        ys = torch.cat(ys, dim=0) # [H]
+
+        # construct grid
+        grid_x, grid_y = torch.meshgrid(xs, ys, indexing='xy') # [H, W]
+        grid = torch.stack([grid_x, grid_y], dim=-1) # [H, W, 2]
+
+        grids.append(grid)
+    
+    grids = torch.stack(grids, dim=0).to(images.device) # [B, H, W, 2]
+
+    # grid sample
+    images = F.grid_sample(images, grids, align_corners=False)
+
+    return images

lgm/infer.py

@@ -0,0 +1,226 @@
+
+import os
+import tyro
+import glob
+import imageio
+import numpy as np
+import tqdm
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from safetensors.torch import load_file
+import rembg
+
+import kiui
+from kiui.op import recenter
+from kiui.cam import orbit_camera
+
+from core.options import AllConfigs, Options
+from core.models import LGM
+from mvdream.pipeline_mvdream import MVDreamPipeline
+import cv2
+
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
+
+opt = tyro.cli(AllConfigs)
+
+# model
+model = LGM(opt)
+
+# resume pretrained checkpoint
+if opt.resume is not None:
+    if opt.resume.endswith('safetensors'):
+        ckpt = load_file(opt.resume, device='cpu')
+    else:
+        ckpt = torch.load(opt.resume, map_location='cpu')
+    model.load_state_dict(ckpt, strict=False)
+    print(f'[INFO] Loaded checkpoint from {opt.resume}')
+else:
+    print(f'[WARN] model randomly initialized, are you sure?')
+
+# device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = model.half().to(device)
+model.eval()
+
+rays_embeddings = model.prepare_default_rays(device)
+
+tan_half_fov = np.tan(0.5 * np.deg2rad(opt.fovy))
+proj_matrix = torch.zeros(4, 4, dtype=torch.float32, device=device)
+proj_matrix[0, 0] = 1 / tan_half_fov
+proj_matrix[1, 1] = 1 / tan_half_fov
+proj_matrix[2, 2] = (opt.zfar + opt.znear) / (opt.zfar - opt.znear)
+proj_matrix[3, 2] = - (opt.zfar * opt.znear) / (opt.zfar - opt.znear)
+proj_matrix[2, 3] = 1
+
+# load image dream
+pipe = MVDreamPipeline.from_pretrained(
+    "ashawkey/imagedream-ipmv-diffusers", # remote weights
+    torch_dtype=torch.float16,
+    trust_remote_code=True,
+    # local_files_only=True,
+)
+pipe = pipe.to(device)
+
+# load rembg
+bg_remover = rembg.new_session()
+
+# process function
+def process(opt: Options, path):
+    name = os.path.splitext(os.path.basename(path))[0]
+    if 'CONSISTENT4D' in path:
+            name = path.split('/')[-2]
+    print(f'[INFO] Processing {path} --> {name}')
+    os.makedirs('vis_data', exist_ok=True)
+    os.makedirs('logs', exist_ok=True)
+
+    input_image = kiui.read_image(path, mode='uint8')
+
+    # bg removal
+    carved_image = rembg.remove(input_image, session=bg_remover) # [H, W, 4]
+    mask = carved_image[..., -1] > 0
+
+    # recenter
+    image = recenter(carved_image, mask, border_ratio=0.2)
+    
+    # generate mv
+    image = image.astype(np.float32) / 255.0
+
+    # rgba to rgb white bg
+    if image.shape[-1] == 4:
+        image = image[..., :3] * image[..., 3:4] + (1 - image[..., 3:4])
+
+    mv_image = pipe('', image, guidance_scale=5.0, num_inference_steps=30, elevation=0)
+    mv_image = np.stack([mv_image[1], mv_image[2], mv_image[3], mv_image[0]], axis=0) # [4, 256, 256, 3], float32
+
+    # generate gaussians
+    input_image = torch.from_numpy(mv_image).permute(0, 3, 1, 2).float().to(device) # [4, 3, 256, 256]
+    input_image = F.interpolate(input_image, size=(opt.input_size, opt.input_size), mode='bilinear', align_corners=False)
+    input_image = TF.normalize(input_image, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
+
+    input_image = torch.cat([input_image, rays_embeddings], dim=1).unsqueeze(0) # [1, 4, 9, H, W]
+
+    with torch.no_grad():
+        ############## align azimuth #####################
+        with torch.autocast(device_type='cuda', dtype=torch.float16):
+            # generate gaussians
+            gaussians = model.forward_gaussians(input_image)
+
+        best_azi = 0
+        best_diff = 1e8
+        for v, azi in enumerate(np.arange(-180, 180, 1)):            
+            cam_poses = torch.from_numpy(orbit_camera(0, azi, radius=opt.cam_radius, opengl=True)).unsqueeze(0).to(device)
+
+            cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
+            
+            # cameras needed by gaussian rasterizer
+            cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
+            cam_view_proj = cam_view @ proj_matrix # [V, 4, 4]
+            cam_pos = - cam_poses[:, :3, 3] # [V, 3]
+
+            # scale = min(azi / 360, 1)
+            scale = 1
+
+
+            result = model.gs.render(gaussians, cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0), scale_modifier=scale)
+            rendered_image = result['image']
+
+            rendered_image = rendered_image.squeeze(1).permute(0,2,3,1).squeeze(0).contiguous().float().cpu().numpy()
+            rendered_image = cv2.resize(rendered_image, (image.shape[0], image.shape[1]), interpolation=cv2.INTER_AREA)
+
+            diff = np.mean((rendered_image- image) ** 2)
+
+            if diff < best_diff:
+                best_diff = diff
+                best_azi = azi
+        print("Best aligned azimuth: ", best_azi)
+
+        mv_image = []
+        for v, azi in enumerate([0, 90, 180, 270]):
+            cam_poses = torch.from_numpy(orbit_camera(0, azi + best_azi, radius=opt.cam_radius, opengl=True)).unsqueeze(0).to(device)
+
+            cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
+            
+            # cameras needed by gaussian rasterizer
+            cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
+            cam_view_proj = cam_view @ proj_matrix # [V, 4, 4]
+            cam_pos = - cam_poses[:, :3, 3] # [V, 3]
+
+            # scale = min(azi / 360, 1)
+            scale = 1
+
+
+            result = model.gs.render(gaussians, cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0), scale_modifier=scale)
+            rendered_image = result['image'] 
+            rendered_image = rendered_image.squeeze(1)
+            rendered_image = F.interpolate(rendered_image, (256, 256))
+            rendered_image = rendered_image.permute(0,2,3,1).contiguous().float().cpu().numpy()
+            mv_image.append(rendered_image)
+        mv_image = np.concatenate(mv_image, axis=0)
+
+        input_image = torch.from_numpy(mv_image).permute(0, 3, 1, 2).float().to(device) # [4, 3, 256, 256]
+        input_image = F.interpolate(input_image, size=(opt.input_size, opt.input_size), mode='bilinear', align_corners=False)
+        input_image = TF.normalize(input_image, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
+
+        input_image = torch.cat([input_image, rays_embeddings], dim=1).unsqueeze(0) # [1, 4, 9, H, W]
+
+        ################################
+
+        with torch.autocast(device_type='cuda', dtype=torch.float16):
+            # generate gaussians
+            gaussians = model.forward_gaussians(input_image)
+        
+        # save gaussians
+        model.gs.save_ply(gaussians, os.path.join('logs', name + '_model.ply'))
+
+        # render 360 video 
+        images = []
+        elevation = 0
+
+        if opt.fancy_video:
+
+            azimuth = np.arange(0, 720, 4, dtype=np.int32)
+            for azi in tqdm.tqdm(azimuth):
+                
+                cam_poses = torch.from_numpy(orbit_camera(elevation, azi, radius=opt.cam_radius, opengl=True)).unsqueeze(0).to(device)
+
+                cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
+                
+                # cameras needed by gaussian rasterizer
+                cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
+                cam_view_proj = cam_view @ proj_matrix # [V, 4, 4]
+                cam_pos = - cam_poses[:, :3, 3] # [V, 3]
+
+                scale = min(azi / 360, 1)
+
+                image = model.gs.render(gaussians, cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0), scale_modifier=scale)['image']
+                images.append((image.squeeze(1).permute(0,2,3,1).contiguous().float().cpu().numpy() * 255).astype(np.uint8))
+        else:
+            azimuth = np.arange(0, 360, 2, dtype=np.int32)
+            for azi in tqdm.tqdm(azimuth):
+                
+                cam_poses = torch.from_numpy(orbit_camera(elevation, azi, radius=opt.cam_radius, opengl=True)).unsqueeze(0).to(device)
+
+                cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
+                
+                # cameras needed by gaussian rasterizer
+                cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
+                cam_view_proj = cam_view @ proj_matrix # [V, 4, 4]
+                cam_pos = - cam_poses[:, :3, 3] # [V, 3]
+
+                image = model.gs.render(gaussians, cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0), scale_modifier=1)['image']
+                images.append((image.squeeze(1).permute(0,2,3,1).contiguous().float().cpu().numpy() * 255).astype(np.uint8))
+
+        images = np.concatenate(images, axis=0)
+        imageio.mimwrite(os.path.join('vis_data', name + '_static.mp4'), images, fps=30)
+
+
+assert opt.test_path is not None
+if os.path.isdir(opt.test_path):
+    file_paths = glob.glob(os.path.join(opt.test_path, "*"))
+else:
+    file_paths = [opt.test_path]
+for path in file_paths:
+    process(opt, path)

lgm/mvdream/mv_unet.py

@@ -0,0 +1,1005 @@
+import math
+import numpy as np
+from inspect import isfunction
+from typing import Optional, Any, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+
+from diffusers.configuration_utils import ConfigMixin
+from diffusers.models.modeling_utils import ModelMixin
+
+# require xformers!
+import xformers
+import xformers.ops
+
+from kiui.cam import orbit_camera
+
+def get_camera(
+    num_frames, elevation=0, azimuth_start=0, azimuth_span=360, blender_coord=True, extra_view=False,
+):
+    angle_gap = azimuth_span / num_frames
+    cameras = []
+    for azimuth in np.arange(azimuth_start, azimuth_span + azimuth_start, angle_gap):
+        
+        pose = orbit_camera(elevation, azimuth, radius=1) # [4, 4]
+
+        # opengl to blender
+        if blender_coord:
+            pose[2] *= -1
+            pose[[1, 2]] = pose[[2, 1]]
+
+        cameras.append(pose.flatten())
+
+    if extra_view:
+        cameras.append(np.zeros_like(cameras[0]))
+
+    return torch.from_numpy(np.stack(cameras, axis=0)).float() # [num_frames, 16]
+
+
+def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
+    """
+    Create sinusoidal timestep embeddings.
+    :param timesteps: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an [N x dim] Tensor of positional embeddings.
+    """
+    if not repeat_only:
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period)
+            * torch.arange(start=0, end=half, dtype=torch.float32)
+            / half
+        ).to(device=timesteps.device)
+        args = timesteps[:, None] * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat(
+                [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
+            )
+    else:
+        embedding = repeat(timesteps, "b -> b d", d=dim)
+    # import pdb; pdb.set_trace()
+    return embedding
+
+
+def zero_module(module):
+    """
+    Zero out the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+def conv_nd(dims, *args, **kwargs):
+    """
+    Create a 1D, 2D, or 3D convolution module.
+    """
+    if dims == 1:
+        return nn.Conv1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.Conv2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.Conv3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def avg_pool_nd(dims, *args, **kwargs):
+    """
+    Create a 1D, 2D, or 3D average pooling module.
+    """
+    if dims == 1:
+        return nn.AvgPool1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.AvgPool2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.AvgPool3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def default(val, d):
+    if val is not None:
+        return val
+    return d() if isfunction(d) else d
+
+
+class GEGLU(nn.Module):
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def forward(self, x):
+        x, gate = self.proj(x).chunk(2, dim=-1)
+        return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = default(dim_out, dim)
+        project_in = (
+            nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())
+            if not glu
+            else GEGLU(dim, inner_dim)
+        )
+
+        self.net = nn.Sequential(
+            project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class MemoryEfficientCrossAttention(nn.Module):
+    # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
+    def __init__(
+            self, 
+            query_dim, 
+            context_dim=None, 
+            heads=8, 
+            dim_head=64, 
+            dropout=0.0,
+            ip_dim=0,
+            ip_weight=1,
+        ):
+        super().__init__()
+        
+        inner_dim = dim_head * heads
+        context_dim = default(context_dim, query_dim)
+
+        self.heads = heads
+        self.dim_head = dim_head
+
+        self.ip_dim = ip_dim
+        self.ip_weight = ip_weight
+
+        if self.ip_dim > 0:
+            self.to_k_ip = nn.Linear(context_dim, inner_dim, bias=False)
+            self.to_v_ip = nn.Linear(context_dim, inner_dim, bias=False)
+
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)
+        )
+        self.attention_op: Optional[Any] = None
+
+    def forward(self, x, context=None):
+        q = self.to_q(x)
+        context = default(context, x)
+
+        if self.ip_dim > 0:
+            # context： [B, 77 + 16(ip), 1024]
+            token_len = context.shape[1]
+            context_ip = context[:, -self.ip_dim :, :]
+            k_ip = self.to_k_ip(context_ip)
+            v_ip = self.to_v_ip(context_ip)
+            context = context[:, : (token_len - self.ip_dim), :]
+
+        k = self.to_k(context)
+        v = self.to_v(context)
+
+        b, _, _ = q.shape
+        q, k, v = map(
+            lambda t: t.unsqueeze(3)
+            .reshape(b, t.shape[1], self.heads, self.dim_head)
+            .permute(0, 2, 1, 3)
+            .reshape(b * self.heads, t.shape[1], self.dim_head)
+            .contiguous(),
+            (q, k, v),
+        )
+
+        # actually compute the attention, what we cannot get enough of
+        out = xformers.ops.memory_efficient_attention(
+            q, k, v, attn_bias=None, op=self.attention_op
+        )
+
+        if self.ip_dim > 0:
+            k_ip, v_ip = map(
+                lambda t: t.unsqueeze(3)
+                .reshape(b, t.shape[1], self.heads, self.dim_head)
+                .permute(0, 2, 1, 3)
+                .reshape(b * self.heads, t.shape[1], self.dim_head)
+                .contiguous(),
+                (k_ip, v_ip),
+            )
+            # actually compute the attention, what we cannot get enough of
+            out_ip = xformers.ops.memory_efficient_attention(
+                q, k_ip, v_ip, attn_bias=None, op=self.attention_op
+            )
+            out = out + self.ip_weight * out_ip
+
+        out = (
+            out.unsqueeze(0)
+            .reshape(b, self.heads, out.shape[1], self.dim_head)
+            .permute(0, 2, 1, 3)
+            .reshape(b, out.shape[1], self.heads * self.dim_head)
+        )
+        return self.to_out(out)
+
+
+class BasicTransformerBlock3D(nn.Module):
+    
+    def __init__(
+        self,
+        dim,
+        n_heads,
+        d_head,
+        context_dim,
+        dropout=0.0,
+        gated_ff=True,
+        ip_dim=0,
+        ip_weight=1,
+    ):
+        super().__init__()
+
+        self.attn1 = MemoryEfficientCrossAttention(
+            query_dim=dim,
+            context_dim=None, # self-attention
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+        )
+        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+        self.attn2 = MemoryEfficientCrossAttention(
+            query_dim=dim,
+            context_dim=context_dim,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+            # ip only applies to cross-attention
+            ip_dim=ip_dim,
+            ip_weight=ip_weight,
+        ) 
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.norm3 = nn.LayerNorm(dim)
+
+    def forward(self, x, context=None, num_frames=1):
+        x = rearrange(x, "(b f) l c -> b (f l) c", f=num_frames).contiguous()
+        x = self.attn1(self.norm1(x), context=None) + x
+        x = rearrange(x, "b (f l) c -> (b f) l c", f=num_frames).contiguous()
+        x = self.attn2(self.norm2(x), context=context) + x
+        x = self.ff(self.norm3(x)) + x
+        return x
+
+
+class SpatialTransformer3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        n_heads,
+        d_head,
+        context_dim, # cross attention input dim
+        depth=1,
+        dropout=0.0,
+        ip_dim=0,
+        ip_weight=1,
+    ):
+        super().__init__()
+
+        if not isinstance(context_dim, list):
+            context_dim = [context_dim]
+
+        self.in_channels = in_channels
+
+        inner_dim = n_heads * d_head
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.proj_in = nn.Linear(in_channels, inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock3D(
+                    inner_dim,
+                    n_heads,
+                    d_head,
+                    context_dim=context_dim[d],
+                    dropout=dropout,
+                    ip_dim=ip_dim,
+                    ip_weight=ip_weight,
+                )
+                for d in range(depth)
+            ]
+        )
+        
+        self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
+        
+
+    def forward(self, x, context=None, num_frames=1):
+        # note: if no context is given, cross-attention defaults to self-attention
+        if not isinstance(context, list):
+            context = [context]
+        b, c, h, w = x.shape
+        x_in = x
+        x = self.norm(x)
+        x = rearrange(x, "b c h w -> b (h w) c").contiguous()
+        x = self.proj_in(x)
+        for i, block in enumerate(self.transformer_blocks):
+            x = block(x, context=context[i], num_frames=num_frames)
+        x = self.proj_out(x)
+        x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
+        
+        return x + x_in
+
+
+class PerceiverAttention(nn.Module):
+    def __init__(self, *, dim, dim_head=64, heads=8):
+        super().__init__()
+        self.scale = dim_head ** -0.5
+        self.dim_head = dim_head
+        self.heads = heads
+        inner_dim = dim_head * heads
+
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias=False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
+        self.to_out = nn.Linear(inner_dim, dim, bias=False)
+
+    def forward(self, x, latents):
+        """
+        Args:
+            x (torch.Tensor): image features
+                shape (b, n1, D)
+            latent (torch.Tensor): latent features
+                shape (b, n2, D)
+        """
+        x = self.norm1(x)
+        latents = self.norm2(latents)
+
+        b, l, _ = latents.shape
+
+        q = self.to_q(latents)
+        kv_input = torch.cat((x, latents), dim=-2)
+        k, v = self.to_kv(kv_input).chunk(2, dim=-1)
+
+        q, k, v = map(
+            lambda t: t.reshape(b, t.shape[1], self.heads, -1)
+            .transpose(1, 2)
+            .reshape(b, self.heads, t.shape[1], -1)
+            .contiguous(),
+            (q, k, v),
+        )
+
+        # attention
+        scale = 1 / math.sqrt(math.sqrt(self.dim_head))
+        weight = (q * scale) @ (k * scale).transpose(-2, -1)  # More stable with f16 than dividing afterwards
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        out = weight @ v
+
+        out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
+
+        return self.to_out(out)
+
+
+class Resampler(nn.Module):
+    def __init__(
+        self,
+        dim=1024,
+        depth=8,
+        dim_head=64,
+        heads=16,
+        num_queries=8,
+        embedding_dim=768,
+        output_dim=1024,
+        ff_mult=4,
+    ):
+        super().__init__()
+        self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim ** 0.5)
+        self.proj_in = nn.Linear(embedding_dim, dim)
+        self.proj_out = nn.Linear(dim, output_dim)
+        self.norm_out = nn.LayerNorm(output_dim)
+
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
+                        nn.Sequential(
+                            nn.LayerNorm(dim),
+                            nn.Linear(dim, dim * ff_mult, bias=False),
+                            nn.GELU(),
+                            nn.Linear(dim * ff_mult, dim, bias=False),
+                        )
+                    ]
+                )
+            )
+
+    def forward(self, x):
+        latents = self.latents.repeat(x.size(0), 1, 1)
+        x = self.proj_in(x)
+        for attn, ff in self.layers:
+            latents = attn(x, latents) + latents
+            latents = ff(latents) + latents
+
+        latents = self.proj_out(latents)
+        return self.norm_out(latents)
+
+
+class CondSequential(nn.Sequential):
+    """
+    A sequential module that passes timestep embeddings to the children that
+    support it as an extra input.
+    """
+
+    def forward(self, x, emb, context=None, num_frames=1):
+        for layer in self:
+            if isinstance(layer, ResBlock):
+                x = layer(x, emb)
+            elif isinstance(layer, SpatialTransformer3D):
+                x = layer(x, context, num_frames=num_frames)
+            else:
+                x = layer(x)
+        return x
+
+
+class Upsample(nn.Module):
+    """
+    An upsampling layer with an optional convolution.
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+                 upsampling occurs in the inner-two dimensions.
+    """
+
+    def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        if use_conv:
+            self.conv = conv_nd(
+                dims, self.channels, self.out_channels, 3, padding=padding
+            )
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        if self.dims == 3:
+            x = F.interpolate(
+                x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
+            )
+        else:
+            x = F.interpolate(x, scale_factor=2, mode="nearest")
+        if self.use_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    """
+    A downsampling layer with an optional convolution.
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+                 downsampling occurs in the inner-two dimensions.
+    """
+
+    def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        stride = 2 if dims != 3 else (1, 2, 2)
+        if use_conv:
+            self.op = conv_nd(
+                dims,
+                self.channels,
+                self.out_channels,
+                3,
+                stride=stride,
+                padding=padding,
+            )
+        else:
+            assert self.channels == self.out_channels
+            self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        return self.op(x)
+
+
+class ResBlock(nn.Module):
+    """
+    A residual block that can optionally change the number of channels.
+    :param channels: the number of input channels.
+    :param emb_channels: the number of timestep embedding channels.
+    :param dropout: the rate of dropout.
+    :param out_channels: if specified, the number of out channels.
+    :param use_conv: if True and out_channels is specified, use a spatial
+        convolution instead of a smaller 1x1 convolution to change the
+        channels in the skip connection.
+    :param dims: determines if the signal is 1D, 2D, or 3D.
+    :param up: if True, use this block for upsampling.
+    :param down: if True, use this block for downsampling.
+    """
+
+    def __init__(
+        self,
+        channels,
+        emb_channels,
+        dropout,
+        out_channels=None,
+        use_conv=False,
+        use_scale_shift_norm=False,
+        dims=2,
+        up=False,
+        down=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.emb_channels = emb_channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_scale_shift_norm = use_scale_shift_norm
+
+        self.in_layers = nn.Sequential(
+            nn.GroupNorm(32, channels),
+            nn.SiLU(),
+            conv_nd(dims, channels, self.out_channels, 3, padding=1),
+        )
+
+        self.updown = up or down
+
+        if up:
+            self.h_upd = Upsample(channels, False, dims)
+            self.x_upd = Upsample(channels, False, dims)
+        elif down:
+            self.h_upd = Downsample(channels, False, dims)
+            self.x_upd = Downsample(channels, False, dims)
+        else:
+            self.h_upd = self.x_upd = nn.Identity()
+
+        self.emb_layers = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(
+                emb_channels,
+                2 * self.out_channels if use_scale_shift_norm else self.out_channels,
+            ),
+        )
+        self.out_layers = nn.Sequential(
+            nn.GroupNorm(32, self.out_channels),
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            zero_module(
+                conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
+            ),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        elif use_conv:
+            self.skip_connection = conv_nd(
+                dims, channels, self.out_channels, 3, padding=1
+            )
+        else:
+            self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
+
+    def forward(self, x, emb):
+        if self.updown:
+            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+            h = in_rest(x)
+            h = self.h_upd(h)
+            x = self.x_upd(x)
+            h = in_conv(h)
+        else:
+            h = self.in_layers(x)
+        emb_out = self.emb_layers(emb).type(h.dtype)
+        while len(emb_out.shape) < len(h.shape):
+            emb_out = emb_out[..., None]
+        if self.use_scale_shift_norm:
+            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
+            scale, shift = torch.chunk(emb_out, 2, dim=1)
+            h = out_norm(h) * (1 + scale) + shift
+            h = out_rest(h)
+        else:
+            h = h + emb_out
+            h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+
+class MultiViewUNetModel(ModelMixin, ConfigMixin):
+    """
+    The full multi-view UNet model with attention, timestep embedding and camera embedding.
+    :param in_channels: channels in the input Tensor.
+    :param model_channels: base channel count for the model.
+    :param out_channels: channels in the output Tensor.
+    :param num_res_blocks: number of residual blocks per downsample.
+    :param attention_resolutions: a collection of downsample rates at which
+        attention will take place. May be a set, list, or tuple.
+        For example, if this contains 4, then at 4x downsampling, attention
+        will be used.
+    :param dropout: the dropout probability.
+    :param channel_mult: channel multiplier for each level of the UNet.
+    :param conv_resample: if True, use learned convolutions for upsampling and
+        downsampling.
+    :param dims: determines if the signal is 1D, 2D, or 3D.
+    :param num_classes: if specified (as an int), then this model will be
+        class-conditional with `num_classes` classes.
+    :param num_heads: the number of attention heads in each attention layer.
+    :param num_heads_channels: if specified, ignore num_heads and instead use
+                               a fixed channel width per attention head.
+    :param num_heads_upsample: works with num_heads to set a different number
+                               of heads for upsampling. Deprecated.
+    :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
+    :param resblock_updown: use residual blocks for up/downsampling.
+    :param use_new_attention_order: use a different attention pattern for potentially
+                                    increased efficiency.
+    :param camera_dim: dimensionality of camera input.
+    """
+
+    def __init__(
+        self,
+        image_size,
+        in_channels,
+        model_channels,
+        out_channels,
+        num_res_blocks,
+        attention_resolutions,
+        dropout=0,
+        channel_mult=(1, 2, 4, 8),
+        conv_resample=True,
+        dims=2,
+        num_classes=None,
+        num_heads=-1,
+        num_head_channels=-1,
+        num_heads_upsample=-1,
+        use_scale_shift_norm=False,
+        resblock_updown=False,
+        transformer_depth=1,
+        context_dim=None,
+        n_embed=None,
+        num_attention_blocks=None,
+        adm_in_channels=None,
+        camera_dim=None,
+        ip_dim=0, # imagedream uses ip_dim > 0
+        ip_weight=1.0,
+        **kwargs,
+    ):
+        super().__init__()
+        assert context_dim is not None
+        
+        if num_heads_upsample == -1:
+            num_heads_upsample = num_heads
+
+        if num_heads == -1:
+            assert (
+                num_head_channels != -1
+            ), "Either num_heads or num_head_channels has to be set"
+
+        if num_head_channels == -1:
+            assert (
+                num_heads != -1
+            ), "Either num_heads or num_head_channels has to be set"
+
+        self.image_size = image_size
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        if isinstance(num_res_blocks, int):
+            self.num_res_blocks = len(channel_mult) * [num_res_blocks]
+        else:
+            if len(num_res_blocks) != len(channel_mult):
+                raise ValueError(
+                    "provide num_res_blocks either as an int (globally constant) or "
+                    "as a list/tuple (per-level) with the same length as channel_mult"
+                )
+            self.num_res_blocks = num_res_blocks
+        
+        if num_attention_blocks is not None:
+            assert len(num_attention_blocks) == len(self.num_res_blocks)
+            assert all(
+                map(
+                    lambda i: self.num_res_blocks[i] >= num_attention_blocks[i],
+                    range(len(num_attention_blocks)),
+                )
+            )
+            print(
+                f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
+                f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
+                f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
+                f"attention will still not be set."
+            )
+
+        self.attention_resolutions = attention_resolutions
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.num_classes = num_classes
+        self.num_heads = num_heads
+        self.num_head_channels = num_head_channels
+        self.num_heads_upsample = num_heads_upsample
+        self.predict_codebook_ids = n_embed is not None
+
+        self.ip_dim = ip_dim
+        self.ip_weight = ip_weight
+
+        if self.ip_dim > 0:
+            self.image_embed = Resampler(
+                dim=context_dim,
+                depth=4,
+                dim_head=64,
+                heads=12,
+                num_queries=ip_dim,  # num token
+                embedding_dim=1280,
+                output_dim=context_dim,
+                ff_mult=4,
+            )
+
+        time_embed_dim = model_channels * 4
+        self.time_embed = nn.Sequential(
+            nn.Linear(model_channels, time_embed_dim),
+            nn.SiLU(),
+            nn.Linear(time_embed_dim, time_embed_dim),
+        )
+
+        if camera_dim is not None:
+            time_embed_dim = model_channels * 4
+            self.camera_embed = nn.Sequential(
+                nn.Linear(camera_dim, time_embed_dim),
+                nn.SiLU(),
+                nn.Linear(time_embed_dim, time_embed_dim),
+            )
+
+        if self.num_classes is not None:
+            if isinstance(self.num_classes, int):
+                self.label_emb = nn.Embedding(self.num_classes, time_embed_dim)
+            elif self.num_classes == "continuous":
+                # print("setting up linear c_adm embedding layer")
+                self.label_emb = nn.Linear(1, time_embed_dim)
+            elif self.num_classes == "sequential":
+                assert adm_in_channels is not None
+                self.label_emb = nn.Sequential(
+                    nn.Sequential(
+                        nn.Linear(adm_in_channels, time_embed_dim),
+                        nn.SiLU(),
+                        nn.Linear(time_embed_dim, time_embed_dim),
+                    )
+                )
+            else:
+                raise ValueError()
+
+        self.input_blocks = nn.ModuleList(
+            [
+                CondSequential(
+                    conv_nd(dims, in_channels, model_channels, 3, padding=1)
+                )
+            ]
+        )
+        self._feature_size = model_channels
+        input_block_chans = [model_channels]
+        ch = model_channels
+        ds = 1
+        for level, mult in enumerate(channel_mult):
+            for nr in range(self.num_res_blocks[level]):
+                layers: List[Any] = [
+                    ResBlock(
+                        ch,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=mult * model_channels,
+                        dims=dims,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = mult * model_channels
+                if ds in attention_resolutions:
+                    if num_head_channels == -1:
+                        dim_head = ch // num_heads
+                    else:
+                        num_heads = ch // num_head_channels
+                        dim_head = num_head_channels
+
+                    if num_attention_blocks is None or nr < num_attention_blocks[level]:
+                        layers.append(
+                            SpatialTransformer3D(
+                                ch,
+                                num_heads,
+                                dim_head,
+                                context_dim=context_dim,
+                                depth=transformer_depth,
+                                ip_dim=self.ip_dim,
+                                ip_weight=self.ip_weight,
+                            )
+                        )
+                self.input_blocks.append(CondSequential(*layers))
+                self._feature_size += ch
+                input_block_chans.append(ch)
+            if level != len(channel_mult) - 1:
+                out_ch = ch
+                self.input_blocks.append(
+                    CondSequential(
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            down=True,
+                        )
+                        if resblock_updown
+                        else Downsample(
+                            ch, conv_resample, dims=dims, out_channels=out_ch
+                        )
+                    )
+                )
+                ch = out_ch
+                input_block_chans.append(ch)
+                ds *= 2
+                self._feature_size += ch
+
+        if num_head_channels == -1:
+            dim_head = ch // num_heads
+        else:
+            num_heads = ch // num_head_channels
+            dim_head = num_head_channels
+        
+        self.middle_block = CondSequential(
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+            SpatialTransformer3D(
+                ch,
+                num_heads,
+                dim_head,
+                context_dim=context_dim,
+                depth=transformer_depth,
+                ip_dim=self.ip_dim,
+                ip_weight=self.ip_weight,
+            ), 
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+        )
+        self._feature_size += ch
+
+        self.output_blocks = nn.ModuleList([])
+        for level, mult in list(enumerate(channel_mult))[::-1]:
+            for i in range(self.num_res_blocks[level] + 1):
+                ich = input_block_chans.pop()
+                layers = [
+                    ResBlock(
+                        ch + ich,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=model_channels * mult,
+                        dims=dims,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = model_channels * mult
+                if ds in attention_resolutions:
+                    if num_head_channels == -1:
+                        dim_head = ch // num_heads
+                    else:
+                        num_heads = ch // num_head_channels
+                        dim_head = num_head_channels
+
+                    if num_attention_blocks is None or i < num_attention_blocks[level]:
+                        layers.append(
+                            SpatialTransformer3D(
+                                ch,
+                                num_heads,
+                                dim_head,
+                                context_dim=context_dim,
+                                depth=transformer_depth,
+                                ip_dim=self.ip_dim,
+                                ip_weight=self.ip_weight,
+                            )
+                        )
+                if level and i == self.num_res_blocks[level]:
+                    out_ch = ch
+                    layers.append(
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            up=True,
+                        )
+                        if resblock_updown
+                        else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
+                    )
+                    ds //= 2
+                self.output_blocks.append(CondSequential(*layers))
+                self._feature_size += ch
+
+        self.out = nn.Sequential(
+            nn.GroupNorm(32, ch),
+            nn.SiLU(),
+            zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
+        )
+        if self.predict_codebook_ids:
+            self.id_predictor = nn.Sequential(
+                nn.GroupNorm(32, ch),
+                conv_nd(dims, model_channels, n_embed, 1),
+                # nn.LogSoftmax(dim=1)  # change to cross_entropy and produce non-normalized logits
+            )
+
+    def forward(
+        self,
+        x,
+        timesteps=None,
+        context=None,
+        y=None,
+        camera=None,
+        num_frames=1,
+        ip=None,
+        ip_img=None,
+        **kwargs,
+    ):
+        """
+        Apply the model to an input batch.
+        :param x: an [(N x F) x C x ...] Tensor of inputs. F is the number of frames (views).
+        :param timesteps: a 1-D batch of timesteps.
+        :param context: conditioning plugged in via crossattn
+        :param y: an [N] Tensor of labels, if class-conditional.
+        :param num_frames: a integer indicating number of frames for tensor reshaping.
+        :return: an [(N x F) x C x ...] Tensor of outputs. F is the number of frames (views).
+        """
+        assert (
+            x.shape[0] % num_frames == 0
+        ), "input batch size must be dividable by num_frames!"
+        assert (y is not None) == (
+            self.num_classes is not None
+        ), "must specify y if and only if the model is class-conditional"
+
+        hs = []
+
+        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
+
+        emb = self.time_embed(t_emb)
+
+        if self.num_classes is not None:
+            assert y is not None
+            assert y.shape[0] == x.shape[0]
+            emb = emb + self.label_emb(y)
+
+        # Add camera embeddings
+        if camera is not None:
+            emb = emb + self.camera_embed(camera)
+        
+        # imagedream variant
+        if self.ip_dim > 0:
+            x[(num_frames - 1) :: num_frames, :, :, :] = ip_img # place at [4, 9]
+            ip_emb = self.image_embed(ip)
+            context = torch.cat((context, ip_emb), 1)
+
+        h = x
+        for module in self.input_blocks:
+            h = module(h, emb, context, num_frames=num_frames)
+            hs.append(h)
+        h = self.middle_block(h, emb, context, num_frames=num_frames)
+        for module in self.output_blocks:
+            h = torch.cat([h, hs.pop()], dim=1)
+            h = module(h, emb, context, num_frames=num_frames)
+        h = h.type(x.dtype)
+        if self.predict_codebook_ids:
+            return self.id_predictor(h)
+        else:
+            return self.out(h)

lgm/mvdream/pipeline_mvdream.py

@@ -0,0 +1,559 @@
+import torch
+import torch.nn.functional as F
+import inspect
+import numpy as np
+from typing import Callable, List, Optional, Union
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModel, CLIPImageProcessor
+from diffusers import AutoencoderKL, DiffusionPipeline
+from diffusers.utils import (
+    deprecate,
+    is_accelerate_available,
+    is_accelerate_version,
+    logging,
+)
+from diffusers.configuration_utils import FrozenDict
+from diffusers.schedulers import DDIMScheduler
+from diffusers.utils.torch_utils import randn_tensor
+
+from mvdream.mv_unet import MultiViewUNetModel, get_camera
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class MVDreamPipeline(DiffusionPipeline):
+
+    _optional_components = ["feature_extractor", "image_encoder"]
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        unet: MultiViewUNetModel,
+        tokenizer: CLIPTokenizer,
+        text_encoder: CLIPTextModel,
+        scheduler: DDIMScheduler,
+        # imagedream variant
+        feature_extractor: CLIPImageProcessor,
+        image_encoder: CLIPVisionModel,
+        requires_safety_checker: bool = False,
+    ):
+        super().__init__()
+
+        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:  # type: ignore
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
+                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "  # type: ignore
+                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
+                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
+                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
+                " file"
+            )
+            deprecate(
+                "steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False
+            )
+            new_config = dict(scheduler.config)
+            new_config["steps_offset"] = 1
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:  # type: ignore
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
+                " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
+                " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
+                " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
+                " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
+            )
+            deprecate(
+                "clip_sample not set", "1.0.0", deprecation_message, standard_warn=False
+            )
+            new_config = dict(scheduler.config)
+            new_config["clip_sample"] = False
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        self.register_modules(
+            vae=vae,
+            unet=unet,
+            scheduler=scheduler,
+            tokenizer=tokenizer,
+            text_encoder=text_encoder,
+            feature_extractor=feature_extractor,
+            image_encoder=image_encoder,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.register_to_config(requires_safety_checker=requires_safety_checker)
+
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding.
+
+        When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several
+        steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_slicing()
+
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding.
+
+        When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in
+        several steps. This is useful to save a large amount of memory and to allow the processing of larger images.
+        """
+        self.vae.enable_tiling()
+
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously invoked, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_tiling()
+
+    def enable_sequential_cpu_offload(self, gpu_id=0):
+        r"""
+        Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
+        text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
+        `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
+        Note that offloading happens on a submodule basis. Memory savings are higher than with
+        `enable_model_cpu_offload`, but performance is lower.
+        """
+        if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):
+            from accelerate import cpu_offload
+        else:
+            raise ImportError(
+                "`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher"
+            )
+
+        device = torch.device(f"cuda:{gpu_id}")
+
+        if self.device.type != "cpu":
+            self.to("cpu", silence_dtype_warnings=True)
+            torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
+
+        for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
+            cpu_offload(cpu_offloaded_model, device)
+
+    def enable_model_cpu_offload(self, gpu_id=0):
+        r"""
+        Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
+        to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
+        method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
+        `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
+        """
+        if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
+            from accelerate import cpu_offload_with_hook
+        else:
+            raise ImportError(
+                "`enable_model_offload` requires `accelerate v0.17.0` or higher."
+            )
+
+        device = torch.device(f"cuda:{gpu_id}")
+
+        if self.device.type != "cpu":
+            self.to("cpu", silence_dtype_warnings=True)
+            torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
+
+        hook = None
+        for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:
+            _, hook = cpu_offload_with_hook(
+                cpu_offloaded_model, device, prev_module_hook=hook
+            )
+
+        # We'll offload the last model manually.
+        self.final_offload_hook = hook
+
+    @property
+    def _execution_device(self):
+        r"""
+        Returns the device on which the pipeline's models will be executed. After calling
+        `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
+        hooks.
+        """
+        if not hasattr(self.unet, "_hf_hook"):
+            return self.device
+        for module in self.unet.modules():
+            if (
+                hasattr(module, "_hf_hook")
+                and hasattr(module._hf_hook, "execution_device")
+                and module._hf_hook.execution_device is not None
+            ):
+                return torch.device(module._hf_hook.execution_device)
+        return self.device
+
+    def _encode_prompt(
+        self,
+        prompt,
+        device,
+        num_images_per_prompt,
+        do_classifier_free_guidance: bool,
+        negative_prompt=None,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+             prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
+                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+        """
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            raise ValueError(
+                f"`prompt` should be either a string or a list of strings, but got {type(prompt)}."
+            )
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer(
+            prompt, padding="longest", return_tensors="pt"
+        ).input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+            text_input_ids, untruncated_ids
+        ):
+            removed_text = self.tokenizer.batch_decode(
+                untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
+            )
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {self.tokenizer.model_max_length} tokens: {removed_text}"
+            )
+
+        if (
+            hasattr(self.text_encoder.config, "use_attention_mask")
+            and self.text_encoder.config.use_attention_mask
+        ):
+            attention_mask = text_inputs.attention_mask.to(device)
+        else:
+            attention_mask = None
+
+        prompt_embeds = self.text_encoder(
+            text_input_ids.to(device),
+            attention_mask=attention_mask,
+        )
+        prompt_embeds = prompt_embeds[0]
+
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(
+            bs_embed * num_images_per_prompt, seq_len, -1
+        )
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance:
+            uncond_tokens: List[str]
+            if negative_prompt is None:
+                uncond_tokens = [""] * batch_size
+            elif type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif isinstance(negative_prompt, str):
+                uncond_tokens = [negative_prompt]
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+            else:
+                uncond_tokens = negative_prompt
+
+            max_length = prompt_embeds.shape[1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+
+            if (
+                hasattr(self.text_encoder.config, "use_attention_mask")
+                and self.text_encoder.config.use_attention_mask
+            ):
+                attention_mask = uncond_input.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            negative_prompt_embeds = self.text_encoder(
+                uncond_input.input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            negative_prompt_embeds = negative_prompt_embeds[0]
+
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(
+                dtype=self.text_encoder.dtype, device=device
+            )
+
+            negative_prompt_embeds = negative_prompt_embeds.repeat(
+                1, num_images_per_prompt, 1
+            )
+            negative_prompt_embeds = negative_prompt_embeds.view(
+                batch_size * num_images_per_prompt, seq_len, -1
+            )
+
+            # For classifier free guidance, we need to do two forward passes.
+            # Here we concatenate the unconditional and text embeddings into a single batch
+            # to avoid doing two forward passes
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+        return prompt_embeds
+
+    def decode_latents(self, latents):
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        return image
+
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(
+            inspect.signature(self.scheduler.step).parameters.keys()
+        )
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(
+            inspect.signature(self.scheduler.step).parameters.keys()
+        )
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        shape = (
+            batch_size,
+            num_channels_latents,
+            height // self.vae_scale_factor,
+            width // self.vae_scale_factor,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(
+                shape, generator=generator, device=device, dtype=dtype
+            )
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    def encode_image(self, image, device, num_images_per_prompt):
+        dtype = next(self.image_encoder.parameters()).dtype
+
+        if image.dtype == np.float32:
+            image = (image * 255).astype(np.uint8)
+            
+        image = self.feature_extractor(image, return_tensors="pt").pixel_values
+        image = image.to(device=device, dtype=dtype)
+        
+        image_embeds = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
+        image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
+
+        return torch.zeros_like(image_embeds), image_embeds
+
+    def encode_image_latents(self, image, device, num_images_per_prompt):
+        
+        dtype = next(self.image_encoder.parameters()).dtype
+
+        image = torch.from_numpy(image).unsqueeze(0).permute(0, 3, 1, 2).to(device=device) # [1, 3, H, W]
+        image = 2 * image - 1
+        image = F.interpolate(image, (256, 256), mode='bilinear', align_corners=False)
+        image = image.to(dtype=dtype)
+
+        posterior = self.vae.encode(image).latent_dist
+        latents = posterior.sample() * self.vae.config.scaling_factor # [B, C, H, W]
+        latents = latents.repeat_interleave(num_images_per_prompt, dim=0)
+
+        return torch.zeros_like(latents), latents
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: str = "",
+        image: Optional[np.ndarray] = None,
+        height: int = 256,
+        width: int = 256,
+        elevation: float = 0,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.0,
+        negative_prompt: str = "",
+        num_images_per_prompt: int = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        output_type: Optional[str] = "numpy", # pil, numpy, latents
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: int = 1,
+        num_frames: int = 4,
+        device=torch.device("cuda:0"),
+    ):
+        self.unet = self.unet.to(device=device)
+        self.vae = self.vae.to(device=device)
+        self.text_encoder = self.text_encoder.to(device=device)
+
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # imagedream variant
+        if image is not None:
+            assert isinstance(image, np.ndarray) and image.dtype == np.float32
+            self.image_encoder = self.image_encoder.to(device=device)
+            image_embeds_neg, image_embeds_pos = self.encode_image(image, device, num_images_per_prompt)
+            image_latents_neg, image_latents_pos = self.encode_image_latents(image, device, num_images_per_prompt)
+            
+        _prompt_embeds = self._encode_prompt(
+            prompt=prompt,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+        )  # type: ignore
+        prompt_embeds_neg, prompt_embeds_pos = _prompt_embeds.chunk(2)
+
+        # Prepare latent variables
+        actual_num_frames = num_frames if image is None else num_frames + 1
+        latents: torch.Tensor = self.prepare_latents(
+            actual_num_frames * num_images_per_prompt,
+            4,
+            height,
+            width,
+            prompt_embeds_pos.dtype,
+            device,
+            generator,
+            None,
+        )
+
+        if image is not None:
+            camera = get_camera(num_frames, elevation=elevation, extra_view=True).to(dtype=latents.dtype, device=device)
+        else:
+            camera = get_camera(num_frames, elevation=elevation, extra_view=False).to(dtype=latents.dtype, device=device)
+        camera = camera.repeat_interleave(num_images_per_prompt, dim=0)
+
+        # Prepare extra step kwargs.
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                multiplier = 2 if do_classifier_free_guidance else 1
+                latent_model_input = torch.cat([latents] * multiplier)
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                unet_inputs = {
+                    'x': latent_model_input,
+                    'timesteps': torch.tensor([t] * actual_num_frames * multiplier, dtype=latent_model_input.dtype, device=device),
+                    'context': torch.cat([prompt_embeds_neg] * actual_num_frames + [prompt_embeds_pos] * actual_num_frames),
+                    'num_frames': actual_num_frames,
+                    'camera': torch.cat([camera] * multiplier),
+                }
+
+                if image is not None:
+                    unet_inputs['ip'] = torch.cat([image_embeds_neg] * actual_num_frames + [image_embeds_pos] * actual_num_frames)
+                    unet_inputs['ip_img'] = torch.cat([image_latents_neg] + [image_latents_pos]) # no repeat
+                
+                # predict the noise residual
+                noise_pred = self.unet.forward(**unet_inputs)
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (
+                        noise_pred_text - noise_pred_uncond
+                    )
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents: torch.Tensor = self.scheduler.step(
+                    noise_pred, t, latents, **extra_step_kwargs, return_dict=False
+                )[0]
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)  # type: ignore
+
+        # Post-processing
+        if output_type == "latent":
+            image = latents
+        elif output_type == "pil":
+            image = self.decode_latents(latents)
+            image = self.numpy_to_pil(image)
+        else: # numpy
+            image = self.decode_latents(latents)
+
+        # Offload last model to CPU
+        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+            self.final_offload_hook.offload()
+
+        return image

main_4d.py

@@ -0,0 +1,601 @@
+import os
+import cv2
+import time
+import tqdm
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+
+import rembg
+
+from cam_utils import orbit_camera, OrbitCamera
+from gs_renderer_4d import Renderer, MiniCam
+
+from grid_put import mipmap_linear_grid_put_2d
+import imageio
+
+import copy
+
+
+class GUI:
+    def __init__(self, opt):
+        self.opt = opt  # shared with the trainer's opt to support in-place modification of rendering parameters.
+        self.gui = opt.gui # enable gui
+        self.W = opt.W
+        self.H = opt.H
+        self.cam = OrbitCamera(opt.W, opt.H, r=opt.radius, fovy=opt.fovy)
+
+        self.mode = "image"
+        # self.seed = "random"
+        self.seed = 888
+
+        self.buffer_image = np.ones((self.W, self.H, 3), dtype=np.float32)
+        self.need_update = True  # update buffer_image
+
+        # models
+        self.device = torch.device("cuda")
+        self.bg_remover = None
+
+        self.guidance_sd = None
+        self.guidance_zero123 = None
+        self.guidance_svd = None
+
+
+        self.enable_sd = False
+        self.enable_zero123 = False
+        self.enable_svd = False
+
+
+        # renderer
+        self.renderer = Renderer(self.opt, sh_degree=self.opt.sh_degree)
+        self.gaussain_scale_factor = 1
+
+        # input image
+        self.input_img = None
+        self.input_mask = None
+        self.input_img_torch = None
+        self.input_mask_torch = None
+        self.overlay_input_img = False
+        self.overlay_input_img_ratio = 0.5
+
+        self.input_img_list = None
+        self.input_mask_list = None
+        self.input_img_torch_list = None
+        self.input_mask_torch_list = None
+
+        # input text
+        self.prompt = ""
+        self.negative_prompt = ""
+
+        # training stuff
+        self.training = False
+        self.optimizer = None
+        self.step = 0
+        self.train_steps = 1  # steps per rendering loop
+        
+        # load input data from cmdline
+        if self.opt.input is not None: # True
+            self.load_input(self.opt.input) # load imgs, if has bg, then rm bg; or just load imgs
+        
+        # override prompt from cmdline
+        if self.opt.prompt is not None: # None
+            self.prompt = self.opt.prompt
+
+        # override if provide a checkpoint
+        if self.opt.load is not None: # not None
+            self.renderer.initialize(self.opt.load)  
+            # self.renderer.gaussians.load_model(opt.outdir, opt.save_path)             
+        else:
+            # initialize gaussians to a blob
+            self.renderer.initialize(num_pts=self.opt.num_pts)
+
+        self.seed_everything()
+
+    def seed_everything(self):
+        try:
+            seed = int(self.seed)
+        except:
+            seed = np.random.randint(0, 1000000)
+
+        print(f'Seed: {seed:d}')
+        os.environ["PYTHONHASHSEED"] = str(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = True
+
+        self.last_seed = seed
+
+    def prepare_train(self):
+
+        self.step = 0
+
+        # setup training
+        self.renderer.gaussians.training_setup(self.opt)
+
+        # # do not do progressive sh-level
+        self.renderer.gaussians.active_sh_degree = self.renderer.gaussians.max_sh_degree
+        self.optimizer = self.renderer.gaussians.optimizer
+
+        # default camera
+        if self.opt.mvdream or self.opt.imagedream:
+            # the second view is the front view for mvdream/imagedream.
+            pose = orbit_camera(self.opt.elevation, 90, self.opt.radius)
+        else:
+            pose = orbit_camera(self.opt.elevation, 0, self.opt.radius)
+        self.fixed_cam = MiniCam(
+            pose,
+            self.opt.ref_size,
+            self.opt.ref_size,
+            self.cam.fovy,
+            self.cam.fovx,
+            self.cam.near,
+            self.cam.far,
+        )
+
+        self.enable_sd = self.opt.lambda_sd > 0
+        self.enable_zero123 = self.opt.lambda_zero123 > 0
+        self.enable_svd = self.opt.lambda_svd > 0 and self.input_img is not None
+
+        # lazy load guidance model
+        if self.guidance_sd is None and self.enable_sd:
+            if self.opt.mvdream:
+                print(f"[INFO] loading MVDream...")
+                from guidance.mvdream_utils import MVDream
+                self.guidance_sd = MVDream(self.device)
+                print(f"[INFO] loaded MVDream!")
+            elif self.opt.imagedream:
+                print(f"[INFO] loading ImageDream...")
+                from guidance.imagedream_utils import ImageDream
+                self.guidance_sd = ImageDream(self.device)
+                print(f"[INFO] loaded ImageDream!")
+            else:
+                print(f"[INFO] loading SD...")
+                from guidance.sd_utils import StableDiffusion
+                self.guidance_sd = StableDiffusion(self.device)
+                print(f"[INFO] loaded SD!")
+
+        if self.guidance_zero123 is None and self.enable_zero123:
+            print(f"[INFO] loading zero123...")
+            from guidance.zero123_utils import Zero123
+            if self.opt.stable_zero123:
+                self.guidance_zero123 = Zero123(self.device, model_key='ashawkey/stable-zero123-diffusers')
+            else:
+                self.guidance_zero123 = Zero123(self.device, model_key='ashawkey/zero123-xl-diffusers')
+            print(f"[INFO] loaded zero123!")
+
+        if self.guidance_svd is None and self.enable_svd: # False
+            print(f"[INFO] loading SVD...")
+            from guidance.svd_utils import StableVideoDiffusion
+            self.guidance_svd = StableVideoDiffusion(self.device)
+            print(f"[INFO] loaded SVD!")
+
+        # input image
+        if self.input_img is not None:
+            self.input_img_torch = torch.from_numpy(self.input_img).permute(2, 0, 1).unsqueeze(0).to(self.device)
+            self.input_img_torch = F.interpolate(self.input_img_torch, (self.opt.ref_size, self.opt.ref_size), mode="bilinear", align_corners=False)
+
+            self.input_mask_torch = torch.from_numpy(self.input_mask).permute(2, 0, 1).unsqueeze(0).to(self.device)
+            self.input_mask_torch = F.interpolate(self.input_mask_torch, (self.opt.ref_size, self.opt.ref_size), mode="bilinear", align_corners=False)
+
+        if self.input_img_list is not None:
+            self.input_img_torch_list = [torch.from_numpy(input_img).permute(2, 0, 1).unsqueeze(0).to(self.device) for input_img in self.input_img_list]
+            self.input_img_torch_list = [F.interpolate(input_img_torch, (self.opt.ref_size, self.opt.ref_size), mode="bilinear", align_corners=False) for input_img_torch in self.input_img_torch_list]
+            
+            self.input_mask_torch_list = [torch.from_numpy(input_mask).permute(2, 0, 1).unsqueeze(0).to(self.device) for input_mask in self.input_mask_list]
+            self.input_mask_torch_list = [F.interpolate(input_mask_torch, (self.opt.ref_size, self.opt.ref_size), mode="bilinear", align_corners=False) for input_mask_torch in self.input_mask_torch_list]
+        # prepare embeddings
+        with torch.no_grad():
+
+            if self.enable_sd:
+                if self.opt.imagedream:
+                    img_pos_list, img_neg_list, ip_pos_list, ip_neg_list, emb_pos_list, emb_neg_list = [], [], [], [], [], []
+                    for _ in range(self.opt.n_views):
+                        for input_img_torch in self.input_img_torch_list:
+                            img_pos, img_neg, ip_pos, ip_neg, emb_pos, emb_neg = self.guidance_sd.get_image_text_embeds(input_img_torch, [self.prompt], [self.negative_prompt])
+                            img_pos_list.append(img_pos)
+                            img_neg_list.append(img_neg)
+                            ip_pos_list.append(ip_pos)
+                            ip_neg_list.append(ip_neg)
+                            emb_pos_list.append(emb_pos)
+                            emb_neg_list.append(emb_neg)
+                    self.guidance_sd.image_embeddings['pos'] = torch.cat(img_pos_list, 0)
+                    self.guidance_sd.image_embeddings['neg'] = torch.cat(img_pos_list, 0)
+                    self.guidance_sd.image_embeddings['ip_img'] = torch.cat(ip_pos_list, 0)
+                    self.guidance_sd.image_embeddings['neg_ip_img'] = torch.cat(ip_neg_list, 0)
+                    self.guidance_sd.embeddings['pos'] = torch.cat(emb_pos_list, 0)
+                    self.guidance_sd.embeddings['neg'] = torch.cat(emb_neg_list, 0)
+                else:
+                    self.guidance_sd.get_text_embeds([self.prompt], [self.negative_prompt])
+
+            if self.enable_zero123:
+                c_list, v_list = [], []
+                for _ in range(self.opt.n_views):
+                    for input_img_torch in self.input_img_torch_list:
+                        c, v = self.guidance_zero123.get_img_embeds(input_img_torch)
+                        c_list.append(c)
+                        v_list.append(v)
+                self.guidance_zero123.embeddings = [torch.cat(c_list, 0), torch.cat(v_list, 0)]
+            
+            if self.enable_svd:
+                self.guidance_svd.get_img_embeds(self.input_img)
+
+    def train_step(self):
+        starter = torch.cuda.Event(enable_timing=True)
+        ender = torch.cuda.Event(enable_timing=True)
+        starter.record()
+
+        for _ in range(self.train_steps): # 1
+
+            self.step += 1 # self.step starts from 0
+            step_ratio = min(1, self.step / self.opt.iters) # 1, step / 500
+
+            # update lr
+            self.renderer.gaussians.update_learning_rate(self.step)
+
+            loss = 0
+
+            self.renderer.prepare_render()
+        
+            ### known view
+            if not self.opt.imagedream:
+                for b_idx in range(self.opt.batch_size):
+                    cur_cam = copy.deepcopy(self.fixed_cam)
+                    cur_cam.time = b_idx
+                    out = self.renderer.render(cur_cam)
+
+                    # rgb loss
+                    image = out["image"].unsqueeze(0) # [1, 3, H, W] in [0, 1]
+                    loss = loss + 10000 * step_ratio * F.mse_loss(image, self.input_img_torch_list[b_idx]) / self.opt.batch_size
+
+                    # mask loss
+                    mask = out["alpha"].unsqueeze(0) # [1, 1, H, W] in [0, 1]
+                    loss = loss + 1000 * step_ratio * F.mse_loss(mask, self.input_mask_torch_list[b_idx]) / self.opt.batch_size
+
+            ### novel view (manual batch)
+            render_resolution = 128 if step_ratio < 0.3 else (256 if step_ratio < 0.6 else 512)
+            # render_resolution = 512
+            images = []
+            poses = []
+            vers, hors, radii = [], [], []
+            # avoid too large elevation (> 80 or < -80), and make sure it always cover [-30, 30]
+            min_ver = max(min(self.opt.min_ver, self.opt.min_ver - self.opt.elevation), -80 - self.opt.elevation)
+            max_ver = min(max(self.opt.max_ver, self.opt.max_ver - self.opt.elevation), 80 - self.opt.elevation)
+
+            for _ in range(self.opt.n_views):
+                for b_idx in range(self.opt.batch_size):
+
+                    # render random view
+                    ver = np.random.randint(min_ver, max_ver)
+                    hor = np.random.randint(-180, 180)
+                    radius = 0
+
+                    vers.append(ver)
+                    hors.append(hor)
+                    radii.append(radius)
+
+                    pose = orbit_camera(self.opt.elevation + ver, hor, self.opt.radius + radius)
+                    poses.append(pose)
+
+                    cur_cam = MiniCam(pose, render_resolution, render_resolution, self.cam.fovy, self.cam.fovx, self.cam.near, self.cam.far, time=b_idx)
+
+                    bg_color = torch.tensor([1, 1, 1] if np.random.rand() > self.opt.invert_bg_prob else [0, 0, 0], dtype=torch.float32, device="cuda")
+                    out = self.renderer.render(cur_cam, bg_color=bg_color)
+
+                    image = out["image"].unsqueeze(0) # [1, 3, H, W] in [0, 1]
+                    images.append(image)
+
+                    # enable mvdream training
+                    if self.opt.mvdream or self.opt.imagedream: # False
+                        for view_i in range(1, 4):
+                            pose_i = orbit_camera(self.opt.elevation + ver, hor + 90 * view_i, self.opt.radius + radius)
+                            poses.append(pose_i)
+
+                            cur_cam_i = MiniCam(pose_i, render_resolution, render_resolution, self.cam.fovy, self.cam.fovx, self.cam.near, self.cam.far)
+
+                            # bg_color = torch.tensor([0.5, 0.5, 0.5], dtype=torch.float32, device="cuda")
+                            out_i = self.renderer.render(cur_cam_i, bg_color=bg_color)
+
+                            image = out_i["image"].unsqueeze(0) # [1, 3, H, W] in [0, 1]
+                            images.append(image)
+
+
+
+            images = torch.cat(images, dim=0)
+            poses = torch.from_numpy(np.stack(poses, axis=0)).to(self.device)
+
+            # guidance loss
+            if self.enable_sd:
+                if self.opt.mvdream or self.opt.imagedream:
+                    loss = loss + self.opt.lambda_sd * self.guidance_sd.train_step(images, poses, step_ratio)
+                else:
+                    loss = loss + self.opt.lambda_sd * self.guidance_sd.train_step(images, step_ratio)
+
+            if self.enable_zero123:
+                loss = loss + self.opt.lambda_zero123 * self.guidance_zero123.train_step(images, vers, hors, radii, step_ratio) / (self.opt.batch_size * self.opt.n_views)
+
+            if self.enable_svd:
+                loss = loss + self.opt.lambda_svd * self.guidance_svd.train_step(images, step_ratio)
+            
+            # optimize step
+            loss.backward()
+            self.optimizer.step()
+            self.optimizer.zero_grad()
+
+            # densify and prune
+            if self.step >= self.opt.density_start_iter and self.step <= self.opt.density_end_iter:
+                viewspace_point_tensor, visibility_filter, radii = out["viewspace_points"], out["visibility_filter"], out["radii"]
+                self.renderer.gaussians.max_radii2D[visibility_filter] = torch.max(self.renderer.gaussians.max_radii2D[visibility_filter], radii[visibility_filter])
+                self.renderer.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter)
+
+                if self.step % self.opt.densification_interval == 0:
+                    # size_threshold = 20 if self.step > self.opt.opacity_reset_interval else None
+                    self.renderer.gaussians.densify_and_prune(self.opt.densify_grad_threshold, min_opacity=0.01, extent=0.5, max_screen_size=1)
+                
+                if self.step % self.opt.opacity_reset_interval == 0:
+                    self.renderer.gaussians.reset_opacity()
+
+        ender.record()
+        torch.cuda.synchronize()
+        t = starter.elapsed_time(ender)
+
+        self.need_update = True
+
+    
+    def load_input(self, file):
+        if self.opt.data_mode == 'c4d':
+            file_list = [os.path.join(file, f'{x * self.opt.downsample_rate}.png') for x in range(self.opt.batch_size)] 
+        elif self.opt.data_mode == 'svd':
+            # file_list = [file.replace('.png', f'_frames/{x* self.opt.downsample_rate:03d}_rgba.png') for x in range(self.opt.batch_size)]
+            # file_list = [x if os.path.exists(x) else (x.replace('_rgba.png', '.png')) for x in file_list]
+            file_list = [file.replace('.png', f'_frames/{x* self.opt.downsample_rate:03d}.png') for x in range(self.opt.batch_size)]
+        else:
+            raise NotImplementedError
+        self.input_img_list, self.input_mask_list = [], []
+        for file in file_list:
+            # load image
+            print(f'[INFO] load image from {file}...')
+            img = cv2.imread(file, cv2.IMREAD_UNCHANGED)
+            if img.shape[-1] == 3:
+                if self.bg_remover is None:
+                    self.bg_remover = rembg.new_session()
+                img = rembg.remove(img, session=self.bg_remover)
+                # cv2.imwrite(file.replace('.png', '_rgba.png'), img) 
+            img = cv2.resize(img, (self.W, self.H), interpolation=cv2.INTER_AREA)
+            img = img.astype(np.float32) / 255.0
+            input_mask = img[..., 3:]
+            # white bg
+            input_img = img[..., :3] * input_mask + (1 - input_mask)
+            # bgr to rgb
+            input_img = input_img[..., ::-1].copy()
+            self.input_img_list.append(input_img)
+            self.input_mask_list.append(input_mask)
+
+    @torch.no_grad()
+    def save_model(self, mode='geo', texture_size=1024, interp=1):
+        os.makedirs(self.opt.outdir, exist_ok=True)
+        if mode == 'geo':
+            path = f'logs/{opt.save_path}_mesh_{t:03d}.ply'
+            mesh = self.renderer.gaussians.extract_mesh_t(path, self.opt.density_thresh, t=t)
+            mesh.write_ply(path)
+
+        elif mode == 'geo+tex':
+            from mesh import Mesh, safe_normalize
+            os.makedirs(os.path.join(self.opt.outdir, self.opt.save_path+'_meshes'), exist_ok=True)
+            for t in range(self.opt.batch_size):
+                path = os.path.join(self.opt.outdir, self.opt.save_path+'_meshes', f'{t:03d}.obj')
+                mesh = self.renderer.gaussians.extract_mesh_t(path, self.opt.density_thresh, t=t)
+
+                # perform texture extraction
+                print(f"[INFO] unwrap uv...")
+                h = w = texture_size
+                mesh.auto_uv()
+                mesh.auto_normal()
+
+                albedo = torch.zeros((h, w, 3), device=self.device, dtype=torch.float32)
+                cnt = torch.zeros((h, w, 1), device=self.device, dtype=torch.float32)
+
+                vers = [0] * 8 + [-45] * 8 + [45] * 8 + [-89.9, 89.9]
+                hors = [0, 45, -45, 90, -90, 135, -135, 180] * 3 + [0, 0]
+
+                render_resolution = 512
+
+                import nvdiffrast.torch as dr
+
+                if not self.opt.force_cuda_rast and (not self.opt.gui or os.name == 'nt'):
+                    glctx = dr.RasterizeGLContext()
+                else:
+                    glctx = dr.RasterizeCudaContext()
+
+                for ver, hor in zip(vers, hors):
+                    # render image
+                    pose = orbit_camera(ver, hor, self.cam.radius)
+
+                    cur_cam = MiniCam(
+                        pose,
+                        render_resolution,
+                        render_resolution,
+                        self.cam.fovy,
+                        self.cam.fovx,
+                        self.cam.near,
+                        self.cam.far,
+                        time=t
+                    )
+                    
+                    cur_out = self.renderer.render(cur_cam)
+
+                    rgbs = cur_out["image"].unsqueeze(0) # [1, 3, H, W] in [0, 1]
+                        
+                    # get coordinate in texture image
+                    pose = torch.from_numpy(pose.astype(np.float32)).to(self.device)
+                    proj = torch.from_numpy(self.cam.perspective.astype(np.float32)).to(self.device)
+
+                    v_cam = torch.matmul(F.pad(mesh.v, pad=(0, 1), mode='constant', value=1.0), torch.inverse(pose).T).float().unsqueeze(0)
+                    v_clip = v_cam @ proj.T
+                    rast, rast_db = dr.rasterize(glctx, v_clip, mesh.f, (render_resolution, render_resolution))
+
+                    depth, _ = dr.interpolate(-v_cam[..., [2]], rast, mesh.f) # [1, H, W, 1]
+                    depth = depth.squeeze(0) # [H, W, 1]
+
+                    alpha = (rast[0, ..., 3:] > 0).float()
+
+                    uvs, _ = dr.interpolate(mesh.vt.unsqueeze(0), rast, mesh.ft)  # [1, 512, 512, 2] in [0, 1]
+
+                    # use normal to produce a back-project mask
+                    normal, _ = dr.interpolate(mesh.vn.unsqueeze(0).contiguous(), rast, mesh.fn)
+                    normal = safe_normalize(normal[0])
+
+                    # rotated normal (where [0, 0, 1] always faces camera)
+                    rot_normal = normal @ pose[:3, :3]
+                    viewcos = rot_normal[..., [2]]
+
+                    mask = (alpha > 0) & (viewcos > 0.5)  # [H, W, 1]
+                    mask = mask.view(-1)
+
+                    uvs = uvs.view(-1, 2).clamp(0, 1)[mask]
+                    rgbs = rgbs.view(3, -1).permute(1, 0)[mask].contiguous()
+                    
+                    # update texture image
+                    cur_albedo, cur_cnt = mipmap_linear_grid_put_2d(
+                        h, w,
+                        uvs[..., [1, 0]] * 2 - 1,
+                        rgbs,
+                        min_resolution=256,
+                        return_count=True,
+                    )
+                    
+                    mask = cnt.squeeze(-1) < 0.1
+                    albedo[mask] += cur_albedo[mask]
+                    cnt[mask] += cur_cnt[mask]
+
+                mask = cnt.squeeze(-1) > 0
+                albedo[mask] = albedo[mask] / cnt[mask].repeat(1, 3)
+
+                mask = mask.view(h, w)
+
+                albedo = albedo.detach().cpu().numpy()
+                mask = mask.detach().cpu().numpy()
+
+                # dilate texture
+                from sklearn.neighbors import NearestNeighbors
+                from scipy.ndimage import binary_dilation, binary_erosion
+
+                inpaint_region = binary_dilation(mask, iterations=32)
+                inpaint_region[mask] = 0
+
+                search_region = mask.copy()
+                not_search_region = binary_erosion(search_region, iterations=3)
+                search_region[not_search_region] = 0
+
+                search_coords = np.stack(np.nonzero(search_region), axis=-1)
+                inpaint_coords = np.stack(np.nonzero(inpaint_region), axis=-1)
+
+                knn = NearestNeighbors(n_neighbors=1, algorithm="kd_tree").fit(
+                    search_coords
+                )
+                _, indices = knn.kneighbors(inpaint_coords)
+
+                albedo[tuple(inpaint_coords.T)] = albedo[tuple(search_coords[indices[:, 0]].T)]
+
+                mesh.albedo = torch.from_numpy(albedo).to(self.device)
+                mesh.write(path)
+
+            
+        elif mode == 'frames':
+            os.makedirs(os.path.join(self.opt.outdir, self.opt.save_path+'_frames'), exist_ok=True)
+            for t in range(self.opt.batch_size * interp):
+                tt = t / interp
+                path = os.path.join(self.opt.outdir, self.opt.save_path+'_frames', f'{t:03d}.ply')
+                self.renderer.gaussians.save_frame_ply(path, tt)
+        else:
+            path = os.path.join(self.opt.outdir, self.opt.save_path + '_4d_model.ply')
+            self.renderer.gaussians.save_ply(path)
+            self.renderer.gaussians.save_deformation(self.opt.outdir, self.opt.save_path)
+
+        print(f"[INFO] save model to {path}.")
+
+    # no gui mode
+    def train(self, iters=500, ui=False):
+        if self.gui:
+            from visualizer.visergui import ViserViewer
+            self.viser_gui = ViserViewer(device="cuda", viewer_port=8080)
+        if iters > 0:
+            self.prepare_train()
+            if self.gui:
+                self.viser_gui.set_renderer(self.renderer, self.fixed_cam)
+            
+            for i in tqdm.trange(iters):
+                self.train_step()
+                if self.gui:
+                    self.viser_gui.update()
+        if self.opt.mesh_format == 'frames':
+            self.save_model(mode='frames', interp=4)
+        elif self.opt.mesh_format == 'obj':
+            self.save_model(mode='geo+tex')
+        
+        if self.opt.save_model:
+            self.save_model(mode='model')
+
+        # render eval
+        image_list =[]
+        nframes = self.opt.batch_size * 7 + 15 * 7
+        hor = 180
+        delta_hor = 45 / 15
+        delta_time = 1
+        for i in range(8):
+            time = 0
+            for j in range(self.opt.batch_size + 15):
+                pose = orbit_camera(self.opt.elevation, hor-180, self.opt.radius)
+                cur_cam = MiniCam(
+                    pose,
+                    512,
+                    512,
+                    self.cam.fovy,
+                    self.cam.fovx,
+                    self.cam.near,
+                    self.cam.far,
+                    time=time
+                )
+                with torch.no_grad():
+                    outputs = self.renderer.render(cur_cam)
+
+                out = outputs["image"].cpu().detach().numpy().astype(np.float32)
+                out = np.transpose(out, (1, 2, 0))
+                out = np.uint8(out*255)
+                image_list.append(out)
+
+                time = (time + delta_time) % self.opt.batch_size
+                if j >= self.opt.batch_size:
+                    hor = (hor+delta_hor) % 360
+
+
+        imageio.mimwrite(f'vis_data/{opt.save_path}.mp4', image_list, fps=7)
+
+        if self.gui:
+            while True:
+                self.viser_gui.update()
+
+if __name__ == "__main__":
+    import argparse
+    from omegaconf import OmegaConf
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", required=True, help="path to the yaml config file")
+    args, extras = parser.parse_known_args()
+
+    # override default config from cli
+    opt = OmegaConf.merge(OmegaConf.load(args.config), OmegaConf.from_cli(extras))
+    opt.save_path = os.path.splitext(os.path.basename(opt.input))[0] if opt.save_path == '' else opt.save_path
+
+
+    # auto find mesh from stage 1
+    opt.load = os.path.join(opt.outdir, opt.save_path + '_model.ply')
+
+    gui = GUI(opt)
+
+    gui.train(opt.iters)
+
+
+# python main_4d.py  --config configs/4d_low.yaml input=data/CONSISTENT4D_DATA/in-the-wild/blooming_rose

requirements.txt

@@ -0,0 +1,35 @@
+tqdm
+rich
+ninja
+numpy
+pandas
+scipy
+scikit-learn
+matplotlib
+opencv-python
+imageio
+imageio-ffmpeg
+omegaconf
+
+torch==2.1.0 --index-url https://download.pytorch.org/whl/cu118
+torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu118
+torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118
+xformer --index-url https://download.pytorch.org/whl/cu118 --no-deps
+einops
+plyfile
+pygltflib
+torchvision
+
+# for stable-diffusion
+huggingface_hub
+diffusers
+accelerate
+transformers
+
+rembg[gpu,cli]
+
+# gradio demo
+gradio
+gradio-model4dgs
+
+-e git+https://github.com/ashawkey/kiuikit.git@main#egg=kiui

scene/deformation.py

@@ -0,0 +1,241 @@
+import functools
+import math
+import os
+import time
+from tkinter import W
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.cpp_extension import load
+import torch.nn.init as init
+from scene.hexplane import HexPlaneField
+
+
+class Linear_Res(nn.Module):
+    def __init__(self, W):
+        super(Linear_Res, self).__init__()
+        self.main_stream = nn.Linear(W, W)
+
+    def forward(self, x):
+        x = F.relu(x)
+        return x + self.main_stream(x)
+
+
+class Head_Res_Net(nn.Module):
+    def __init__(self, W, H):
+        super(Head_Res_Net, self).__init__()
+        self.W = W
+        self.H = H
+
+        self.feature_out = [Linear_Res(self.W)]
+        self.feature_out.append(nn.Linear(W, self.H))
+        self.feature_out = nn.Sequential(*self.feature_out)
+    
+    def initialize_weights(self,):
+        for m in self.feature_out.modules():
+            if isinstance(m, nn.Linear):
+                init.constant_(m.weight, 0)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        return self.feature_out(x)
+
+
+
+class Deformation(nn.Module):
+    def __init__(self, D=8, W=256, input_ch=27, input_ch_time=9, skips=[], args=None, use_res=False):
+        super(Deformation, self).__init__()
+        self.D = D
+        self.W = W
+        self.input_ch = input_ch
+        self.input_ch_time = input_ch_time
+        self.skips = skips
+
+        self.no_grid = args.no_grid
+        self.grid = HexPlaneField(args.bounds, args.kplanes_config, args.multires)
+
+        self.use_res = use_res
+        if not self.use_res:
+            self.pos_deform, self.scales_deform, self.rotations_deform, self.opacity_deform = self.create_net()
+        else:
+            self.pos_deform, self.scales_deform, self.rotations_deform, self.opacity_deform = self.create_res_net()
+        self.args = args
+    
+    def create_net(self):
+        
+        mlp_out_dim = 0
+        if self.no_grid:
+            self.feature_out = [nn.Linear(4,self.W)]
+        else:
+            self.feature_out = [nn.Linear(mlp_out_dim + self.grid.feat_dim ,self.W)]
+        
+        for i in range(self.D-1):
+            self.feature_out.append(nn.ReLU())
+            self.feature_out.append(nn.Linear(self.W,self.W))
+        self.feature_out = nn.Sequential(*self.feature_out)
+        output_dim = self.W
+        return  \
+            nn.Sequential(nn.ReLU(),nn.Linear(self.W,self.W),nn.ReLU(),nn.Linear(self.W, 3)),\
+            nn.Sequential(nn.ReLU(),nn.Linear(self.W,self.W),nn.ReLU(),nn.Linear(self.W, 3)),\
+            nn.Sequential(nn.ReLU(),nn.Linear(self.W,self.W),nn.ReLU(),nn.Linear(self.W, 4)), \
+            nn.Sequential(nn.ReLU(),nn.Linear(self.W,self.W),nn.ReLU(),nn.Linear(self.W, 1))
+    
+    def create_res_net(self,):
+        
+        mlp_out_dim = 0
+
+        if self.no_grid:
+            self.feature_out = [nn.Linear(4,self.W)]
+        else:
+            self.feature_out = [nn.Linear(mlp_out_dim + self.grid.feat_dim ,self.W)]
+        
+        for i in range(self.D-1):
+            self.feature_out.append(nn.ReLU())
+            self.feature_out.append(nn.Linear(self.W,self.W))
+        self.feature_out = nn.Sequential(*self.feature_out)
+
+        output_dim = self.W
+        return  \
+            Head_Res_Net(self.W, 3), \
+            Head_Res_Net(self.W, 3), \
+            Head_Res_Net(self.W, 4), \
+            Head_Res_Net(self.W, 1) 
+
+    
+    def query_time(self, rays_pts_emb, scales_emb, rotations_emb, time_emb):
+        if self.args.no_mlp:
+            assert not self.no_grid
+            grid_feature = self.grid(rays_pts_emb[:,:3], time_emb[:,:1])
+            h = grid_feature
+        elif not self.use_res:
+            if self.no_grid:
+                h = torch.cat([rays_pts_emb[:,:3],time_emb[:,:1]],-1)
+            else:
+                grid_feature = self.grid(rays_pts_emb[:,:3], time_emb[:,:1])
+
+                h = grid_feature
+            
+            h = self.feature_out(h)
+        else:
+            if self.no_grid:
+                h = torch.cat([rays_pts_emb[:,:3],time_emb[:,:1]],-1)
+                h = self.feature_out(h)
+            else:
+                grid_feature = self.grid(rays_pts_emb[:,:3], time_emb[:,:1])
+                h = self.feature_out(grid_feature)
+        return h
+
+    def forward(self, rays_pts_emb, scales_emb=None, rotations_emb=None, opacity = None, time_emb=None):
+        if time_emb is None:
+            return self.forward_static(rays_pts_emb[:,:3])
+        else:
+            return self.forward_dynamic(rays_pts_emb, scales_emb, rotations_emb, opacity, time_emb)
+
+    def forward_static(self, rays_pts_emb):
+        grid_feature = self.grid(rays_pts_emb[:,:3])
+        dx = self.static_mlp(grid_feature)
+        return rays_pts_emb[:, :3] + dx
+
+    def forward_dynamic(self,rays_pts_emb, scales_emb, rotations_emb, opacity_emb, time_emb):
+        hidden = self.query_time(rays_pts_emb, scales_emb, rotations_emb, time_emb).float()
+        if self.args.no_mlp:
+            return hidden[:, :3], hidden[:, 3:6], hidden[:, 6:10], hidden[:, 10:11]
+        dx = self.pos_deform(hidden)
+        pts = dx
+        if self.args.no_ds:
+            scales = scales_emb[:,:3]
+        else:
+            ds = self.scales_deform(hidden)
+            scales = ds
+        if self.args.no_dr:
+            rotations = rotations_emb[:,:4]
+        else:
+            dr = self.rotations_deform(hidden)
+            rotations = dr
+        if self.args.no_do:
+            opacity = opacity_emb[:,:1] 
+        else:
+            do = self.opacity_deform(hidden) 
+            opacity = do
+
+        return pts, scales, rotations, opacity
+    def get_mlp_parameters(self):
+        parameter_list = []
+        for name, param in self.named_parameters():
+            if  "grid" not in name:
+                parameter_list.append(param)
+        return parameter_list
+    def get_grid_parameters(self):
+        return list(self.grid.parameters() ) 
+
+
+class deform_network(nn.Module):
+    def __init__(self, args) :
+        super(deform_network, self).__init__()
+        net_width = args.net_width
+        timebase_pe = args.timebase_pe
+        defor_depth= args.defor_depth
+        posbase_pe= args.posebase_pe
+        scale_rotation_pe = args.scale_rotation_pe
+        opacity_pe = args.opacity_pe
+        timenet_width = args.timenet_width
+        timenet_output = args.timenet_output
+        times_ch = 2*timebase_pe+1
+        self.timenet = nn.Sequential(
+        nn.Linear(times_ch, timenet_width), nn.ReLU(),
+        nn.Linear(timenet_width, timenet_output))
+        
+        self.use_res = args.use_res
+        if self.use_res:
+            print("Using zero-init and residual")
+        self.deformation_net = Deformation(W=net_width, D=defor_depth, input_ch=(4+3)+((4+3)*scale_rotation_pe)*2, input_ch_time=timenet_output, args=args, use_res=self.use_res)
+        self.register_buffer('time_poc', torch.FloatTensor([(2**i) for i in range(timebase_pe)]))
+        self.register_buffer('pos_poc', torch.FloatTensor([(2**i) for i in range(posbase_pe)]))
+        self.register_buffer('rotation_scaling_poc', torch.FloatTensor([(2**i) for i in range(scale_rotation_pe)]))
+        self.register_buffer('opacity_poc', torch.FloatTensor([(2**i) for i in range(opacity_pe)]))
+        self.apply(initialize_weights)
+
+        if self.use_res:
+            self.deformation_net.pos_deform.initialize_weights()
+            self.deformation_net.scales_deform.initialize_weights()
+            self.deformation_net.rotations_deform.initialize_weights()
+            self.deformation_net.opacity_deform.initialize_weights()
+
+
+    def forward(self, point, scales=None, rotations=None, opacity=None, times_sel=None):
+        if times_sel is not None:
+            return self.forward_dynamic(point, scales, rotations, opacity, times_sel)
+        else:
+            return self.forward_static(point)
+
+        
+    def forward_static(self, points):
+        points = self.deformation_net(points)
+        return points
+    def forward_dynamic(self, point, scales=None, rotations=None, opacity=None, times_sel=None):
+        means3D, scales, rotations, opacity = self.deformation_net( point,
+                                                  scales,
+                                                rotations,
+                                                opacity,
+                                                times_sel)
+        return means3D, scales, rotations, opacity
+    def get_mlp_parameters(self):
+        return self.deformation_net.get_mlp_parameters() + list(self.timenet.parameters())
+    def get_grid_parameters(self):
+        return self.deformation_net.get_grid_parameters()
+
+
+def initialize_weights(m):
+    if isinstance(m, nn.Linear):
+        init.xavier_uniform_(m.weight,gain=1)
+        if m.bias is not None:
+            init.xavier_uniform_(m.weight,gain=1)
+
+def initialize_zeros_weights(m):
+    if isinstance(m, nn.Linear):
+        init.constant_(m.weight, 0)
+        if m.bias is not None:
+            init.constant_(m.bias, 0)

scene/hexplane.py

@@ -0,0 +1,182 @@
+import itertools
+import logging as log
+from typing import Optional, Union, List, Dict, Sequence, Iterable, Collection, Callable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def get_normalized_directions(directions):
+    """SH encoding must be in the range [0, 1]
+
+    Args:
+        directions: batch of directions
+    """
+    return (directions + 1.0) / 2.0
+
+
+def normalize_aabb(pts, aabb):
+    return (pts - aabb[0]) * (2.0 / (aabb[1] - aabb[0])) - 1.0
+def grid_sample_wrapper(grid: torch.Tensor, coords: torch.Tensor, align_corners: bool = True) -> torch.Tensor:
+    grid_dim = coords.shape[-1]
+
+    if grid.dim() == grid_dim + 1:
+        # no batch dimension present, need to add it
+        grid = grid.unsqueeze(0)
+    if coords.dim() == 2:
+        coords = coords.unsqueeze(0)
+
+    if grid_dim == 2 or grid_dim == 3:
+        grid_sampler = F.grid_sample
+    else:
+        raise NotImplementedError(f"Grid-sample was called with {grid_dim}D data but is only "
+                                  f"implemented for 2 and 3D data.")
+
+    coords = coords.view([coords.shape[0]] + [1] * (grid_dim - 1) + list(coords.shape[1:]))
+    B, feature_dim = grid.shape[:2]
+    n = coords.shape[-2]
+    interp = grid_sampler(
+        grid,  # [B, feature_dim, reso, ...]
+        coords,  # [B, 1, ..., n, grid_dim]
+        align_corners=align_corners,
+        mode='bilinear', padding_mode='border')
+    interp = interp.view(B, feature_dim, n).transpose(-1, -2)  # [B, n, feature_dim]
+    interp = interp.squeeze()  # [B?, n, feature_dim?]
+    return interp
+
+def init_grid_param(
+        grid_nd: int,
+        in_dim: int,
+        out_dim: int,
+        reso: Sequence[int],
+        a: float = 0.1,
+        b: float = 0.5):
+    assert in_dim == len(reso), "Resolution must have same number of elements as input-dimension"
+    has_time_planes = in_dim == 4
+    assert grid_nd <= in_dim
+    coo_combs = list(itertools.combinations(range(in_dim), grid_nd))
+    grid_coefs = nn.ParameterList()
+    for ci, coo_comb in enumerate(coo_combs):
+        new_grid_coef = nn.Parameter(torch.empty(
+            [1, out_dim] + [reso[cc] for cc in coo_comb[::-1]]
+        ))
+        if has_time_planes and 3 in coo_comb:  # Initialize time planes to 1
+            nn.init.ones_(new_grid_coef)
+        else:
+            nn.init.uniform_(new_grid_coef, a=a, b=b)
+        grid_coefs.append(new_grid_coef)
+
+    return grid_coefs
+
+
+def interpolate_ms_features(pts: torch.Tensor,
+                            ms_grids: Collection[Iterable[nn.Module]],
+                            grid_dimensions: int,
+                            concat_features: bool,
+                            num_levels: Optional[int],
+                            ) -> torch.Tensor:
+    coo_combs = list(itertools.combinations(
+        range(pts.shape[-1]), grid_dimensions)
+    )
+    if num_levels is None:
+        num_levels = len(ms_grids)
+    multi_scale_interp = [] if concat_features else 0.
+    grid: nn.ParameterList
+    for scale_id,  grid in enumerate(ms_grids[:num_levels]):
+        interp_space = 1.
+        for ci, coo_comb in enumerate(coo_combs):
+            # interpolate in plane
+            feature_dim = grid[ci].shape[1]  # shape of grid[ci]: 1, out_dim, *reso
+            interp_out_plane = (
+                grid_sample_wrapper(grid[ci], pts[..., coo_comb])
+                .view(-1, feature_dim)
+            )
+            # compute product over planes
+            interp_space = interp_space * interp_out_plane
+
+        # combine over scales
+        if concat_features:
+            multi_scale_interp.append(interp_space)
+        else:
+            multi_scale_interp = multi_scale_interp + interp_space
+
+    if concat_features:
+        multi_scale_interp = torch.cat(multi_scale_interp, dim=-1)
+    return multi_scale_interp
+
+
+class HexPlaneField(nn.Module):
+    def __init__(
+        self,
+        
+        bounds,
+        planeconfig,
+        multires
+    ) -> None:
+        super().__init__()
+        aabb = torch.tensor([[bounds,bounds,bounds],
+                             [-bounds,-bounds,-bounds]])
+        self.aabb = nn.Parameter(aabb, requires_grad=False)
+        self.grid_config =  [planeconfig]
+        self.multiscale_res_multipliers = multires
+        self.concat_features = True
+
+        # 1. Init planes
+        self.grids = nn.ModuleList()
+        self.feat_dim = 0
+        for res in self.multiscale_res_multipliers:
+            # initialize coordinate grid
+            config = self.grid_config[0].copy()
+            # Resolution fix: multi-res only on spatial planes
+            config["resolution"] = [
+                r * res for r in config["resolution"][:3]
+            ] + config["resolution"][3:]
+            gp = init_grid_param(
+                grid_nd=config["grid_dimensions"],
+                in_dim=config["input_coordinate_dim"],
+                out_dim=config["output_coordinate_dim"],
+                reso=config["resolution"],
+            )
+            # shape[1] is out-dim - Concatenate over feature len for each scale
+            if self.concat_features:
+                self.feat_dim += gp[-1].shape[1]
+            else:
+                self.feat_dim = gp[-1].shape[1]
+            self.grids.append(gp)
+        # print(f"Initialized model grids: {self.grids}")
+        print("feature_dim:",self.feat_dim)
+
+
+    def set_aabb(self,xyz_max, xyz_min):
+        aabb = torch.tensor([
+            xyz_max,
+            xyz_min
+        ])
+        self.aabb = nn.Parameter(aabb,requires_grad=True)
+        print("Voxel Plane: set aabb=",self.aabb)
+
+    def get_density(self, pts: torch.Tensor, timestamps: Optional[torch.Tensor] = None):
+        """Computes and returns the densities."""
+
+        pts = normalize_aabb(pts, self.aabb)
+        pts = torch.cat((pts, timestamps), dim=-1)  # [n_rays, n_samples, 4]
+
+        pts = pts.reshape(-1, pts.shape[-1])
+        features = interpolate_ms_features(
+            pts, ms_grids=self.grids,  # noqa
+            grid_dimensions=self.grid_config[0]["grid_dimensions"],
+            concat_features=self.concat_features, num_levels=None)
+        if len(features) < 1:
+            features = torch.zeros((0, 1)).to(features.device)
+
+
+        return features
+
+    def forward(self,
+                pts: torch.Tensor,
+                timestamps: Optional[torch.Tensor] = None):
+
+        features = self.get_density(pts, timestamps)
+
+        return features

scene/regulation.py

@@ -0,0 +1,176 @@
+import abc
+import os
+from typing import Sequence
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.optim.lr_scheduler
+from torch import nn
+
+
+
+def compute_plane_tv(t):
+    batch_size, c, h, w = t.shape
+    count_h = batch_size * c * (h - 1) * w
+    count_w = batch_size * c * h * (w - 1)
+    h_tv = torch.square(t[..., 1:, :] - t[..., :h-1, :]).sum()
+    w_tv = torch.square(t[..., :, 1:] - t[..., :, :w-1]).sum()
+    return 2 * (h_tv / count_h + w_tv / count_w)  # This is summing over batch and c instead of avg
+
+
+def compute_plane_smoothness(t):
+    batch_size, c, h, w = t.shape
+    # Convolve with a second derivative filter, in the time dimension which is dimension 2
+    first_difference = t[..., 1:, :] - t[..., :h-1, :]  # [batch, c, h-1, w]
+    second_difference = first_difference[..., 1:, :] - first_difference[..., :h-2, :]  # [batch, c, h-2, w]
+    # Take the L2 norm of the result
+    return torch.square(second_difference).mean()
+
+
+class Regularizer():
+    def __init__(self, reg_type, initialization):
+        self.reg_type = reg_type
+        self.initialization = initialization
+        self.weight = float(self.initialization)
+        self.last_reg = None
+
+    def step(self, global_step):
+        pass
+
+    def report(self, d):
+        if self.last_reg is not None:
+            d[self.reg_type].update(self.last_reg.item())
+
+    def regularize(self, *args, **kwargs) -> torch.Tensor:
+        out = self._regularize(*args, **kwargs) * self.weight
+        self.last_reg = out.detach()
+        return out
+
+    @abc.abstractmethod
+    def _regularize(self, *args, **kwargs) -> torch.Tensor:
+        raise NotImplementedError()
+
+    def __str__(self):
+        return f"Regularizer({self.reg_type}, weight={self.weight})"
+
+
+class PlaneTV(Regularizer):
+    def __init__(self, initial_value, what: str = 'field'):
+        if what not in {'field', 'proposal_network'}:
+            raise ValueError(f'what must be one of "field" or "proposal_network" '
+                             f'but {what} was passed.')
+        name = f'planeTV-{what[:2]}'
+        super().__init__(name, initial_value)
+        self.what = what
+
+    def step(self, global_step):
+        pass
+
+    def _regularize(self, model, **kwargs):
+        multi_res_grids: Sequence[nn.ParameterList]
+        if self.what == 'field':
+            multi_res_grids = model.field.grids
+        elif self.what == 'proposal_network':
+            multi_res_grids = [p.grids for p in model.proposal_networks]
+        else:
+            raise NotImplementedError(self.what)
+        total = 0
+        # Note: input to compute_plane_tv should be of shape [batch_size, c, h, w]
+        for grids in multi_res_grids:
+            if len(grids) == 3:
+                spatial_grids = [0, 1, 2]
+            else:
+                spatial_grids = [0, 1, 3]  # These are the spatial grids; the others are spatiotemporal
+            for grid_id in spatial_grids:
+                total += compute_plane_tv(grids[grid_id])
+            for grid in grids:
+                # grid: [1, c, h, w]
+                total += compute_plane_tv(grid)
+        return total
+
+
+class TimeSmoothness(Regularizer):
+    def __init__(self, initial_value, what: str = 'field'):
+        if what not in {'field', 'proposal_network'}:
+            raise ValueError(f'what must be one of "field" or "proposal_network" '
+                             f'but {what} was passed.')
+        name = f'time-smooth-{what[:2]}'
+        super().__init__(name, initial_value)
+        self.what = what
+
+    def _regularize(self, model, **kwargs) -> torch.Tensor:
+        multi_res_grids: Sequence[nn.ParameterList]
+        if self.what == 'field':
+            multi_res_grids = model.field.grids
+        elif self.what == 'proposal_network':
+            multi_res_grids = [p.grids for p in model.proposal_networks]
+        else:
+            raise NotImplementedError(self.what)
+        total = 0
+        # model.grids is 6 x [1, rank * F_dim, reso, reso]
+        for grids in multi_res_grids:
+            if len(grids) == 3:
+                time_grids = []
+            else:
+                time_grids = [2, 4, 5]
+            for grid_id in time_grids:
+                total += compute_plane_smoothness(grids[grid_id])
+        return torch.as_tensor(total)
+
+
+
+class L1ProposalNetwork(Regularizer):
+    def __init__(self, initial_value):
+        super().__init__('l1-proposal-network', initial_value)
+
+    def _regularize(self, model, **kwargs) -> torch.Tensor:
+        grids = [p.grids for p in model.proposal_networks]
+        total = 0.0
+        for pn_grids in grids:
+            for grid in pn_grids:
+                total += torch.abs(grid).mean()
+        return torch.as_tensor(total)
+
+
+class DepthTV(Regularizer):
+    def __init__(self, initial_value):
+        super().__init__('tv-depth', initial_value)
+
+    def _regularize(self, model, model_out, **kwargs) -> torch.Tensor:
+        depth = model_out['depth']
+        tv = compute_plane_tv(
+            depth.reshape(64, 64)[None, None, :, :]
+        )
+        return tv
+
+
+class L1TimePlanes(Regularizer):
+    def __init__(self, initial_value, what='field'):
+        if what not in {'field', 'proposal_network'}:
+            raise ValueError(f'what must be one of "field" or "proposal_network" '
+                             f'but {what} was passed.')
+        super().__init__(f'l1-time-{what[:2]}', initial_value)
+        self.what = what
+
+    def _regularize(self, model, **kwargs) -> torch.Tensor:
+        # model.grids is 6 x [1, rank * F_dim, reso, reso]
+        multi_res_grids: Sequence[nn.ParameterList]
+        if self.what == 'field':
+            multi_res_grids = model.field.grids
+        elif self.what == 'proposal_network':
+            multi_res_grids = [p.grids for p in model.proposal_networks]
+        else:
+            raise NotImplementedError(self.what)
+
+        total = 0.0
+        for grids in multi_res_grids:
+            if len(grids) == 3:
+                continue
+            else:
+                # These are the spatiotemporal grids
+                spatiotemporal_grids = [2, 4, 5]
+            for grid_id in spatiotemporal_grids:
+                total += torch.abs(1 - grids[grid_id]).mean()
+        return torch.as_tensor(total)
+

scene/utils.py

@@ -0,0 +1,429 @@
+import copy
+import json
+import math
+import os
+import pathlib
+from typing import Any, Callable, List, Optional, Text, Tuple, Union
+
+import numpy as np
+import scipy.signal
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+PRNGKey = Any
+Shape = Tuple[int]
+Dtype = Any  # this could be a real type?
+Array = Any
+Activation = Callable[[Array], Array]
+Initializer = Callable[[PRNGKey, Shape, Dtype], Array]
+Normalizer = Callable[[], Callable[[Array], Array]]
+PathType = Union[Text, pathlib.PurePosixPath]
+
+from pathlib import PurePosixPath as GPath
+
+
+def _compute_residual_and_jacobian(
+    x: np.ndarray,
+    y: np.ndarray,
+    xd: np.ndarray,
+    yd: np.ndarray,
+    k1: float = 0.0,
+    k2: float = 0.0,
+    k3: float = 0.0,
+    p1: float = 0.0,
+    p2: float = 0.0,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray,
+           np.ndarray]:
+  """Auxiliary function of radial_and_tangential_undistort()."""
+
+  r = x * x + y * y
+  d = 1.0 + r * (k1 + r * (k2 + k3 * r))
+
+  fx = d * x + 2 * p1 * x * y + p2 * (r + 2 * x * x) - xd
+  fy = d * y + 2 * p2 * x * y + p1 * (r + 2 * y * y) - yd
+
+  # Compute derivative of d over [x, y]
+  d_r = (k1 + r * (2.0 * k2 + 3.0 * k3 * r))
+  d_x = 2.0 * x * d_r
+  d_y = 2.0 * y * d_r
+
+  # Compute derivative of fx over x and y.
+  fx_x = d + d_x * x + 2.0 * p1 * y + 6.0 * p2 * x
+  fx_y = d_y * x + 2.0 * p1 * x + 2.0 * p2 * y
+
+  # Compute derivative of fy over x and y.
+  fy_x = d_x * y + 2.0 * p2 * y + 2.0 * p1 * x
+  fy_y = d + d_y * y + 2.0 * p2 * x + 6.0 * p1 * y
+
+  return fx, fy, fx_x, fx_y, fy_x, fy_y
+
+
+def _radial_and_tangential_undistort(
+    xd: np.ndarray,
+    yd: np.ndarray,
+    k1: float = 0,
+    k2: float = 0,
+    k3: float = 0,
+    p1: float = 0,
+    p2: float = 0,
+    eps: float = 1e-9,
+    max_iterations=10) -> Tuple[np.ndarray, np.ndarray]:
+  """Computes undistorted (x, y) from (xd, yd)."""
+  # Initialize from the distorted point.
+  x = xd.copy()
+  y = yd.copy()
+
+  for _ in range(max_iterations):
+    fx, fy, fx_x, fx_y, fy_x, fy_y = _compute_residual_and_jacobian(
+        x=x, y=y, xd=xd, yd=yd, k1=k1, k2=k2, k3=k3, p1=p1, p2=p2)
+    denominator = fy_x * fx_y - fx_x * fy_y
+    x_numerator = fx * fy_y - fy * fx_y
+    y_numerator = fy * fx_x - fx * fy_x
+    step_x = np.where(
+        np.abs(denominator) > eps, x_numerator / denominator,
+        np.zeros_like(denominator))
+    step_y = np.where(
+        np.abs(denominator) > eps, y_numerator / denominator,
+        np.zeros_like(denominator))
+
+    x = x + step_x
+    y = y + step_y
+
+  return x, y
+
+
+class Camera:
+  """Class to handle camera geometry."""
+
+  def __init__(self,
+               orientation: np.ndarray,
+               position: np.ndarray,
+               focal_length: Union[np.ndarray, float],
+               principal_point: np.ndarray,
+               image_size: np.ndarray,
+               skew: Union[np.ndarray, float] = 0.0,
+               pixel_aspect_ratio: Union[np.ndarray, float] = 1.0,
+               radial_distortion: Optional[np.ndarray] = None,
+               tangential_distortion: Optional[np.ndarray] = None,
+               dtype=np.float32):
+    """Constructor for camera class."""
+    if radial_distortion is None:
+      radial_distortion = np.array([0.0, 0.0, 0.0], dtype)
+    if tangential_distortion is None:
+      tangential_distortion = np.array([0.0, 0.0], dtype)
+
+    self.orientation = np.array(orientation, dtype)
+    self.position = np.array(position, dtype)
+    self.focal_length = np.array(focal_length, dtype)
+    self.principal_point = np.array(principal_point, dtype)
+    self.skew = np.array(skew, dtype)
+    self.pixel_aspect_ratio = np.array(pixel_aspect_ratio, dtype)
+    self.radial_distortion = np.array(radial_distortion, dtype)
+    self.tangential_distortion = np.array(tangential_distortion, dtype)
+    self.image_size = np.array(image_size, np.uint32)
+    self.dtype = dtype
+
+  @classmethod
+  def from_json(cls, path: PathType):
+    """Loads a JSON camera into memory."""
+    path = GPath(path)
+    # with path.open('r') as fp:
+    with open(path, 'r') as fp:
+      camera_json = json.load(fp)
+
+    # Fix old camera JSON.
+    if 'tangential' in camera_json:
+      camera_json['tangential_distortion'] = camera_json['tangential']
+
+    return cls(
+        orientation=np.asarray(camera_json['orientation']),
+        position=np.asarray(camera_json['position']),
+        focal_length=camera_json['focal_length'],
+        principal_point=np.asarray(camera_json['principal_point']),
+        skew=camera_json['skew'],
+        pixel_aspect_ratio=camera_json['pixel_aspect_ratio'],
+        radial_distortion=np.asarray(camera_json['radial_distortion']),
+        tangential_distortion=np.asarray(camera_json['tangential_distortion']),
+        image_size=np.asarray(camera_json['image_size']),
+    )
+
+  def to_json(self):
+    return {
+        k: (v.tolist() if hasattr(v, 'tolist') else v)
+        for k, v in self.get_parameters().items()
+    }
+
+  def get_parameters(self):
+    return {
+        'orientation': self.orientation,
+        'position': self.position,
+        'focal_length': self.focal_length,
+        'principal_point': self.principal_point,
+        'skew': self.skew,
+        'pixel_aspect_ratio': self.pixel_aspect_ratio,
+        'radial_distortion': self.radial_distortion,
+        'tangential_distortion': self.tangential_distortion,
+        'image_size': self.image_size,
+    }
+
+  @property
+  def scale_factor_x(self):
+    return self.focal_length
+
+  @property
+  def scale_factor_y(self):
+    return self.focal_length * self.pixel_aspect_ratio
+
+  @property
+  def principal_point_x(self):
+    return self.principal_point[0]
+
+  @property
+  def principal_point_y(self):
+    return self.principal_point[1]
+
+  @property
+  def has_tangential_distortion(self):
+    return any(self.tangential_distortion != 0.0)
+
+  @property
+  def has_radial_distortion(self):
+    return any(self.radial_distortion != 0.0)
+
+  @property
+  def image_size_y(self):
+    return self.image_size[1]
+
+  @property
+  def image_size_x(self):
+    return self.image_size[0]
+
+  @property
+  def image_shape(self):
+    return self.image_size_y, self.image_size_x
+
+  @property
+  def optical_axis(self):
+    return self.orientation[2, :]
+
+  @property
+  def translation(self):
+    return -np.matmul(self.orientation, self.position)
+
+  def pixel_to_local_rays(self, pixels: np.ndarray):
+    """Returns the local ray directions for the provided pixels."""
+    y = ((pixels[..., 1] - self.principal_point_y) / self.scale_factor_y)
+    x = ((pixels[..., 0] - self.principal_point_x - y * self.skew) /
+         self.scale_factor_x)
+
+    if self.has_radial_distortion or self.has_tangential_distortion:
+      x, y = _radial_and_tangential_undistort(
+          x,
+          y,
+          k1=self.radial_distortion[0],
+          k2=self.radial_distortion[1],
+          k3=self.radial_distortion[2],
+          p1=self.tangential_distortion[0],
+          p2=self.tangential_distortion[1])
+
+    dirs = np.stack([x, y, np.ones_like(x)], axis=-1)
+    return dirs / np.linalg.norm(dirs, axis=-1, keepdims=True)
+
+  def pixels_to_rays(self, pixels: np.ndarray) -> np.ndarray:
+    """Returns the rays for the provided pixels.
+
+    Args:
+      pixels: [A1, ..., An, 2] tensor or np.array containing 2d pixel positions.
+
+    Returns:
+        An array containing the normalized ray directions in world coordinates.
+    """
+    if pixels.shape[-1] != 2:
+      raise ValueError('The last dimension of pixels must be 2.')
+    if pixels.dtype != self.dtype:
+      raise ValueError(f'pixels dtype ({pixels.dtype!r}) must match camera '
+                       f'dtype ({self.dtype!r})')
+
+    batch_shape = pixels.shape[:-1]
+    pixels = np.reshape(pixels, (-1, 2))
+
+    local_rays_dir = self.pixel_to_local_rays(pixels)
+    rays_dir = np.matmul(self.orientation.T, local_rays_dir[..., np.newaxis])
+    rays_dir = np.squeeze(rays_dir, axis=-1)
+
+    # Normalize rays.
+    rays_dir /= np.linalg.norm(rays_dir, axis=-1, keepdims=True)
+    rays_dir = rays_dir.reshape((*batch_shape, 3))
+    return rays_dir
+
+  def pixels_to_points(self, pixels: np.ndarray, depth: np.ndarray):
+    rays_through_pixels = self.pixels_to_rays(pixels)
+    cosa = np.matmul(rays_through_pixels, self.optical_axis)
+    points = (
+        rays_through_pixels * depth[..., np.newaxis] / cosa[..., np.newaxis] +
+        self.position)
+    return points
+
+  def points_to_local_points(self, points: np.ndarray):
+    translated_points = points - self.position
+    local_points = (np.matmul(self.orientation, translated_points.T)).T
+    return local_points
+
+  def project(self, points: np.ndarray):
+    """Projects a 3D point (x,y,z) to a pixel position (x,y)."""
+    batch_shape = points.shape[:-1]
+    points = points.reshape((-1, 3))
+    local_points = self.points_to_local_points(points)
+
+    # Get normalized local pixel positions.
+    x = local_points[..., 0] / local_points[..., 2]
+    y = local_points[..., 1] / local_points[..., 2]
+    r2 = x**2 + y**2
+
+    # Apply radial distortion.
+    distortion = 1.0 + r2 * (
+        self.radial_distortion[0] + r2 *
+        (self.radial_distortion[1] + self.radial_distortion[2] * r2))
+
+    # Apply tangential distortion.
+    x_times_y = x * y
+    x = (
+        x * distortion + 2.0 * self.tangential_distortion[0] * x_times_y +
+        self.tangential_distortion[1] * (r2 + 2.0 * x**2))
+    y = (
+        y * distortion + 2.0 * self.tangential_distortion[1] * x_times_y +
+        self.tangential_distortion[0] * (r2 + 2.0 * y**2))
+
+    # Map the distorted ray to the image plane and return the depth.
+    pixel_x = self.focal_length * x + self.skew * y + self.principal_point_x
+    pixel_y = (self.focal_length * self.pixel_aspect_ratio * y
+               + self.principal_point_y)
+
+    pixels = np.stack([pixel_x, pixel_y], axis=-1)
+    return pixels.reshape((*batch_shape, 2))
+
+  def get_pixel_centers(self):
+    """Returns the pixel centers."""
+    xx, yy = np.meshgrid(np.arange(self.image_size_x, dtype=self.dtype),
+                         np.arange(self.image_size_y, dtype=self.dtype))
+    return np.stack([xx, yy], axis=-1) + 0.5
+
+  def scale(self, scale: float):
+    """Scales the camera."""
+    if scale <= 0:
+      raise ValueError('scale needs to be positive.')
+
+    new_camera = Camera(
+        orientation=self.orientation.copy(),
+        position=self.position.copy(),
+        focal_length=self.focal_length * scale,
+        principal_point=self.principal_point.copy() * scale,
+        skew=self.skew,
+        pixel_aspect_ratio=self.pixel_aspect_ratio,
+        radial_distortion=self.radial_distortion.copy(),
+        tangential_distortion=self.tangential_distortion.copy(),
+        image_size=np.array((int(round(self.image_size[0] * scale)),
+                             int(round(self.image_size[1] * scale)))),
+    )
+    return new_camera
+
+  def look_at(self, position, look_at, up, eps=1e-6):
+    """Creates a copy of the camera which looks at a given point.
+
+    Copies the provided vision_sfm camera and returns a new camera that is
+    positioned at `camera_position` while looking at `look_at_position`.
+    Camera intrinsics are copied by this method. A common value for the
+    up_vector is (0, 1, 0).
+
+    Args:
+      position: A (3,) numpy array representing the position of the camera.
+      look_at: A (3,) numpy array representing the location the camera
+        looks at.
+      up: A (3,) numpy array representing the up direction, whose
+        projection is parallel to the y-axis of the image plane.
+      eps: a small number to prevent divides by zero.
+
+    Returns:
+      A new camera that is copied from the original but is positioned and
+        looks at the provided coordinates.
+
+    Raises:
+      ValueError: If the camera position and look at position are very close
+        to each other or if the up-vector is parallel to the requested optical
+        axis.
+    """
+
+    look_at_camera = self.copy()
+    optical_axis = look_at - position
+    norm = np.linalg.norm(optical_axis)
+    if norm < eps:
+      raise ValueError('The camera center and look at position are too close.')
+    optical_axis /= norm
+
+    right_vector = np.cross(optical_axis, up)
+    norm = np.linalg.norm(right_vector)
+    if norm < eps:
+      raise ValueError('The up-vector is parallel to the optical axis.')
+    right_vector /= norm
+
+    # The three directions here are orthogonal to each other and form a right
+    # handed coordinate system.
+    camera_rotation = np.identity(3)
+    camera_rotation[0, :] = right_vector
+    camera_rotation[1, :] = np.cross(optical_axis, right_vector)
+    camera_rotation[2, :] = optical_axis
+
+    look_at_camera.position = position
+    look_at_camera.orientation = camera_rotation
+    return look_at_camera
+
+  def crop_image_domain(
+      self, left: int = 0, right: int = 0, top: int = 0, bottom: int = 0):
+    """Returns a copy of the camera with adjusted image bounds.
+
+    Args:
+      left: number of pixels by which to reduce (or augment, if negative) the
+        image domain at the associated boundary.
+      right: likewise.
+      top: likewise.
+      bottom: likewise.
+
+    The crop parameters may not cause the camera image domain dimensions to
+    become non-positive.
+
+    Returns:
+      A camera with adjusted image dimensions.  The focal length is unchanged,
+      and the principal point is updated to preserve the original principal
+      axis.
+    """
+
+    crop_left_top = np.array([left, top])
+    crop_right_bottom = np.array([right, bottom])
+    new_resolution = self.image_size - crop_left_top - crop_right_bottom
+    new_principal_point = self.principal_point - crop_left_top
+    if np.any(new_resolution <= 0):
+      raise ValueError('Crop would result in non-positive image dimensions.')
+
+    new_camera = self.copy()
+    new_camera.image_size = np.array([int(new_resolution[0]),
+                                      int(new_resolution[1])])
+    new_camera.principal_point = np.array([new_principal_point[0],
+                                           new_principal_point[1]])
+    return new_camera
+
+  def copy(self):
+    return copy.deepcopy(self)
+
+
+''' Misc
+'''
+mse2psnr = lambda x : -10. * torch.log10(x)
+to8b = lambda x : (255*np.clip(x,0,1)).astype(np.uint8)
+
+
+
+''' Checkpoint utils
+'''

scripts/add_bg_to_gt.py

@@ -0,0 +1,18 @@
+import os
+import cv2
+
+os.makedirs('data/CONSISTENT4D_DATA/test_dataset/eval_gt_rgb', exist_ok=True)
+file_list = []
+for img_name in ['aurorus', 'crocodile', 'guppie', 'monster', 'pistol', 'skull', 'trump']:
+    os.makedirs(f'data/CONSISTENT4D_DATA/test_dataset/eval_gt_rgb/{img_name}', exist_ok=True)
+    for view in range(4):
+        os.makedirs(f'datdata/CONSISTENT4D_DATAa/test_dataset/eval_gt_rgb/{img_name}/eval_{view}', exist_ok=True)
+        for t in range(32):
+            file_list.append(f'data/CONSISTENT4D_DATA/test_dataset/eval_gt/{img_name}/eval_{view}/{t}.png')
+for file in file_list:
+    img = cv2.imread(file, cv2.IMREAD_UNCHANGED)
+    input_mask = img[..., 3:]
+    input_mask = input_mask / 255.
+    input_img = img[..., :3] * input_mask + (1 - input_mask) * 255
+    fpath = file.replace('eval_gt', 'eval_gt_rgb')
+    cv2.imwrite(fpath, input_img) 

scripts/convert_obj_to_video.py

@@ -0,0 +1,20 @@
+import os
+import glob
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--dir', default='logs', type=str, help='Directory where obj files are stored')
+parser.add_argument('--out', default='videos', type=str, help='Directory where videos will be saved')
+args = parser.parse_args()
+
+out = args.out
+os.makedirs(out, exist_ok=True)
+
+files = glob.glob(f'{args.dir}/*.obj')
+for f in files:
+    name = os.path.basename(f)
+    # first stage model, ignore
+    if name.endswith('_mesh.obj'): 
+        continue
+    print(f'[INFO] process {name}')
+    os.system(f"python -m kiui.render {f} --save_video {os.path.join(out, name.replace('.obj', '.mp4'))} ")

scripts/gen_vid.py

@@ -0,0 +1,117 @@
+import torch
+
+from diffusers import StableVideoDiffusionPipeline
+
+from PIL import Image
+import numpy as np
+
+import cv2
+import rembg
+
+import argparse
+import imageio
+import os
+
+def add_margin(pil_img, top, right, bottom, left, color):
+    width, height = pil_img.size
+    new_width = width + right + left
+    new_height = height + top + bottom
+    result = Image.new(pil_img.mode, (new_width, new_height), color)
+    result.paste(pil_img, (left, top))
+    return result
+
+def resize_image(image, output_size=(1024, 576)):
+    image = image.resize((output_size[1],output_size[1]))
+    pad_size = (output_size[0]-output_size[1]) //2
+    image = add_margin(image, 0, pad_size, 0, pad_size, tuple(np.array(image)[0,0]))
+    return image
+
+
+def load_image(file, W, H, bg='white'):
+    # load image
+    print(f'[INFO] load image from {file}...')
+    img = cv2.imread(file, cv2.IMREAD_UNCHANGED)
+    bg_remover = rembg.new_session()
+    img = rembg.remove(img, session=bg_remover)
+    img = cv2.resize(img, (W, H), interpolation=cv2.INTER_AREA)
+    img = img.astype(np.float32) / 255.0
+    input_mask = img[..., 3:]
+    # white bg
+    if bg == 'white':
+        input_img = img[..., :3] * input_mask + (1 - input_mask)
+    elif bg == 'black':
+        input_img = img[..., :3]
+    else:
+        raise NotImplementedError
+    # bgr to rgb
+    input_img = input_img[..., ::-1].copy()
+    input_img = Image.fromarray(np.uint8(input_img*255))
+    return input_img
+
+def load_image_w_bg(file, W, H):
+    # load image
+    print(f'[INFO] load image from {file}...')
+    img = cv2.imread(file, cv2.IMREAD_UNCHANGED)
+    img = cv2.resize(img, (W, H), interpolation=cv2.INTER_AREA)
+    img = img.astype(np.float32) / 255.0
+    input_img = img[..., :3]
+    # bgr to rgb
+    input_img = input_img[..., ::-1].copy()
+    input_img = Image.fromarray(np.uint8(input_img*255))
+    return input_img
+
+def gen_vid(input_path, seed, bg, is_pad):
+    name = input_path.split('/')[-1].split('.')[0]
+    input_dir = os.path.dirname(input_path)
+    pipe = StableVideoDiffusionPipeline.from_pretrained(
+        "stabilityai/stable-video-diffusion-img2vid", torch_dtype=torch.float16, variant="fp16"
+    )
+    # pipe.enable_model_cpu_offload()
+    pipe.to("cuda")
+
+    if is_pad:
+        height, width = 576, 1024
+    else:
+        height, width = 512, 512
+
+    if seed is None:
+        for bg in ['white', 'black', 'orig']:
+            if bg == 'orig':
+                if 'rgba' in name:
+                    continue
+                image = load_image_w_bg(input_path, width, height)
+            else:
+                image = load_image(input_path, width, height, bg)
+            if is_pad:
+                image = resize_image(image, output_size=(width, height))
+            for seed in range(20):
+                generator = torch.manual_seed(seed)
+                frames = pipe(image, height, width, generator=generator).frames[0]
+                imageio.mimwrite(f"{input_dir}/videos/{name}_{bg}_{seed:03}.mp4", frames, fps=7)
+    else:
+        if bg == 'orig':
+            if 'rgba' in name:
+                raise ValueError
+            image = load_image_w_bg(input_path, width, height)
+        else:
+            image = load_image(input_path, width, height, bg)
+        if is_pad:
+            image = resize_image(image, output_size=(width, height))
+        generator = torch.manual_seed(seed)
+        frames = pipe(image, height, width, generator=generator).frames[0]
+
+        imageio.mimwrite(f"{input_dir}/{name}_generated.mp4", frames, fps=7)
+        os.makedirs(f"{input_dir}/{name}_frames", exist_ok=True)
+        for idx, img in enumerate(frames):
+            if is_pad:
+                img = img.crop(((width-height) //2, 0, width - (width-height) //2, height))
+            img.save(f"{input_dir}/{name}_frames/{idx:03}.png")
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--path", type=str, required=True)
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument("--bg", type=str, default='white')
+    parser.add_argument("--is_pad", type=bool, default=False)
+    args, extras = parser.parse_known_args()
+    gen_vid(args.path, args.seed, args.bg, args.is_pad)

scripts/process.py

@@ -0,0 +1,92 @@
+import os
+import glob
+import sys
+import cv2
+import argparse
+import numpy as np
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms
+from PIL import Image
+import rembg
+
+class BLIP2():
+    def __init__(self, device='cuda'):
+        self.device = device
+        from transformers import AutoProcessor, Blip2ForConditionalGeneration
+        self.processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
+        self.model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16).to(device)
+
+    @torch.no_grad()
+    def __call__(self, image):
+        image = Image.fromarray(image)
+        inputs = self.processor(image, return_tensors="pt").to(self.device, torch.float16)
+
+        generated_ids = self.model.generate(**inputs, max_new_tokens=20)
+        generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+
+        return generated_text
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('path', type=str, help="path to image (png, jpeg, etc.)")
+    parser.add_argument('--model', default='u2net', type=str, help="rembg model, see https://github.com/danielgatis/rembg#models")
+    parser.add_argument('--size', default=256, type=int, help="output resolution")
+    parser.add_argument('--border_ratio', default=0.2, type=float, help="output border ratio")
+    parser.add_argument('--recenter', type=bool, default=True, help="recenter, potentially not helpful for multiview zero123")    
+    opt = parser.parse_args()
+
+    session = rembg.new_session(model_name=opt.model)
+
+    if os.path.isdir(opt.path):
+        print(f'[INFO] processing directory {opt.path}...')
+        files = glob.glob(f'{opt.path}/*')
+        out_dir = opt.path
+    else: # isfile
+        files = [opt.path]
+        out_dir = os.path.dirname(opt.path)
+    
+    for file in files:
+
+        out_base = os.path.basename(file).split('.')[0]
+        out_rgba = os.path.join(out_dir, out_base + '_rgba.png')
+
+        # load image
+        print(f'[INFO] loading image {file}...')
+        image = cv2.imread(file, cv2.IMREAD_UNCHANGED)
+        
+        # carve background
+        print(f'[INFO] background removal...')
+        carved_image = rembg.remove(image, session=session) # [H, W, 4]
+        mask = carved_image[..., -1] > 0
+
+        # recenter
+        if opt.recenter:
+            print(f'[INFO] recenter...')
+            final_rgba = np.zeros((opt.size, opt.size, 4), dtype=np.uint8)
+            
+            coords = np.nonzero(mask)
+            x_min, x_max = coords[0].min(), coords[0].max()
+            y_min, y_max = coords[1].min(), coords[1].max()
+            h = x_max - x_min
+            w = y_max - y_min
+            desired_size = int(opt.size * (1 - opt.border_ratio))
+            scale = desired_size / max(h, w)
+            h2 = int(h * scale)
+            w2 = int(w * scale)
+            x2_min = (opt.size - h2) // 2
+            x2_max = x2_min + h2
+            y2_min = (opt.size - w2) // 2
+            y2_max = y2_min + w2
+            final_rgba[x2_min:x2_max, y2_min:y2_max] = cv2.resize(carved_image[x_min:x_max, y_min:y_max], (w2, h2), interpolation=cv2.INTER_AREA)
+            
+        else:
+            final_rgba = carved_image
+        
+        # write image
+        cv2.imwrite(out_rgba, final_rgba)

scripts/runall.py

@@ -0,0 +1,48 @@
+import os
+import glob
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--dir', default='data', type=str, help='Directory where processed images are stored')
+parser.add_argument('--out', default='logs', type=str, help='Directory where obj files will be saved')
+parser.add_argument('--video-out', default='videos', type=str, help='Directory where videos will be saved')
+parser.add_argument('--gpu', default=0, type=int, help='ID of GPU to use')
+parser.add_argument('--elevation', default=0, type=int, help='Elevation angle of view in degrees')
+parser.add_argument('--config', default='configs', type=str, help='Path to config directory, which contains image.yaml')
+args = parser.parse_args()
+
+files = glob.glob(f'{args.dir}/*_rgba.png')
+configs_dir = args.config
+
+# check if image.yaml exists
+if not os.path.exists(os.path.join(configs_dir, 'image.yaml')):
+    raise FileNotFoundError(
+        f'image.yaml not found in {configs_dir} directory. Please check if the directory is correct.'
+    )
+
+# create output directories if not exists
+out_dir = args.out
+os.makedirs(out_dir, exist_ok=True)
+video_dir = args.video_out
+os.makedirs(video_dir, exist_ok=True)
+
+
+for file in files:
+    name = os.path.basename(file).replace("_rgba.png", "")
+    print(f'======== processing {name} ========')
+    # first stage
+    os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main.py '
+              f'--config {configs_dir}/image.yaml '
+              f'input={file} '
+              f'save_path={name} elevation={args.elevation}')
+    # second stage
+    os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main2.py '
+              f'--config {configs_dir}/image.yaml '
+              f'input={file} '
+              f'save_path={name} elevation={args.elevation}')
+    # export video
+    mesh_path = os.path.join(out_dir, f'{name}.obj')
+    os.system(f'python -m kiui.render {mesh_path} '
+              f'--save_video {video_dir}/{name}.mp4 '
+              f'--wogui '
+              f'--elevation {args.elevation}')

scripts/runall_mvdream.py

@@ -0,0 +1,44 @@
+import os
+import glob
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--gpu', default=0, type=int)
+args = parser.parse_args()
+
+prompts = [
+    # ('butterfly', 'a beautiful, intricate butterfly'),
+    # ('boy', 'a nendoroid of a chibi cute boy'),
+    # ('axe', 'a viking axe, fantasy, blender'),
+    # ('dog_rocket', 'corgi riding a rocket'),
+    ('teapot', 'a chinese teapot'),
+    ('squirrel_guitar', 'a DSLR photo of a squirrel playing guitar'),
+    # ('house', 'fisherman house, cute, cartoon, blender, stylized'),
+    # ('ship', 'Higly detailed, majestic royal tall ship, realistic painting'),
+    ('einstein', 'Albert Einstein with grey suit is riding a bicycle'),
+    # ('angle', 'a statue of an angle'),
+    ('lion', 'A 3D model of Simba, the lion cub from The Lion King, standing majestically on Pride Rock, character'),
+    # ('paris', 'mini Paris, highly detailed 3d model'),
+    # ('pig_backpack', 'a pig wearing a backpack'),
+    ('pisa_tower', 'Picture of the Leaning Tower of Pisa, featuring its tilted structure and marble facade'),
+    # ('robot', 'a human-like full body robot'),
+    ('coin', 'a golden coin'),
+    # ('cake', 'a delicious and beautiful cake'),
+    # ('horse', 'a DSLR photo of a horse'),
+    # ('cat', 'a photo of a cat'),
+    ('cat_hat', 'a photo of a cat wearing a wizard hat'),
+    # ('cat_ball', 'a photo of a cat playing with a red ball'),
+    # ('nendoroid', 'a nendoroid of a chibi girl'),
+
+]
+
+for name, prompt in prompts:
+    print(f'======== processing {name} ========')
+    # first stage
+    os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main.py --config configs/text_mv.yaml prompt="{prompt}" save_path={name}')
+    # second stage
+    os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main2.py --config configs/text_mv.yaml  prompt="{prompt}" save_path={name}')
+    # export video
+    mesh_path = os.path.join('logs', f'{name}.obj')
+    os.makedirs('videos', exist_ok=True)
+    os.system(f'python -m kiui.render {mesh_path} --save_video videos/{name}.mp4 --wogui')

scripts/runall_sd.py

@@ -0,0 +1,45 @@
+import os
+import glob
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--gpu', default=0, type=int)
+args = parser.parse_args()
+
+prompts = [
+    ('strawberry', 'a ripe strawberry'),
+    ('cactus_pot', 'a small saguaro cactus planted in a clay pot'),
+    ('hamburger', 'a delicious hamburger'),
+    ('icecream', 'an icecream'),
+    ('tulip', 'a blue tulip'),
+    ('pineapple', 'a ripe pineapple'),
+    ('goblet', 'a golden goblet'),
+    # ('squitopus', 'a squirrel-octopus hybrid'),
+    # ('astronaut', 'Michelangelo style statue of an astronaut'),
+    # ('teddy_bear', 'a teddy bear'),
+    # ('corgi_nurse', 'a plush toy of a corgi nurse'),
+    # ('teapot', 'a blue and white porcelain teapot'),
+    # ('skull', "a human skull"),
+    # ('penguin', 'a penguin'),
+    # ('campfire', 'a campfire'),
+    # ('donut', 'a donut with pink icing'),
+    # ('cupcake', 'a birthday cupcake'),
+    # ('pie', 'shepherds pie'),
+    # ('cone', 'a traffic cone'),
+    # ('schoolbus', 'a schoolbus'),
+    # ('avocado_chair', 'a chair that looks like an avocado'),
+    # ('glasses', 'a pair of sunglasses')
+    # ('potion', 'a bottle of green potion'),
+    # ('chalice', 'a delicate chalice'),
+]
+
+for name, prompt in prompts:
+    print(f'======== processing {name} ========')
+    # first stage
+    os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main.py --config configs/text.yaml prompt="{prompt}" save_path={name}')
+    # second stage
+    os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main2.py --config configs/text.yaml  prompt="{prompt}" save_path={name}')
+    # export video
+    mesh_path = os.path.join('logs', f'{name}.obj')
+    os.makedirs('videos', exist_ok=True)
+    os.system(f'python -m kiui.render {mesh_path} --save_video videos/{name}.mp4 --wogui')

sh_utils.py

@@ -0,0 +1,118 @@
+#  Copyright 2021 The PlenOctree Authors.
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  1. Redistributions of source code must retain the above copyright notice,
+#  this list of conditions and the following disclaimer.
+#
+#  2. Redistributions in binary form must reproduce the above copyright notice,
+#  this list of conditions and the following disclaimer in the documentation
+#  and/or other materials provided with the distribution.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+#  POSSIBILITY OF SUCH DAMAGE.
+
+import torch
+
+C0 = 0.28209479177387814
+C1 = 0.4886025119029199
+C2 = [
+    1.0925484305920792,
+    -1.0925484305920792,
+    0.31539156525252005,
+    -1.0925484305920792,
+    0.5462742152960396
+]
+C3 = [
+    -0.5900435899266435,
+    2.890611442640554,
+    -0.4570457994644658,
+    0.3731763325901154,
+    -0.4570457994644658,
+    1.445305721320277,
+    -0.5900435899266435
+]
+C4 = [
+    2.5033429417967046,
+    -1.7701307697799304,
+    0.9461746957575601,
+    -0.6690465435572892,
+    0.10578554691520431,
+    -0.6690465435572892,
+    0.47308734787878004,
+    -1.7701307697799304,
+    0.6258357354491761,
+]   
+
+
+def eval_sh(deg, sh, dirs):
+    """
+    Evaluate spherical harmonics at unit directions
+    using hardcoded SH polynomials.
+    Works with torch/np/jnp.
+    ... Can be 0 or more batch dimensions.
+    Args:
+        deg: int SH deg. Currently, 0-3 supported
+        sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2]
+        dirs: jnp.ndarray unit directions [..., 3]
+    Returns:
+        [..., C]
+    """
+    assert deg <= 4 and deg >= 0
+    coeff = (deg + 1) ** 2
+    assert sh.shape[-1] >= coeff
+
+    result = C0 * sh[..., 0]
+    if deg > 0:
+        x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3]
+        result = (result -
+                C1 * y * sh[..., 1] +
+                C1 * z * sh[..., 2] -
+                C1 * x * sh[..., 3])
+
+        if deg > 1:
+            xx, yy, zz = x * x, y * y, z * z
+            xy, yz, xz = x * y, y * z, x * z
+            result = (result +
+                    C2[0] * xy * sh[..., 4] +
+                    C2[1] * yz * sh[..., 5] +
+                    C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] +
+                    C2[3] * xz * sh[..., 7] +
+                    C2[4] * (xx - yy) * sh[..., 8])
+
+            if deg > 2:
+                result = (result +
+                C3[0] * y * (3 * xx - yy) * sh[..., 9] +
+                C3[1] * xy * z * sh[..., 10] +
+                C3[2] * y * (4 * zz - xx - yy)* sh[..., 11] +
+                C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] +
+                C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] +
+                C3[5] * z * (xx - yy) * sh[..., 14] +
+                C3[6] * x * (xx - 3 * yy) * sh[..., 15])
+
+                if deg > 3:
+                    result = (result + C4[0] * xy * (xx - yy) * sh[..., 16] +
+                            C4[1] * yz * (3 * xx - yy) * sh[..., 17] +
+                            C4[2] * xy * (7 * zz - 1) * sh[..., 18] +
+                            C4[3] * yz * (7 * zz - 3) * sh[..., 19] +
+                            C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] +
+                            C4[5] * xz * (7 * zz - 3) * sh[..., 21] +
+                            C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] +
+                            C4[7] * xz * (xx - 3 * yy) * sh[..., 23] +
+                            C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) * sh[..., 24])
+    return result
+
+def RGB2SH(rgb):
+    return (rgb - 0.5) / C0
+
+def SH2RGB(sh):
+    return sh * C0 + 0.5

utils/camera_utils.py

@@ -0,0 +1,65 @@
+#
+# Copyright (C) 2023, Inria
+# GRAPHDECO research group, https://team.inria.fr/graphdeco
+# All rights reserved.
+#
+# This software is free for non-commercial, research and evaluation use 
+# under the terms of the LICENSE.md file.
+#
+# For inquiries contact  [email protected]
+#
+
+from scene.cameras import Camera
+import numpy as np
+from utils.general_utils import PILtoTorch
+from utils.graphics_utils import fov2focal
+
+WARNED = False
+
+def loadCam(args, id, cam_info, resolution_scale):
+
+
+    # resized_image_rgb = PILtoTorch(cam_info.image, resolution)
+
+    # gt_image = resized_image_rgb[:3, ...]
+    # loaded_mask = None
+
+    # if resized_image_rgb.shape[1] == 4:
+        # loaded_mask = resized_image_rgb[3:4, ...]
+
+    return Camera(colmap_id=cam_info.uid, R=cam_info.R, T=cam_info.T, 
+                  FoVx=cam_info.FovX, FoVy=cam_info.FovY, 
+                  image=cam_info.image, gt_alpha_mask=None,
+                  image_name=cam_info.image_name, uid=id, data_device=args.data_device, 
+                  time = cam_info.time,
+)
+
+def cameraList_from_camInfos(cam_infos, resolution_scale, args):
+    camera_list = []
+
+    for id, c in enumerate(cam_infos):
+        camera_list.append(loadCam(args, id, c, resolution_scale))
+
+    return camera_list
+
+def camera_to_JSON(id, camera : Camera):
+    Rt = np.zeros((4, 4))
+    Rt[:3, :3] = camera.R.transpose()
+    Rt[:3, 3] = camera.T
+    Rt[3, 3] = 1.0
+
+    W2C = np.linalg.inv(Rt)
+    pos = W2C[:3, 3]
+    rot = W2C[:3, :3]
+    serializable_array_2d = [x.tolist() for x in rot]
+    camera_entry = {
+        'id' : id,
+        'img_name' : camera.image_name,
+        'width' : camera.width,
+        'height' : camera.height,
+        'position': pos.tolist(),
+        'rotation': serializable_array_2d,
+        'fy' : fov2focal(camera.FovY, camera.height),
+        'fx' : fov2focal(camera.FovX, camera.width)
+    }
+    return camera_entry

utils/general_utils.py

@@ -0,0 +1,136 @@
+#
+# Copyright (C) 2023, Inria
+# GRAPHDECO research group, https://team.inria.fr/graphdeco
+# All rights reserved.
+#
+# This software is free for non-commercial, research and evaluation use 
+# under the terms of the LICENSE.md file.
+#
+# For inquiries contact  [email protected]
+#
+
+import torch
+import sys
+from datetime import datetime
+import numpy as np
+import random
+
+def inverse_sigmoid(x):
+    return torch.log(x/(1-x))
+
+def PILtoTorch(pil_image, resolution):
+    if resolution is not None:
+        resized_image_PIL = pil_image.resize(resolution)
+    else:
+        resized_image_PIL = pil_image
+    resized_image = torch.from_numpy(np.array(resized_image_PIL)) / 255.0
+    if len(resized_image.shape) == 3:
+        return resized_image.permute(2, 0, 1)
+    else:
+        return resized_image.unsqueeze(dim=-1).permute(2, 0, 1)
+
+def get_expon_lr_func(
+    lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000
+):
+    """
+    Copied from Plenoxels
+
+    Continuous learning rate decay function. Adapted from JaxNeRF
+    The returned rate is lr_init when step=0 and lr_final when step=max_steps, and
+    is log-linearly interpolated elsewhere (equivalent to exponential decay).
+    If lr_delay_steps>0 then the learning rate will be scaled by some smooth
+    function of lr_delay_mult, such that the initial learning rate is
+    lr_init*lr_delay_mult at the beginning of optimization but will be eased back
+    to the normal learning rate when steps>lr_delay_steps.
+    :param conf: config subtree 'lr' or similar
+    :param max_steps: int, the number of steps during optimization.
+    :return HoF which takes step as input
+    """
+
+    def helper(step):
+        if step < 0 or (lr_init == 0.0 and lr_final == 0.0):
+            # Disable this parameter
+            return 0.0
+        if lr_delay_steps > 0:
+            # A kind of reverse cosine decay.
+            delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin(
+                0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1)
+            )
+        else:
+            delay_rate = 1.0
+        t = np.clip(step / max_steps, 0, 1)
+        log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t)
+        return delay_rate * log_lerp
+
+    return helper
+
+def strip_lowerdiag(L):
+    uncertainty = torch.zeros((L.shape[0], 6), dtype=torch.float, device="cuda")
+
+    uncertainty[:, 0] = L[:, 0, 0]
+    uncertainty[:, 1] = L[:, 0, 1]
+    uncertainty[:, 2] = L[:, 0, 2]
+    uncertainty[:, 3] = L[:, 1, 1]
+    uncertainty[:, 4] = L[:, 1, 2]
+    uncertainty[:, 5] = L[:, 2, 2]
+    return uncertainty
+
+def strip_symmetric(sym):
+    return strip_lowerdiag(sym)
+
+def build_rotation(r):
+    norm = torch.sqrt(r[:,0]*r[:,0] + r[:,1]*r[:,1] + r[:,2]*r[:,2] + r[:,3]*r[:,3])
+
+    q = r / norm[:, None]
+
+    R = torch.zeros((q.size(0), 3, 3), device='cuda')
+
+    w = q[:, 0]
+    x = q[:, 1]
+    y = q[:, 2]
+    z = q[:, 3]
+
+    R[:, 0, 0] = 1 - 2 * (y*y + z*z)
+    R[:, 0, 1] = 2 * (x*y - w*z)
+    R[:, 0, 2] = 2 * (x*z + w*y)
+    R[:, 1, 0] = 2 * (x*y + w*z)
+    R[:, 1, 1] = 1 - 2 * (x*x + z*z)
+    R[:, 1, 2] = 2 * (y*z - w*x)
+    R[:, 2, 0] = 2 * (x*z - w*y)
+    R[:, 2, 1] = 2 * (y*z + w*x)
+    R[:, 2, 2] = 1 - 2 * (x*x + y*y)
+    return R
+
+def build_scaling_rotation(s, r):
+    L = torch.zeros((s.shape[0], 3, 3), dtype=torch.float, device="cuda")
+    R = build_rotation(r)
+
+    L[:,0,0] = s[:,0]
+    L[:,1,1] = s[:,1]
+    L[:,2,2] = s[:,2]
+
+    L = R @ L
+    return L
+
+def safe_state(silent):
+    old_f = sys.stdout
+    class F:
+        def __init__(self, silent):
+            self.silent = silent
+
+        def write(self, x):
+            if not self.silent:
+                if x.endswith("\n"):
+                    old_f.write(x.replace("\n", " [{}]\n".format(str(datetime.now().strftime("%d/%m %H:%M:%S")))))
+                else:
+                    old_f.write(x)
+
+        def flush(self):
+            old_f.flush()
+
+    sys.stdout = F(silent)
+
+    random.seed(0)
+    np.random.seed(0)
+    torch.manual_seed(0)
+    torch.cuda.set_device(torch.device("cuda:0"))