jiaweir commited on
Commit
21c4e64
β€’
1 Parent(s): e59b351
This view is limited to 50 files because it contains too many changes. Β  See raw diff
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ *.png* filter=lfs diff=lfs merge=lfs -text
.gitignore ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ __pycache__/
2
+ build/
3
+ *.egg-info/
4
+ *.so
5
+ venv_*/
6
+ .vs/
7
+ .vscode/
8
+ .idea/
9
+
10
+ tmp_*
11
+ data?
12
+ data??
13
+ scripts2
14
+
15
+ model_cache
16
+
17
+ logs
18
+ videos
19
+ images
20
+ *.mp4
21
+
22
+ vis_data*/
23
+ logs*/
24
+ data*/
25
+ eval_data*/
26
+
27
+
28
+ *.sh
29
+ *.out
30
+ batchscript*
31
+
32
+ pretrained/
33
+ diff-gaussian-rasterization/
34
+
35
+ tmp_data/
.gitmodules ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ [submodule "diff-gaussian-rasterization"]
2
+ path = diff-gaussian-rasterization
3
+ url = https://github.com/ashawkey/diff-gaussian-rasterization
app.py ADDED
@@ -0,0 +1,135 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import os
3
+ from PIL import Image
4
+ import subprocess
5
+ from gradio_model4dgs import Model4DGS
6
+ import numpy
7
+ import hashlib
8
+
9
+ os.system('pip install -e ./simple-knn')
10
+ os.system('pip install -e ./diff-gaussian-rasterization')
11
+
12
+ from huggingface_hub import hf_hub_download
13
+ ckpt_path = hf_hub_download(repo_id="ashawkey/LGM", filename="model_fp16_fixrot.safetensors")
14
+
15
+ js_func = """
16
+ function refresh() {
17
+ const url = new URL(window.location);
18
+
19
+ if (url.searchParams.get('__theme') !== 'light') {
20
+ url.searchParams.set('__theme', 'light');
21
+ window.location.href = url.href;
22
+ }
23
+ }
24
+ """
25
+
26
+ # check if there is a picture uploaded or selected
27
+ def check_img_input(control_image):
28
+ if control_image is None:
29
+ raise gr.Error("Please select or upload an input image")
30
+
31
+ # check if there is a picture uploaded or selected
32
+ def check_video_input(image_block: Image.Image):
33
+ img_hash = hashlib.sha256(image_block.tobytes()).hexdigest()
34
+ if not os.path.exists(os.path.join('tmp_data', f'{img_hash}_rgba_generated.mp4')):
35
+ raise gr.Error("Please generate a video first")
36
+
37
+
38
+ def optimize_stage_1(image_block: Image.Image, preprocess_chk: bool, seed_slider: int):
39
+ if not os.path.exists('tmp_data'):
40
+ os.makedirs('tmp_data')
41
+ img_hash = hashlib.sha256(image_block.tobytes()).hexdigest()
42
+ if preprocess_chk:
43
+ # save image to a designated path
44
+ image_block.save(os.path.join('tmp_data', f'{img_hash}.png'))
45
+
46
+ # preprocess image
47
+ print(f'python scripts/process.py {os.path.join("tmp_data", f"{img_hash}.png")}')
48
+ subprocess.run(f'python scripts/process.py {os.path.join("tmp_data", f"{img_hash}.png")}', shell=True)
49
+ else:
50
+ image_block.save(os.path.join('tmp_data', f'{img_hash}_rgba.png'))
51
+
52
+ # stage 1
53
+ 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)
54
+
55
+ # return [os.path.join('logs', 'tmp_rgba_model.ply')]
56
+ return os.path.join('tmp_data', f'{img_hash}_rgba_generated.mp4')
57
+
58
+
59
+ def optimize_stage_2(image_block: Image.Image, seed_slider: int):
60
+ img_hash = hashlib.sha256(image_block.tobytes()).hexdigest()
61
+ subprocess.run(f'python lgm/infer.py big --resume {ckpt_path} --test_path tmp_data/{img_hash}_rgba.png', shell=True)
62
+ # stage 2
63
+ 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)
64
+ # os.rename(os.path.join('logs', f'{img_hash}_rgba_frames'), os.path.join('logs', f'{img_hash}_{seed_slider:03d}_rgba_frames'))
65
+ image_dir = os.path.join('logs', f'{img_hash}_rgba_frames')
66
+ # return 'vis_data/tmp_rgba.mp4', [os.path.join(image_dir, file) for file in os.listdir(image_dir) if file.endswith('.ply')]
67
+ return [image_dir+f'/{t:03d}.ply' for t in range(28)]
68
+
69
+
70
+ if __name__ == "__main__":
71
+ _TITLE = '''DreamGaussian4D: Generative 4D Gaussian Splatting'''
72
+
73
+ _DESCRIPTION = '''
74
+ <div>
75
+ <a style="display:inline-block" href="https://jiawei-ren.github.io/projects/dreamgaussian4d/"><img src='https://img.shields.io/badge/public_website-8A2BE2'></a>
76
+ <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>
77
+ <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>
78
+ </div>
79
+ We present DreamGausssion4D, an efficient 4D generation framework that builds on Gaussian Splatting.
80
+ '''
81
+ _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**."
82
+
83
+ # load images in 'data' folder as examples
84
+ example_folder = os.path.join(os.path.dirname(__file__), 'data')
85
+ example_fns = os.listdir(example_folder)
86
+ example_fns.sort()
87
+ examples_full = [os.path.join(example_folder, x) for x in example_fns if x.endswith('.png')]
88
+
89
+ # Compose demo layout & data flow
90
+ with gr.Blocks(title=_TITLE, theme=gr.themes.Soft(), js=js_func) as demo:
91
+ with gr.Row():
92
+ with gr.Column(scale=1):
93
+ gr.Markdown('# ' + _TITLE)
94
+ gr.Markdown(_DESCRIPTION)
95
+
96
+ # Image-to-3D
97
+ with gr.Row(variant='panel'):
98
+ with gr.Column(scale=4):
99
+ image_block = gr.Image(type='pil', image_mode='RGBA', height=290, label='Input image')
100
+
101
+ # elevation_slider = gr.Slider(-90, 90, value=0, step=1, label='Estimated elevation angle')
102
+ seed_slider = gr.Slider(0, 100000, value=0, step=1, label='Random Seed')
103
+ gr.Markdown(
104
+ "random seed for video generation.")
105
+
106
+ preprocess_chk = gr.Checkbox(True,
107
+ label='Preprocess image automatically (remove background and recenter object)')
108
+
109
+ gr.Examples(
110
+ examples=examples_full, # NOTE: elements must match inputs list!
111
+ inputs=[image_block],
112
+ outputs=[image_block],
113
+ cache_examples=False,
114
+ label='Examples (click one of the images below to start)',
115
+ examples_per_page=40
116
+ )
117
+ img_run_btn = gr.Button("Generate Video")
118
+ fourd_run_btn = gr.Button("Generate 4D")
119
+ img_guide_text = gr.Markdown(_IMG_USER_GUIDE, visible=True)
120
+
121
+ with gr.Column(scale=5):
122
+ obj3d = gr.Video(label="video",height=290)
123
+ obj4d = Model4DGS(label="4D Model", height=500, fps=14)
124
+
125
+ img_run_btn.click(check_img_input, inputs=[image_block], queue=False).success(optimize_stage_1,
126
+ inputs=[image_block,
127
+ preprocess_chk,
128
+ seed_slider],
129
+ outputs=[
130
+ obj3d])
131
+ fourd_run_btn.click(check_video_input, inputs=[image_block], queue=False).success(optimize_stage_2, inputs=[image_block, seed_slider], outputs=[obj4d])
132
+
133
+ # demo.queue().launch(share=True)
134
+ demo.queue(max_size=10) # <-- Sets up a queue with default parameters
135
+ demo.launch(share=True)
cam_utils.py ADDED
@@ -0,0 +1,146 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ from scipy.spatial.transform import Rotation as R
3
+
4
+ import torch
5
+
6
+ def dot(x, y):
7
+ if isinstance(x, np.ndarray):
8
+ return np.sum(x * y, -1, keepdims=True)
9
+ else:
10
+ return torch.sum(x * y, -1, keepdim=True)
11
+
12
+
13
+ def length(x, eps=1e-20):
14
+ if isinstance(x, np.ndarray):
15
+ return np.sqrt(np.maximum(np.sum(x * x, axis=-1, keepdims=True), eps))
16
+ else:
17
+ return torch.sqrt(torch.clamp(dot(x, x), min=eps))
18
+
19
+
20
+ def safe_normalize(x, eps=1e-20):
21
+ return x / length(x, eps)
22
+
23
+
24
+ def look_at(campos, target, opengl=True):
25
+ # campos: [N, 3], camera/eye position
26
+ # target: [N, 3], object to look at
27
+ # return: [N, 3, 3], rotation matrix
28
+ if not opengl:
29
+ # camera forward aligns with -z
30
+ forward_vector = safe_normalize(target - campos)
31
+ up_vector = np.array([0, 1, 0], dtype=np.float32)
32
+ right_vector = safe_normalize(np.cross(forward_vector, up_vector))
33
+ up_vector = safe_normalize(np.cross(right_vector, forward_vector))
34
+ else:
35
+ # camera forward aligns with +z
36
+ forward_vector = safe_normalize(campos - target)
37
+ up_vector = np.array([0, 1, 0], dtype=np.float32)
38
+ right_vector = safe_normalize(np.cross(up_vector, forward_vector))
39
+ up_vector = safe_normalize(np.cross(forward_vector, right_vector))
40
+ R = np.stack([right_vector, up_vector, forward_vector], axis=1)
41
+ return R
42
+
43
+
44
+ # elevation & azimuth to pose (cam2world) matrix
45
+ def orbit_camera(elevation, azimuth, radius=1, is_degree=True, target=None, opengl=True):
46
+ # radius: scalar
47
+ # elevation: scalar, in (-90, 90), from +y to -y is (-90, 90)
48
+ # azimuth: scalar, in (-180, 180), from +z to +x is (0, 90)
49
+ # return: [4, 4], camera pose matrix
50
+ if is_degree:
51
+ elevation = np.deg2rad(elevation)
52
+ azimuth = np.deg2rad(azimuth)
53
+ x = radius * np.cos(elevation) * np.sin(azimuth)
54
+ y = - radius * np.sin(elevation)
55
+ z = radius * np.cos(elevation) * np.cos(azimuth)
56
+ if target is None:
57
+ target = np.zeros([3], dtype=np.float32)
58
+ campos = np.array([x, y, z]) + target # [3]
59
+ T = np.eye(4, dtype=np.float32)
60
+ T[:3, :3] = look_at(campos, target, opengl)
61
+ T[:3, 3] = campos
62
+ return T
63
+
64
+
65
+ class OrbitCamera:
66
+ def __init__(self, W, H, r=2, fovy=60, near=0.01, far=100):
67
+ self.W = W
68
+ self.H = H
69
+ self.radius = r # camera distance from center
70
+ self.fovy = np.deg2rad(fovy) # deg 2 rad
71
+ self.near = near
72
+ self.far = far
73
+ self.center = np.array([0, 0, 0], dtype=np.float32) # look at this point
74
+ self.rot = R.from_matrix(np.eye(3))
75
+ self.up = np.array([0, 1, 0], dtype=np.float32) # need to be normalized!
76
+
77
+ @property
78
+ def fovx(self):
79
+ return 2 * np.arctan(np.tan(self.fovy / 2) * self.W / self.H)
80
+
81
+ @property
82
+ def campos(self):
83
+ return self.pose[:3, 3]
84
+
85
+ # pose (c2w)
86
+ @property
87
+ def pose(self):
88
+ # first move camera to radius
89
+ res = np.eye(4, dtype=np.float32)
90
+ res[2, 3] = self.radius # opengl convention...
91
+ # rotate
92
+ rot = np.eye(4, dtype=np.float32)
93
+ rot[:3, :3] = self.rot.as_matrix()
94
+ res = rot @ res
95
+ # translate
96
+ res[:3, 3] -= self.center
97
+ return res
98
+
99
+ # view (w2c)
100
+ @property
101
+ def view(self):
102
+ return np.linalg.inv(self.pose)
103
+
104
+ # projection (perspective)
105
+ @property
106
+ def perspective(self):
107
+ y = np.tan(self.fovy / 2)
108
+ aspect = self.W / self.H
109
+ return np.array(
110
+ [
111
+ [1 / (y * aspect), 0, 0, 0],
112
+ [0, -1 / y, 0, 0],
113
+ [
114
+ 0,
115
+ 0,
116
+ -(self.far + self.near) / (self.far - self.near),
117
+ -(2 * self.far * self.near) / (self.far - self.near),
118
+ ],
119
+ [0, 0, -1, 0],
120
+ ],
121
+ dtype=np.float32,
122
+ )
123
+
124
+ # intrinsics
125
+ @property
126
+ def intrinsics(self):
127
+ focal = self.H / (2 * np.tan(self.fovy / 2))
128
+ return np.array([focal, focal, self.W // 2, self.H // 2], dtype=np.float32)
129
+
130
+ @property
131
+ def mvp(self):
132
+ return self.perspective @ np.linalg.inv(self.pose) # [4, 4]
133
+
134
+ def orbit(self, dx, dy):
135
+ # rotate along camera up/side axis!
136
+ side = self.rot.as_matrix()[:3, 0]
137
+ rotvec_x = self.up * np.radians(-0.05 * dx)
138
+ rotvec_y = side * np.radians(-0.05 * dy)
139
+ self.rot = R.from_rotvec(rotvec_x) * R.from_rotvec(rotvec_y) * self.rot
140
+
141
+ def scale(self, delta):
142
+ self.radius *= 1.1 ** (-delta)
143
+
144
+ def pan(self, dx, dy, dz=0):
145
+ # pan in camera coordinate system (careful on the sensitivity!)
146
+ self.center += 0.0005 * self.rot.as_matrix()[:3, :3] @ np.array([-dx, -dy, dz])
configs/4d.yaml ADDED
@@ -0,0 +1,121 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ### Input
2
+ # input rgba image path (default to None, can be load in GUI too)
3
+ input:
4
+ # input text prompt (default to None, can be input in GUI too)
5
+ prompt:
6
+ # input mesh for stage 2 (auto-search from stage 1 output path if None)
7
+ mesh:
8
+ # estimated elevation angle for input image
9
+ elevation: 0
10
+ # reference image resolution
11
+ ref_size: 256
12
+ # density thresh for mesh extraction
13
+ density_thresh: 0.5
14
+
15
+ ### Output
16
+ outdir: logs
17
+ mesh_format: frames
18
+ save_path: ''
19
+ save_model: False
20
+
21
+ ### Training
22
+ # guidance loss weights (0 to disable)
23
+ mvdream: False
24
+ imagedream: False
25
+ lambda_sd: 0
26
+ lambda_zero123: 1
27
+ # use stable-zero123 instead of zero123-xl
28
+ stable_zero123: True
29
+ lambda_svd: 0
30
+ # training batch size per iter
31
+ batch_size: 14
32
+ # training iterations for stage 1
33
+ iters: 500
34
+ # training iterations for stage 2
35
+ iters_refine: 50
36
+ # training camera radius
37
+ radius: 1.5
38
+ # training camera fovy
39
+ fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
40
+ # training camera min elevation
41
+ min_ver: -30
42
+ # training camera max elevation
43
+ max_ver: 30
44
+ # checkpoint to load for stage 1 (should be a ply file)
45
+ load:
46
+ # whether allow geom training in stage 2
47
+ train_geo: False
48
+ # prob to invert background color during training (0 = always black, 1 = always white)
49
+ invert_bg_prob: 0.
50
+ n_views: 4
51
+ t_max: 0.5
52
+
53
+
54
+ ### GUI
55
+ gui: False
56
+ force_cuda_rast: False
57
+ # GUI resolution
58
+ H: 800
59
+ W: 800
60
+
61
+ ### Gaussian splatting
62
+ optimize_gaussians: True
63
+ position_lr_init: 0.001
64
+ position_lr_final: 0.00002
65
+ position_lr_delay_mult: 0.02
66
+ position_lr_max_steps: 500
67
+ feature_lr: 0.01
68
+ opacity_lr: 0.05
69
+ scaling_lr: 0.005
70
+ rotation_lr: 0.005
71
+
72
+ num_pts: 5000
73
+ sh_degree: 0
74
+ percent_dense: 0.1
75
+ density_start_iter: 3000
76
+ density_end_iter: 3000
77
+ densification_interval: 100
78
+ opacity_reset_interval: 700
79
+ densify_grad_threshold: 0.05
80
+
81
+ # deformation field
82
+ deformation_lr_init: 0.00064
83
+ deformation_lr_final: 0.00064
84
+ deformation_lr_delay_mult: 0.01
85
+ grid_lr_init: 0.0064
86
+ grid_lr_final: 0.0064
87
+
88
+ ### Textured Mesh
89
+ geom_lr: 0.0001
90
+ texture_lr: 0.2
91
+
92
+ deformation:
93
+ net_width: 64
94
+ timebase_pe: 4
95
+ defor_depth: 1
96
+ posebase_pe: 10
97
+ scale_rotation_pe: 2
98
+ opacity_pe: 2
99
+ timenet_width: 64
100
+ timenet_output: 32
101
+ bounds: 1.6
102
+ plane_tv_weight: 0.0001
103
+ time_smoothness_weight: 0.01
104
+ l1_time_planes: 0.0001
105
+ kplanes_config:
106
+ grid_dimensions: 2
107
+ input_coordinate_dim: 4
108
+ output_coordinate_dim: 32
109
+ resolution: [32, 32, 32, 12]
110
+ multires: [1]
111
+ no_grid: False
112
+ no_mlp: False
113
+ no_ds: False
114
+ no_dr: False
115
+ no_do: True
116
+ use_res: True
117
+
118
+ data_mode: svd
119
+ downsample_rate: 1
120
+ # data_mode: c4d
121
+ # downsample_rate: 2
configs/4d_c4d.yaml ADDED
@@ -0,0 +1,119 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ### Input
2
+ # input rgba image path (default to None, can be load in GUI too)
3
+ input:
4
+ # input text prompt (default to None, can be input in GUI too)
5
+ prompt:
6
+ # input mesh for stage 2 (auto-search from stage 1 output path if None)
7
+ mesh:
8
+ # estimated elevation angle for input image
9
+ elevation: 0
10
+ # reference image resolution
11
+ ref_size: 256
12
+ # density thresh for mesh extraction
13
+ density_thresh: 0.5
14
+
15
+ ### Output
16
+ outdir: logs
17
+ mesh_format: frames
18
+ save_path: ''
19
+ save_model: False
20
+
21
+ ### Training
22
+ # guidance loss weights (0 to disable)
23
+ mvdream: False
24
+ imagedream: False
25
+ lambda_sd: 0
26
+ lambda_zero123: 1
27
+ # use stable-zero123 instead of zero123-xl
28
+ stable_zero123: True
29
+ lambda_svd: 0
30
+ # training batch size per iter
31
+ batch_size: 32
32
+ # training iterations for stage 1
33
+ iters: 500
34
+ # training iterations for stage 2
35
+ iters_refine: 50
36
+ # training camera radius
37
+ radius: 1.5
38
+ # training camera fovy
39
+ fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
40
+ # training camera min elevation
41
+ min_ver: -30
42
+ # training camera max elevation
43
+ max_ver: 30
44
+ # checkpoint to load for stage 1 (should be a ply file)
45
+ load:
46
+ # whether allow geom training in stage 2
47
+ train_geo: False
48
+ # prob to invert background color during training (0 = always black, 1 = always white)
49
+ invert_bg_prob: 0.
50
+ n_views: 4
51
+ t_max: 0.5
52
+
53
+
54
+ ### GUI
55
+ gui: False
56
+ force_cuda_rast: False
57
+ # GUI resolution
58
+ H: 800
59
+ W: 800
60
+
61
+ ### Gaussian splatting
62
+ optimize_gaussians: True
63
+ position_lr_init: 0.001
64
+ position_lr_final: 0.00002
65
+ position_lr_delay_mult: 0.02
66
+ position_lr_max_steps: 500
67
+ feature_lr: 0.01
68
+ opacity_lr: 0.05
69
+ scaling_lr: 0.005
70
+ rotation_lr: 0.005
71
+
72
+ num_pts: 5000
73
+ sh_degree: 0
74
+ percent_dense: 0.1
75
+ density_start_iter: 3000
76
+ density_end_iter: 3000
77
+ densification_interval: 100
78
+ opacity_reset_interval: 700
79
+ densify_grad_threshold: 0.05
80
+
81
+ # deformation field
82
+ deformation_lr_init: 0.00064
83
+ deformation_lr_final: 0.00064
84
+ deformation_lr_delay_mult: 0.01
85
+ grid_lr_init: 0.0064
86
+ grid_lr_final: 0.0064
87
+
88
+ ### Textured Mesh
89
+ geom_lr: 0.0001
90
+ texture_lr: 0.2
91
+
92
+ deformation:
93
+ net_width: 64
94
+ timebase_pe: 4
95
+ defor_depth: 1
96
+ posebase_pe: 10
97
+ scale_rotation_pe: 2
98
+ opacity_pe: 2
99
+ timenet_width: 64
100
+ timenet_output: 32
101
+ bounds: 1.6
102
+ plane_tv_weight: 0.0001
103
+ time_smoothness_weight: 0.01
104
+ l1_time_planes: 0.0001
105
+ kplanes_config:
106
+ grid_dimensions: 2
107
+ input_coordinate_dim: 4
108
+ output_coordinate_dim: 32
109
+ resolution: [32, 32, 32, 32]
110
+ multires: [1]
111
+ no_grid: False
112
+ no_mlp: False
113
+ no_ds: False
114
+ no_dr: False
115
+ no_do: True
116
+ use_res: True
117
+
118
+ data_mode: c4d
119
+ downsample_rate: 1
configs/4d_c4d_low.yaml ADDED
@@ -0,0 +1,119 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ### Input
2
+ # input rgba image path (default to None, can be load in GUI too)
3
+ input:
4
+ # input text prompt (default to None, can be input in GUI too)
5
+ prompt:
6
+ # input mesh for stage 2 (auto-search from stage 1 output path if None)
7
+ mesh:
8
+ # estimated elevation angle for input image
9
+ elevation: 0
10
+ # reference image resolution
11
+ ref_size: 256
12
+ # density thresh for mesh extraction
13
+ density_thresh: 0.5
14
+
15
+ ### Output
16
+ outdir: logs
17
+ mesh_format: frames
18
+ save_path: ''
19
+ save_model: False
20
+
21
+ ### Training
22
+ # guidance loss weights (0 to disable)
23
+ mvdream: False
24
+ imagedream: False
25
+ lambda_sd: 0
26
+ lambda_zero123: 1
27
+ # use stable-zero123 instead of zero123-xl
28
+ stable_zero123: True
29
+ lambda_svd: 0
30
+ # training batch size per iter
31
+ batch_size: 8
32
+ # training iterations for stage 1
33
+ iters: 500
34
+ # training iterations for stage 2
35
+ iters_refine: 50
36
+ # training camera radius
37
+ radius: 1.5
38
+ # training camera fovy
39
+ fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
40
+ # training camera min elevation
41
+ min_ver: -30
42
+ # training camera max elevation
43
+ max_ver: 30
44
+ # checkpoint to load for stage 1 (should be a ply file)
45
+ load:
46
+ # whether allow geom training in stage 2
47
+ train_geo: False
48
+ # prob to invert background color during training (0 = always black, 1 = always white)
49
+ invert_bg_prob: 0.
50
+ n_views: 1
51
+ t_max: 0.5
52
+
53
+
54
+ ### GUI
55
+ gui: False
56
+ force_cuda_rast: False
57
+ # GUI resolution
58
+ H: 800
59
+ W: 800
60
+
61
+ ### Gaussian splatting
62
+ optimize_gaussians: True
63
+ position_lr_init: 0.001
64
+ position_lr_final: 0.00002
65
+ position_lr_delay_mult: 0.02
66
+ position_lr_max_steps: 500
67
+ feature_lr: 0.01
68
+ opacity_lr: 0.05
69
+ scaling_lr: 0.005
70
+ rotation_lr: 0.005
71
+
72
+ num_pts: 5000
73
+ sh_degree: 0
74
+ percent_dense: 0.1
75
+ density_start_iter: 3000
76
+ density_end_iter: 3000
77
+ densification_interval: 100
78
+ opacity_reset_interval: 700
79
+ densify_grad_threshold: 0.05
80
+
81
+ # deformation field
82
+ deformation_lr_init: 0.00064
83
+ deformation_lr_final: 0.00064
84
+ deformation_lr_delay_mult: 0.01
85
+ grid_lr_init: 0.0064
86
+ grid_lr_final: 0.0064
87
+
88
+ ### Textured Mesh
89
+ geom_lr: 0.0001
90
+ texture_lr: 0.2
91
+
92
+ deformation:
93
+ net_width: 64
94
+ timebase_pe: 4
95
+ defor_depth: 1
96
+ posebase_pe: 10
97
+ scale_rotation_pe: 2
98
+ opacity_pe: 2
99
+ timenet_width: 64
100
+ timenet_output: 32
101
+ bounds: 1.6
102
+ plane_tv_weight: 0.0001
103
+ time_smoothness_weight: 0.01
104
+ l1_time_planes: 0.0001
105
+ kplanes_config:
106
+ grid_dimensions: 2
107
+ input_coordinate_dim: 4
108
+ output_coordinate_dim: 32
109
+ resolution: [32, 32, 32, 12]
110
+ multires: [1]
111
+ no_grid: False
112
+ no_mlp: False
113
+ no_ds: False
114
+ no_dr: False
115
+ no_do: True
116
+ use_res: True
117
+
118
+ data_mode: c4d
119
+ downsample_rate: 4
configs/4d_demo.yaml ADDED
@@ -0,0 +1,121 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ### Input
2
+ # input rgba image path (default to None, can be load in GUI too)
3
+ input:
4
+ # input text prompt (default to None, can be input in GUI too)
5
+ prompt:
6
+ # input mesh for stage 2 (auto-search from stage 1 output path if None)
7
+ mesh:
8
+ # estimated elevation angle for input image
9
+ elevation: 0
10
+ # reference image resolution
11
+ ref_size: 256
12
+ # density thresh for mesh extraction
13
+ density_thresh: 0.5
14
+
15
+ ### Output
16
+ outdir: logs
17
+ mesh_format: frames
18
+ save_path: ''
19
+ save_model: False
20
+
21
+ ### Training
22
+ # guidance loss weights (0 to disable)
23
+ mvdream: False
24
+ imagedream: False
25
+ lambda_sd: 0
26
+ lambda_zero123: 1
27
+ # use stable-zero123 instead of zero123-xl
28
+ stable_zero123: True
29
+ lambda_svd: 0
30
+ # training batch size per iter
31
+ batch_size: 7
32
+ # training iterations for stage 1
33
+ iters: 500
34
+ # training iterations for stage 2
35
+ iters_refine: 50
36
+ # training camera radius
37
+ radius: 1.5
38
+ # training camera fovy
39
+ fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
40
+ # training camera min elevation
41
+ min_ver: -30
42
+ # training camera max elevation
43
+ max_ver: 30
44
+ # checkpoint to load for stage 1 (should be a ply file)
45
+ load:
46
+ # whether allow geom training in stage 2
47
+ train_geo: False
48
+ # prob to invert background color during training (0 = always black, 1 = always white)
49
+ invert_bg_prob: 0.
50
+ n_views: 1
51
+ t_max: 0.5
52
+
53
+
54
+ ### GUI
55
+ gui: False
56
+ force_cuda_rast: False
57
+ # GUI resolution
58
+ H: 800
59
+ W: 800
60
+
61
+ ### Gaussian splatting
62
+ optimize_gaussians: True
63
+ position_lr_init: 0.001
64
+ position_lr_final: 0.00002
65
+ position_lr_delay_mult: 0.02
66
+ position_lr_max_steps: 500
67
+ feature_lr: 0.01
68
+ opacity_lr: 0.05
69
+ scaling_lr: 0.005
70
+ rotation_lr: 0.005
71
+
72
+ num_pts: 5000
73
+ sh_degree: 0
74
+ percent_dense: 0.1
75
+ density_start_iter: 3000
76
+ density_end_iter: 3000
77
+ densification_interval: 100
78
+ opacity_reset_interval: 700
79
+ densify_grad_threshold: 0.05
80
+
81
+ # deformation field
82
+ deformation_lr_init: 0.00064
83
+ deformation_lr_final: 0.00064
84
+ deformation_lr_delay_mult: 0.01
85
+ grid_lr_init: 0.0064
86
+ grid_lr_final: 0.0064
87
+
88
+ ### Textured Mesh
89
+ geom_lr: 0.0001
90
+ texture_lr: 0.2
91
+
92
+ deformation:
93
+ net_width: 64
94
+ timebase_pe: 4
95
+ defor_depth: 1
96
+ posebase_pe: 10
97
+ scale_rotation_pe: 2
98
+ opacity_pe: 2
99
+ timenet_width: 64
100
+ timenet_output: 32
101
+ bounds: 1.6
102
+ plane_tv_weight: 0.0001
103
+ time_smoothness_weight: 0.01
104
+ l1_time_planes: 0.0001
105
+ kplanes_config:
106
+ grid_dimensions: 2
107
+ input_coordinate_dim: 4
108
+ output_coordinate_dim: 32
109
+ resolution: [32, 32, 32, 12]
110
+ multires: [1]
111
+ no_grid: False
112
+ no_mlp: False
113
+ no_ds: False
114
+ no_dr: False
115
+ no_do: True
116
+ use_res: True
117
+
118
+ data_mode: svd
119
+ downsample_rate: 2
120
+ # data_mode: c4d
121
+ # downsample_rate: 2
configs/4d_low.yaml ADDED
@@ -0,0 +1,121 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ### Input
2
+ # input rgba image path (default to None, can be load in GUI too)
3
+ input:
4
+ # input text prompt (default to None, can be input in GUI too)
5
+ prompt:
6
+ # input mesh for stage 2 (auto-search from stage 1 output path if None)
7
+ mesh:
8
+ # estimated elevation angle for input image
9
+ elevation: 0
10
+ # reference image resolution
11
+ ref_size: 256
12
+ # density thresh for mesh extraction
13
+ density_thresh: 0.5
14
+
15
+ ### Output
16
+ outdir: logs
17
+ mesh_format: frames
18
+ save_path: ''
19
+ save_model: False
20
+
21
+ ### Training
22
+ # guidance loss weights (0 to disable)
23
+ mvdream: False
24
+ imagedream: False
25
+ lambda_sd: 0
26
+ lambda_zero123: 1
27
+ # use stable-zero123 instead of zero123-xl
28
+ stable_zero123: True
29
+ lambda_svd: 0
30
+ # training batch size per iter
31
+ batch_size: 14
32
+ # training iterations for stage 1
33
+ iters: 500
34
+ # training iterations for stage 2
35
+ iters_refine: 50
36
+ # training camera radius
37
+ radius: 1.5
38
+ # training camera fovy
39
+ fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
40
+ # training camera min elevation
41
+ min_ver: -30
42
+ # training camera max elevation
43
+ max_ver: 30
44
+ # checkpoint to load for stage 1 (should be a ply file)
45
+ load:
46
+ # whether allow geom training in stage 2
47
+ train_geo: False
48
+ # prob to invert background color during training (0 = always black, 1 = always white)
49
+ invert_bg_prob: 0.
50
+ n_views: 1
51
+ t_max: 0.5
52
+
53
+
54
+ ### GUI
55
+ gui: False
56
+ force_cuda_rast: False
57
+ # GUI resolution
58
+ H: 800
59
+ W: 800
60
+
61
+ ### Gaussian splatting
62
+ optimize_gaussians: True
63
+ position_lr_init: 0.001
64
+ position_lr_final: 0.00002
65
+ position_lr_delay_mult: 0.02
66
+ position_lr_max_steps: 500
67
+ feature_lr: 0.01
68
+ opacity_lr: 0.05
69
+ scaling_lr: 0.005
70
+ rotation_lr: 0.005
71
+
72
+ num_pts: 5000
73
+ sh_degree: 0
74
+ percent_dense: 0.1
75
+ density_start_iter: 3000
76
+ density_end_iter: 3000
77
+ densification_interval: 100
78
+ opacity_reset_interval: 700
79
+ densify_grad_threshold: 0.05
80
+
81
+ # deformation field
82
+ deformation_lr_init: 0.00064
83
+ deformation_lr_final: 0.00064
84
+ deformation_lr_delay_mult: 0.01
85
+ grid_lr_init: 0.0064
86
+ grid_lr_final: 0.0064
87
+
88
+ ### Textured Mesh
89
+ geom_lr: 0.0001
90
+ texture_lr: 0.2
91
+
92
+ deformation:
93
+ net_width: 64
94
+ timebase_pe: 4
95
+ defor_depth: 1
96
+ posebase_pe: 10
97
+ scale_rotation_pe: 2
98
+ opacity_pe: 2
99
+ timenet_width: 64
100
+ timenet_output: 32
101
+ bounds: 1.6
102
+ plane_tv_weight: 0.0001
103
+ time_smoothness_weight: 0.01
104
+ l1_time_planes: 0.0001
105
+ kplanes_config:
106
+ grid_dimensions: 2
107
+ input_coordinate_dim: 4
108
+ output_coordinate_dim: 32
109
+ resolution: [32, 32, 32, 22]
110
+ multires: [1]
111
+ no_grid: False
112
+ no_mlp: False
113
+ no_ds: False
114
+ no_dr: False
115
+ no_do: True
116
+ use_res: True
117
+
118
+ data_mode: svd
119
+ downsample_rate: 1
120
+ # data_mode: c4d
121
+ # downsample_rate: 2
configs/dg.yaml ADDED
@@ -0,0 +1,85 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ### Input
2
+ # input rgba image path (default to None, can be load in GUI too)
3
+ input:
4
+ # input text prompt (default to None, can be input in GUI too)
5
+ prompt:
6
+ negative_prompt:
7
+ # input mesh for stage 2 (auto-search from stage 1 output path if None)
8
+ mesh:
9
+ # estimated elevation angle for input image
10
+ elevation: 0
11
+ # reference image resolution
12
+ ref_size: 256
13
+ # density thresh for mesh extraction
14
+ density_thresh: 1
15
+
16
+ ### Output
17
+ outdir: logs
18
+ mesh_format: obj
19
+ save_path: ''
20
+
21
+ ### Training
22
+ # use mvdream instead of sd 2.1
23
+ mvdream: False
24
+ # use imagedream
25
+ imagedream: False
26
+ # use stable-zero123 instead of zero123-xl
27
+ stable_zero123: False
28
+ # guidance loss weights (0 to disable)
29
+ lambda_sd: 0
30
+ lambda_zero123: 1
31
+ # warmup rgb supervision for image-to-3d
32
+ warmup_rgb_loss: True
33
+ # training batch size per iter
34
+ batch_size: 1
35
+ # training iterations for stage 1
36
+ iters: 500
37
+ # whether to linearly anneal timestep
38
+ anneal_timestep: True
39
+ # training iterations for stage 2
40
+ iters_refine: 50
41
+ # training camera radius
42
+ radius: 2
43
+ # training camera fovy
44
+ fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
45
+ # training camera min elevation
46
+ min_ver: -30
47
+ # training camera max elevation
48
+ max_ver: 30
49
+ # checkpoint to load for stage 1 (should be a ply file)
50
+ load:
51
+ # whether allow geom training in stage 2
52
+ train_geo: False
53
+ # prob to invert background color during training (0 = always black, 1 = always white)
54
+ invert_bg_prob: 0.5
55
+
56
+
57
+ ### GUI
58
+ gui: False
59
+ force_cuda_rast: False
60
+ # GUI resolution
61
+ H: 800
62
+ W: 800
63
+
64
+ ### Gaussian splatting
65
+ num_pts: 5000
66
+ sh_degree: 0
67
+ position_lr_init: 0.001
68
+ position_lr_final: 0.00002
69
+ position_lr_delay_mult: 0.02
70
+ position_lr_max_steps: 500
71
+ feature_lr: 0.01
72
+ opacity_lr: 0.05
73
+ scaling_lr: 0.005
74
+ rotation_lr: 0.005
75
+ percent_dense: 0.1
76
+ density_start_iter: 100
77
+ density_end_iter: 3000
78
+ densification_interval: 100
79
+ opacity_reset_interval: 700
80
+ densify_grad_threshold: 0.5
81
+
82
+
83
+ ### Textured Mesh
84
+ geom_lr: 0.0001
85
+ texture_lr: 0.2
configs/dghd.yaml ADDED
@@ -0,0 +1,72 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ### Input
2
+ # input rgba image path (default to None, can be load in GUI too)
3
+ input:
4
+ # input text prompt (default to None, can be input in GUI too)
5
+ prompt:
6
+ # input mesh for stage 2 (auto-search from stage 1 output path if None)
7
+ mesh:
8
+ # estimated elevation angle for input image
9
+ elevation: 0
10
+ # reference image resolution
11
+ ref_size: 256
12
+ # density thresh for mesh extraction
13
+ density_thresh: 1
14
+
15
+ ### Output
16
+ outdir: logs
17
+ mesh_format: obj
18
+ save_path: ''
19
+
20
+ ### Training
21
+ # guidance loss weights (0 to disable)
22
+ lambda_sd: 0
23
+ mvdream: False
24
+ lambda_zero123: 1
25
+ # use stable-zero123 instead of zero123-xl
26
+ stable_zero123: False
27
+ # training batch size per iter
28
+ batch_size: 16
29
+ # training iterations for stage 1
30
+ iters: 500
31
+ # training iterations for stage 2
32
+ iters_refine: 50
33
+ # training camera radius
34
+ radius: 2
35
+ # training camera fovy
36
+ fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
37
+ # checkpoint to load for stage 1 (should be a ply file)
38
+ load:
39
+ # whether allow geom training in stage 2
40
+ train_geo: False
41
+ # prob to invert background color during training (0 = always black, 1 = always white)
42
+ invert_bg_prob: 0.
43
+
44
+
45
+ ### GUI
46
+ gui: False
47
+ force_cuda_rast: False
48
+ # GUI resolution
49
+ H: 800
50
+ W: 800
51
+
52
+ ### Gaussian splatting
53
+ num_pts: 5000
54
+ sh_degree: 0
55
+ position_lr_init: 0.001
56
+ position_lr_final: 0.00002
57
+ position_lr_delay_mult: 0.02
58
+ position_lr_max_steps: 500
59
+ feature_lr: 0.01
60
+ opacity_lr: 0.05
61
+ scaling_lr: 0.005
62
+ rotation_lr: 0.005
63
+ percent_dense: 0.1
64
+ density_start_iter: 0
65
+ density_end_iter: 3000
66
+ densification_interval: 100
67
+ opacity_reset_interval: 700
68
+ densify_grad_threshold: 0.05
69
+
70
+ ### Textured Mesh
71
+ geom_lr: 0.0001
72
+ texture_lr: 0.2
configs/refine.yaml ADDED
@@ -0,0 +1,79 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ### Input
2
+ # input rgba image path (default to None, can be load in GUI too)
3
+ input:
4
+ # input text prompt (default to None, can be input in GUI too)
5
+ prompt:
6
+ # input mesh for stage 2 (auto-search from stage 1 output path if None)
7
+ mesh:
8
+ # estimated elevation angle for input image
9
+ elevation: 0
10
+ # reference image resolution
11
+ ref_size: 256
12
+ # density thresh for mesh extraction
13
+ density_thresh: 1
14
+
15
+ ### Output
16
+ outdir: logs
17
+ mesh_format: obj
18
+ save_path: ''
19
+
20
+ ### Training
21
+ # guidance loss weights (0 to disable)
22
+ lambda_sd: 0
23
+ mvdream: False
24
+ lambda_zero123: 0
25
+ # use stable-zero123 instead of zero123-xl
26
+ stable_zero123: False
27
+ lambda_svd: 1
28
+ # training batch size per iter
29
+ batch_size: 1
30
+ # training iterations for stage 1
31
+ iters: 500
32
+ # training iterations for stage 2
33
+ iters_refine: 50
34
+ # training camera radius
35
+ radius: 1.5
36
+ # training camera fovy
37
+ fovy: 49.1 # align with zero123 rendering setting (ref: https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py#L61
38
+ # checkpoint to load for stage 1 (should be a ply file)
39
+ load:
40
+ # whether allow geom training in stage 2
41
+ train_geo: False
42
+ # prob to invert background color during training (0 = always black, 1 = always white)
43
+ invert_bg_prob: 0.5
44
+
45
+
46
+ ### GUI
47
+ gui: False
48
+ force_cuda_rast: False
49
+ # GUI resolution
50
+ H: 800
51
+ W: 800
52
+
53
+ ### Gaussian splatting
54
+ num_pts: 5000
55
+ sh_degree: 0
56
+ position_lr_init: 0.001
57
+ position_lr_final: 0.00002
58
+ position_lr_delay_mult: 0.02
59
+ position_lr_max_steps: 500
60
+ feature_lr: 0.01
61
+ opacity_lr: 0.05
62
+ scaling_lr: 0.005
63
+ rotation_lr: 0.005
64
+ percent_dense: 0.1
65
+ density_start_iter: 100
66
+ density_end_iter: 3000
67
+ densification_interval: 100
68
+ opacity_reset_interval: 700
69
+ densify_grad_threshold: 0.5
70
+
71
+ ### Textured Mesh
72
+ geom_lr: 0.0001
73
+ texture_lr: 0.2
74
+
75
+ static_model: lgm
76
+ data_mode: svd
77
+ downsample_rate: 2
78
+
79
+ oom_hack: False
data/anya_rgba.png ADDED
data/catstatue_rgba.png ADDED
data/csm_luigi_rgba.png ADDED
data/zelda_rgba.png ADDED
gaussian_model_4d.py ADDED
@@ -0,0 +1,773 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #
2
+ # Copyright (C) 2023, Inria
3
+ # GRAPHDECO research group, https://team.inria.fr/graphdeco
4
+ # All rights reserved.
5
+ #
6
+ # This software is free for non-commercial, research and evaluation use
7
+ # under the terms of the LICENSE.md file.
8
+ #
9
+ # For inquiries contact george.drettakis@inria.fr
10
+ #
11
+
12
+ import torch
13
+ import numpy as np
14
+ from utils.general_utils import inverse_sigmoid, get_expon_lr_func, build_rotation
15
+ from torch import nn
16
+ import os
17
+ from utils.system_utils import mkdir_p
18
+ from plyfile import PlyData, PlyElement
19
+ from random import randint
20
+ from utils.sh_utils import RGB2SH
21
+ from simple_knn._C import distCUDA2
22
+ from utils.graphics_utils import BasicPointCloud
23
+ from utils.general_utils import strip_symmetric, build_scaling_rotation
24
+ from scene.deformation import deform_network
25
+ from scene.regulation import compute_plane_smoothness
26
+
27
+
28
+ def gaussian_3d_coeff(xyzs, covs):
29
+ # xyzs: [N, 3]
30
+ # covs: [N, 6]
31
+ x, y, z = xyzs[:, 0], xyzs[:, 1], xyzs[:, 2]
32
+ a, b, c, d, e, f = covs[:, 0], covs[:, 1], covs[:, 2], covs[:, 3], covs[:, 4], covs[:, 5]
33
+
34
+ # eps must be small enough !!!
35
+ inv_det = 1 / (a * d * f + 2 * e * c * b - e**2 * a - c**2 * d - b**2 * f + 1e-24)
36
+ inv_a = (d * f - e**2) * inv_det
37
+ inv_b = (e * c - b * f) * inv_det
38
+ inv_c = (e * b - c * d) * inv_det
39
+ inv_d = (a * f - c**2) * inv_det
40
+ inv_e = (b * c - e * a) * inv_det
41
+ inv_f = (a * d - b**2) * inv_det
42
+
43
+ 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
44
+
45
+ power[power > 0] = -1e10 # abnormal values... make weights 0
46
+
47
+ return torch.exp(power)
48
+
49
+ class GaussianModel:
50
+
51
+ def setup_functions(self):
52
+ def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation):
53
+ L = build_scaling_rotation(scaling_modifier * scaling, rotation)
54
+ actual_covariance = L @ L.transpose(1, 2)
55
+ symm = strip_symmetric(actual_covariance)
56
+ return symm
57
+
58
+ self.scaling_activation = torch.exp
59
+ self.scaling_inverse_activation = torch.log
60
+
61
+ self.covariance_activation = build_covariance_from_scaling_rotation
62
+
63
+ self.opacity_activation = torch.sigmoid
64
+ self.inverse_opacity_activation = inverse_sigmoid
65
+
66
+ self.rotation_activation = torch.nn.functional.normalize
67
+
68
+
69
+ def __init__(self, sh_degree : int, args):
70
+ self.active_sh_degree = 0
71
+ self.max_sh_degree = sh_degree
72
+ self._xyz = torch.empty(0)
73
+ # self._deformation = torch.empty(0)
74
+ self._deformation = deform_network(args)
75
+ # self.grid = TriPlaneGrid()
76
+ self._features_dc = torch.empty(0)
77
+ self._features_rest = torch.empty(0)
78
+ self._scaling = torch.empty(0)
79
+ self._rotation = torch.empty(0)
80
+ self._opacity = torch.empty(0)
81
+ self.max_radii2D = torch.empty(0)
82
+ self.xyz_gradient_accum = torch.empty(0)
83
+ self.denom = torch.empty(0)
84
+ self.optimizer = None
85
+ self.percent_dense = 0
86
+ self.spatial_lr_scale = 0
87
+ self._deformation_table = torch.empty(0)
88
+ self.setup_functions()
89
+
90
+ def capture(self):
91
+ return (
92
+ self.active_sh_degree,
93
+ self._xyz,
94
+ self._deformation.state_dict(),
95
+ self._deformation_table,
96
+ # self.grid,
97
+ self._features_dc,
98
+ self._features_rest,
99
+ self._scaling,
100
+ self._rotation,
101
+ self._opacity,
102
+ self.max_radii2D,
103
+ self.xyz_gradient_accum,
104
+ self.denom,
105
+ self.optimizer.state_dict(),
106
+ self.spatial_lr_scale,
107
+ )
108
+
109
+ def restore(self, model_args, training_args):
110
+ (self.active_sh_degree,
111
+ self._xyz,
112
+ self._deformation_table,
113
+ self._deformation,
114
+ # self.grid,
115
+ self._features_dc,
116
+ self._features_rest,
117
+ self._scaling,
118
+ self._rotation,
119
+ self._opacity,
120
+ self.max_radii2D,
121
+ xyz_gradient_accum,
122
+ denom,
123
+ opt_dict,
124
+ self.spatial_lr_scale) = model_args
125
+ self.training_setup(training_args)
126
+ self.xyz_gradient_accum = xyz_gradient_accum
127
+ self.denom = denom
128
+ self.optimizer.load_state_dict(opt_dict)
129
+
130
+ @property
131
+ def get_scaling(self):
132
+ return self.scaling_activation(self._scaling)
133
+
134
+ @property
135
+ def get_rotation(self):
136
+ return self.rotation_activation(self._rotation)
137
+
138
+ @property
139
+ def get_xyz(self):
140
+ return self._xyz
141
+
142
+ @property
143
+ def get_features(self):
144
+ features_dc = self._features_dc
145
+ features_rest = self._features_rest
146
+ return torch.cat((features_dc, features_rest), dim=1)
147
+
148
+ @property
149
+ def get_opacity(self):
150
+ return self.opacity_activation(self._opacity)
151
+
152
+
153
+ def get_deformed_everything(self, time):
154
+ means3D = self.get_xyz
155
+ time = torch.tensor(time).to(means3D.device).repeat(means3D.shape[0],1)
156
+ time = ((time.float() / self.T) - 0.5) * 2
157
+
158
+ opacity = self._opacity
159
+ scales = self._scaling
160
+ rotations = self._rotation
161
+
162
+ deformation_point = self._deformation_table
163
+ means3D_deform, scales_deform, rotations_deform, opacity_deform = self._deformation(means3D[deformation_point], scales[deformation_point],
164
+ rotations[deformation_point], opacity[deformation_point],
165
+ time[deformation_point])
166
+
167
+ means3D_final = means3D + means3D_deform
168
+ rotations_final = rotations + rotations_deform
169
+ scales_final = scales + scales_deform
170
+ opacity_final = opacity
171
+
172
+ return means3D_final, rotations_final, scales_final, opacity_final
173
+
174
+
175
+
176
+ @torch.no_grad()
177
+ def extract_fields_t(self, resolution=128, num_blocks=16, relax_ratio=1.5, t=0):
178
+ # resolution: resolution of field
179
+
180
+ block_size = 2 / num_blocks
181
+
182
+ assert resolution % block_size == 0
183
+ split_size = resolution // num_blocks
184
+
185
+ xyzs, rotation, scale, opacities = self.get_deformed_everything(t)
186
+
187
+ scale = self.scaling_activation(scale)
188
+ opacities = self.opacity_activation(opacities)
189
+
190
+ # pre-filter low opacity gaussians to save computation
191
+ mask = (opacities > 0.005).squeeze(1)
192
+
193
+ opacities = opacities[mask]
194
+ xyzs = xyzs[mask]
195
+ stds = scale[mask]
196
+
197
+ # normalize to ~ [-1, 1]
198
+ mn, mx = xyzs.amin(0), xyzs.amax(0)
199
+ self.center = (mn + mx) / 2
200
+ self.scale = 1.8 / (mx - mn).amax().item()
201
+
202
+ xyzs = (xyzs - self.center) * self.scale
203
+ stds = stds * self.scale
204
+
205
+ covs = self.covariance_activation(stds, 1, rotation[mask])
206
+
207
+ # tile
208
+ device = opacities.device
209
+ occ = torch.zeros([resolution] * 3, dtype=torch.float32, device=device)
210
+
211
+ X = torch.linspace(-1, 1, resolution).split(split_size)
212
+ Y = torch.linspace(-1, 1, resolution).split(split_size)
213
+ Z = torch.linspace(-1, 1, resolution).split(split_size)
214
+
215
+
216
+ # loop blocks (assume max size of gaussian is small than relax_ratio * block_size !!!)
217
+ for xi, xs in enumerate(X):
218
+ for yi, ys in enumerate(Y):
219
+ for zi, zs in enumerate(Z):
220
+ xx, yy, zz = torch.meshgrid(xs, ys, zs)
221
+ # sample points [M, 3]
222
+ pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1).to(device)
223
+ # in-tile gaussians mask
224
+ vmin, vmax = pts.amin(0), pts.amax(0)
225
+ vmin -= block_size * relax_ratio
226
+ vmax += block_size * relax_ratio
227
+ mask = (xyzs < vmax).all(-1) & (xyzs > vmin).all(-1)
228
+ # if hit no gaussian, continue to next block
229
+ if not mask.any():
230
+ continue
231
+ mask_xyzs = xyzs[mask] # [L, 3]
232
+ mask_covs = covs[mask] # [L, 6]
233
+ mask_opas = opacities[mask].view(1, -1) # [L, 1] --> [1, L]
234
+
235
+ # query per point-gaussian pair.
236
+ g_pts = pts.unsqueeze(1).repeat(1, mask_covs.shape[0], 1) - mask_xyzs.unsqueeze(0) # [M, L, 3]
237
+ g_covs = mask_covs.unsqueeze(0).repeat(pts.shape[0], 1, 1) # [M, L, 6]
238
+
239
+ # batch on gaussian to avoid OOM
240
+ batch_g = 1024
241
+ val = 0
242
+ for start in range(0, g_covs.shape[1], batch_g):
243
+ end = min(start + batch_g, g_covs.shape[1])
244
+ 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]
245
+ val += (mask_opas[:, start:end] * w).sum(-1)
246
+
247
+ # kiui.lo(val, mask_opas, w)
248
+
249
+ occ[xi * split_size: xi * split_size + len(xs),
250
+ yi * split_size: yi * split_size + len(ys),
251
+ zi * split_size: zi * split_size + len(zs)] = val.reshape(len(xs), len(ys), len(zs))
252
+ return occ
253
+
254
+ def extract_mesh_t(self, path, density_thresh=1, t=0, resolution=128, decimate_target=1e5):
255
+ from mesh import Mesh
256
+ from mesh_utils import decimate_mesh, clean_mesh
257
+
258
+ os.makedirs(os.path.dirname(path), exist_ok=True)
259
+
260
+ occ = self.extract_fields_t(resolution, t=t).detach().cpu().numpy()
261
+
262
+ import mcubes
263
+ vertices, triangles = mcubes.marching_cubes(occ, density_thresh)
264
+ vertices = vertices / (resolution - 1.0) * 2 - 1
265
+
266
+ # transform back to the original space
267
+ vertices = vertices / self.scale + self.center.detach().cpu().numpy()
268
+
269
+ vertices, triangles = clean_mesh(vertices, triangles, remesh=True, remesh_size=0.015)
270
+ if decimate_target > 0 and triangles.shape[0] > decimate_target:
271
+ vertices, triangles = decimate_mesh(vertices, triangles, decimate_target)
272
+
273
+ v = torch.from_numpy(vertices.astype(np.float32)).contiguous().cuda()
274
+ f = torch.from_numpy(triangles.astype(np.int32)).contiguous().cuda()
275
+
276
+ print(
277
+ f"[INFO] marching cubes result: {v.shape} ({v.min().item()}-{v.max().item()}), {f.shape}"
278
+ )
279
+
280
+ mesh = Mesh(v=v, f=f, device='cuda')
281
+
282
+ return mesh
283
+
284
+ def get_covariance(self, scaling_modifier = 1):
285
+ return self.covariance_activation(self.get_scaling, scaling_modifier, self._rotation)
286
+
287
+ def oneupSHdegree(self):
288
+ if self.active_sh_degree < self.max_sh_degree:
289
+ self.active_sh_degree += 1
290
+
291
+ def create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : float, time_line: int):
292
+ self.spatial_lr_scale = spatial_lr_scale
293
+ fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda()
294
+ fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda())
295
+ features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda()
296
+ features[:, :3, 0 ] = fused_color
297
+ features[:, 3:, 1:] = 0.0
298
+
299
+ print("Number of points at initialisation : ", fused_point_cloud.shape[0])
300
+
301
+ dist2 = torch.clamp_min(distCUDA2(torch.from_numpy(np.asarray(pcd.points)).float().cuda()), 0.0000001)
302
+ scales = torch.log(torch.sqrt(dist2))[...,None].repeat(1, 3)
303
+ rots = torch.zeros((fused_point_cloud.shape[0], 4), device="cuda")
304
+ rots[:, 0] = 1
305
+
306
+ opacities = inverse_sigmoid(0.1 * torch.ones((fused_point_cloud.shape[0], 1), dtype=torch.float, device="cuda"))
307
+
308
+ self._xyz = nn.Parameter(fused_point_cloud.requires_grad_(True))
309
+ self._deformation = self._deformation.to("cuda")
310
+ # self.grid = self.grid.to("cuda")
311
+ self._features_dc = nn.Parameter(features[:,:,0:1].transpose(1, 2).contiguous().requires_grad_(True))
312
+ self._features_rest = nn.Parameter(features[:,:,1:].transpose(1, 2).contiguous().requires_grad_(True))
313
+ self._scaling = nn.Parameter(scales.requires_grad_(True))
314
+ self._rotation = nn.Parameter(rots.requires_grad_(True))
315
+ self._opacity = nn.Parameter(opacities.requires_grad_(True))
316
+ self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
317
+ self._deformation_table = torch.gt(torch.ones((self.get_xyz.shape[0]),device="cuda"),0)
318
+
319
+ def training_setup(self, training_args):
320
+ self.percent_dense = training_args.percent_dense
321
+ self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
322
+ self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
323
+ self._deformation_accum = torch.zeros((self.get_xyz.shape[0],3),device="cuda")
324
+ self.T = training_args.batch_size
325
+
326
+ if training_args.optimize_gaussians:
327
+ l = [
328
+ {'params': [self._xyz], 'lr': training_args.position_lr_init * self.spatial_lr_scale, "name": "xyz"},
329
+ {'params': list(self._deformation.get_mlp_parameters()), 'lr': training_args.deformation_lr_init * self.spatial_lr_scale, "name": "deformation"},
330
+ {'params': list(self._deformation.get_grid_parameters()), 'lr': training_args.grid_lr_init * self.spatial_lr_scale, "name": "grid"},
331
+ {'params': [self._features_dc], 'lr': training_args.feature_lr, "name": "f_dc"},
332
+ {'params': [self._features_rest], 'lr': training_args.feature_lr / 20.0, "name": "f_rest"},
333
+ {'params': [self._opacity], 'lr': training_args.opacity_lr, "name": "opacity"},
334
+ {'params': [self._scaling], 'lr': training_args.scaling_lr, "name": "scaling"},
335
+ {'params': [self._rotation], 'lr': training_args.rotation_lr, "name": "rotation"}
336
+ ]
337
+ else:
338
+ l = [
339
+ {'params': list(self._deformation.get_mlp_parameters()), 'lr': training_args.deformation_lr_init * self.spatial_lr_scale, "name": "deformation"},
340
+ {'params': list(self._deformation.get_grid_parameters()), 'lr': training_args.grid_lr_init * self.spatial_lr_scale, "name": "grid"},
341
+ ]
342
+
343
+ self.optimizer = torch.optim.Adam(l, lr=0.0, eps=1e-15)
344
+ self.xyz_scheduler_args = get_expon_lr_func(lr_init=training_args.position_lr_init*self.spatial_lr_scale,
345
+ lr_final=training_args.position_lr_final*self.spatial_lr_scale,
346
+ lr_delay_mult=training_args.position_lr_delay_mult,
347
+ max_steps=training_args.position_lr_max_steps)
348
+ self.deformation_scheduler_args = get_expon_lr_func(lr_init=training_args.deformation_lr_init*self.spatial_lr_scale,
349
+ lr_final=training_args.deformation_lr_final*self.spatial_lr_scale,
350
+ lr_delay_mult=training_args.deformation_lr_delay_mult,
351
+ max_steps=training_args.position_lr_max_steps)
352
+ self.grid_scheduler_args = get_expon_lr_func(lr_init=training_args.grid_lr_init*self.spatial_lr_scale,
353
+ lr_final=training_args.grid_lr_final*self.spatial_lr_scale,
354
+ lr_delay_mult=training_args.deformation_lr_delay_mult,
355
+ max_steps=training_args.position_lr_max_steps)
356
+
357
+ def update_learning_rate(self, iteration):
358
+ ''' Learning rate scheduling per step '''
359
+ for param_group in self.optimizer.param_groups:
360
+ if param_group["name"] == "xyz":
361
+ lr = self.xyz_scheduler_args(iteration)
362
+ param_group['lr'] = lr
363
+ # return lr
364
+ if "grid" in param_group["name"]:
365
+ lr = self.grid_scheduler_args(iteration)
366
+ param_group['lr'] = lr
367
+ # return lr
368
+ elif param_group["name"] == "deformation":
369
+ lr = self.deformation_scheduler_args(iteration)
370
+ param_group['lr'] = lr
371
+ # return lr
372
+
373
+ def construct_list_of_attributes(self):
374
+ l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
375
+ # All channels except the 3 DC
376
+ for i in range(self._features_dc.shape[1]*self._features_dc.shape[2]):
377
+ l.append('f_dc_{}'.format(i))
378
+ for i in range(self._features_rest.shape[1]*self._features_rest.shape[2]):
379
+ l.append('f_rest_{}'.format(i))
380
+ l.append('opacity')
381
+ for i in range(self._scaling.shape[1]):
382
+ l.append('scale_{}'.format(i))
383
+ for i in range(self._rotation.shape[1]):
384
+ l.append('rot_{}'.format(i))
385
+ return l
386
+ def compute_deformation(self,time):
387
+
388
+ deform = self._deformation[:,:,:time].sum(dim=-1)
389
+ xyz = self._xyz + deform
390
+ return xyz
391
+
392
+ def load_model(self, path, name):
393
+ print("loading model from exists{}".format(path))
394
+ weight_dict = torch.load(os.path.join(path, name+"_deformation.pth"),map_location="cuda")
395
+ self._deformation.load_state_dict(weight_dict)
396
+ self._deformation = self._deformation.to("cuda")
397
+ self._deformation_table = torch.gt(torch.ones((self.get_xyz.shape[0]),device="cuda"),0)
398
+ self._deformation_accum = torch.zeros((self.get_xyz.shape[0],3),device="cuda")
399
+ if os.path.exists(os.path.join(path, name+"_deformation_table.pth")):
400
+ self._deformation_table = torch.load(os.path.join(path, name+"_deformation_table.pth"),map_location="cuda")
401
+ if os.path.exists(os.path.join(path,name+"_deformation_accum.pth")):
402
+ self._deformation_accum = torch.load(os.path.join(path, name+"_deformation_accum.pth"),map_location="cuda")
403
+ self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
404
+
405
+ def save_deformation(self, path, name):
406
+ torch.save(self._deformation.state_dict(),os.path.join(path, name+"_deformation.pth"))
407
+ torch.save(self._deformation_table,os.path.join(path, name+"_deformation_table.pth"))
408
+ torch.save(self._deformation_accum,os.path.join(path, name+"_deformation_accum.pth"))
409
+
410
+ def save_ply(self, path):
411
+ mkdir_p(os.path.dirname(path))
412
+
413
+ xyz = self._xyz.detach().cpu().numpy()
414
+ normals = np.zeros_like(xyz)
415
+ f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
416
+ f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
417
+ opacities = self._opacity.detach().cpu().numpy()
418
+ scale = self._scaling.detach().cpu().numpy()
419
+ rotation = self._rotation.detach().cpu().numpy()
420
+
421
+ dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()]
422
+
423
+ elements = np.empty(xyz.shape[0], dtype=dtype_full)
424
+ attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1)
425
+ elements[:] = list(map(tuple, attributes))
426
+ el = PlyElement.describe(elements, 'vertex')
427
+ PlyData([el]).write(path)
428
+
429
+ def save_frame_ply(self, path, t):
430
+ mkdir_p(os.path.dirname(path))
431
+
432
+ xyzs, rotation, scale, opacities = self.get_deformed_everything(t)
433
+
434
+ xyz = xyzs.detach().cpu().numpy()
435
+ normals = np.zeros_like(xyz)
436
+ f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
437
+ f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
438
+ opacities = opacities.detach().cpu().numpy()
439
+ scale = scale.detach().cpu().numpy()
440
+ rotation = rotation.detach().cpu().numpy()
441
+
442
+ dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()]
443
+
444
+ elements = np.empty(xyz.shape[0], dtype=dtype_full)
445
+ attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1)
446
+ elements[:] = list(map(tuple, attributes))
447
+ el = PlyElement.describe(elements, 'vertex')
448
+ PlyData([el]).write(path)
449
+ # def save_frame_ply(self, path, t):
450
+ # mkdir_p(os.path.dirname(path))
451
+
452
+ # xyz = self._xyz.detach().cpu().numpy()
453
+ # normals = np.zeros_like(xyz)
454
+ # f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
455
+ # f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
456
+ # opacities = self._opacity.detach().cpu().numpy()
457
+ # scale = self._scaling.detach().cpu().numpy()
458
+ # rotation = self._rotation.detach().cpu().numpy()
459
+
460
+ # dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()]
461
+
462
+ # elements = np.empty(xyz.shape[0], dtype=dtype_full)
463
+ # attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1)
464
+ # elements[:] = list(map(tuple, attributes))
465
+ # el = PlyElement.describe(elements, 'vertex')
466
+ # PlyData([el]).write(path)
467
+
468
+ def reset_opacity(self):
469
+ opacities_new = inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01))
470
+ optimizable_tensors = self.replace_tensor_to_optimizer(opacities_new, "opacity")
471
+ self._opacity = optimizable_tensors["opacity"]
472
+
473
+ def load_ply(self, path):
474
+ self.spatial_lr_scale = 1
475
+ plydata = PlyData.read(path)
476
+
477
+ xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
478
+ np.asarray(plydata.elements[0]["y"]),
479
+ np.asarray(plydata.elements[0]["z"])), axis=1)
480
+ opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]
481
+
482
+ features_dc = np.zeros((xyz.shape[0], 3, 1))
483
+ features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
484
+ features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"])
485
+ features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"])
486
+
487
+ extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")]
488
+ extra_f_names = sorted(extra_f_names, key = lambda x: int(x.split('_')[-1]))
489
+ assert len(extra_f_names)==3*(self.max_sh_degree + 1) ** 2 - 3
490
+ features_extra = np.zeros((xyz.shape[0], len(extra_f_names)))
491
+ for idx, attr_name in enumerate(extra_f_names):
492
+ features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name])
493
+ # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC)
494
+ features_extra = features_extra.reshape((features_extra.shape[0], 3, (self.max_sh_degree + 1) ** 2 - 1))
495
+
496
+ scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
497
+ scale_names = sorted(scale_names, key = lambda x: int(x.split('_')[-1]))
498
+ scales = np.zeros((xyz.shape[0], len(scale_names)))
499
+ for idx, attr_name in enumerate(scale_names):
500
+ scales[:, idx] = np.asarray(plydata.elements[0][attr_name])
501
+
502
+ rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")]
503
+ rot_names = sorted(rot_names, key = lambda x: int(x.split('_')[-1]))
504
+ rots = np.zeros((xyz.shape[0], len(rot_names)))
505
+ for idx, attr_name in enumerate(rot_names):
506
+ rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
507
+
508
+ self._xyz = nn.Parameter(torch.tensor(xyz, dtype=torch.float, device="cuda").requires_grad_(True))
509
+ self._features_dc = nn.Parameter(torch.tensor(features_dc, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True))
510
+ self._features_rest = nn.Parameter(torch.tensor(features_extra, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True))
511
+ self._opacity = nn.Parameter(torch.tensor(opacities, dtype=torch.float, device="cuda").requires_grad_(True))
512
+ self._scaling = nn.Parameter(torch.tensor(scales, dtype=torch.float, device="cuda").requires_grad_(True))
513
+ self._rotation = nn.Parameter(torch.tensor(rots, dtype=torch.float, device="cuda").requires_grad_(True))
514
+ self.active_sh_degree = self.max_sh_degree
515
+ self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
516
+ self._deformation = self._deformation.to("cuda")
517
+ self._deformation_table = torch.gt(torch.ones((self.get_xyz.shape[0]),device="cuda"),0) # everything deformed
518
+
519
+ print(self._xyz.shape)
520
+
521
+
522
+ def replace_tensor_to_optimizer(self, tensor, name):
523
+ optimizable_tensors = {}
524
+ for group in self.optimizer.param_groups:
525
+ if group["name"] == name:
526
+ stored_state = self.optimizer.state.get(group['params'][0], None)
527
+ stored_state["exp_avg"] = torch.zeros_like(tensor)
528
+ stored_state["exp_avg_sq"] = torch.zeros_like(tensor)
529
+
530
+ del self.optimizer.state[group['params'][0]]
531
+ group["params"][0] = nn.Parameter(tensor.requires_grad_(True))
532
+ self.optimizer.state[group['params'][0]] = stored_state
533
+
534
+ optimizable_tensors[group["name"]] = group["params"][0]
535
+ return optimizable_tensors
536
+
537
+ def _prune_optimizer(self, mask):
538
+ optimizable_tensors = {}
539
+ for group in self.optimizer.param_groups:
540
+ if len(group["params"]) > 1:
541
+ continue
542
+ stored_state = self.optimizer.state.get(group['params'][0], None)
543
+ if stored_state is not None:
544
+ stored_state["exp_avg"] = stored_state["exp_avg"][mask]
545
+ stored_state["exp_avg_sq"] = stored_state["exp_avg_sq"][mask]
546
+
547
+ del self.optimizer.state[group['params'][0]]
548
+ group["params"][0] = nn.Parameter((group["params"][0][mask].requires_grad_(True)))
549
+ self.optimizer.state[group['params'][0]] = stored_state
550
+
551
+ optimizable_tensors[group["name"]] = group["params"][0]
552
+ else:
553
+ group["params"][0] = nn.Parameter(group["params"][0][mask].requires_grad_(True))
554
+ optimizable_tensors[group["name"]] = group["params"][0]
555
+ return optimizable_tensors
556
+
557
+ def prune_points(self, mask):
558
+ valid_points_mask = ~mask
559
+ optimizable_tensors = self._prune_optimizer(valid_points_mask)
560
+
561
+ self._xyz = optimizable_tensors["xyz"]
562
+ self._features_dc = optimizable_tensors["f_dc"]
563
+ self._features_rest = optimizable_tensors["f_rest"]
564
+ self._opacity = optimizable_tensors["opacity"]
565
+ self._scaling = optimizable_tensors["scaling"]
566
+ self._rotation = optimizable_tensors["rotation"]
567
+ self._deformation_accum = self._deformation_accum[valid_points_mask]
568
+ self.xyz_gradient_accum = self.xyz_gradient_accum[valid_points_mask]
569
+ self._deformation_table = self._deformation_table[valid_points_mask]
570
+ self.denom = self.denom[valid_points_mask]
571
+ self.max_radii2D = self.max_radii2D[valid_points_mask]
572
+
573
+ def cat_tensors_to_optimizer(self, tensors_dict):
574
+ optimizable_tensors = {}
575
+ for group in self.optimizer.param_groups:
576
+ if len(group["params"])>1:continue
577
+ assert len(group["params"]) == 1
578
+ extension_tensor = tensors_dict[group["name"]]
579
+ stored_state = self.optimizer.state.get(group['params'][0], None)
580
+ if stored_state is not None:
581
+
582
+ stored_state["exp_avg"] = torch.cat((stored_state["exp_avg"], torch.zeros_like(extension_tensor)), dim=0)
583
+ stored_state["exp_avg_sq"] = torch.cat((stored_state["exp_avg_sq"], torch.zeros_like(extension_tensor)), dim=0)
584
+
585
+ del self.optimizer.state[group['params'][0]]
586
+ group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))
587
+ self.optimizer.state[group['params'][0]] = stored_state
588
+
589
+ optimizable_tensors[group["name"]] = group["params"][0]
590
+ else:
591
+ group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))
592
+ optimizable_tensors[group["name"]] = group["params"][0]
593
+
594
+ return optimizable_tensors
595
+
596
+ def densification_postfix(self, new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation, new_deformation_table):
597
+ d = {"xyz": new_xyz,
598
+ "f_dc": new_features_dc,
599
+ "f_rest": new_features_rest,
600
+ "opacity": new_opacities,
601
+ "scaling" : new_scaling,
602
+ "rotation" : new_rotation,
603
+ # "deformation": new_deformation
604
+ }
605
+
606
+ optimizable_tensors = self.cat_tensors_to_optimizer(d)
607
+ self._xyz = optimizable_tensors["xyz"]
608
+ self._features_dc = optimizable_tensors["f_dc"]
609
+ self._features_rest = optimizable_tensors["f_rest"]
610
+ self._opacity = optimizable_tensors["opacity"]
611
+ self._scaling = optimizable_tensors["scaling"]
612
+ self._rotation = optimizable_tensors["rotation"]
613
+ # self._deformation = optimizable_tensors["deformation"]
614
+
615
+ self._deformation_table = torch.cat([self._deformation_table,new_deformation_table],-1)
616
+ self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
617
+ self._deformation_accum = torch.zeros((self.get_xyz.shape[0], 3), device="cuda")
618
+ self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
619
+ self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
620
+
621
+ def densify_and_split(self, grads, grad_threshold, scene_extent, N=2):
622
+ n_init_points = self.get_xyz.shape[0]
623
+ # Extract points that satisfy the gradient condition
624
+ padded_grad = torch.zeros((n_init_points), device="cuda")
625
+ padded_grad[:grads.shape[0]] = grads.squeeze()
626
+ selected_pts_mask = torch.where(padded_grad >= grad_threshold, True, False)
627
+ selected_pts_mask = torch.logical_and(selected_pts_mask,
628
+ torch.max(self.get_scaling, dim=1).values > self.percent_dense*scene_extent)
629
+ if not selected_pts_mask.any():
630
+ return
631
+ stds = self.get_scaling[selected_pts_mask].repeat(N,1)
632
+ means =torch.zeros((stds.size(0), 3),device="cuda")
633
+ samples = torch.normal(mean=means, std=stds)
634
+ rots = build_rotation(self._rotation[selected_pts_mask]).repeat(N,1,1)
635
+ new_xyz = torch.bmm(rots, samples.unsqueeze(-1)).squeeze(-1) + self.get_xyz[selected_pts_mask].repeat(N, 1)
636
+ new_scaling = self.scaling_inverse_activation(self.get_scaling[selected_pts_mask].repeat(N,1) / (0.8*N))
637
+ new_rotation = self._rotation[selected_pts_mask].repeat(N,1)
638
+ new_features_dc = self._features_dc[selected_pts_mask].repeat(N,1,1)
639
+ new_features_rest = self._features_rest[selected_pts_mask].repeat(N,1,1)
640
+ new_opacity = self._opacity[selected_pts_mask].repeat(N,1)
641
+ new_deformation_table = self._deformation_table[selected_pts_mask].repeat(N)
642
+ self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacity, new_scaling, new_rotation, new_deformation_table)
643
+
644
+ prune_filter = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device="cuda", dtype=bool)))
645
+ self.prune_points(prune_filter)
646
+
647
+ def densify_and_clone(self, grads, grad_threshold, scene_extent):
648
+ # Extract points that satisfy the gradient condition
649
+ selected_pts_mask = torch.where(torch.norm(grads, dim=-1) >= grad_threshold, True, False)
650
+ selected_pts_mask = torch.logical_and(selected_pts_mask,
651
+ torch.max(self.get_scaling, dim=1).values <= self.percent_dense*scene_extent)
652
+
653
+ new_xyz = self._xyz[selected_pts_mask]
654
+ # - 0.001 * self._xyz.grad[selected_pts_mask]
655
+ new_features_dc = self._features_dc[selected_pts_mask]
656
+ new_features_rest = self._features_rest[selected_pts_mask]
657
+ new_opacities = self._opacity[selected_pts_mask]
658
+ new_scaling = self._scaling[selected_pts_mask]
659
+ new_rotation = self._rotation[selected_pts_mask]
660
+ new_deformation_table = self._deformation_table[selected_pts_mask]
661
+
662
+ self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation, new_deformation_table)
663
+
664
+ def prune(self, min_opacity, extent, max_screen_size):
665
+ prune_mask = (self.get_opacity < min_opacity).squeeze()
666
+ # 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"))
667
+ # prune_mask = torch.logical_or(prune_mask, prune_mask_2)
668
+ # deformation_sum = abs(self._deformation).sum(dim=-1).mean(dim=-1)
669
+ # deformation_mask = (deformation_sum < torch.quantile(deformation_sum, torch.tensor([0.5]).to("cuda")))
670
+ # prune_mask = prune_mask & deformation_mask
671
+ if max_screen_size:
672
+ big_points_vs = self.max_radii2D > max_screen_size
673
+ big_points_ws = self.get_scaling.max(dim=1).values > 0.1 * extent
674
+ prune_mask = torch.logical_or(prune_mask, big_points_vs)
675
+
676
+ prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws)
677
+ self.prune_points(prune_mask)
678
+
679
+ torch.cuda.empty_cache()
680
+ def densify(self, max_grad, min_opacity, extent, max_screen_size):
681
+ grads = self.xyz_gradient_accum / self.denom
682
+ grads[grads.isnan()] = 0.0
683
+
684
+ self.densify_and_clone(grads, max_grad, extent)
685
+ self.densify_and_split(grads, max_grad, extent)
686
+ def standard_constaint(self):
687
+
688
+ means3D = self._xyz.detach()
689
+ scales = self._scaling.detach()
690
+ rotations = self._rotation.detach()
691
+ opacity = self._opacity.detach()
692
+ time = torch.tensor(0).to("cuda").repeat(means3D.shape[0],1)
693
+ means3D_deform, scales_deform, rotations_deform, _ = self._deformation(means3D, scales, rotations, opacity, time)
694
+ position_error = (means3D_deform - means3D)**2
695
+ rotation_error = (rotations_deform - rotations)**2
696
+ scaling_erorr = (scales_deform - scales)**2
697
+ return position_error.mean() + rotation_error.mean() + scaling_erorr.mean()
698
+
699
+
700
+ def add_densification_stats(self, viewspace_point_tensor, update_filter):
701
+ self.xyz_gradient_accum[update_filter] += torch.norm(viewspace_point_tensor[update_filter,:2], dim=-1, keepdim=True)
702
+ self.denom[update_filter] += 1
703
+ @torch.no_grad()
704
+ def update_deformation_table(self,threshold):
705
+ # print("origin deformation point nums:",self._deformation_table.sum())
706
+ self._deformation_table = torch.gt(self._deformation_accum.max(dim=-1).values/100,threshold)
707
+ def print_deformation_weight_grad(self):
708
+ for name, weight in self._deformation.named_parameters():
709
+ if weight.requires_grad:
710
+ if weight.grad is None:
711
+
712
+ print(name," :",weight.grad)
713
+ else:
714
+ if weight.grad.mean() != 0:
715
+ print(name," :",weight.grad.mean(), weight.grad.min(), weight.grad.max())
716
+ print("-"*50)
717
+ def _plane_regulation(self):
718
+ multi_res_grids = self._deformation.deformation_net.grid.grids
719
+ total = 0
720
+ # model.grids is 6 x [1, rank * F_dim, reso, reso]
721
+ for grids in multi_res_grids:
722
+ if len(grids) == 3:
723
+ time_grids = []
724
+ else:
725
+ time_grids = [0,1,3]
726
+ for grid_id in time_grids:
727
+ total += compute_plane_smoothness(grids[grid_id])
728
+ return total
729
+ def _time_regulation(self):
730
+ multi_res_grids = self._deformation.deformation_net.grid.grids
731
+ total = 0
732
+ # model.grids is 6 x [1, rank * F_dim, reso, reso]
733
+ for grids in multi_res_grids:
734
+ if len(grids) == 3:
735
+ time_grids = []
736
+ else:
737
+ time_grids =[2, 4, 5]
738
+ for grid_id in time_grids:
739
+ total += compute_plane_smoothness(grids[grid_id])
740
+ return total
741
+ def _l1_regulation(self):
742
+ # model.grids is 6 x [1, rank * F_dim, reso, reso]
743
+ multi_res_grids = self._deformation.deformation_net.grid.grids
744
+
745
+ total = 0.0
746
+ for grids in multi_res_grids:
747
+ if len(grids) == 3:
748
+ continue
749
+ else:
750
+ # These are the spatiotemporal grids
751
+ spatiotemporal_grids = [2, 4, 5]
752
+ for grid_id in spatiotemporal_grids:
753
+ total += torch.abs(1 - grids[grid_id]).mean()
754
+ return total
755
+ def compute_regulation(self, time_smoothness_weight, l1_time_planes_weight, plane_tv_weight):
756
+ return plane_tv_weight * self._plane_regulation() + time_smoothness_weight * self._time_regulation() + l1_time_planes_weight * self._l1_regulation()
757
+
758
+
759
+ def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_size):
760
+ grads = self.xyz_gradient_accum / self.denom
761
+ grads[grads.isnan()] = 0.0
762
+
763
+ self.densify_and_clone(grads, max_grad, extent)
764
+ self.densify_and_split(grads, max_grad, extent)
765
+
766
+ prune_mask = (self.get_opacity < min_opacity).squeeze()
767
+ if max_screen_size:
768
+ big_points_vs = self.max_radii2D > max_screen_size
769
+ big_points_ws = self.get_scaling.max(dim=1).values > 0.1 * extent
770
+ prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws)
771
+ self.prune_points(prune_mask)
772
+
773
+ torch.cuda.empty_cache()
grid_put.py ADDED
@@ -0,0 +1,300 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn.functional as F
3
+
4
+ def stride_from_shape(shape):
5
+ stride = [1]
6
+ for x in reversed(shape[1:]):
7
+ stride.append(stride[-1] * x)
8
+ return list(reversed(stride))
9
+
10
+
11
+ def scatter_add_nd(input, indices, values):
12
+ # input: [..., C], D dimension + C channel
13
+ # indices: [N, D], long
14
+ # values: [N, C]
15
+
16
+ D = indices.shape[-1]
17
+ C = input.shape[-1]
18
+ size = input.shape[:-1]
19
+ stride = stride_from_shape(size)
20
+
21
+ assert len(size) == D
22
+
23
+ input = input.view(-1, C) # [HW, C]
24
+ flatten_indices = (indices * torch.tensor(stride, dtype=torch.long, device=indices.device)).sum(-1) # [N]
25
+
26
+ input.scatter_add_(0, flatten_indices.unsqueeze(1).repeat(1, C), values)
27
+
28
+ return input.view(*size, C)
29
+
30
+
31
+ def scatter_add_nd_with_count(input, count, indices, values, weights=None):
32
+ # input: [..., C], D dimension + C channel
33
+ # count: [..., 1], D dimension
34
+ # indices: [N, D], long
35
+ # values: [N, C]
36
+
37
+ D = indices.shape[-1]
38
+ C = input.shape[-1]
39
+ size = input.shape[:-1]
40
+ stride = stride_from_shape(size)
41
+
42
+ assert len(size) == D
43
+
44
+ input = input.view(-1, C) # [HW, C]
45
+ count = count.view(-1, 1)
46
+
47
+ flatten_indices = (indices * torch.tensor(stride, dtype=torch.long, device=indices.device)).sum(-1) # [N]
48
+
49
+ if weights is None:
50
+ weights = torch.ones_like(values[..., :1])
51
+
52
+ input.scatter_add_(0, flatten_indices.unsqueeze(1).repeat(1, C), values)
53
+ count.scatter_add_(0, flatten_indices.unsqueeze(1), weights)
54
+
55
+ return input.view(*size, C), count.view(*size, 1)
56
+
57
+ def nearest_grid_put_2d(H, W, coords, values, return_count=False):
58
+ # coords: [N, 2], float in [-1, 1]
59
+ # values: [N, C]
60
+
61
+ C = values.shape[-1]
62
+
63
+ indices = (coords * 0.5 + 0.5) * torch.tensor(
64
+ [H - 1, W - 1], dtype=torch.float32, device=coords.device
65
+ )
66
+ indices = indices.round().long() # [N, 2]
67
+
68
+ result = torch.zeros(H, W, C, device=values.device, dtype=values.dtype) # [H, W, C]
69
+ count = torch.zeros(H, W, 1, device=values.device, dtype=values.dtype) # [H, W, 1]
70
+ weights = torch.ones_like(values[..., :1]) # [N, 1]
71
+
72
+ result, count = scatter_add_nd_with_count(result, count, indices, values, weights)
73
+
74
+ if return_count:
75
+ return result, count
76
+
77
+ mask = (count.squeeze(-1) > 0)
78
+ result[mask] = result[mask] / count[mask].repeat(1, C)
79
+
80
+ return result
81
+
82
+
83
+ def linear_grid_put_2d(H, W, coords, values, return_count=False):
84
+ # coords: [N, 2], float in [-1, 1]
85
+ # values: [N, C]
86
+
87
+ C = values.shape[-1]
88
+
89
+ indices = (coords * 0.5 + 0.5) * torch.tensor(
90
+ [H - 1, W - 1], dtype=torch.float32, device=coords.device
91
+ )
92
+ indices_00 = indices.floor().long() # [N, 2]
93
+ indices_00[:, 0].clamp_(0, H - 2)
94
+ indices_00[:, 1].clamp_(0, W - 2)
95
+ indices_01 = indices_00 + torch.tensor(
96
+ [0, 1], dtype=torch.long, device=indices.device
97
+ )
98
+ indices_10 = indices_00 + torch.tensor(
99
+ [1, 0], dtype=torch.long, device=indices.device
100
+ )
101
+ indices_11 = indices_00 + torch.tensor(
102
+ [1, 1], dtype=torch.long, device=indices.device
103
+ )
104
+
105
+ h = indices[..., 0] - indices_00[..., 0].float()
106
+ w = indices[..., 1] - indices_00[..., 1].float()
107
+ w_00 = (1 - h) * (1 - w)
108
+ w_01 = (1 - h) * w
109
+ w_10 = h * (1 - w)
110
+ w_11 = h * w
111
+
112
+ result = torch.zeros(H, W, C, device=values.device, dtype=values.dtype) # [H, W, C]
113
+ count = torch.zeros(H, W, 1, device=values.device, dtype=values.dtype) # [H, W, 1]
114
+ weights = torch.ones_like(values[..., :1]) # [N, 1]
115
+
116
+ result, count = scatter_add_nd_with_count(result, count, indices_00, values * w_00.unsqueeze(1), weights* w_00.unsqueeze(1))
117
+ result, count = scatter_add_nd_with_count(result, count, indices_01, values * w_01.unsqueeze(1), weights* w_01.unsqueeze(1))
118
+ result, count = scatter_add_nd_with_count(result, count, indices_10, values * w_10.unsqueeze(1), weights* w_10.unsqueeze(1))
119
+ result, count = scatter_add_nd_with_count(result, count, indices_11, values * w_11.unsqueeze(1), weights* w_11.unsqueeze(1))
120
+
121
+ if return_count:
122
+ return result, count
123
+
124
+ mask = (count.squeeze(-1) > 0)
125
+ result[mask] = result[mask] / count[mask].repeat(1, C)
126
+
127
+ return result
128
+
129
+ def mipmap_linear_grid_put_2d(H, W, coords, values, min_resolution=32, return_count=False):
130
+ # coords: [N, 2], float in [-1, 1]
131
+ # values: [N, C]
132
+
133
+ C = values.shape[-1]
134
+
135
+ result = torch.zeros(H, W, C, device=values.device, dtype=values.dtype) # [H, W, C]
136
+ count = torch.zeros(H, W, 1, device=values.device, dtype=values.dtype) # [H, W, 1]
137
+
138
+ cur_H, cur_W = H, W
139
+
140
+ while min(cur_H, cur_W) > min_resolution:
141
+
142
+ # try to fill the holes
143
+ mask = (count.squeeze(-1) == 0)
144
+ if not mask.any():
145
+ break
146
+
147
+ cur_result, cur_count = linear_grid_put_2d(cur_H, cur_W, coords, values, return_count=True)
148
+ 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]
149
+ 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]
150
+ cur_H //= 2
151
+ cur_W //= 2
152
+
153
+ if return_count:
154
+ return result, count
155
+
156
+ mask = (count.squeeze(-1) > 0)
157
+ result[mask] = result[mask] / count[mask].repeat(1, C)
158
+
159
+ return result
160
+
161
+ def nearest_grid_put_3d(H, W, D, coords, values, return_count=False):
162
+ # coords: [N, 3], float in [-1, 1]
163
+ # values: [N, C]
164
+
165
+ C = values.shape[-1]
166
+
167
+ indices = (coords * 0.5 + 0.5) * torch.tensor(
168
+ [H - 1, W - 1, D - 1], dtype=torch.float32, device=coords.device
169
+ )
170
+ indices = indices.round().long() # [N, 2]
171
+
172
+ result = torch.zeros(H, W, D, C, device=values.device, dtype=values.dtype) # [H, W, C]
173
+ count = torch.zeros(H, W, D, 1, device=values.device, dtype=values.dtype) # [H, W, 1]
174
+ weights = torch.ones_like(values[..., :1]) # [N, 1]
175
+
176
+ result, count = scatter_add_nd_with_count(result, count, indices, values, weights)
177
+
178
+ if return_count:
179
+ return result, count
180
+
181
+ mask = (count.squeeze(-1) > 0)
182
+ result[mask] = result[mask] / count[mask].repeat(1, C)
183
+
184
+ return result
185
+
186
+
187
+ def linear_grid_put_3d(H, W, D, coords, values, return_count=False):
188
+ # coords: [N, 3], float in [-1, 1]
189
+ # values: [N, C]
190
+
191
+ C = values.shape[-1]
192
+
193
+ indices = (coords * 0.5 + 0.5) * torch.tensor(
194
+ [H - 1, W - 1, D - 1], dtype=torch.float32, device=coords.device
195
+ )
196
+ indices_000 = indices.floor().long() # [N, 3]
197
+ indices_000[:, 0].clamp_(0, H - 2)
198
+ indices_000[:, 1].clamp_(0, W - 2)
199
+ indices_000[:, 2].clamp_(0, D - 2)
200
+
201
+ indices_001 = indices_000 + torch.tensor([0, 0, 1], dtype=torch.long, device=indices.device)
202
+ indices_010 = indices_000 + torch.tensor([0, 1, 0], dtype=torch.long, device=indices.device)
203
+ indices_011 = indices_000 + torch.tensor([0, 1, 1], dtype=torch.long, device=indices.device)
204
+ indices_100 = indices_000 + torch.tensor([1, 0, 0], dtype=torch.long, device=indices.device)
205
+ indices_101 = indices_000 + torch.tensor([1, 0, 1], dtype=torch.long, device=indices.device)
206
+ indices_110 = indices_000 + torch.tensor([1, 1, 0], dtype=torch.long, device=indices.device)
207
+ indices_111 = indices_000 + torch.tensor([1, 1, 1], dtype=torch.long, device=indices.device)
208
+
209
+ h = indices[..., 0] - indices_000[..., 0].float()
210
+ w = indices[..., 1] - indices_000[..., 1].float()
211
+ d = indices[..., 2] - indices_000[..., 2].float()
212
+
213
+ w_000 = (1 - h) * (1 - w) * (1 - d)
214
+ w_001 = (1 - h) * w * (1 - d)
215
+ w_010 = h * (1 - w) * (1 - d)
216
+ w_011 = h * w * (1 - d)
217
+ w_100 = (1 - h) * (1 - w) * d
218
+ w_101 = (1 - h) * w * d
219
+ w_110 = h * (1 - w) * d
220
+ w_111 = h * w * d
221
+
222
+ result = torch.zeros(H, W, D, C, device=values.device, dtype=values.dtype) # [H, W, D, C]
223
+ count = torch.zeros(H, W, D, 1, device=values.device, dtype=values.dtype) # [H, W, D, 1]
224
+ weights = torch.ones_like(values[..., :1]) # [N, 1]
225
+
226
+ result, count = scatter_add_nd_with_count(result, count, indices_000, values * w_000.unsqueeze(1), weights * w_000.unsqueeze(1))
227
+ result, count = scatter_add_nd_with_count(result, count, indices_001, values * w_001.unsqueeze(1), weights * w_001.unsqueeze(1))
228
+ result, count = scatter_add_nd_with_count(result, count, indices_010, values * w_010.unsqueeze(1), weights * w_010.unsqueeze(1))
229
+ result, count = scatter_add_nd_with_count(result, count, indices_011, values * w_011.unsqueeze(1), weights * w_011.unsqueeze(1))
230
+ result, count = scatter_add_nd_with_count(result, count, indices_100, values * w_100.unsqueeze(1), weights * w_100.unsqueeze(1))
231
+ result, count = scatter_add_nd_with_count(result, count, indices_101, values * w_101.unsqueeze(1), weights * w_101.unsqueeze(1))
232
+ result, count = scatter_add_nd_with_count(result, count, indices_110, values * w_110.unsqueeze(1), weights * w_110.unsqueeze(1))
233
+ result, count = scatter_add_nd_with_count(result, count, indices_111, values * w_111.unsqueeze(1), weights * w_111.unsqueeze(1))
234
+
235
+ if return_count:
236
+ return result, count
237
+
238
+ mask = (count.squeeze(-1) > 0)
239
+ result[mask] = result[mask] / count[mask].repeat(1, C)
240
+
241
+ return result
242
+
243
+ def mipmap_linear_grid_put_3d(H, W, D, coords, values, min_resolution=32, return_count=False):
244
+ # coords: [N, 3], float in [-1, 1]
245
+ # values: [N, C]
246
+
247
+ C = values.shape[-1]
248
+
249
+ result = torch.zeros(H, W, D, C, device=values.device, dtype=values.dtype) # [H, W, D, C]
250
+ count = torch.zeros(H, W, D, 1, device=values.device, dtype=values.dtype) # [H, W, D, 1]
251
+ cur_H, cur_W, cur_D = H, W, D
252
+
253
+ while min(min(cur_H, cur_W), cur_D) > min_resolution:
254
+
255
+ # try to fill the holes
256
+ mask = (count.squeeze(-1) == 0)
257
+ if not mask.any():
258
+ break
259
+
260
+ cur_result, cur_count = linear_grid_put_3d(cur_H, cur_W, cur_D, coords, values, return_count=True)
261
+ 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]
262
+ 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]
263
+ cur_H //= 2
264
+ cur_W //= 2
265
+ cur_D //= 2
266
+
267
+ if return_count:
268
+ return result, count
269
+
270
+ mask = (count.squeeze(-1) > 0)
271
+ result[mask] = result[mask] / count[mask].repeat(1, C)
272
+
273
+ return result
274
+
275
+
276
+ def grid_put(shape, coords, values, mode='linear-mipmap', min_resolution=32, return_raw=False):
277
+ # shape: [D], list/tuple
278
+ # coords: [N, D], float in [-1, 1]
279
+ # values: [N, C]
280
+
281
+ D = len(shape)
282
+ assert D in [2, 3], f'only support D == 2 or 3, but got D == {D}'
283
+
284
+ if mode == 'nearest':
285
+ if D == 2:
286
+ return nearest_grid_put_2d(*shape, coords, values, return_raw)
287
+ else:
288
+ return nearest_grid_put_3d(*shape, coords, values, return_raw)
289
+ elif mode == 'linear':
290
+ if D == 2:
291
+ return linear_grid_put_2d(*shape, coords, values, return_raw)
292
+ else:
293
+ return linear_grid_put_3d(*shape, coords, values, return_raw)
294
+ elif mode == 'linear-mipmap':
295
+ if D == 2:
296
+ return mipmap_linear_grid_put_2d(*shape, coords, values, min_resolution, return_raw)
297
+ else:
298
+ return mipmap_linear_grid_put_3d(*shape, coords, values, min_resolution, return_raw)
299
+ else:
300
+ raise NotImplementedError(f"got mode {mode}")
gs_renderer_4d.py ADDED
@@ -0,0 +1,277 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import numpy as np
3
+
4
+ import torch
5
+
6
+ from diff_gaussian_rasterization import (
7
+ GaussianRasterizationSettings,
8
+ GaussianRasterizer,
9
+ )
10
+
11
+ from sh_utils import eval_sh, SH2RGB, RGB2SH
12
+
13
+ from gaussian_model_4d import GaussianModel, BasicPointCloud
14
+
15
+ def getProjectionMatrix(znear, zfar, fovX, fovY):
16
+ tanHalfFovY = math.tan((fovY / 2))
17
+ tanHalfFovX = math.tan((fovX / 2))
18
+
19
+ P = torch.zeros(4, 4)
20
+
21
+ z_sign = 1.0
22
+
23
+ P[0, 0] = 1 / tanHalfFovX
24
+ P[1, 1] = 1 / tanHalfFovY
25
+ P[3, 2] = z_sign
26
+ P[2, 2] = z_sign * zfar / (zfar - znear)
27
+ P[2, 3] = -(zfar * znear) / (zfar - znear)
28
+ return P
29
+
30
+
31
+ class MiniCam:
32
+ def __init__(self, c2w, width, height, fovy, fovx, znear, zfar, time=0, gs_convention=True):
33
+ # c2w (pose) should be in NeRF convention.
34
+
35
+ self.image_width = width
36
+ self.image_height = height
37
+ self.FoVy = fovy
38
+ self.FoVx = fovx
39
+ self.znear = znear
40
+ self.zfar = zfar
41
+
42
+ w2c = np.linalg.inv(c2w)
43
+
44
+ if gs_convention:
45
+ # rectify...
46
+ w2c[1:3, :3] *= -1
47
+ w2c[:3, 3] *= -1
48
+
49
+ self.world_view_transform = torch.tensor(w2c).transpose(0, 1).cuda()
50
+ self.projection_matrix = (
51
+ getProjectionMatrix(
52
+ znear=self.znear, zfar=self.zfar, fovX=self.FoVx, fovY=self.FoVy
53
+ )
54
+ .transpose(0, 1)
55
+ .cuda()
56
+ )
57
+ self.full_proj_transform = self.world_view_transform @ self.projection_matrix
58
+ self.camera_center = -torch.tensor(c2w[:3, 3]).cuda()
59
+
60
+ self.time = time
61
+
62
+
63
+ class Renderer:
64
+ def __init__(self, opt, sh_degree=3, white_background=True, radius=1):
65
+
66
+ self.sh_degree = sh_degree
67
+ self.white_background = white_background
68
+ self.radius = radius
69
+ self.opt = opt
70
+ self.T = self.opt.batch_size
71
+
72
+ self.gaussians = GaussianModel(sh_degree, opt.deformation)
73
+
74
+ self.bg_color = torch.tensor(
75
+ [1, 1, 1] if white_background else [0, 0, 0],
76
+ dtype=torch.float32,
77
+ device="cuda",
78
+ )
79
+ self.means3D_deform_T = None
80
+ self.opacity_deform_T = None
81
+ self.scales_deform_T = None
82
+ self.rotations_deform_T = None
83
+
84
+
85
+
86
+ def initialize(self, input=None, num_pts=5000, radius=0.5):
87
+ # load checkpoint
88
+ if input is None:
89
+ # init from random point cloud
90
+
91
+ phis = np.random.random((num_pts,)) * 2 * np.pi
92
+ costheta = np.random.random((num_pts,)) * 2 - 1
93
+ thetas = np.arccos(costheta)
94
+ mu = np.random.random((num_pts,))
95
+ radius = radius * np.cbrt(mu)
96
+ x = radius * np.sin(thetas) * np.cos(phis)
97
+ y = radius * np.sin(thetas) * np.sin(phis)
98
+ z = radius * np.cos(thetas)
99
+ xyz = np.stack((x, y, z), axis=1)
100
+ # xyz = np.random.random((num_pts, 3)) * 2.6 - 1.3
101
+
102
+ shs = np.random.random((num_pts, 3)) / 255.0
103
+ pcd = BasicPointCloud(
104
+ points=xyz, colors=SH2RGB(shs), normals=np.zeros((num_pts, 3))
105
+ )
106
+ # self.gaussians.create_from_pcd(pcd, 10)
107
+ self.gaussians.create_from_pcd(pcd, 10, 1)
108
+ elif isinstance(input, BasicPointCloud):
109
+ # load from a provided pcd
110
+ self.gaussians.create_from_pcd(input, 1)
111
+ else:
112
+ # load from saved ply
113
+ self.gaussians.load_ply(input)
114
+
115
+ def prepare_render(
116
+ self,
117
+ ):
118
+ means3D = self.gaussians.get_xyz
119
+ opacity = self.gaussians._opacity
120
+ scales = self.gaussians._scaling
121
+ rotations = self.gaussians._rotation
122
+
123
+ means3D_T = []
124
+ opacity_T = []
125
+ scales_T = []
126
+ rotations_T = []
127
+ time_T = []
128
+
129
+ for t in range(self.T):
130
+ time = torch.tensor(t).to(means3D.device).repeat(means3D.shape[0],1)
131
+ time = ((time.float() / self.T) - 0.5) * 2
132
+
133
+ means3D_T.append(means3D)
134
+ opacity_T.append(opacity)
135
+ scales_T.append(scales)
136
+ rotations_T.append(rotations)
137
+ time_T.append(time)
138
+
139
+ means3D_T = torch.cat(means3D_T)
140
+ opacity_T = torch.cat(opacity_T)
141
+ scales_T = torch.cat(scales_T)
142
+ rotations_T = torch.cat(rotations_T)
143
+ time_T = torch.cat(time_T)
144
+
145
+
146
+ means3D_deform_T, scales_deform_T, rotations_deform_T, opacity_deform_T = self.gaussians._deformation(means3D_T, scales_T,
147
+ rotations_T, opacity_T,
148
+ time_T) # time is not none
149
+ self.means3D_deform_T = means3D_deform_T.reshape([self.T, means3D_deform_T.shape[0]//self.T, -1])
150
+ self.opacity_deform_T = opacity_deform_T.reshape([self.T, means3D_deform_T.shape[0]//self.T, -1])
151
+ self.scales_deform_T = scales_deform_T.reshape([self.T, means3D_deform_T.shape[0]//self.T, -1])
152
+ self.rotations_deform_T = rotations_deform_T.reshape([self.T, means3D_deform_T.shape[0]//self.T, -1])
153
+
154
+
155
+ def render(
156
+ self,
157
+ viewpoint_camera,
158
+ scaling_modifier=1.0,
159
+ bg_color=None,
160
+ override_color=None,
161
+ compute_cov3D_python=False,
162
+ convert_SHs_python=False,
163
+ ):
164
+ # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
165
+ screenspace_points = (
166
+ torch.zeros_like(
167
+ self.gaussians.get_xyz,
168
+ dtype=self.gaussians.get_xyz.dtype,
169
+ requires_grad=True,
170
+ device="cuda",
171
+ )
172
+ + 0
173
+ )
174
+ try:
175
+ screenspace_points.retain_grad()
176
+ except:
177
+ pass
178
+
179
+ # Set up rasterization configuration
180
+ tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
181
+ tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)
182
+
183
+ raster_settings = GaussianRasterizationSettings(
184
+ image_height=int(viewpoint_camera.image_height),
185
+ image_width=int(viewpoint_camera.image_width),
186
+ tanfovx=tanfovx,
187
+ tanfovy=tanfovy,
188
+ bg=self.bg_color if bg_color is None else bg_color,
189
+ scale_modifier=scaling_modifier,
190
+ viewmatrix=viewpoint_camera.world_view_transform,
191
+ projmatrix=viewpoint_camera.full_proj_transform,
192
+ sh_degree=self.gaussians.active_sh_degree,
193
+ campos=viewpoint_camera.camera_center,
194
+ prefiltered=False,
195
+ debug=False,
196
+ )
197
+
198
+ rasterizer = GaussianRasterizer(raster_settings=raster_settings)
199
+
200
+ means3D = self.gaussians.get_xyz
201
+ time = torch.tensor(viewpoint_camera.time).to(means3D.device).repeat(means3D.shape[0],1)
202
+ time = ((time.float() / self.T) - 0.5) * 2
203
+
204
+ means2D = screenspace_points
205
+ opacity = self.gaussians._opacity
206
+
207
+ # If precomputed 3d covariance is provided, use it. If not, then it will be computed from
208
+ # scaling / rotation by the rasterizer.
209
+ scales = None
210
+ rotations = None
211
+ cov3D_precomp = None
212
+ if compute_cov3D_python:
213
+ cov3D_precomp = self.gaussians.get_covariance(scaling_modifier)
214
+ else:
215
+ scales = self.gaussians._scaling
216
+ rotations = self.gaussians._rotation
217
+
218
+ 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]
219
+
220
+
221
+ means3D_final = means3D + means3D_deform
222
+ rotations_final = rotations + rotations_deform
223
+ scales_final = scales + scales_deform
224
+ opacity_final = opacity + opacity_deform
225
+
226
+
227
+
228
+ scales_final = self.gaussians.scaling_activation(scales_final)
229
+ rotations_final = self.gaussians.rotation_activation(rotations_final)
230
+ opacity = self.gaussians.opacity_activation(opacity)
231
+
232
+
233
+ # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
234
+ # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
235
+ shs = None
236
+ colors_precomp = None
237
+ if colors_precomp is None:
238
+ if convert_SHs_python:
239
+ shs_view = self.gaussians.get_features.transpose(1, 2).view(
240
+ -1, 3, (self.gaussians.max_sh_degree + 1) ** 2
241
+ )
242
+ dir_pp = self.gaussians.get_xyz - viewpoint_camera.camera_center.repeat(
243
+ self.gaussians.get_features.shape[0], 1
244
+ )
245
+ dir_pp_normalized = dir_pp / dir_pp.norm(dim=1, keepdim=True)
246
+ sh2rgb = eval_sh(
247
+ self.gaussians.active_sh_degree, shs_view, dir_pp_normalized
248
+ )
249
+ colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)
250
+ else:
251
+ shs = self.gaussians.get_features
252
+ else:
253
+ colors_precomp = override_color
254
+
255
+ rendered_image, radii, rendered_depth, rendered_alpha = rasterizer(
256
+ means3D = means3D_final,
257
+ means2D = means2D,
258
+ shs = shs,
259
+ colors_precomp = colors_precomp,
260
+ opacities = opacity,
261
+ scales = scales_final,
262
+ rotations = rotations_final,
263
+ cov3D_precomp = cov3D_precomp)
264
+
265
+
266
+ rendered_image = rendered_image.clamp(0, 1)
267
+
268
+ # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
269
+ # They will be excluded from value updates used in the splitting criteria.
270
+ return {
271
+ "image": rendered_image,
272
+ "depth": rendered_depth,
273
+ "alpha": rendered_alpha,
274
+ "viewspace_points": screenspace_points,
275
+ "visibility_filter": radii > 0,
276
+ "radii": radii,
277
+ }
guidance/imagedream_utils.py ADDED
@@ -0,0 +1,326 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import torch
3
+ import torch.nn as nn
4
+ import torch.nn.functional as F
5
+ import torchvision.transforms.functional as TF
6
+
7
+ from imagedream.camera_utils import get_camera, convert_opengl_to_blender, normalize_camera
8
+ from imagedream.model_zoo import build_model
9
+ from imagedream.ldm.models.diffusion.ddim import DDIMSampler
10
+
11
+ from diffusers import DDIMScheduler
12
+
13
+ class ImageDream(nn.Module):
14
+ def __init__(
15
+ self,
16
+ device,
17
+ model_name='sd-v2.1-base-4view-ipmv',
18
+ ckpt_path=None,
19
+ t_range=[0.02, 0.98],
20
+ ):
21
+ super().__init__()
22
+
23
+ self.device = device
24
+ self.model_name = model_name
25
+ self.ckpt_path = ckpt_path
26
+
27
+ self.model = build_model(self.model_name, ckpt_path=self.ckpt_path).eval().to(self.device)
28
+ self.model.device = device
29
+ for p in self.model.parameters():
30
+ p.requires_grad_(False)
31
+
32
+ self.dtype = torch.float32
33
+
34
+ self.num_train_timesteps = 1000
35
+ self.min_step = int(self.num_train_timesteps * t_range[0])
36
+ self.max_step = int(self.num_train_timesteps * t_range[1])
37
+
38
+ self.image_embeddings = {}
39
+ self.embeddings = {}
40
+
41
+ self.scheduler = DDIMScheduler.from_pretrained(
42
+ "stabilityai/stable-diffusion-2-1-base", subfolder="scheduler", torch_dtype=self.dtype
43
+ )
44
+
45
+ @torch.no_grad()
46
+ def get_image_text_embeds(self, image, prompts, negative_prompts):
47
+
48
+ image = F.interpolate(image, (256, 256), mode='bilinear', align_corners=False)
49
+ image_pil = TF.to_pil_image(image[0])
50
+ image_embeddings = self.model.get_learned_image_conditioning(image_pil).repeat(5,1,1) # [5, 257, 1280]
51
+ self.image_embeddings['pos'] = image_embeddings
52
+ self.image_embeddings['neg'] = torch.zeros_like(image_embeddings)
53
+
54
+ self.image_embeddings['ip_img'] = self.encode_imgs(image)
55
+ self.image_embeddings['neg_ip_img'] = torch.zeros_like(self.image_embeddings['ip_img'])
56
+
57
+ pos_embeds = self.encode_text(prompts).repeat(5,1,1)
58
+ neg_embeds = self.encode_text(negative_prompts).repeat(5,1,1)
59
+ self.embeddings['pos'] = pos_embeds
60
+ self.embeddings['neg'] = neg_embeds
61
+
62
+ 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']
63
+
64
+ def encode_text(self, prompt):
65
+ # prompt: [str]
66
+ embeddings = self.model.get_learned_conditioning(prompt).to(self.device)
67
+ return embeddings
68
+
69
+ @torch.no_grad()
70
+ def refine(self, pred_rgb, camera,
71
+ guidance_scale=5, steps=50, strength=0.8,
72
+ ):
73
+
74
+ batch_size = pred_rgb.shape[0]
75
+ real_batch_size = batch_size // 4
76
+ pred_rgb_256 = F.interpolate(pred_rgb, (256, 256), mode='bilinear', align_corners=False)
77
+ latents = self.encode_imgs(pred_rgb_256.to(self.dtype))
78
+
79
+ self.scheduler.set_timesteps(steps)
80
+ init_step = int(steps * strength)
81
+ latents = self.scheduler.add_noise(latents, torch.randn_like(latents), self.scheduler.timesteps[init_step])
82
+
83
+ camera = camera[:, [0, 2, 1, 3]] # to blender convention (flip y & z axis)
84
+ camera[:, 1] *= -1
85
+ camera = normalize_camera(camera).view(batch_size, 16)
86
+
87
+ # extra view
88
+ camera = camera.view(real_batch_size, 4, 16)
89
+ camera = torch.cat([camera, torch.zeros_like(camera[:, :1])], dim=1) # [rB, 5, 16]
90
+ camera = camera.view(real_batch_size * 5, 16)
91
+
92
+ camera = camera.repeat(2, 1)
93
+ embeddings = torch.cat([self.embeddings['neg'].repeat(real_batch_size, 1, 1), self.embeddings['pos'].repeat(real_batch_size, 1, 1)], dim=0)
94
+ 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)
95
+ 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)
96
+
97
+ context = {
98
+ "context": embeddings,
99
+ "ip": image_embeddings,
100
+ "ip_img": ip_img_embeddings,
101
+ "camera": camera,
102
+ "num_frames": 4 + 1
103
+ }
104
+
105
+ for i, t in enumerate(self.scheduler.timesteps[init_step:]):
106
+
107
+ # extra view
108
+
109
+ latents = latents.view(real_batch_size, 4, 4, 32, 32)
110
+ latents = torch.cat([latents, torch.zeros_like(latents[:, :1])], dim=1).view(-1, 4, 32, 32)
111
+ latent_model_input = torch.cat([latents] * 2)
112
+
113
+ tt = torch.cat([t.unsqueeze(0).repeat(real_batch_size * 5)] * 2).to(self.device)
114
+
115
+ noise_pred = self.model.apply_model(latent_model_input, tt, context)
116
+
117
+ noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
118
+
119
+ # remove extra view
120
+ noise_pred_uncond = noise_pred_uncond.reshape(real_batch_size, 5, 4, 32, 32)[:, :-1].reshape(-1, 4, 32, 32)
121
+ noise_pred_cond = noise_pred_cond.reshape(real_batch_size, 5, 4, 32, 32)[:, :-1].reshape(-1, 4, 32, 32)
122
+ latents = latents.reshape(real_batch_size, 5, 4, 32, 32)[:, :-1].reshape(-1, 4, 32, 32)
123
+
124
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
125
+
126
+ latents = self.scheduler.step(noise_pred, t, latents).prev_sample
127
+
128
+ imgs = self.decode_latents(latents) # [1, 3, 512, 512]
129
+ return imgs
130
+
131
+ def train_step(
132
+ self,
133
+ pred_rgb, # [B, C, H, W]
134
+ camera, # [B, 4, 4]
135
+ step_ratio=None,
136
+ guidance_scale=5,
137
+ as_latent=False,
138
+ ):
139
+
140
+ batch_size = pred_rgb.shape[0]
141
+ real_batch_size = batch_size // 4
142
+ pred_rgb = pred_rgb.to(self.dtype)
143
+
144
+ if as_latent:
145
+ latents = F.interpolate(pred_rgb, (32, 32), mode="bilinear", align_corners=False) * 2 - 1
146
+ else:
147
+ # interp to 256x256 to be fed into vae.
148
+ pred_rgb_256 = F.interpolate(pred_rgb, (256, 256), mode="bilinear", align_corners=False)
149
+ # encode image into latents with vae, requires grad!
150
+ latents = self.encode_imgs(pred_rgb_256)
151
+
152
+ if step_ratio is not None:
153
+ # dreamtime-like
154
+ # t = self.max_step - (self.max_step - self.min_step) * np.sqrt(step_ratio)
155
+ t = np.round((1 - step_ratio) * self.num_train_timesteps).clip(self.min_step, self.max_step)
156
+ t = torch.full((batch_size,), t, dtype=torch.long, device=self.device)
157
+ else:
158
+ t = torch.randint(self.min_step, self.max_step + 1, (real_batch_size,), dtype=torch.long, device=self.device).repeat(4)
159
+
160
+ camera = camera[:, [0, 2, 1, 3]] # to blender convention (flip y & z axis)
161
+ camera[:, 1] *= -1
162
+ camera = normalize_camera(camera).view(batch_size, 16)
163
+
164
+ # extra view
165
+ camera = camera.view(real_batch_size, 4, 16)
166
+ camera = torch.cat([camera, torch.zeros_like(camera[:, :1])], dim=1) # [rB, 5, 16]
167
+ camera = camera.view(real_batch_size * 5, 16)
168
+
169
+ camera = camera.repeat(2, 1)
170
+ embeddings = torch.cat([self.embeddings['neg'].repeat(real_batch_size, 1, 1), self.embeddings['pos'].repeat(real_batch_size, 1, 1)], dim=0)
171
+ 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)
172
+ 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)
173
+
174
+ context = {
175
+ "context": embeddings,
176
+ "ip": image_embeddings,
177
+ "ip_img": ip_img_embeddings,
178
+ "camera": camera,
179
+ "num_frames": 4 + 1
180
+ }
181
+
182
+ # predict the noise residual with unet, NO grad!
183
+ with torch.no_grad():
184
+ # add noise
185
+ noise = torch.randn_like(latents)
186
+ latents_noisy = self.model.q_sample(latents, t, noise) # [B=4, 4, 32, 32]
187
+ # extra view
188
+ t = t.view(real_batch_size, 4)
189
+ t = torch.cat([t, t[:, :1]], dim=1).view(-1)
190
+ latents_noisy = latents_noisy.view(real_batch_size, 4, 4, 32, 32)
191
+ latents_noisy = torch.cat([latents_noisy, torch.zeros_like(latents_noisy[:, :1])], dim=1).view(-1, 4, 32, 32)
192
+ # pred noise
193
+ latent_model_input = torch.cat([latents_noisy] * 2)
194
+ tt = torch.cat([t] * 2)
195
+
196
+ # import kiui
197
+ # kiui.lo(latent_model_input, t, context['context'], context['camera'])
198
+
199
+ noise_pred = self.model.apply_model(latent_model_input, tt, context)
200
+
201
+ # perform guidance (high scale from paper!)
202
+ noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
203
+
204
+ # remove extra view
205
+ noise_pred_uncond = noise_pred_uncond.reshape(real_batch_size, 5, 4, 32, 32)[:, :-1].reshape(-1, 4, 32, 32)
206
+ noise_pred_cond = noise_pred_cond.reshape(real_batch_size, 5, 4, 32, 32)[:, :-1].reshape(-1, 4, 32, 32)
207
+
208
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
209
+
210
+ grad = (noise_pred - noise)
211
+ grad = torch.nan_to_num(grad)
212
+
213
+ target = (latents - grad).detach()
214
+ loss = 0.5 * F.mse_loss(latents.float(), target, reduction='sum') / latents.shape[0]
215
+
216
+ return loss
217
+
218
+ def decode_latents(self, latents):
219
+ imgs = self.model.decode_first_stage(latents)
220
+ imgs = ((imgs + 1) / 2).clamp(0, 1)
221
+ return imgs
222
+
223
+ def encode_imgs(self, imgs):
224
+ # imgs: [B, 3, 256, 256]
225
+ imgs = 2 * imgs - 1
226
+ latents = self.model.get_first_stage_encoding(self.model.encode_first_stage(imgs))
227
+ return latents # [B, 4, 32, 32]
228
+
229
+ @torch.no_grad()
230
+ def prompt_to_img(
231
+ self,
232
+ image,
233
+ prompts,
234
+ negative_prompts="",
235
+ height=256,
236
+ width=256,
237
+ num_inference_steps=50,
238
+ guidance_scale=5.0,
239
+ latents=None,
240
+ elevation=0,
241
+ azimuth_start=0,
242
+ ):
243
+ if isinstance(prompts, str):
244
+ prompts = [prompts]
245
+
246
+ if isinstance(negative_prompts, str):
247
+ negative_prompts = [negative_prompts]
248
+
249
+ real_batch_size = len(prompts)
250
+ batch_size = len(prompts) * 5
251
+
252
+ # Text embeds -> img latents
253
+ sampler = DDIMSampler(self.model)
254
+ shape = [4, height // 8, width // 8]
255
+
256
+ c_ = {"context": self.encode_text(prompts).repeat(5,1,1)}
257
+ uc_ = {"context": self.encode_text(negative_prompts).repeat(5,1,1)}
258
+
259
+ # image embeddings
260
+ image = F.interpolate(image, (256, 256), mode='bilinear', align_corners=False)
261
+ image_pil = TF.to_pil_image(image[0])
262
+ image_embeddings = self.model.get_learned_image_conditioning(image_pil).repeat(5,1,1).to(self.device)
263
+ c_["ip"] = image_embeddings
264
+ uc_["ip"] = torch.zeros_like(image_embeddings)
265
+
266
+ ip_img = self.encode_imgs(image)
267
+ c_["ip_img"] = ip_img
268
+ uc_["ip_img"] = torch.zeros_like(ip_img)
269
+
270
+ camera = get_camera(4, elevation=elevation, azimuth_start=azimuth_start, extra_view=True)
271
+ camera = camera.repeat(real_batch_size, 1).to(self.device)
272
+
273
+ c_["camera"] = uc_["camera"] = camera
274
+ c_["num_frames"] = uc_["num_frames"] = 5
275
+
276
+ kiui.lo(image_embeddings, ip_img, camera)
277
+
278
+ latents, _ = sampler.sample(S=num_inference_steps, conditioning=c_,
279
+ batch_size=batch_size, shape=shape,
280
+ verbose=False,
281
+ unconditional_guidance_scale=guidance_scale,
282
+ unconditional_conditioning=uc_,
283
+ eta=0, x_T=None)
284
+
285
+ # Img latents -> imgs
286
+ imgs = self.decode_latents(latents) # [4, 3, 256, 256]
287
+
288
+ kiui.lo(latents, imgs)
289
+
290
+ # Img to Numpy
291
+ imgs = imgs.detach().cpu().permute(0, 2, 3, 1).numpy()
292
+ imgs = (imgs * 255).round().astype("uint8")
293
+
294
+ return imgs
295
+
296
+
297
+ if __name__ == "__main__":
298
+ import argparse
299
+ import matplotlib.pyplot as plt
300
+ import kiui
301
+
302
+ parser = argparse.ArgumentParser()
303
+ parser.add_argument("image", type=str)
304
+ parser.add_argument("prompt", type=str)
305
+ parser.add_argument("--negative", default="", type=str)
306
+ parser.add_argument("--steps", type=int, default=30)
307
+ opt = parser.parse_args()
308
+
309
+ device = torch.device("cuda")
310
+
311
+ sd = ImageDream(device)
312
+
313
+ image = kiui.read_image(opt.image, mode='tensor')
314
+ image = image.permute(2, 0, 1).unsqueeze(0).to(device)
315
+
316
+ while True:
317
+ imgs = sd.prompt_to_img(image, opt.prompt, opt.negative, num_inference_steps=opt.steps)
318
+
319
+ grid = np.concatenate([
320
+ np.concatenate([imgs[0], imgs[1]], axis=1),
321
+ np.concatenate([imgs[2], imgs[3]], axis=1),
322
+ ], axis=0)
323
+
324
+ # visualize image
325
+ plt.imshow(grid)
326
+ plt.show()
guidance/mvdream_utils.py ADDED
@@ -0,0 +1,271 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import torch
3
+ import torch.nn as nn
4
+ import torch.nn.functional as F
5
+
6
+ from mvdream.camera_utils import get_camera, convert_opengl_to_blender, normalize_camera
7
+ from mvdream.model_zoo import build_model
8
+ from mvdream.ldm.models.diffusion.ddim import DDIMSampler
9
+
10
+ from diffusers import DDIMScheduler
11
+
12
+ class MVDream(nn.Module):
13
+ def __init__(
14
+ self,
15
+ device,
16
+ model_name='sd-v2.1-base-4view',
17
+ ckpt_path=None,
18
+ t_range=[0.02, 0.98],
19
+ ):
20
+ super().__init__()
21
+
22
+ self.device = device
23
+ self.model_name = model_name
24
+ self.ckpt_path = ckpt_path
25
+
26
+ self.model = build_model(self.model_name, ckpt_path=self.ckpt_path).eval().to(self.device)
27
+ self.model.device = device
28
+ for p in self.model.parameters():
29
+ p.requires_grad_(False)
30
+
31
+ self.dtype = torch.float32
32
+
33
+ self.num_train_timesteps = 1000
34
+ self.min_step = int(self.num_train_timesteps * t_range[0])
35
+ self.max_step = int(self.num_train_timesteps * t_range[1])
36
+
37
+ self.embeddings = None
38
+
39
+ self.scheduler = DDIMScheduler.from_pretrained(
40
+ "stabilityai/stable-diffusion-2-1-base", subfolder="scheduler", torch_dtype=self.dtype
41
+ )
42
+
43
+ @torch.no_grad()
44
+ def get_text_embeds(self, prompts, negative_prompts):
45
+ pos_embeds = self.encode_text(prompts).repeat(4,1,1) # [1, 77, 768]
46
+ neg_embeds = self.encode_text(negative_prompts).repeat(4,1,1)
47
+ self.embeddings = torch.cat([neg_embeds, pos_embeds], dim=0) # [2, 77, 768]
48
+
49
+ def encode_text(self, prompt):
50
+ # prompt: [str]
51
+ embeddings = self.model.get_learned_conditioning(prompt).to(self.device)
52
+ return embeddings
53
+
54
+ @torch.no_grad()
55
+ def refine(self, pred_rgb, camera,
56
+ guidance_scale=100, steps=50, strength=0.8,
57
+ ):
58
+
59
+ batch_size = pred_rgb.shape[0]
60
+ pred_rgb_256 = F.interpolate(pred_rgb, (256, 256), mode='bilinear', align_corners=False)
61
+ latents = self.encode_imgs(pred_rgb_256.to(self.dtype))
62
+ # latents = torch.randn((1, 4, 64, 64), device=self.device, dtype=self.dtype)
63
+
64
+ self.scheduler.set_timesteps(steps)
65
+ init_step = int(steps * strength)
66
+ latents = self.scheduler.add_noise(latents, torch.randn_like(latents), self.scheduler.timesteps[init_step])
67
+
68
+ camera = camera[:, [0, 2, 1, 3]] # to blender convention (flip y & z axis)
69
+ camera[:, 1] *= -1
70
+ camera = normalize_camera(camera).view(batch_size, 16)
71
+ camera = camera.repeat(2, 1)
72
+ context = {"context": self.embeddings, "camera": camera, "num_frames": 4}
73
+
74
+ for i, t in enumerate(self.scheduler.timesteps[init_step:]):
75
+
76
+ latent_model_input = torch.cat([latents] * 2)
77
+
78
+ tt = torch.cat([t.unsqueeze(0).repeat(batch_size)] * 2).to(self.device)
79
+
80
+ noise_pred = self.model.apply_model(latent_model_input, tt, context)
81
+
82
+ noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
83
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
84
+
85
+ latents = self.scheduler.step(noise_pred, t, latents).prev_sample
86
+
87
+ imgs = self.decode_latents(latents) # [1, 3, 512, 512]
88
+ return imgs
89
+
90
+ def train_step(
91
+ self,
92
+ pred_rgb, # [B, C, H, W], B is multiples of 4
93
+ camera, # [B, 4, 4]
94
+ step_ratio=None,
95
+ guidance_scale=50,
96
+ as_latent=False,
97
+ ):
98
+
99
+ batch_size = pred_rgb.shape[0]
100
+ pred_rgb = pred_rgb.to(self.dtype)
101
+
102
+ if as_latent:
103
+ latents = F.interpolate(pred_rgb, (32, 32), mode="bilinear", align_corners=False) * 2 - 1
104
+ else:
105
+ # interp to 256x256 to be fed into vae.
106
+ pred_rgb_256 = F.interpolate(pred_rgb, (256, 256), mode="bilinear", align_corners=False)
107
+ # encode image into latents with vae, requires grad!
108
+ latents = self.encode_imgs(pred_rgb_256)
109
+
110
+ if step_ratio is not None:
111
+ # dreamtime-like
112
+ # t = self.max_step - (self.max_step - self.min_step) * np.sqrt(step_ratio)
113
+ t = np.round((1 - step_ratio) * self.num_train_timesteps).clip(self.min_step, self.max_step)
114
+ t = torch.full((batch_size,), t, dtype=torch.long, device=self.device)
115
+ else:
116
+ t = torch.randint(self.min_step, self.max_step + 1, (batch_size,), dtype=torch.long, device=self.device)
117
+
118
+ # camera = convert_opengl_to_blender(camera)
119
+ # flip_yz = torch.tensor([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]).unsqueeze(0)
120
+ # camera = torch.matmul(flip_yz.to(camera), camera)
121
+ camera = camera[:, [0, 2, 1, 3]] # to blender convention (flip y & z axis)
122
+ camera[:, 1] *= -1
123
+ camera = normalize_camera(camera).view(batch_size, 16)
124
+
125
+ ###############
126
+ # sampler = DDIMSampler(self.model)
127
+ # shape = [4, 32, 32]
128
+ # c_ = {"context": self.embeddings[4:]}
129
+ # uc_ = {"context": self.embeddings[:4]}
130
+
131
+ # # print(camera)
132
+
133
+ # # camera = get_camera(4, elevation=0, azimuth_start=0)
134
+ # # camera = camera.repeat(batch_size // 4, 1).to(self.device)
135
+
136
+ # # print(camera)
137
+
138
+ # c_["camera"] = uc_["camera"] = camera
139
+ # c_["num_frames"] = uc_["num_frames"] = 4
140
+
141
+ # latents_, _ = sampler.sample(S=30, conditioning=c_,
142
+ # batch_size=batch_size, shape=shape,
143
+ # verbose=False,
144
+ # unconditional_guidance_scale=guidance_scale,
145
+ # unconditional_conditioning=uc_,
146
+ # eta=0, x_T=None)
147
+
148
+ # # Img latents -> imgs
149
+ # imgs = self.decode_latents(latents_) # [4, 3, 256, 256]
150
+ # import kiui
151
+ # kiui.vis.plot_image(imgs)
152
+ ###############
153
+
154
+ camera = camera.repeat(2, 1)
155
+ context = {"context": self.embeddings, "camera": camera, "num_frames": 4}
156
+
157
+ # predict the noise residual with unet, NO grad!
158
+ with torch.no_grad():
159
+ # add noise
160
+ noise = torch.randn_like(latents)
161
+ latents_noisy = self.model.q_sample(latents, t, noise)
162
+ # pred noise
163
+ latent_model_input = torch.cat([latents_noisy] * 2)
164
+ tt = torch.cat([t] * 2)
165
+
166
+ # import kiui
167
+ # kiui.lo(latent_model_input, t, context['context'], context['camera'])
168
+
169
+ noise_pred = self.model.apply_model(latent_model_input, tt, context)
170
+
171
+ # perform guidance (high scale from paper!)
172
+ noise_pred_uncond, noise_pred_pos = noise_pred.chunk(2)
173
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_pos - noise_pred_uncond)
174
+
175
+ grad = (noise_pred - noise)
176
+ grad = torch.nan_to_num(grad)
177
+
178
+ # seems important to avoid NaN...
179
+ # grad = grad.clamp(-1, 1)
180
+
181
+ target = (latents - grad).detach()
182
+ loss = 0.5 * F.mse_loss(latents.float(), target, reduction='sum') / latents.shape[0]
183
+
184
+ return loss
185
+
186
+ def decode_latents(self, latents):
187
+ imgs = self.model.decode_first_stage(latents)
188
+ imgs = ((imgs + 1) / 2).clamp(0, 1)
189
+ return imgs
190
+
191
+ def encode_imgs(self, imgs):
192
+ # imgs: [B, 3, 256, 256]
193
+ imgs = 2 * imgs - 1
194
+ latents = self.model.get_first_stage_encoding(self.model.encode_first_stage(imgs))
195
+ return latents # [B, 4, 32, 32]
196
+
197
+ @torch.no_grad()
198
+ def prompt_to_img(
199
+ self,
200
+ prompts,
201
+ negative_prompts="",
202
+ height=256,
203
+ width=256,
204
+ num_inference_steps=50,
205
+ guidance_scale=7.5,
206
+ latents=None,
207
+ elevation=0,
208
+ azimuth_start=0,
209
+ ):
210
+ if isinstance(prompts, str):
211
+ prompts = [prompts]
212
+
213
+ if isinstance(negative_prompts, str):
214
+ negative_prompts = [negative_prompts]
215
+
216
+ batch_size = len(prompts) * 4
217
+
218
+ # Text embeds -> img latents
219
+ sampler = DDIMSampler(self.model)
220
+ shape = [4, height // 8, width // 8]
221
+ c_ = {"context": self.encode_text(prompts).repeat(4,1,1)}
222
+ uc_ = {"context": self.encode_text(negative_prompts).repeat(4,1,1)}
223
+
224
+ camera = get_camera(4, elevation=elevation, azimuth_start=azimuth_start)
225
+ camera = camera.repeat(batch_size // 4, 1).to(self.device)
226
+
227
+ c_["camera"] = uc_["camera"] = camera
228
+ c_["num_frames"] = uc_["num_frames"] = 4
229
+
230
+ latents, _ = sampler.sample(S=num_inference_steps, conditioning=c_,
231
+ batch_size=batch_size, shape=shape,
232
+ verbose=False,
233
+ unconditional_guidance_scale=guidance_scale,
234
+ unconditional_conditioning=uc_,
235
+ eta=0, x_T=None)
236
+
237
+ # Img latents -> imgs
238
+ imgs = self.decode_latents(latents) # [4, 3, 256, 256]
239
+
240
+ # Img to Numpy
241
+ imgs = imgs.detach().cpu().permute(0, 2, 3, 1).numpy()
242
+ imgs = (imgs * 255).round().astype("uint8")
243
+
244
+ return imgs
245
+
246
+
247
+ if __name__ == "__main__":
248
+ import argparse
249
+ import matplotlib.pyplot as plt
250
+
251
+ parser = argparse.ArgumentParser()
252
+ parser.add_argument("prompt", type=str)
253
+ parser.add_argument("--negative", default="", type=str)
254
+ parser.add_argument("--steps", type=int, default=30)
255
+ opt = parser.parse_args()
256
+
257
+ device = torch.device("cuda")
258
+
259
+ sd = MVDream(device)
260
+
261
+ while True:
262
+ imgs = sd.prompt_to_img(opt.prompt, opt.negative, num_inference_steps=opt.steps)
263
+
264
+ grid = np.concatenate([
265
+ np.concatenate([imgs[0], imgs[1]], axis=1),
266
+ np.concatenate([imgs[2], imgs[3]], axis=1),
267
+ ], axis=0)
268
+
269
+ # visualize image
270
+ plt.imshow(grid)
271
+ plt.show()
guidance/sd_utils.py ADDED
@@ -0,0 +1,334 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import CLIPTextModel, CLIPTokenizer, logging
2
+ from diffusers import (
3
+ AutoencoderKL,
4
+ UNet2DConditionModel,
5
+ PNDMScheduler,
6
+ DDIMScheduler,
7
+ StableDiffusionPipeline,
8
+ )
9
+ from diffusers.utils.import_utils import is_xformers_available
10
+
11
+ # suppress partial model loading warning
12
+ logging.set_verbosity_error()
13
+
14
+ import numpy as np
15
+ import torch
16
+ import torch.nn as nn
17
+ import torch.nn.functional as F
18
+
19
+
20
+ def seed_everything(seed):
21
+ torch.manual_seed(seed)
22
+ torch.cuda.manual_seed(seed)
23
+ # torch.backends.cudnn.deterministic = True
24
+ # torch.backends.cudnn.benchmark = True
25
+
26
+
27
+ class StableDiffusion(nn.Module):
28
+ def __init__(
29
+ self,
30
+ device,
31
+ fp16=True,
32
+ vram_O=False,
33
+ sd_version="2.1",
34
+ hf_key=None,
35
+ t_range=[0.02, 0.98],
36
+ ):
37
+ super().__init__()
38
+
39
+ self.device = device
40
+ self.sd_version = sd_version
41
+
42
+ if hf_key is not None:
43
+ print(f"[INFO] using hugging face custom model key: {hf_key}")
44
+ model_key = hf_key
45
+ elif self.sd_version == "2.1":
46
+ model_key = "stabilityai/stable-diffusion-2-1-base"
47
+ elif self.sd_version == "2.0":
48
+ model_key = "stabilityai/stable-diffusion-2-base"
49
+ elif self.sd_version == "1.5":
50
+ model_key = "runwayml/stable-diffusion-v1-5"
51
+ else:
52
+ raise ValueError(
53
+ f"Stable-diffusion version {self.sd_version} not supported."
54
+ )
55
+
56
+ self.dtype = torch.float16 if fp16 else torch.float32
57
+
58
+ # Create model
59
+ pipe = StableDiffusionPipeline.from_pretrained(
60
+ model_key, torch_dtype=self.dtype
61
+ )
62
+
63
+ if vram_O:
64
+ pipe.enable_sequential_cpu_offload()
65
+ pipe.enable_vae_slicing()
66
+ pipe.unet.to(memory_format=torch.channels_last)
67
+ pipe.enable_attention_slicing(1)
68
+ # pipe.enable_model_cpu_offload()
69
+ else:
70
+ pipe.to(device)
71
+
72
+ self.vae = pipe.vae
73
+ self.tokenizer = pipe.tokenizer
74
+ self.text_encoder = pipe.text_encoder
75
+ self.unet = pipe.unet
76
+
77
+ self.scheduler = DDIMScheduler.from_pretrained(
78
+ model_key, subfolder="scheduler", torch_dtype=self.dtype
79
+ )
80
+
81
+ del pipe
82
+
83
+ self.num_train_timesteps = self.scheduler.config.num_train_timesteps
84
+ self.min_step = int(self.num_train_timesteps * t_range[0])
85
+ self.max_step = int(self.num_train_timesteps * t_range[1])
86
+ self.alphas = self.scheduler.alphas_cumprod.to(self.device) # for convenience
87
+
88
+ self.embeddings = None
89
+
90
+ @torch.no_grad()
91
+ def get_text_embeds(self, prompts, negative_prompts):
92
+ pos_embeds = self.encode_text(prompts) # [1, 77, 768]
93
+ neg_embeds = self.encode_text(negative_prompts)
94
+ self.embeddings = torch.cat([neg_embeds, pos_embeds], dim=0) # [2, 77, 768]
95
+
96
+ def encode_text(self, prompt):
97
+ # prompt: [str]
98
+ inputs = self.tokenizer(
99
+ prompt,
100
+ padding="max_length",
101
+ max_length=self.tokenizer.model_max_length,
102
+ return_tensors="pt",
103
+ )
104
+ embeddings = self.text_encoder(inputs.input_ids.to(self.device))[0]
105
+ return embeddings
106
+
107
+ @torch.no_grad()
108
+ def refine(self, pred_rgb,
109
+ guidance_scale=100, steps=50, strength=0.8,
110
+ ):
111
+
112
+ batch_size = pred_rgb.shape[0]
113
+ pred_rgb_512 = F.interpolate(pred_rgb, (512, 512), mode='bilinear', align_corners=False)
114
+ latents = self.encode_imgs(pred_rgb_512.to(self.dtype))
115
+ # latents = torch.randn((1, 4, 64, 64), device=self.device, dtype=self.dtype)
116
+
117
+ self.scheduler.set_timesteps(steps)
118
+ init_step = int(steps * strength)
119
+ latents = self.scheduler.add_noise(latents, torch.randn_like(latents), self.scheduler.timesteps[init_step])
120
+
121
+ for i, t in enumerate(self.scheduler.timesteps[init_step:]):
122
+
123
+ latent_model_input = torch.cat([latents] * 2)
124
+
125
+ noise_pred = self.unet(
126
+ latent_model_input, t, encoder_hidden_states=self.embeddings,
127
+ ).sample
128
+
129
+ noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
130
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
131
+
132
+ latents = self.scheduler.step(noise_pred, t, latents).prev_sample
133
+
134
+ imgs = self.decode_latents(latents) # [1, 3, 512, 512]
135
+ return imgs
136
+
137
+ def train_step(
138
+ self,
139
+ pred_rgb,
140
+ step_ratio=None,
141
+ guidance_scale=100,
142
+ as_latent=False,
143
+ ):
144
+
145
+ batch_size = pred_rgb.shape[0]
146
+ pred_rgb = pred_rgb.to(self.dtype)
147
+
148
+ if as_latent:
149
+ latents = F.interpolate(pred_rgb, (64, 64), mode="bilinear", align_corners=False) * 2 - 1
150
+ else:
151
+ # interp to 512x512 to be fed into vae.
152
+ pred_rgb_512 = F.interpolate(pred_rgb, (512, 512), mode="bilinear", align_corners=False)
153
+ # encode image into latents with vae, requires grad!
154
+ latents = self.encode_imgs(pred_rgb_512)
155
+
156
+ if step_ratio is not None:
157
+ # dreamtime-like
158
+ # t = self.max_step - (self.max_step - self.min_step) * np.sqrt(step_ratio)
159
+ t = np.round((1 - step_ratio) * self.num_train_timesteps).clip(self.min_step, self.max_step)
160
+ t = torch.full((batch_size,), t, dtype=torch.long, device=self.device)
161
+ else:
162
+ t = torch.randint(self.min_step, self.max_step + 1, (batch_size,), dtype=torch.long, device=self.device)
163
+
164
+ # w(t), sigma_t^2
165
+ w = (1 - self.alphas[t]).view(batch_size, 1, 1, 1)
166
+
167
+ # predict the noise residual with unet, NO grad!
168
+ with torch.no_grad():
169
+ # add noise
170
+ noise = torch.randn_like(latents)
171
+ latents_noisy = self.scheduler.add_noise(latents, noise, t)
172
+ # pred noise
173
+ latent_model_input = torch.cat([latents_noisy] * 2)
174
+ tt = torch.cat([t] * 2)
175
+
176
+ noise_pred = self.unet(
177
+ latent_model_input, tt, encoder_hidden_states=self.embeddings.repeat(batch_size, 1, 1)
178
+ ).sample
179
+
180
+ # perform guidance (high scale from paper!)
181
+ noise_pred_uncond, noise_pred_pos = noise_pred.chunk(2)
182
+ noise_pred = noise_pred_uncond + guidance_scale * (
183
+ noise_pred_pos - noise_pred_uncond
184
+ )
185
+
186
+ grad = w * (noise_pred - noise)
187
+ grad = torch.nan_to_num(grad)
188
+
189
+ # seems important to avoid NaN...
190
+ # grad = grad.clamp(-1, 1)
191
+
192
+ target = (latents - grad).detach()
193
+ loss = 0.5 * F.mse_loss(latents.float(), target, reduction='sum') / latents.shape[0]
194
+
195
+ return loss
196
+
197
+ @torch.no_grad()
198
+ def produce_latents(
199
+ self,
200
+ height=512,
201
+ width=512,
202
+ num_inference_steps=50,
203
+ guidance_scale=7.5,
204
+ latents=None,
205
+ ):
206
+ if latents is None:
207
+ latents = torch.randn(
208
+ (
209
+ self.embeddings.shape[0] // 2,
210
+ self.unet.in_channels,
211
+ height // 8,
212
+ width // 8,
213
+ ),
214
+ device=self.device,
215
+ )
216
+
217
+ self.scheduler.set_timesteps(num_inference_steps)
218
+
219
+ for i, t in enumerate(self.scheduler.timesteps):
220
+ # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
221
+ latent_model_input = torch.cat([latents] * 2)
222
+ # predict the noise residual
223
+ noise_pred = self.unet(
224
+ latent_model_input, t, encoder_hidden_states=self.embeddings
225
+ ).sample
226
+
227
+ # perform guidance
228
+ noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
229
+ noise_pred = noise_pred_uncond + guidance_scale * (
230
+ noise_pred_cond - noise_pred_uncond
231
+ )
232
+
233
+ # compute the previous noisy sample x_t -> x_t-1
234
+ latents = self.scheduler.step(noise_pred, t, latents).prev_sample
235
+
236
+ return latents
237
+
238
+ def decode_latents(self, latents):
239
+ latents = 1 / self.vae.config.scaling_factor * latents
240
+
241
+ imgs = self.vae.decode(latents).sample
242
+ imgs = (imgs / 2 + 0.5).clamp(0, 1)
243
+
244
+ return imgs
245
+
246
+ def encode_imgs(self, imgs):
247
+ # imgs: [B, 3, H, W]
248
+
249
+ imgs = 2 * imgs - 1
250
+
251
+ posterior = self.vae.encode(imgs).latent_dist
252
+ latents = posterior.sample() * self.vae.config.scaling_factor
253
+
254
+ return latents
255
+
256
+ def prompt_to_img(
257
+ self,
258
+ prompts,
259
+ negative_prompts="",
260
+ height=512,
261
+ width=512,
262
+ num_inference_steps=50,
263
+ guidance_scale=7.5,
264
+ latents=None,
265
+ ):
266
+ if isinstance(prompts, str):
267
+ prompts = [prompts]
268
+
269
+ if isinstance(negative_prompts, str):
270
+ negative_prompts = [negative_prompts]
271
+
272
+ # Prompts -> text embeds
273
+ self.get_text_embeds(prompts, negative_prompts)
274
+
275
+ # Text embeds -> img latents
276
+ latents = self.produce_latents(
277
+ height=height,
278
+ width=width,
279
+ latents=latents,
280
+ num_inference_steps=num_inference_steps,
281
+ guidance_scale=guidance_scale,
282
+ ) # [1, 4, 64, 64]
283
+
284
+ # Img latents -> imgs
285
+ imgs = self.decode_latents(latents) # [1, 3, 512, 512]
286
+
287
+ # Img to Numpy
288
+ imgs = imgs.detach().cpu().permute(0, 2, 3, 1).numpy()
289
+ imgs = (imgs * 255).round().astype("uint8")
290
+
291
+ return imgs
292
+
293
+
294
+ if __name__ == "__main__":
295
+ import argparse
296
+ import matplotlib.pyplot as plt
297
+
298
+ parser = argparse.ArgumentParser()
299
+ parser.add_argument("prompt", type=str)
300
+ parser.add_argument("--negative", default="", type=str)
301
+ parser.add_argument(
302
+ "--sd_version",
303
+ type=str,
304
+ default="2.1",
305
+ choices=["1.5", "2.0", "2.1"],
306
+ help="stable diffusion version",
307
+ )
308
+ parser.add_argument(
309
+ "--hf_key",
310
+ type=str,
311
+ default=None,
312
+ help="hugging face Stable diffusion model key",
313
+ )
314
+ parser.add_argument("--fp16", action="store_true", help="use float16 for training")
315
+ parser.add_argument(
316
+ "--vram_O", action="store_true", help="optimization for low VRAM usage"
317
+ )
318
+ parser.add_argument("-H", type=int, default=512)
319
+ parser.add_argument("-W", type=int, default=512)
320
+ parser.add_argument("--seed", type=int, default=0)
321
+ parser.add_argument("--steps", type=int, default=50)
322
+ opt = parser.parse_args()
323
+
324
+ seed_everything(opt.seed)
325
+
326
+ device = torch.device("cuda")
327
+
328
+ sd = StableDiffusion(device, opt.fp16, opt.vram_O, opt.sd_version, opt.hf_key)
329
+
330
+ imgs = sd.prompt_to_img(opt.prompt, opt.negative, opt.H, opt.W, opt.steps)
331
+
332
+ # visualize image
333
+ plt.imshow(imgs[0])
334
+ plt.show()
guidance/svd_utils.py ADDED
@@ -0,0 +1,221 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+ from svd import StableVideoDiffusionPipeline
4
+ from diffusers import DDIMScheduler
5
+
6
+ from PIL import Image
7
+ import numpy as np
8
+
9
+ import torch.nn as nn
10
+ import torch.nn.functional as F
11
+
12
+
13
+ class StableVideoDiffusion:
14
+ def __init__(
15
+ self,
16
+ device,
17
+ fp16=True,
18
+ t_range=[0.02, 0.98],
19
+ ):
20
+ super().__init__()
21
+
22
+ self.guidance_type = [
23
+ 'sds',
24
+ 'pixel reconstruction',
25
+ 'latent reconstruction'
26
+ ][1]
27
+
28
+ self.device = device
29
+ self.dtype = torch.float16 if fp16 else torch.float32
30
+
31
+ # Create model
32
+ pipe = StableVideoDiffusionPipeline.from_pretrained(
33
+ "stabilityai/stable-video-diffusion-img2vid", torch_dtype=torch.float16, variant="fp16"
34
+ )
35
+ pipe.to(device)
36
+
37
+ self.pipe = pipe
38
+
39
+ self.num_train_timesteps = self.pipe.scheduler.config.num_train_timesteps if self.guidance_type == 'sds' else 25
40
+ self.pipe.scheduler.set_timesteps(self.num_train_timesteps, device=device) # set sigma for euler discrete scheduling
41
+
42
+ self.min_step = int(self.num_train_timesteps * t_range[0])
43
+ self.max_step = int(self.num_train_timesteps * t_range[1])
44
+ self.alphas = self.pipe.scheduler.alphas_cumprod.to(self.device) # for convenience
45
+
46
+ self.embeddings = None
47
+ self.image = None
48
+ self.target_cache = None
49
+
50
+ @torch.no_grad()
51
+ def get_img_embeds(self, image):
52
+ self.image = Image.fromarray(np.uint8(image*255))
53
+
54
+ def encode_image(self, image):
55
+ image = image * 2 -1
56
+ latents = self.pipe._encode_vae_image(image, self.device, num_videos_per_prompt=1, do_classifier_free_guidance=False)
57
+ latents = self.pipe.vae.config.scaling_factor * latents
58
+ return latents
59
+
60
+ def refine(self,
61
+ pred_rgb,
62
+ steps=25, strength=0.8,
63
+ min_guidance_scale: float = 1.0,
64
+ max_guidance_scale: float = 3.0,
65
+ ):
66
+ # strength = 0.8
67
+ batch_size = pred_rgb.shape[0]
68
+ pred_rgb = pred_rgb.to(self.dtype)
69
+
70
+ # interp to 512x512 to be fed into vae.
71
+ pred_rgb_512 = F.interpolate(pred_rgb, (512, 512), mode="bilinear", align_corners=False)
72
+ # encode image into latents with vae, requires grad!
73
+
74
+ # latents = []
75
+ # for i in range(batch_size):
76
+ # latent = self.encode_image(pred_rgb_512[i:i+1])
77
+ # latents.append(latent)
78
+ # latents = torch.cat(latents, 0)
79
+
80
+ latents = self.encode_image(pred_rgb_512)
81
+ latents = latents.unsqueeze(0)
82
+
83
+ if strength == 0:
84
+ init_step = 0
85
+ latents = torch.randn_like(latents)
86
+ else:
87
+ init_step = int(steps * strength)
88
+ latents = self.pipe.scheduler.add_noise(latents, torch.randn_like(latents), self.pipe.scheduler.timesteps[init_step:init_step+1])
89
+
90
+ target = self.pipe(
91
+ image=self.image,
92
+ height=512,
93
+ width=512,
94
+ latents=latents,
95
+ denoise_beg=init_step,
96
+ denoise_end=steps,
97
+ output_type='frame',
98
+ num_frames=batch_size,
99
+ min_guidance_scale=min_guidance_scale,
100
+ max_guidance_scale=max_guidance_scale,
101
+ num_inference_steps=steps,
102
+ decode_chunk_size=1
103
+ ).frames[0]
104
+ target = (target + 1) * 0.5
105
+ target = target.permute(1,0,2,3)
106
+ return target
107
+
108
+ # frames = self.pipe(
109
+ # image=self.image,
110
+ # height=512,
111
+ # width=512,
112
+ # latents=latents,
113
+ # denoise_beg=init_step,
114
+ # denoise_end=steps,
115
+ # num_frames=batch_size,
116
+ # min_guidance_scale=min_guidance_scale,
117
+ # max_guidance_scale=max_guidance_scale,
118
+ # num_inference_steps=steps,
119
+ # decode_chunk_size=1
120
+ # ).frames[0]
121
+ # export_to_gif(frames, f"tmp.gif")
122
+ # raise
123
+
124
+ def train_step(
125
+ self,
126
+ pred_rgb,
127
+ step_ratio=None,
128
+ min_guidance_scale: float = 1.0,
129
+ max_guidance_scale: float = 3.0,
130
+ ):
131
+
132
+ batch_size = pred_rgb.shape[0]
133
+ pred_rgb = pred_rgb.to(self.dtype)
134
+
135
+ # interp to 512x512 to be fed into vae.
136
+ pred_rgb_512 = F.interpolate(pred_rgb, (512, 512), mode="bilinear", align_corners=False)
137
+ # encode image into latents with vae, requires grad!
138
+ # latents = self.pipe._encode_image(pred_rgb_512, self.device, num_videos_per_prompt=1, do_classifier_free_guidance=True)
139
+ latents = self.encode_image(pred_rgb_512)
140
+ latents = latents.unsqueeze(0)
141
+
142
+ if step_ratio is not None:
143
+ # dreamtime-like
144
+ # t = self.max_step - (self.max_step - self.min_step) * np.sqrt(step_ratio)
145
+ t = np.round((1 - step_ratio) * self.num_train_timesteps).clip(self.min_step, self.max_step)
146
+ t = torch.full((1,), t, dtype=torch.long, device=self.device)
147
+ else:
148
+ t = torch.randint(self.min_step, self.max_step + 1, (1,), dtype=torch.long, device=self.device)
149
+ # print(t)
150
+
151
+ w = (1 - self.alphas[t]).view(1, 1, 1, 1)
152
+
153
+
154
+ if self.guidance_type == 'sds':
155
+ # predict the noise residual with unet, NO grad!
156
+ with torch.no_grad():
157
+ t = self.num_train_timesteps - t.item()
158
+ # add noise
159
+ noise = torch.randn_like(latents)
160
+ latents_noisy = self.pipe.scheduler.add_noise(latents, noise, self.pipe.scheduler.timesteps[t:t+1]) # t=0 noise;t=999 clean
161
+ noise_pred = self.pipe(
162
+ image=self.image,
163
+ # image_embeddings=self.embeddings,
164
+ height=512,
165
+ width=512,
166
+ latents=latents_noisy,
167
+ output_type='noise',
168
+ denoise_beg=t,
169
+ denoise_end=t + 1,
170
+ min_guidance_scale=min_guidance_scale,
171
+ max_guidance_scale=max_guidance_scale,
172
+ num_frames=batch_size,
173
+ num_inference_steps=self.num_train_timesteps
174
+ ).frames[0]
175
+
176
+ grad = w * (noise_pred - noise)
177
+ grad = torch.nan_to_num(grad)
178
+
179
+ target = (latents - grad).detach()
180
+ loss = 0.5 * F.mse_loss(latents.float(), target, reduction='sum') / latents.shape[1]
181
+ print(loss.item())
182
+ return loss
183
+
184
+ elif self.guidance_type == 'pixel reconstruction':
185
+ # pixel space reconstruction
186
+ if self.target_cache is None:
187
+ with torch.no_grad():
188
+ self.target_cache = self.pipe(
189
+ image=self.image,
190
+ height=512,
191
+ width=512,
192
+ output_type='frame',
193
+ num_frames=batch_size,
194
+ num_inference_steps=self.num_train_timesteps,
195
+ decode_chunk_size=1
196
+ ).frames[0]
197
+ self.target_cache = (self.target_cache + 1) * 0.5
198
+ self.target_cache = self.target_cache.permute(1,0,2,3)
199
+
200
+ loss = 0.5 * F.mse_loss(pred_rgb_512.float(), self.target_cache.detach().float(), reduction='sum') / latents.shape[1]
201
+ print(loss.item())
202
+
203
+ return loss
204
+
205
+ elif self.guidance_type == 'latent reconstruction':
206
+ # latent space reconstruction
207
+ if self.target_cache is None:
208
+ with torch.no_grad():
209
+ self.target_cache = self.pipe(
210
+ image=self.image,
211
+ height=512,
212
+ width=512,
213
+ output_type='latent',
214
+ num_frames=batch_size,
215
+ num_inference_steps=self.num_train_timesteps,
216
+ ).frames[0]
217
+
218
+ loss = 0.5 * F.mse_loss(latents.float(), self.target_cache.detach().float(), reduction='sum') / latents.shape[1]
219
+ print(loss.item())
220
+
221
+ return loss
guidance/zero123_utils.py ADDED
@@ -0,0 +1,237 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusers import DDIMScheduler
2
+ import torchvision.transforms.functional as TF
3
+
4
+ import numpy as np
5
+ import torch
6
+ import torch.nn as nn
7
+ import torch.nn.functional as F
8
+
9
+ import sys
10
+ sys.path.append('./')
11
+
12
+ from zero123 import Zero123Pipeline
13
+
14
+
15
+ class Zero123(nn.Module):
16
+ def __init__(self, device, fp16=True, t_range=[0.02, 0.98], model_key="ashawkey/zero123-xl-diffusers"):
17
+ super().__init__()
18
+
19
+ self.device = device
20
+ self.fp16 = fp16
21
+ self.dtype = torch.float16 if fp16 else torch.float32
22
+
23
+ assert self.fp16, 'Only zero123 fp16 is supported for now.'
24
+
25
+ # model_key = "ashawkey/zero123-xl-diffusers"
26
+ # model_key = './model_cache/stable_zero123_diffusers'
27
+
28
+ self.pipe = Zero123Pipeline.from_pretrained(
29
+ model_key,
30
+ torch_dtype=self.dtype,
31
+ trust_remote_code=True,
32
+ ).to(self.device)
33
+
34
+ # stable-zero123 has a different camera embedding
35
+ self.use_stable_zero123 = 'stable' in model_key
36
+
37
+ self.pipe.image_encoder.eval()
38
+ self.pipe.vae.eval()
39
+ self.pipe.unet.eval()
40
+ self.pipe.clip_camera_projection.eval()
41
+
42
+ self.vae = self.pipe.vae
43
+ self.unet = self.pipe.unet
44
+
45
+ self.pipe.set_progress_bar_config(disable=True)
46
+
47
+ self.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config)
48
+ self.num_train_timesteps = self.scheduler.config.num_train_timesteps
49
+
50
+ self.min_step = int(self.num_train_timesteps * t_range[0])
51
+ self.max_step = int(self.num_train_timesteps * t_range[1])
52
+ self.alphas = self.scheduler.alphas_cumprod.to(self.device) # for convenience
53
+
54
+ self.embeddings = None
55
+
56
+ @torch.no_grad()
57
+ def get_img_embeds(self, x):
58
+ # x: image tensor in [0, 1]
59
+ x = F.interpolate(x, (256, 256), mode='bilinear', align_corners=False)
60
+ x_pil = [TF.to_pil_image(image) for image in x]
61
+ x_clip = self.pipe.feature_extractor(images=x_pil, return_tensors="pt").pixel_values.to(device=self.device, dtype=self.dtype)
62
+ c = self.pipe.image_encoder(x_clip).image_embeds
63
+ v = self.encode_imgs(x.to(self.dtype)) / self.vae.config.scaling_factor
64
+ self.embeddings = [c, v]
65
+ return c, v
66
+
67
+ def get_cam_embeddings(self, elevation, azimuth, radius, default_elevation=0):
68
+ if self.use_stable_zero123:
69
+ 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)
70
+ else:
71
+ # original zero123 camera embedding
72
+ T = np.stack([np.deg2rad(elevation), np.sin(np.deg2rad(azimuth)), np.cos(np.deg2rad(azimuth)), radius], axis=-1)
73
+ T = torch.from_numpy(T).unsqueeze(1).to(dtype=self.dtype, device=self.device) # [8, 1, 4]
74
+ return T
75
+
76
+ @torch.no_grad()
77
+ def refine(self, pred_rgb, elevation, azimuth, radius,
78
+ guidance_scale=5, steps=50, strength=0.8, default_elevation=0,
79
+ ):
80
+
81
+ batch_size = pred_rgb.shape[0]
82
+
83
+ self.scheduler.set_timesteps(steps)
84
+
85
+ if strength == 0:
86
+ init_step = 0
87
+ latents = torch.randn((1, 4, 32, 32), device=self.device, dtype=self.dtype)
88
+ else:
89
+ init_step = int(steps * strength)
90
+ pred_rgb_256 = F.interpolate(pred_rgb, (256, 256), mode='bilinear', align_corners=False)
91
+ latents = self.encode_imgs(pred_rgb_256.to(self.dtype))
92
+ latents = self.scheduler.add_noise(latents, torch.randn_like(latents), self.scheduler.timesteps[init_step])
93
+
94
+ T = self.get_cam_embeddings(elevation, azimuth, radius, default_elevation)
95
+ cc_emb = torch.cat([self.embeddings[0].repeat(batch_size, 1, 1), T], dim=-1)
96
+ cc_emb = self.pipe.clip_camera_projection(cc_emb)
97
+ cc_emb = torch.cat([cc_emb, torch.zeros_like(cc_emb)], dim=0)
98
+
99
+ vae_emb = self.embeddings[1].repeat(batch_size, 1, 1, 1)
100
+ vae_emb = torch.cat([vae_emb, torch.zeros_like(vae_emb)], dim=0)
101
+
102
+ for i, t in enumerate(self.scheduler.timesteps[init_step:]):
103
+
104
+ x_in = torch.cat([latents] * 2)
105
+ t_in = t.view(1).to(self.device)
106
+
107
+ noise_pred = self.unet(
108
+ torch.cat([x_in, vae_emb], dim=1),
109
+ t_in.to(self.unet.dtype),
110
+ encoder_hidden_states=cc_emb,
111
+ ).sample
112
+
113
+ noise_pred_cond, noise_pred_uncond = noise_pred.chunk(2)
114
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
115
+
116
+ latents = self.scheduler.step(noise_pred, t, latents).prev_sample
117
+
118
+ imgs = self.decode_latents(latents) # [1, 3, 256, 256]
119
+ return imgs
120
+
121
+ def train_step(self, pred_rgb, elevation, azimuth, radius, step_ratio=None, guidance_scale=5, as_latent=False, default_elevation=0):
122
+ # pred_rgb: tensor [1, 3, H, W] in [0, 1]
123
+
124
+ batch_size = pred_rgb.shape[0]
125
+
126
+ if as_latent:
127
+ latents = F.interpolate(pred_rgb, (32, 32), mode='bilinear', align_corners=False) * 2 - 1
128
+ else:
129
+ pred_rgb_256 = F.interpolate(pred_rgb, (256, 256), mode='bilinear', align_corners=False)
130
+ latents = self.encode_imgs(pred_rgb_256.to(self.dtype))
131
+
132
+ if step_ratio is not None:
133
+ # dreamtime-like
134
+ # t = self.max_step - (self.max_step - self.min_step) * np.sqrt(step_ratio)
135
+ # t = self.max_step - (self.max_step - self.min_step) * (step_ratio ** 2)
136
+ t = np.round((1 - step_ratio) * self.num_train_timesteps).clip(self.min_step, self.max_step)
137
+ t = torch.full((batch_size,), t, dtype=torch.long, device=self.device)
138
+ else:
139
+ t = torch.randint(self.min_step, self.max_step + 1, (batch_size,), dtype=torch.long, device=self.device)
140
+
141
+ w = (1 - self.alphas[t]).view(batch_size, 1, 1, 1)
142
+
143
+ with torch.no_grad():
144
+ noise = torch.randn_like(latents)
145
+ latents_noisy = self.scheduler.add_noise(latents, noise, t)
146
+
147
+ x_in = torch.cat([latents_noisy] * 2)
148
+ t_in = torch.cat([t] * 2)
149
+
150
+ T = self.get_cam_embeddings(elevation, azimuth, radius, default_elevation)
151
+ cc_emb = torch.cat([self.embeddings[0].unsqueeze(1), T], dim=-1)
152
+ cc_emb = self.pipe.clip_camera_projection(cc_emb)
153
+ cc_emb = torch.cat([cc_emb, torch.zeros_like(cc_emb)], dim=0)
154
+
155
+ vae_emb = self.embeddings[1]
156
+ vae_emb = torch.cat([vae_emb, torch.zeros_like(vae_emb)], dim=0)
157
+
158
+ noise_pred = self.unet(
159
+ torch.cat([x_in, vae_emb], dim=1),
160
+ t_in.to(self.unet.dtype),
161
+ encoder_hidden_states=cc_emb,
162
+ ).sample
163
+
164
+ noise_pred_cond, noise_pred_uncond = noise_pred.chunk(2)
165
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
166
+
167
+ grad = w * (noise_pred - noise)
168
+ grad = torch.nan_to_num(grad)
169
+
170
+ target = (latents - grad).detach()
171
+ loss = 0.5 * F.mse_loss(latents.float(), target, reduction='sum')
172
+
173
+ return loss
174
+
175
+ def decode_latents(self, latents):
176
+ latents = 1 / self.vae.config.scaling_factor * latents
177
+
178
+ imgs = self.vae.decode(latents).sample
179
+ imgs = (imgs / 2 + 0.5).clamp(0, 1)
180
+
181
+ return imgs
182
+
183
+ def encode_imgs(self, imgs, mode=False):
184
+ # imgs: [B, 3, H, W]
185
+
186
+ imgs = 2 * imgs - 1
187
+
188
+ posterior = self.vae.encode(imgs).latent_dist
189
+ if mode:
190
+ latents = posterior.mode()
191
+ else:
192
+ latents = posterior.sample()
193
+ latents = latents * self.vae.config.scaling_factor
194
+
195
+ return latents
196
+
197
+
198
+ if __name__ == '__main__':
199
+ import cv2
200
+ import argparse
201
+ import numpy as np
202
+ import matplotlib.pyplot as plt
203
+ import kiui
204
+
205
+ parser = argparse.ArgumentParser()
206
+
207
+ parser.add_argument('input', type=str)
208
+ parser.add_argument('--elevation', type=float, default=0, help='delta elevation angle in [-90, 90]')
209
+ parser.add_argument('--azimuth', type=float, default=0, help='delta azimuth angle in [-180, 180]')
210
+ parser.add_argument('--radius', type=float, default=0, help='delta camera radius multiplier in [-0.5, 0.5]')
211
+ parser.add_argument('--stable', action='store_true')
212
+
213
+ opt = parser.parse_args()
214
+
215
+ device = torch.device('cuda')
216
+
217
+ print(f'[INFO] loading image from {opt.input} ...')
218
+ image = kiui.read_image(opt.input, mode='tensor')
219
+ image = image.permute(2, 0, 1).unsqueeze(0).contiguous().to(device)
220
+ image = F.interpolate(image, (256, 256), mode='bilinear', align_corners=False)
221
+
222
+ print(f'[INFO] loading model ...')
223
+
224
+ if opt.stable:
225
+ zero123 = Zero123(device, model_key='ashawkey/stable-zero123-diffusers')
226
+ else:
227
+ zero123 = Zero123(device, model_key='ashawkey/zero123-xl-diffusers')
228
+
229
+ print(f'[INFO] running model ...')
230
+ zero123.get_img_embeds(image)
231
+
232
+ azimuth = opt.azimuth
233
+ while True:
234
+ outputs = zero123.refine(image, elevation=[opt.elevation], azimuth=[azimuth], radius=[opt.radius], strength=0)
235
+ plt.imshow(outputs.float().cpu().numpy().transpose(0, 2, 3, 1)[0])
236
+ plt.show()
237
+ azimuth = (azimuth + 10) % 360
lgm/core/attention.py ADDED
@@ -0,0 +1,156 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ #
3
+ # This source code is licensed under the Apache License, Version 2.0
4
+ # found in the LICENSE file in the root directory of this source tree.
5
+
6
+ # References:
7
+ # https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
8
+ # https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
9
+
10
+ import os
11
+ import warnings
12
+
13
+ from torch import Tensor
14
+ from torch import nn
15
+
16
+ XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
17
+ try:
18
+ if XFORMERS_ENABLED:
19
+ from xformers.ops import memory_efficient_attention, unbind
20
+
21
+ XFORMERS_AVAILABLE = True
22
+ warnings.warn("xFormers is available (Attention)")
23
+ else:
24
+ warnings.warn("xFormers is disabled (Attention)")
25
+ raise ImportError
26
+ except ImportError:
27
+ XFORMERS_AVAILABLE = False
28
+ warnings.warn("xFormers is not available (Attention)")
29
+
30
+
31
+ class Attention(nn.Module):
32
+ def __init__(
33
+ self,
34
+ dim: int,
35
+ num_heads: int = 8,
36
+ qkv_bias: bool = False,
37
+ proj_bias: bool = True,
38
+ attn_drop: float = 0.0,
39
+ proj_drop: float = 0.0,
40
+ ) -> None:
41
+ super().__init__()
42
+ self.num_heads = num_heads
43
+ head_dim = dim // num_heads
44
+ self.scale = head_dim**-0.5
45
+
46
+ self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
47
+ self.attn_drop = nn.Dropout(attn_drop)
48
+ self.proj = nn.Linear(dim, dim, bias=proj_bias)
49
+ self.proj_drop = nn.Dropout(proj_drop)
50
+
51
+ def forward(self, x: Tensor) -> Tensor:
52
+ B, N, C = x.shape
53
+ qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
54
+
55
+ q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
56
+ attn = q @ k.transpose(-2, -1)
57
+
58
+ attn = attn.softmax(dim=-1)
59
+ attn = self.attn_drop(attn)
60
+
61
+ x = (attn @ v).transpose(1, 2).reshape(B, N, C)
62
+ x = self.proj(x)
63
+ x = self.proj_drop(x)
64
+ return x
65
+
66
+
67
+ class MemEffAttention(Attention):
68
+ def forward(self, x: Tensor, attn_bias=None) -> Tensor:
69
+ if not XFORMERS_AVAILABLE:
70
+ if attn_bias is not None:
71
+ raise AssertionError("xFormers is required for using nested tensors")
72
+ return super().forward(x)
73
+
74
+ B, N, C = x.shape
75
+ qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
76
+
77
+ q, k, v = unbind(qkv, 2)
78
+
79
+ x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
80
+ x = x.reshape([B, N, C])
81
+
82
+ x = self.proj(x)
83
+ x = self.proj_drop(x)
84
+ return x
85
+
86
+
87
+ class CrossAttention(nn.Module):
88
+ def __init__(
89
+ self,
90
+ dim: int,
91
+ dim_q: int,
92
+ dim_k: int,
93
+ dim_v: int,
94
+ num_heads: int = 8,
95
+ qkv_bias: bool = False,
96
+ proj_bias: bool = True,
97
+ attn_drop: float = 0.0,
98
+ proj_drop: float = 0.0,
99
+ ) -> None:
100
+ super().__init__()
101
+ self.dim = dim
102
+ self.num_heads = num_heads
103
+ head_dim = dim // num_heads
104
+ self.scale = head_dim**-0.5
105
+
106
+ self.to_q = nn.Linear(dim_q, dim, bias=qkv_bias)
107
+ self.to_k = nn.Linear(dim_k, dim, bias=qkv_bias)
108
+ self.to_v = nn.Linear(dim_v, dim, bias=qkv_bias)
109
+ self.attn_drop = nn.Dropout(attn_drop)
110
+ self.proj = nn.Linear(dim, dim, bias=proj_bias)
111
+ self.proj_drop = nn.Dropout(proj_drop)
112
+
113
+ def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
114
+ # q: [B, N, Cq]
115
+ # k: [B, M, Ck]
116
+ # v: [B, M, Cv]
117
+ # return: [B, N, C]
118
+
119
+ B, N, _ = q.shape
120
+ M = k.shape[1]
121
+
122
+ 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]
123
+ 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]
124
+ 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]
125
+
126
+ attn = q @ k.transpose(-2, -1) # [B, nh, N, M]
127
+
128
+ attn = attn.softmax(dim=-1) # [B, nh, N, M]
129
+ attn = self.attn_drop(attn)
130
+
131
+ 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]
132
+ x = self.proj(x)
133
+ x = self.proj_drop(x)
134
+ return x
135
+
136
+
137
+ class MemEffCrossAttention(CrossAttention):
138
+ def forward(self, q: Tensor, k: Tensor, v: Tensor, attn_bias=None) -> Tensor:
139
+ if not XFORMERS_AVAILABLE:
140
+ if attn_bias is not None:
141
+ raise AssertionError("xFormers is required for using nested tensors")
142
+ return super().forward(x)
143
+
144
+ B, N, _ = q.shape
145
+ M = k.shape[1]
146
+
147
+ q = self.scale * self.to_q(q).reshape(B, N, self.num_heads, self.dim // self.num_heads) # [B, N, nh, C/nh]
148
+ k = self.to_k(k).reshape(B, M, self.num_heads, self.dim // self.num_heads) # [B, M, nh, C/nh]
149
+ v = self.to_v(v).reshape(B, M, self.num_heads, self.dim // self.num_heads) # [B, M, nh, C/nh]
150
+
151
+ x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
152
+ x = x.reshape(B, N, -1)
153
+
154
+ x = self.proj(x)
155
+ x = self.proj_drop(x)
156
+ return x
lgm/core/gs.py ADDED
@@ -0,0 +1,190 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+
3
+ import torch
4
+ import torch.nn as nn
5
+ import torch.nn.functional as F
6
+
7
+ from diff_gaussian_rasterization import (
8
+ GaussianRasterizationSettings,
9
+ GaussianRasterizer,
10
+ )
11
+
12
+ from core.options import Options
13
+
14
+ import kiui
15
+
16
+ class GaussianRenderer:
17
+ def __init__(self, opt: Options):
18
+
19
+ self.opt = opt
20
+ self.bg_color = torch.tensor([1, 1, 1], dtype=torch.float32, device="cuda")
21
+
22
+ # intrinsics
23
+ self.tan_half_fov = np.tan(0.5 * np.deg2rad(self.opt.fovy))
24
+ self.proj_matrix = torch.zeros(4, 4, dtype=torch.float32)
25
+ self.proj_matrix[0, 0] = 1 / self.tan_half_fov
26
+ self.proj_matrix[1, 1] = 1 / self.tan_half_fov
27
+ self.proj_matrix[2, 2] = (opt.zfar + opt.znear) / (opt.zfar - opt.znear)
28
+ self.proj_matrix[3, 2] = - (opt.zfar * opt.znear) / (opt.zfar - opt.znear)
29
+ self.proj_matrix[2, 3] = 1
30
+
31
+ def render(self, gaussians, cam_view, cam_view_proj, cam_pos, bg_color=None, scale_modifier=1):
32
+ # gaussians: [B, N, 14]
33
+ # cam_view, cam_view_proj: [B, V, 4, 4]
34
+ # cam_pos: [B, V, 3]
35
+
36
+ device = gaussians.device
37
+ B, V = cam_view.shape[:2]
38
+
39
+ # loop of loop...
40
+ images = []
41
+ alphas = []
42
+ for b in range(B):
43
+
44
+ # pos, opacity, scale, rotation, shs
45
+ means3D = gaussians[b, :, 0:3].contiguous().float()
46
+ opacity = gaussians[b, :, 3:4].contiguous().float()
47
+ scales = gaussians[b, :, 4:7].contiguous().float()
48
+ rotations = gaussians[b, :, 7:11].contiguous().float()
49
+ rgbs = gaussians[b, :, 11:].contiguous().float() # [N, 3]
50
+
51
+ for v in range(V):
52
+
53
+ # render novel views
54
+ view_matrix = cam_view[b, v].float()
55
+ view_proj_matrix = cam_view_proj[b, v].float()
56
+ campos = cam_pos[b, v].float()
57
+
58
+ raster_settings = GaussianRasterizationSettings(
59
+ image_height=self.opt.output_size,
60
+ image_width=self.opt.output_size,
61
+ tanfovx=self.tan_half_fov,
62
+ tanfovy=self.tan_half_fov,
63
+ bg=self.bg_color if bg_color is None else bg_color,
64
+ scale_modifier=scale_modifier,
65
+ viewmatrix=view_matrix,
66
+ projmatrix=view_proj_matrix,
67
+ sh_degree=0,
68
+ campos=campos,
69
+ prefiltered=False,
70
+ debug=False,
71
+ )
72
+
73
+ rasterizer = GaussianRasterizer(raster_settings=raster_settings)
74
+
75
+ # Rasterize visible Gaussians to image, obtain their radii (on screen).
76
+ rendered_image, radii, rendered_depth, rendered_alpha = rasterizer(
77
+ means3D=means3D,
78
+ means2D=torch.zeros_like(means3D, dtype=torch.float32, device=device),
79
+ shs=None,
80
+ colors_precomp=rgbs,
81
+ opacities=opacity,
82
+ scales=scales,
83
+ rotations=rotations,
84
+ cov3D_precomp=None,
85
+ )
86
+
87
+ rendered_image = rendered_image.clamp(0, 1)
88
+
89
+ images.append(rendered_image)
90
+ alphas.append(rendered_alpha)
91
+
92
+ images = torch.stack(images, dim=0).view(B, V, 3, self.opt.output_size, self.opt.output_size)
93
+ alphas = torch.stack(alphas, dim=0).view(B, V, 1, self.opt.output_size, self.opt.output_size)
94
+
95
+ return {
96
+ "image": images, # [B, V, 3, H, W]
97
+ "alpha": alphas, # [B, V, 1, H, W]
98
+ }
99
+
100
+
101
+ def save_ply(self, gaussians, path, compatible=True):
102
+ # gaussians: [B, N, 14]
103
+ # compatible: save pre-activated gaussians as in the original paper
104
+
105
+ assert gaussians.shape[0] == 1, 'only support batch size 1'
106
+
107
+ from plyfile import PlyData, PlyElement
108
+
109
+ means3D = gaussians[0, :, 0:3].contiguous().float()
110
+ opacity = gaussians[0, :, 3:4].contiguous().float()
111
+ scales = gaussians[0, :, 4:7].contiguous().float()
112
+ rotations = gaussians[0, :, 7:11].contiguous().float()
113
+ shs = gaussians[0, :, 11:].unsqueeze(1).contiguous().float() # [N, 1, 3]
114
+
115
+ # prune by opacity
116
+ mask = opacity.squeeze(-1) >= 0.005
117
+ means3D = means3D[mask]
118
+ opacity = opacity[mask]
119
+ scales = scales[mask]
120
+ rotations = rotations[mask]
121
+ shs = shs[mask]
122
+
123
+ # invert activation to make it compatible with the original ply format
124
+ if compatible:
125
+ opacity = kiui.op.inverse_sigmoid(opacity)
126
+ scales = torch.log(scales + 1e-8)
127
+ shs = (shs - 0.5) / 0.28209479177387814
128
+
129
+ xyzs = means3D.detach().cpu().numpy()
130
+ f_dc = shs.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
131
+ opacities = opacity.detach().cpu().numpy()
132
+ scales = scales.detach().cpu().numpy()
133
+ rotations = rotations.detach().cpu().numpy()
134
+
135
+ l = ['x', 'y', 'z']
136
+ # All channels except the 3 DC
137
+ for i in range(f_dc.shape[1]):
138
+ l.append('f_dc_{}'.format(i))
139
+ l.append('opacity')
140
+ for i in range(scales.shape[1]):
141
+ l.append('scale_{}'.format(i))
142
+ for i in range(rotations.shape[1]):
143
+ l.append('rot_{}'.format(i))
144
+
145
+ dtype_full = [(attribute, 'f4') for attribute in l]
146
+
147
+ elements = np.empty(xyzs.shape[0], dtype=dtype_full)
148
+ attributes = np.concatenate((xyzs, f_dc, opacities, scales, rotations), axis=1)
149
+ elements[:] = list(map(tuple, attributes))
150
+ el = PlyElement.describe(elements, 'vertex')
151
+
152
+ PlyData([el]).write(path)
153
+
154
+ def load_ply(self, path, compatible=True):
155
+
156
+ from plyfile import PlyData, PlyElement
157
+
158
+ plydata = PlyData.read(path)
159
+
160
+ xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
161
+ np.asarray(plydata.elements[0]["y"]),
162
+ np.asarray(plydata.elements[0]["z"])), axis=1)
163
+ print("Number of points at loading : ", xyz.shape[0])
164
+
165
+ opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]
166
+
167
+ shs = np.zeros((xyz.shape[0], 3))
168
+ shs[:, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
169
+ shs[:, 1] = np.asarray(plydata.elements[0]["f_dc_1"])
170
+ shs[:, 2] = np.asarray(plydata.elements[0]["f_dc_2"])
171
+
172
+ scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
173
+ scales = np.zeros((xyz.shape[0], len(scale_names)))
174
+ for idx, attr_name in enumerate(scale_names):
175
+ scales[:, idx] = np.asarray(plydata.elements[0][attr_name])
176
+
177
+ rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot_")]
178
+ rots = np.zeros((xyz.shape[0], len(rot_names)))
179
+ for idx, attr_name in enumerate(rot_names):
180
+ rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
181
+
182
+ gaussians = np.concatenate([xyz, opacities, scales, rots, shs], axis=1)
183
+ gaussians = torch.from_numpy(gaussians).float() # cpu
184
+
185
+ if compatible:
186
+ gaussians[..., 3:4] = torch.sigmoid(gaussians[..., 3:4])
187
+ gaussians[..., 4:7] = torch.exp(gaussians[..., 4:7])
188
+ gaussians[..., 11:] = 0.28209479177387814 * gaussians[..., 11:] + 0.5
189
+
190
+ return gaussians
lgm/core/models.py ADDED
@@ -0,0 +1,171 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ import torch.nn.functional as F
4
+ import numpy as np
5
+
6
+ import kiui
7
+ from kiui.lpips import LPIPS
8
+
9
+ from core.unet import UNet
10
+ from core.options import Options
11
+ from core.gs import GaussianRenderer
12
+
13
+
14
+ class LGM(nn.Module):
15
+ def __init__(
16
+ self,
17
+ opt: Options,
18
+ ):
19
+ super().__init__()
20
+
21
+ self.opt = opt
22
+
23
+ # unet
24
+ self.unet = UNet(
25
+ 9, 14,
26
+ down_channels=self.opt.down_channels,
27
+ down_attention=self.opt.down_attention,
28
+ mid_attention=self.opt.mid_attention,
29
+ up_channels=self.opt.up_channels,
30
+ up_attention=self.opt.up_attention,
31
+ )
32
+
33
+ # last conv
34
+ self.conv = nn.Conv2d(14, 14, kernel_size=1) # NOTE: maybe remove it if train again
35
+
36
+ # Gaussian Renderer
37
+ self.gs = GaussianRenderer(opt)
38
+
39
+ # activations...
40
+ self.pos_act = lambda x: x.clamp(-1, 1)
41
+ self.scale_act = lambda x: 0.1 * F.softplus(x)
42
+ self.opacity_act = lambda x: torch.sigmoid(x)
43
+ self.rot_act = lambda x: F.normalize(x, dim=-1)
44
+ self.rgb_act = lambda x: 0.5 * torch.tanh(x) + 0.5 # NOTE: may use sigmoid if train again
45
+
46
+ # LPIPS loss
47
+ if self.opt.lambda_lpips > 0:
48
+ self.lpips_loss = LPIPS(net='vgg')
49
+ self.lpips_loss.requires_grad_(False)
50
+
51
+
52
+ def state_dict(self, **kwargs):
53
+ # remove lpips_loss
54
+ state_dict = super().state_dict(**kwargs)
55
+ for k in list(state_dict.keys()):
56
+ if 'lpips_loss' in k:
57
+ del state_dict[k]
58
+ return state_dict
59
+
60
+
61
+ def prepare_default_rays(self, device, elevation=0):
62
+
63
+ from kiui.cam import orbit_camera
64
+ from core.utils import get_rays
65
+
66
+ cam_poses = np.stack([
67
+ orbit_camera(elevation, 0, radius=self.opt.cam_radius),
68
+ orbit_camera(elevation, 90, radius=self.opt.cam_radius),
69
+ orbit_camera(elevation, 180, radius=self.opt.cam_radius),
70
+ orbit_camera(elevation, 270, radius=self.opt.cam_radius),
71
+ ], axis=0) # [4, 4, 4]
72
+ cam_poses = torch.from_numpy(cam_poses)
73
+
74
+ rays_embeddings = []
75
+ for i in range(cam_poses.shape[0]):
76
+ rays_o, rays_d = get_rays(cam_poses[i], self.opt.input_size, self.opt.input_size, self.opt.fovy) # [h, w, 3]
77
+ rays_plucker = torch.cat([torch.cross(rays_o, rays_d, dim=-1), rays_d], dim=-1) # [h, w, 6]
78
+ rays_embeddings.append(rays_plucker)
79
+
80
+ ## visualize rays for plotting figure
81
+ # kiui.vis.plot_image(rays_d * 0.5 + 0.5, save=True)
82
+
83
+ rays_embeddings = torch.stack(rays_embeddings, dim=0).permute(0, 3, 1, 2).contiguous().to(device) # [V, 6, h, w]
84
+
85
+ return rays_embeddings
86
+
87
+
88
+ def forward_gaussians(self, images):
89
+ # images: [B, 4, 9, H, W]
90
+ # return: Gaussians: [B, dim_t]
91
+
92
+ B, V, C, H, W = images.shape
93
+ images = images.view(B*V, C, H, W)
94
+
95
+ x = self.unet(images) # [B*4, 14, h, w]
96
+ x = self.conv(x) # [B*4, 14, h, w]
97
+
98
+ x = x.reshape(B, 4, 14, self.opt.splat_size, self.opt.splat_size)
99
+
100
+ ## visualize multi-view gaussian features for plotting figure
101
+ # tmp_alpha = self.opacity_act(x[0, :, 3:4])
102
+ # tmp_img_rgb = self.rgb_act(x[0, :, 11:]) * tmp_alpha + (1 - tmp_alpha)
103
+ # tmp_img_pos = self.pos_act(x[0, :, 0:3]) * 0.5 + 0.5
104
+ # kiui.vis.plot_image(tmp_img_rgb, save=True)
105
+ # kiui.vis.plot_image(tmp_img_pos, save=True)
106
+
107
+ x = x.permute(0, 1, 3, 4, 2).reshape(B, -1, 14)
108
+
109
+ pos = self.pos_act(x[..., 0:3]) # [B, N, 3]
110
+ opacity = self.opacity_act(x[..., 3:4])
111
+ scale = self.scale_act(x[..., 4:7])
112
+ rotation = self.rot_act(x[..., 7:11])
113
+ rgbs = self.rgb_act(x[..., 11:])
114
+
115
+ gaussians = torch.cat([pos, opacity, scale, rotation, rgbs], dim=-1) # [B, N, 14]
116
+
117
+ return gaussians
118
+
119
+
120
+ def forward(self, data, step_ratio=1):
121
+ # data: output of the dataloader
122
+ # return: loss
123
+
124
+ results = {}
125
+ loss = 0
126
+
127
+ images = data['input'] # [B, 4, 9, h, W], input features
128
+
129
+ # use the first view to predict gaussians
130
+ gaussians = self.forward_gaussians(images) # [B, N, 14]
131
+
132
+ results['gaussians'] = gaussians
133
+
134
+ # always use white bg
135
+ bg_color = torch.ones(3, dtype=torch.float32, device=gaussians.device)
136
+
137
+ # use the other views for rendering and supervision
138
+ results = self.gs.render(gaussians, data['cam_view'], data['cam_view_proj'], data['cam_pos'], bg_color=bg_color)
139
+ pred_images = results['image'] # [B, V, C, output_size, output_size]
140
+ pred_alphas = results['alpha'] # [B, V, 1, output_size, output_size]
141
+
142
+ results['images_pred'] = pred_images
143
+ results['alphas_pred'] = pred_alphas
144
+
145
+ gt_images = data['images_output'] # [B, V, 3, output_size, output_size], ground-truth novel views
146
+ gt_masks = data['masks_output'] # [B, V, 1, output_size, output_size], ground-truth masks
147
+
148
+ gt_images = gt_images * gt_masks + bg_color.view(1, 1, 3, 1, 1) * (1 - gt_masks)
149
+
150
+ loss_mse = F.mse_loss(pred_images, gt_images) + F.mse_loss(pred_alphas, gt_masks)
151
+ loss = loss + loss_mse
152
+
153
+ if self.opt.lambda_lpips > 0:
154
+ loss_lpips = self.lpips_loss(
155
+ # gt_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1,
156
+ # pred_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1,
157
+ # downsampled to at most 256 to reduce memory cost
158
+ F.interpolate(gt_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1, (256, 256), mode='bilinear', align_corners=False),
159
+ F.interpolate(pred_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1, (256, 256), mode='bilinear', align_corners=False),
160
+ ).mean()
161
+ results['loss_lpips'] = loss_lpips
162
+ loss = loss + self.opt.lambda_lpips * loss_lpips
163
+
164
+ results['loss'] = loss
165
+
166
+ # metric
167
+ with torch.no_grad():
168
+ psnr = -10 * torch.log10(torch.mean((pred_images.detach() - gt_images) ** 2))
169
+ results['psnr'] = psnr
170
+
171
+ return results
lgm/core/options.py ADDED
@@ -0,0 +1,120 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import tyro
2
+ from dataclasses import dataclass
3
+ from typing import Tuple, Literal, Dict, Optional
4
+
5
+
6
+ @dataclass
7
+ class Options:
8
+ ### model
9
+ # Unet image input size
10
+ input_size: int = 256
11
+ # Unet definition
12
+ down_channels: Tuple[int, ...] = (64, 128, 256, 512, 1024, 1024)
13
+ down_attention: Tuple[bool, ...] = (False, False, False, True, True, True)
14
+ mid_attention: bool = True
15
+ up_channels: Tuple[int, ...] = (1024, 1024, 512, 256)
16
+ up_attention: Tuple[bool, ...] = (True, True, True, False)
17
+ # Unet output size, dependent on the input_size and U-Net structure!
18
+ splat_size: int = 64
19
+ # gaussian render size
20
+ output_size: int = 256
21
+
22
+ ### dataset
23
+ # data mode (only support s3 now)
24
+ data_mode: Literal['s3'] = 's3'
25
+ # fovy of the dataset
26
+ fovy: float = 49.1
27
+ # camera near plane
28
+ znear: float = 0.5
29
+ # camera far plane
30
+ zfar: float = 2.5
31
+ # number of all views (input + output)
32
+ num_views: int = 12
33
+ # number of views
34
+ num_input_views: int = 4
35
+ # camera radius
36
+ cam_radius: float = 1.5 # to better use [-1, 1]^3 space
37
+ # num workers
38
+ num_workers: int = 8
39
+
40
+ ### training
41
+ # workspace
42
+ workspace: str = './workspace'
43
+ # resume
44
+ resume: Optional[str] = 'pretrained/model_fp16_fixrot.safetensors'
45
+ # batch size (per-GPU)
46
+ batch_size: int = 8
47
+ # gradient accumulation
48
+ gradient_accumulation_steps: int = 1
49
+ # training epochs
50
+ num_epochs: int = 30
51
+ # lpips loss weight
52
+ lambda_lpips: float = 1.0
53
+ # gradient clip
54
+ gradient_clip: float = 1.0
55
+ # mixed precision
56
+ mixed_precision: str = 'bf16'
57
+ # learning rate
58
+ lr: float = 4e-4
59
+ # augmentation prob for grid distortion
60
+ prob_grid_distortion: float = 0.5
61
+ # augmentation prob for camera jitter
62
+ prob_cam_jitter: float = 0.5
63
+
64
+ ### testing
65
+ # test image path
66
+ test_path: Optional[str] = None
67
+
68
+ ### misc
69
+ # nvdiffrast backend setting
70
+ force_cuda_rast: bool = False
71
+ # render fancy video with gaussian scaling effect
72
+ fancy_video: bool = False
73
+
74
+
75
+ # all the default settings
76
+ config_defaults: Dict[str, Options] = {}
77
+ config_doc: Dict[str, str] = {}
78
+
79
+ config_doc['lrm'] = 'the default settings for LGM'
80
+ config_defaults['lrm'] = Options()
81
+
82
+ config_doc['small'] = 'small model with lower resolution Gaussians'
83
+ config_defaults['small'] = Options(
84
+ input_size=256,
85
+ splat_size=64,
86
+ output_size=256,
87
+ batch_size=8,
88
+ gradient_accumulation_steps=1,
89
+ mixed_precision='bf16',
90
+ )
91
+
92
+ config_doc['big'] = 'big model with higher resolution Gaussians'
93
+ config_defaults['big'] = Options(
94
+ input_size=256,
95
+ up_channels=(1024, 1024, 512, 256, 128), # one more decoder
96
+ up_attention=(True, True, True, False, False),
97
+ splat_size=128,
98
+ output_size=512, # render & supervise Gaussians at a higher resolution.
99
+ batch_size=8,
100
+ num_views=8,
101
+ gradient_accumulation_steps=1,
102
+ mixed_precision='bf16',
103
+ )
104
+
105
+ config_doc['tiny'] = 'tiny model for ablation'
106
+ config_defaults['tiny'] = Options(
107
+ input_size=256,
108
+ down_channels=(32, 64, 128, 256, 512),
109
+ down_attention=(False, False, False, False, True),
110
+ up_channels=(512, 256, 128),
111
+ up_attention=(True, False, False, False),
112
+ splat_size=64,
113
+ output_size=256,
114
+ batch_size=16,
115
+ num_views=8,
116
+ gradient_accumulation_steps=1,
117
+ mixed_precision='bf16',
118
+ )
119
+
120
+ AllConfigs = tyro.extras.subcommand_type_from_defaults(config_defaults, config_doc)
lgm/core/unet.py ADDED
@@ -0,0 +1,319 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ import torch.nn.functional as F
4
+
5
+ import numpy as np
6
+ from typing import Tuple, Literal
7
+ from functools import partial
8
+
9
+ from core.attention import MemEffAttention
10
+
11
+ class MVAttention(nn.Module):
12
+ def __init__(
13
+ self,
14
+ dim: int,
15
+ num_heads: int = 8,
16
+ qkv_bias: bool = False,
17
+ proj_bias: bool = True,
18
+ attn_drop: float = 0.0,
19
+ proj_drop: float = 0.0,
20
+ groups: int = 32,
21
+ eps: float = 1e-5,
22
+ residual: bool = True,
23
+ skip_scale: float = 1,
24
+ num_frames: int = 4, # WARN: hardcoded!
25
+ ):
26
+ super().__init__()
27
+
28
+ self.residual = residual
29
+ self.skip_scale = skip_scale
30
+ self.num_frames = num_frames
31
+
32
+ self.norm = nn.GroupNorm(num_groups=groups, num_channels=dim, eps=eps, affine=True)
33
+ self.attn = MemEffAttention(dim, num_heads, qkv_bias, proj_bias, attn_drop, proj_drop)
34
+
35
+ def forward(self, x):
36
+ # x: [B*V, C, H, W]
37
+ BV, C, H, W = x.shape
38
+ B = BV // self.num_frames # assert BV % self.num_frames == 0
39
+
40
+ res = x
41
+ x = self.norm(x)
42
+
43
+ x = x.reshape(B, self.num_frames, C, H, W).permute(0, 1, 3, 4, 2).reshape(B, -1, C)
44
+ x = self.attn(x)
45
+ x = x.reshape(B, self.num_frames, H, W, C).permute(0, 1, 4, 2, 3).reshape(BV, C, H, W)
46
+
47
+ if self.residual:
48
+ x = (x + res) * self.skip_scale
49
+ return x
50
+
51
+ class ResnetBlock(nn.Module):
52
+ def __init__(
53
+ self,
54
+ in_channels: int,
55
+ out_channels: int,
56
+ resample: Literal['default', 'up', 'down'] = 'default',
57
+ groups: int = 32,
58
+ eps: float = 1e-5,
59
+ skip_scale: float = 1, # multiplied to output
60
+ ):
61
+ super().__init__()
62
+
63
+ self.in_channels = in_channels
64
+ self.out_channels = out_channels
65
+ self.skip_scale = skip_scale
66
+
67
+ self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
68
+ self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
69
+
70
+ self.norm2 = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True)
71
+ self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
72
+
73
+ self.act = F.silu
74
+
75
+ self.resample = None
76
+ if resample == 'up':
77
+ self.resample = partial(F.interpolate, scale_factor=2.0, mode="nearest")
78
+ elif resample == 'down':
79
+ self.resample = nn.AvgPool2d(kernel_size=2, stride=2)
80
+
81
+ self.shortcut = nn.Identity()
82
+ if self.in_channels != self.out_channels:
83
+ self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=True)
84
+
85
+
86
+ def forward(self, x):
87
+ res = x
88
+
89
+ x = self.norm1(x)
90
+ x = self.act(x)
91
+
92
+ if self.resample:
93
+ res = self.resample(res)
94
+ x = self.resample(x)
95
+
96
+ x = self.conv1(x)
97
+ x = self.norm2(x)
98
+ x = self.act(x)
99
+ x = self.conv2(x)
100
+
101
+ x = (x + self.shortcut(res)) * self.skip_scale
102
+
103
+ return x
104
+
105
+ class DownBlock(nn.Module):
106
+ def __init__(
107
+ self,
108
+ in_channels: int,
109
+ out_channels: int,
110
+ num_layers: int = 1,
111
+ downsample: bool = True,
112
+ attention: bool = True,
113
+ attention_heads: int = 16,
114
+ skip_scale: float = 1,
115
+ ):
116
+ super().__init__()
117
+
118
+ nets = []
119
+ attns = []
120
+ for i in range(num_layers):
121
+ in_channels = in_channels if i == 0 else out_channels
122
+ nets.append(ResnetBlock(in_channels, out_channels, skip_scale=skip_scale))
123
+ if attention:
124
+ attns.append(MVAttention(out_channels, attention_heads, skip_scale=skip_scale))
125
+ else:
126
+ attns.append(None)
127
+ self.nets = nn.ModuleList(nets)
128
+ self.attns = nn.ModuleList(attns)
129
+
130
+ self.downsample = None
131
+ if downsample:
132
+ self.downsample = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1)
133
+
134
+ def forward(self, x):
135
+ xs = []
136
+
137
+ for attn, net in zip(self.attns, self.nets):
138
+ x = net(x)
139
+ if attn:
140
+ x = attn(x)
141
+ xs.append(x)
142
+
143
+ if self.downsample:
144
+ x = self.downsample(x)
145
+ xs.append(x)
146
+
147
+ return x, xs
148
+
149
+
150
+ class MidBlock(nn.Module):
151
+ def __init__(
152
+ self,
153
+ in_channels: int,
154
+ num_layers: int = 1,
155
+ attention: bool = True,
156
+ attention_heads: int = 16,
157
+ skip_scale: float = 1,
158
+ ):
159
+ super().__init__()
160
+
161
+ nets = []
162
+ attns = []
163
+ # first layer
164
+ nets.append(ResnetBlock(in_channels, in_channels, skip_scale=skip_scale))
165
+ # more layers
166
+ for i in range(num_layers):
167
+ nets.append(ResnetBlock(in_channels, in_channels, skip_scale=skip_scale))
168
+ if attention:
169
+ attns.append(MVAttention(in_channels, attention_heads, skip_scale=skip_scale))
170
+ else:
171
+ attns.append(None)
172
+ self.nets = nn.ModuleList(nets)
173
+ self.attns = nn.ModuleList(attns)
174
+
175
+ def forward(self, x):
176
+ x = self.nets[0](x)
177
+ for attn, net in zip(self.attns, self.nets[1:]):
178
+ if attn:
179
+ x = attn(x)
180
+ x = net(x)
181
+ return x
182
+
183
+
184
+ class UpBlock(nn.Module):
185
+ def __init__(
186
+ self,
187
+ in_channels: int,
188
+ prev_out_channels: int,
189
+ out_channels: int,
190
+ num_layers: int = 1,
191
+ upsample: bool = True,
192
+ attention: bool = True,
193
+ attention_heads: int = 16,
194
+ skip_scale: float = 1,
195
+ ):
196
+ super().__init__()
197
+
198
+ nets = []
199
+ attns = []
200
+ for i in range(num_layers):
201
+ cin = in_channels if i == 0 else out_channels
202
+ cskip = prev_out_channels if (i == num_layers - 1) else out_channels
203
+
204
+ nets.append(ResnetBlock(cin + cskip, out_channels, skip_scale=skip_scale))
205
+ if attention:
206
+ attns.append(MVAttention(out_channels, attention_heads, skip_scale=skip_scale))
207
+ else:
208
+ attns.append(None)
209
+ self.nets = nn.ModuleList(nets)
210
+ self.attns = nn.ModuleList(attns)
211
+
212
+ self.upsample = None
213
+ if upsample:
214
+ self.upsample = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
215
+
216
+ def forward(self, x, xs):
217
+
218
+ for attn, net in zip(self.attns, self.nets):
219
+ res_x = xs[-1]
220
+ xs = xs[:-1]
221
+ x = torch.cat([x, res_x], dim=1)
222
+ x = net(x)
223
+ if attn:
224
+ x = attn(x)
225
+
226
+ if self.upsample:
227
+ x = F.interpolate(x, scale_factor=2.0, mode='nearest')
228
+ x = self.upsample(x)
229
+
230
+ return x
231
+
232
+
233
+ # it could be asymmetric!
234
+ class UNet(nn.Module):
235
+ def __init__(
236
+ self,
237
+ in_channels: int = 3,
238
+ out_channels: int = 3,
239
+ down_channels: Tuple[int, ...] = (64, 128, 256, 512, 1024),
240
+ down_attention: Tuple[bool, ...] = (False, False, False, True, True),
241
+ mid_attention: bool = True,
242
+ up_channels: Tuple[int, ...] = (1024, 512, 256),
243
+ up_attention: Tuple[bool, ...] = (True, True, False),
244
+ layers_per_block: int = 2,
245
+ skip_scale: float = np.sqrt(0.5),
246
+ ):
247
+ super().__init__()
248
+
249
+ # first
250
+ self.conv_in = nn.Conv2d(in_channels, down_channels[0], kernel_size=3, stride=1, padding=1)
251
+
252
+ # down
253
+ down_blocks = []
254
+ cout = down_channels[0]
255
+ for i in range(len(down_channels)):
256
+ cin = cout
257
+ cout = down_channels[i]
258
+
259
+ down_blocks.append(DownBlock(
260
+ cin, cout,
261
+ num_layers=layers_per_block,
262
+ downsample=(i != len(down_channels) - 1), # not final layer
263
+ attention=down_attention[i],
264
+ skip_scale=skip_scale,
265
+ ))
266
+ self.down_blocks = nn.ModuleList(down_blocks)
267
+
268
+ # mid
269
+ self.mid_block = MidBlock(down_channels[-1], attention=mid_attention, skip_scale=skip_scale)
270
+
271
+ # up
272
+ up_blocks = []
273
+ cout = up_channels[0]
274
+ for i in range(len(up_channels)):
275
+ cin = cout
276
+ cout = up_channels[i]
277
+ cskip = down_channels[max(-2 - i, -len(down_channels))] # for assymetric
278
+
279
+ up_blocks.append(UpBlock(
280
+ cin, cskip, cout,
281
+ num_layers=layers_per_block + 1, # one more layer for up
282
+ upsample=(i != len(up_channels) - 1), # not final layer
283
+ attention=up_attention[i],
284
+ skip_scale=skip_scale,
285
+ ))
286
+ self.up_blocks = nn.ModuleList(up_blocks)
287
+
288
+ # last
289
+ self.norm_out = nn.GroupNorm(num_channels=up_channels[-1], num_groups=32, eps=1e-5)
290
+ self.conv_out = nn.Conv2d(up_channels[-1], out_channels, kernel_size=3, stride=1, padding=1)
291
+
292
+
293
+ def forward(self, x):
294
+ # x: [B, Cin, H, W]
295
+
296
+ # first
297
+ x = self.conv_in(x)
298
+
299
+ # down
300
+ xss = [x]
301
+ for block in self.down_blocks:
302
+ x, xs = block(x)
303
+ xss.extend(xs)
304
+
305
+ # mid
306
+ x = self.mid_block(x)
307
+
308
+ # up
309
+ for block in self.up_blocks:
310
+ xs = xss[-len(block.nets):]
311
+ xss = xss[:-len(block.nets)]
312
+ x = block(x, xs)
313
+
314
+ # last
315
+ x = self.norm_out(x)
316
+ x = F.silu(x)
317
+ x = self.conv_out(x) # [B, Cout, H', W']
318
+
319
+ return x
lgm/core/utils.py ADDED
@@ -0,0 +1,108 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+
3
+ import torch
4
+ import torch.nn as nn
5
+ import torch.nn.functional as F
6
+
7
+ import roma
8
+ from kiui.op import safe_normalize
9
+
10
+ def get_rays(pose, h, w, fovy, opengl=True):
11
+
12
+ x, y = torch.meshgrid(
13
+ torch.arange(w, device=pose.device),
14
+ torch.arange(h, device=pose.device),
15
+ indexing="xy",
16
+ )
17
+ x = x.flatten()
18
+ y = y.flatten()
19
+
20
+ cx = w * 0.5
21
+ cy = h * 0.5
22
+
23
+ focal = h * 0.5 / np.tan(0.5 * np.deg2rad(fovy))
24
+
25
+ camera_dirs = F.pad(
26
+ torch.stack(
27
+ [
28
+ (x - cx + 0.5) / focal,
29
+ (y - cy + 0.5) / focal * (-1.0 if opengl else 1.0),
30
+ ],
31
+ dim=-1,
32
+ ),
33
+ (0, 1),
34
+ value=(-1.0 if opengl else 1.0),
35
+ ) # [hw, 3]
36
+
37
+ rays_d = camera_dirs @ pose[:3, :3].transpose(0, 1) # [hw, 3]
38
+ rays_o = pose[:3, 3].unsqueeze(0).expand_as(rays_d) # [hw, 3]
39
+
40
+ rays_o = rays_o.view(h, w, 3)
41
+ rays_d = safe_normalize(rays_d).view(h, w, 3)
42
+
43
+ return rays_o, rays_d
44
+
45
+ def orbit_camera_jitter(poses, strength=0.1):
46
+ # poses: [B, 4, 4], assume orbit camera in opengl format
47
+ # random orbital rotate
48
+
49
+ B = poses.shape[0]
50
+ rotvec_x = poses[:, :3, 1] * strength * np.pi * (torch.rand(B, 1, device=poses.device) * 2 - 1)
51
+ rotvec_y = poses[:, :3, 0] * strength * np.pi / 2 * (torch.rand(B, 1, device=poses.device) * 2 - 1)
52
+
53
+ rot = roma.rotvec_to_rotmat(rotvec_x) @ roma.rotvec_to_rotmat(rotvec_y)
54
+ R = rot @ poses[:, :3, :3]
55
+ T = rot @ poses[:, :3, 3:]
56
+
57
+ new_poses = poses.clone()
58
+ new_poses[:, :3, :3] = R
59
+ new_poses[:, :3, 3:] = T
60
+
61
+ return new_poses
62
+
63
+ def grid_distortion(images, strength=0.5):
64
+ # images: [B, C, H, W]
65
+ # num_steps: int, grid resolution for distortion
66
+ # strength: float in [0, 1], strength of distortion
67
+
68
+ B, C, H, W = images.shape
69
+
70
+ num_steps = np.random.randint(8, 17)
71
+ grid_steps = torch.linspace(-1, 1, num_steps)
72
+
73
+ # have to loop batch...
74
+ grids = []
75
+ for b in range(B):
76
+ # construct displacement
77
+ x_steps = torch.linspace(0, 1, num_steps) # [num_steps], inclusive
78
+ x_steps = (x_steps + strength * (torch.rand_like(x_steps) - 0.5) / (num_steps - 1)).clamp(0, 1) # perturb
79
+ x_steps = (x_steps * W).long() # [num_steps]
80
+ x_steps[0] = 0
81
+ x_steps[-1] = W
82
+ xs = []
83
+ for i in range(num_steps - 1):
84
+ xs.append(torch.linspace(grid_steps[i], grid_steps[i + 1], x_steps[i + 1] - x_steps[i]))
85
+ xs = torch.cat(xs, dim=0) # [W]
86
+
87
+ y_steps = torch.linspace(0, 1, num_steps) # [num_steps], inclusive
88
+ y_steps = (y_steps + strength * (torch.rand_like(y_steps) - 0.5) / (num_steps - 1)).clamp(0, 1) # perturb
89
+ y_steps = (y_steps * H).long() # [num_steps]
90
+ y_steps[0] = 0
91
+ y_steps[-1] = H
92
+ ys = []
93
+ for i in range(num_steps - 1):
94
+ ys.append(torch.linspace(grid_steps[i], grid_steps[i + 1], y_steps[i + 1] - y_steps[i]))
95
+ ys = torch.cat(ys, dim=0) # [H]
96
+
97
+ # construct grid
98
+ grid_x, grid_y = torch.meshgrid(xs, ys, indexing='xy') # [H, W]
99
+ grid = torch.stack([grid_x, grid_y], dim=-1) # [H, W, 2]
100
+
101
+ grids.append(grid)
102
+
103
+ grids = torch.stack(grids, dim=0).to(images.device) # [B, H, W, 2]
104
+
105
+ # grid sample
106
+ images = F.grid_sample(images, grids, align_corners=False)
107
+
108
+ return images
lgm/infer.py ADDED
@@ -0,0 +1,226 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import os
3
+ import tyro
4
+ import glob
5
+ import imageio
6
+ import numpy as np
7
+ import tqdm
8
+ import torch
9
+ import torch.nn as nn
10
+ import torch.nn.functional as F
11
+ import torchvision.transforms.functional as TF
12
+ from safetensors.torch import load_file
13
+ import rembg
14
+
15
+ import kiui
16
+ from kiui.op import recenter
17
+ from kiui.cam import orbit_camera
18
+
19
+ from core.options import AllConfigs, Options
20
+ from core.models import LGM
21
+ from mvdream.pipeline_mvdream import MVDreamPipeline
22
+ import cv2
23
+
24
+ IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
25
+ IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
26
+
27
+ opt = tyro.cli(AllConfigs)
28
+
29
+ # model
30
+ model = LGM(opt)
31
+
32
+ # resume pretrained checkpoint
33
+ if opt.resume is not None:
34
+ if opt.resume.endswith('safetensors'):
35
+ ckpt = load_file(opt.resume, device='cpu')
36
+ else:
37
+ ckpt = torch.load(opt.resume, map_location='cpu')
38
+ model.load_state_dict(ckpt, strict=False)
39
+ print(f'[INFO] Loaded checkpoint from {opt.resume}')
40
+ else:
41
+ print(f'[WARN] model randomly initialized, are you sure?')
42
+
43
+ # device
44
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
45
+ model = model.half().to(device)
46
+ model.eval()
47
+
48
+ rays_embeddings = model.prepare_default_rays(device)
49
+
50
+ tan_half_fov = np.tan(0.5 * np.deg2rad(opt.fovy))
51
+ proj_matrix = torch.zeros(4, 4, dtype=torch.float32, device=device)
52
+ proj_matrix[0, 0] = 1 / tan_half_fov
53
+ proj_matrix[1, 1] = 1 / tan_half_fov
54
+ proj_matrix[2, 2] = (opt.zfar + opt.znear) / (opt.zfar - opt.znear)
55
+ proj_matrix[3, 2] = - (opt.zfar * opt.znear) / (opt.zfar - opt.znear)
56
+ proj_matrix[2, 3] = 1
57
+
58
+ # load image dream
59
+ pipe = MVDreamPipeline.from_pretrained(
60
+ "ashawkey/imagedream-ipmv-diffusers", # remote weights
61
+ torch_dtype=torch.float16,
62
+ trust_remote_code=True,
63
+ # local_files_only=True,
64
+ )
65
+ pipe = pipe.to(device)
66
+
67
+ # load rembg
68
+ bg_remover = rembg.new_session()
69
+
70
+ # process function
71
+ def process(opt: Options, path):
72
+ name = os.path.splitext(os.path.basename(path))[0]
73
+ if 'CONSISTENT4D' in path:
74
+ name = path.split('/')[-2]
75
+ print(f'[INFO] Processing {path} --> {name}')
76
+ os.makedirs('vis_data', exist_ok=True)
77
+ os.makedirs('logs', exist_ok=True)
78
+
79
+ input_image = kiui.read_image(path, mode='uint8')
80
+
81
+ # bg removal
82
+ carved_image = rembg.remove(input_image, session=bg_remover) # [H, W, 4]
83
+ mask = carved_image[..., -1] > 0
84
+
85
+ # recenter
86
+ image = recenter(carved_image, mask, border_ratio=0.2)
87
+
88
+ # generate mv
89
+ image = image.astype(np.float32) / 255.0
90
+
91
+ # rgba to rgb white bg
92
+ if image.shape[-1] == 4:
93
+ image = image[..., :3] * image[..., 3:4] + (1 - image[..., 3:4])
94
+
95
+ mv_image = pipe('', image, guidance_scale=5.0, num_inference_steps=30, elevation=0)
96
+ mv_image = np.stack([mv_image[1], mv_image[2], mv_image[3], mv_image[0]], axis=0) # [4, 256, 256, 3], float32
97
+
98
+ # generate gaussians
99
+ input_image = torch.from_numpy(mv_image).permute(0, 3, 1, 2).float().to(device) # [4, 3, 256, 256]
100
+ input_image = F.interpolate(input_image, size=(opt.input_size, opt.input_size), mode='bilinear', align_corners=False)
101
+ input_image = TF.normalize(input_image, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
102
+
103
+ input_image = torch.cat([input_image, rays_embeddings], dim=1).unsqueeze(0) # [1, 4, 9, H, W]
104
+
105
+ with torch.no_grad():
106
+ ############## align azimuth #####################
107
+ with torch.autocast(device_type='cuda', dtype=torch.float16):
108
+ # generate gaussians
109
+ gaussians = model.forward_gaussians(input_image)
110
+
111
+ best_azi = 0
112
+ best_diff = 1e8
113
+ for v, azi in enumerate(np.arange(-180, 180, 1)):
114
+ cam_poses = torch.from_numpy(orbit_camera(0, azi, radius=opt.cam_radius, opengl=True)).unsqueeze(0).to(device)
115
+
116
+ cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
117
+
118
+ # cameras needed by gaussian rasterizer
119
+ cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
120
+ cam_view_proj = cam_view @ proj_matrix # [V, 4, 4]
121
+ cam_pos = - cam_poses[:, :3, 3] # [V, 3]
122
+
123
+ # scale = min(azi / 360, 1)
124
+ scale = 1
125
+
126
+
127
+ result = model.gs.render(gaussians, cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0), scale_modifier=scale)
128
+ rendered_image = result['image']
129
+
130
+ rendered_image = rendered_image.squeeze(1).permute(0,2,3,1).squeeze(0).contiguous().float().cpu().numpy()
131
+ rendered_image = cv2.resize(rendered_image, (image.shape[0], image.shape[1]), interpolation=cv2.INTER_AREA)
132
+
133
+ diff = np.mean((rendered_image- image) ** 2)
134
+
135
+ if diff < best_diff:
136
+ best_diff = diff
137
+ best_azi = azi
138
+ print("Best aligned azimuth: ", best_azi)
139
+
140
+ mv_image = []
141
+ for v, azi in enumerate([0, 90, 180, 270]):
142
+ cam_poses = torch.from_numpy(orbit_camera(0, azi + best_azi, radius=opt.cam_radius, opengl=True)).unsqueeze(0).to(device)
143
+
144
+ cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
145
+
146
+ # cameras needed by gaussian rasterizer
147
+ cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
148
+ cam_view_proj = cam_view @ proj_matrix # [V, 4, 4]
149
+ cam_pos = - cam_poses[:, :3, 3] # [V, 3]
150
+
151
+ # scale = min(azi / 360, 1)
152
+ scale = 1
153
+
154
+
155
+ result = model.gs.render(gaussians, cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0), scale_modifier=scale)
156
+ rendered_image = result['image']
157
+ rendered_image = rendered_image.squeeze(1)
158
+ rendered_image = F.interpolate(rendered_image, (256, 256))
159
+ rendered_image = rendered_image.permute(0,2,3,1).contiguous().float().cpu().numpy()
160
+ mv_image.append(rendered_image)
161
+ mv_image = np.concatenate(mv_image, axis=0)
162
+
163
+ input_image = torch.from_numpy(mv_image).permute(0, 3, 1, 2).float().to(device) # [4, 3, 256, 256]
164
+ input_image = F.interpolate(input_image, size=(opt.input_size, opt.input_size), mode='bilinear', align_corners=False)
165
+ input_image = TF.normalize(input_image, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
166
+
167
+ input_image = torch.cat([input_image, rays_embeddings], dim=1).unsqueeze(0) # [1, 4, 9, H, W]
168
+
169
+ ################################
170
+
171
+ with torch.autocast(device_type='cuda', dtype=torch.float16):
172
+ # generate gaussians
173
+ gaussians = model.forward_gaussians(input_image)
174
+
175
+ # save gaussians
176
+ model.gs.save_ply(gaussians, os.path.join('logs', name + '_model.ply'))
177
+
178
+ # render 360 video
179
+ images = []
180
+ elevation = 0
181
+
182
+ if opt.fancy_video:
183
+
184
+ azimuth = np.arange(0, 720, 4, dtype=np.int32)
185
+ for azi in tqdm.tqdm(azimuth):
186
+
187
+ cam_poses = torch.from_numpy(orbit_camera(elevation, azi, radius=opt.cam_radius, opengl=True)).unsqueeze(0).to(device)
188
+
189
+ cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
190
+
191
+ # cameras needed by gaussian rasterizer
192
+ cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
193
+ cam_view_proj = cam_view @ proj_matrix # [V, 4, 4]
194
+ cam_pos = - cam_poses[:, :3, 3] # [V, 3]
195
+
196
+ scale = min(azi / 360, 1)
197
+
198
+ image = model.gs.render(gaussians, cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0), scale_modifier=scale)['image']
199
+ images.append((image.squeeze(1).permute(0,2,3,1).contiguous().float().cpu().numpy() * 255).astype(np.uint8))
200
+ else:
201
+ azimuth = np.arange(0, 360, 2, dtype=np.int32)
202
+ for azi in tqdm.tqdm(azimuth):
203
+
204
+ cam_poses = torch.from_numpy(orbit_camera(elevation, azi, radius=opt.cam_radius, opengl=True)).unsqueeze(0).to(device)
205
+
206
+ cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
207
+
208
+ # cameras needed by gaussian rasterizer
209
+ cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
210
+ cam_view_proj = cam_view @ proj_matrix # [V, 4, 4]
211
+ cam_pos = - cam_poses[:, :3, 3] # [V, 3]
212
+
213
+ image = model.gs.render(gaussians, cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0), scale_modifier=1)['image']
214
+ images.append((image.squeeze(1).permute(0,2,3,1).contiguous().float().cpu().numpy() * 255).astype(np.uint8))
215
+
216
+ images = np.concatenate(images, axis=0)
217
+ imageio.mimwrite(os.path.join('vis_data', name + '_static.mp4'), images, fps=30)
218
+
219
+
220
+ assert opt.test_path is not None
221
+ if os.path.isdir(opt.test_path):
222
+ file_paths = glob.glob(os.path.join(opt.test_path, "*"))
223
+ else:
224
+ file_paths = [opt.test_path]
225
+ for path in file_paths:
226
+ process(opt, path)
lgm/mvdream/mv_unet.py ADDED
@@ -0,0 +1,1005 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import numpy as np
3
+ from inspect import isfunction
4
+ from typing import Optional, Any, List
5
+
6
+ import torch
7
+ import torch.nn as nn
8
+ import torch.nn.functional as F
9
+ from einops import rearrange, repeat
10
+
11
+ from diffusers.configuration_utils import ConfigMixin
12
+ from diffusers.models.modeling_utils import ModelMixin
13
+
14
+ # require xformers!
15
+ import xformers
16
+ import xformers.ops
17
+
18
+ from kiui.cam import orbit_camera
19
+
20
+ def get_camera(
21
+ num_frames, elevation=0, azimuth_start=0, azimuth_span=360, blender_coord=True, extra_view=False,
22
+ ):
23
+ angle_gap = azimuth_span / num_frames
24
+ cameras = []
25
+ for azimuth in np.arange(azimuth_start, azimuth_span + azimuth_start, angle_gap):
26
+
27
+ pose = orbit_camera(elevation, azimuth, radius=1) # [4, 4]
28
+
29
+ # opengl to blender
30
+ if blender_coord:
31
+ pose[2] *= -1
32
+ pose[[1, 2]] = pose[[2, 1]]
33
+
34
+ cameras.append(pose.flatten())
35
+
36
+ if extra_view:
37
+ cameras.append(np.zeros_like(cameras[0]))
38
+
39
+ return torch.from_numpy(np.stack(cameras, axis=0)).float() # [num_frames, 16]
40
+
41
+
42
+ def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
43
+ """
44
+ Create sinusoidal timestep embeddings.
45
+ :param timesteps: a 1-D Tensor of N indices, one per batch element.
46
+ These may be fractional.
47
+ :param dim: the dimension of the output.
48
+ :param max_period: controls the minimum frequency of the embeddings.
49
+ :return: an [N x dim] Tensor of positional embeddings.
50
+ """
51
+ if not repeat_only:
52
+ half = dim // 2
53
+ freqs = torch.exp(
54
+ -math.log(max_period)
55
+ * torch.arange(start=0, end=half, dtype=torch.float32)
56
+ / half
57
+ ).to(device=timesteps.device)
58
+ args = timesteps[:, None] * freqs[None]
59
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
60
+ if dim % 2:
61
+ embedding = torch.cat(
62
+ [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
63
+ )
64
+ else:
65
+ embedding = repeat(timesteps, "b -> b d", d=dim)
66
+ # import pdb; pdb.set_trace()
67
+ return embedding
68
+
69
+
70
+ def zero_module(module):
71
+ """
72
+ Zero out the parameters of a module and return it.
73
+ """
74
+ for p in module.parameters():
75
+ p.detach().zero_()
76
+ return module
77
+
78
+
79
+ def conv_nd(dims, *args, **kwargs):
80
+ """
81
+ Create a 1D, 2D, or 3D convolution module.
82
+ """
83
+ if dims == 1:
84
+ return nn.Conv1d(*args, **kwargs)
85
+ elif dims == 2:
86
+ return nn.Conv2d(*args, **kwargs)
87
+ elif dims == 3:
88
+ return nn.Conv3d(*args, **kwargs)
89
+ raise ValueError(f"unsupported dimensions: {dims}")
90
+
91
+
92
+ def avg_pool_nd(dims, *args, **kwargs):
93
+ """
94
+ Create a 1D, 2D, or 3D average pooling module.
95
+ """
96
+ if dims == 1:
97
+ return nn.AvgPool1d(*args, **kwargs)
98
+ elif dims == 2:
99
+ return nn.AvgPool2d(*args, **kwargs)
100
+ elif dims == 3:
101
+ return nn.AvgPool3d(*args, **kwargs)
102
+ raise ValueError(f"unsupported dimensions: {dims}")
103
+
104
+
105
+ def default(val, d):
106
+ if val is not None:
107
+ return val
108
+ return d() if isfunction(d) else d
109
+
110
+
111
+ class GEGLU(nn.Module):
112
+ def __init__(self, dim_in, dim_out):
113
+ super().__init__()
114
+ self.proj = nn.Linear(dim_in, dim_out * 2)
115
+
116
+ def forward(self, x):
117
+ x, gate = self.proj(x).chunk(2, dim=-1)
118
+ return x * F.gelu(gate)
119
+
120
+
121
+ class FeedForward(nn.Module):
122
+ def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):
123
+ super().__init__()
124
+ inner_dim = int(dim * mult)
125
+ dim_out = default(dim_out, dim)
126
+ project_in = (
127
+ nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())
128
+ if not glu
129
+ else GEGLU(dim, inner_dim)
130
+ )
131
+
132
+ self.net = nn.Sequential(
133
+ project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)
134
+ )
135
+
136
+ def forward(self, x):
137
+ return self.net(x)
138
+
139
+
140
+ class MemoryEfficientCrossAttention(nn.Module):
141
+ # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
142
+ def __init__(
143
+ self,
144
+ query_dim,
145
+ context_dim=None,
146
+ heads=8,
147
+ dim_head=64,
148
+ dropout=0.0,
149
+ ip_dim=0,
150
+ ip_weight=1,
151
+ ):
152
+ super().__init__()
153
+
154
+ inner_dim = dim_head * heads
155
+ context_dim = default(context_dim, query_dim)
156
+
157
+ self.heads = heads
158
+ self.dim_head = dim_head
159
+
160
+ self.ip_dim = ip_dim
161
+ self.ip_weight = ip_weight
162
+
163
+ if self.ip_dim > 0:
164
+ self.to_k_ip = nn.Linear(context_dim, inner_dim, bias=False)
165
+ self.to_v_ip = nn.Linear(context_dim, inner_dim, bias=False)
166
+
167
+ self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
168
+ self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
169
+ self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
170
+
171
+ self.to_out = nn.Sequential(
172
+ nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)
173
+ )
174
+ self.attention_op: Optional[Any] = None
175
+
176
+ def forward(self, x, context=None):
177
+ q = self.to_q(x)
178
+ context = default(context, x)
179
+
180
+ if self.ip_dim > 0:
181
+ # context: [B, 77 + 16(ip), 1024]
182
+ token_len = context.shape[1]
183
+ context_ip = context[:, -self.ip_dim :, :]
184
+ k_ip = self.to_k_ip(context_ip)
185
+ v_ip = self.to_v_ip(context_ip)
186
+ context = context[:, : (token_len - self.ip_dim), :]
187
+
188
+ k = self.to_k(context)
189
+ v = self.to_v(context)
190
+
191
+ b, _, _ = q.shape
192
+ q, k, v = map(
193
+ lambda t: t.unsqueeze(3)
194
+ .reshape(b, t.shape[1], self.heads, self.dim_head)
195
+ .permute(0, 2, 1, 3)
196
+ .reshape(b * self.heads, t.shape[1], self.dim_head)
197
+ .contiguous(),
198
+ (q, k, v),
199
+ )
200
+
201
+ # actually compute the attention, what we cannot get enough of
202
+ out = xformers.ops.memory_efficient_attention(
203
+ q, k, v, attn_bias=None, op=self.attention_op
204
+ )
205
+
206
+ if self.ip_dim > 0:
207
+ k_ip, v_ip = map(
208
+ lambda t: t.unsqueeze(3)
209
+ .reshape(b, t.shape[1], self.heads, self.dim_head)
210
+ .permute(0, 2, 1, 3)
211
+ .reshape(b * self.heads, t.shape[1], self.dim_head)
212
+ .contiguous(),
213
+ (k_ip, v_ip),
214
+ )
215
+ # actually compute the attention, what we cannot get enough of
216
+ out_ip = xformers.ops.memory_efficient_attention(
217
+ q, k_ip, v_ip, attn_bias=None, op=self.attention_op
218
+ )
219
+ out = out + self.ip_weight * out_ip
220
+
221
+ out = (
222
+ out.unsqueeze(0)
223
+ .reshape(b, self.heads, out.shape[1], self.dim_head)
224
+ .permute(0, 2, 1, 3)
225
+ .reshape(b, out.shape[1], self.heads * self.dim_head)
226
+ )
227
+ return self.to_out(out)
228
+
229
+
230
+ class BasicTransformerBlock3D(nn.Module):
231
+
232
+ def __init__(
233
+ self,
234
+ dim,
235
+ n_heads,
236
+ d_head,
237
+ context_dim,
238
+ dropout=0.0,
239
+ gated_ff=True,
240
+ ip_dim=0,
241
+ ip_weight=1,
242
+ ):
243
+ super().__init__()
244
+
245
+ self.attn1 = MemoryEfficientCrossAttention(
246
+ query_dim=dim,
247
+ context_dim=None, # self-attention
248
+ heads=n_heads,
249
+ dim_head=d_head,
250
+ dropout=dropout,
251
+ )
252
+ self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
253
+ self.attn2 = MemoryEfficientCrossAttention(
254
+ query_dim=dim,
255
+ context_dim=context_dim,
256
+ heads=n_heads,
257
+ dim_head=d_head,
258
+ dropout=dropout,
259
+ # ip only applies to cross-attention
260
+ ip_dim=ip_dim,
261
+ ip_weight=ip_weight,
262
+ )
263
+ self.norm1 = nn.LayerNorm(dim)
264
+ self.norm2 = nn.LayerNorm(dim)
265
+ self.norm3 = nn.LayerNorm(dim)
266
+
267
+ def forward(self, x, context=None, num_frames=1):
268
+ x = rearrange(x, "(b f) l c -> b (f l) c", f=num_frames).contiguous()
269
+ x = self.attn1(self.norm1(x), context=None) + x
270
+ x = rearrange(x, "b (f l) c -> (b f) l c", f=num_frames).contiguous()
271
+ x = self.attn2(self.norm2(x), context=context) + x
272
+ x = self.ff(self.norm3(x)) + x
273
+ return x
274
+
275
+
276
+ class SpatialTransformer3D(nn.Module):
277
+
278
+ def __init__(
279
+ self,
280
+ in_channels,
281
+ n_heads,
282
+ d_head,
283
+ context_dim, # cross attention input dim
284
+ depth=1,
285
+ dropout=0.0,
286
+ ip_dim=0,
287
+ ip_weight=1,
288
+ ):
289
+ super().__init__()
290
+
291
+ if not isinstance(context_dim, list):
292
+ context_dim = [context_dim]
293
+
294
+ self.in_channels = in_channels
295
+
296
+ inner_dim = n_heads * d_head
297
+ self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
298
+ self.proj_in = nn.Linear(in_channels, inner_dim)
299
+
300
+ self.transformer_blocks = nn.ModuleList(
301
+ [
302
+ BasicTransformerBlock3D(
303
+ inner_dim,
304
+ n_heads,
305
+ d_head,
306
+ context_dim=context_dim[d],
307
+ dropout=dropout,
308
+ ip_dim=ip_dim,
309
+ ip_weight=ip_weight,
310
+ )
311
+ for d in range(depth)
312
+ ]
313
+ )
314
+
315
+ self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
316
+
317
+
318
+ def forward(self, x, context=None, num_frames=1):
319
+ # note: if no context is given, cross-attention defaults to self-attention
320
+ if not isinstance(context, list):
321
+ context = [context]
322
+ b, c, h, w = x.shape
323
+ x_in = x
324
+ x = self.norm(x)
325
+ x = rearrange(x, "b c h w -> b (h w) c").contiguous()
326
+ x = self.proj_in(x)
327
+ for i, block in enumerate(self.transformer_blocks):
328
+ x = block(x, context=context[i], num_frames=num_frames)
329
+ x = self.proj_out(x)
330
+ x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
331
+
332
+ return x + x_in
333
+
334
+
335
+ class PerceiverAttention(nn.Module):
336
+ def __init__(self, *, dim, dim_head=64, heads=8):
337
+ super().__init__()
338
+ self.scale = dim_head ** -0.5
339
+ self.dim_head = dim_head
340
+ self.heads = heads
341
+ inner_dim = dim_head * heads
342
+
343
+ self.norm1 = nn.LayerNorm(dim)
344
+ self.norm2 = nn.LayerNorm(dim)
345
+
346
+ self.to_q = nn.Linear(dim, inner_dim, bias=False)
347
+ self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
348
+ self.to_out = nn.Linear(inner_dim, dim, bias=False)
349
+
350
+ def forward(self, x, latents):
351
+ """
352
+ Args:
353
+ x (torch.Tensor): image features
354
+ shape (b, n1, D)
355
+ latent (torch.Tensor): latent features
356
+ shape (b, n2, D)
357
+ """
358
+ x = self.norm1(x)
359
+ latents = self.norm2(latents)
360
+
361
+ b, l, _ = latents.shape
362
+
363
+ q = self.to_q(latents)
364
+ kv_input = torch.cat((x, latents), dim=-2)
365
+ k, v = self.to_kv(kv_input).chunk(2, dim=-1)
366
+
367
+ q, k, v = map(
368
+ lambda t: t.reshape(b, t.shape[1], self.heads, -1)
369
+ .transpose(1, 2)
370
+ .reshape(b, self.heads, t.shape[1], -1)
371
+ .contiguous(),
372
+ (q, k, v),
373
+ )
374
+
375
+ # attention
376
+ scale = 1 / math.sqrt(math.sqrt(self.dim_head))
377
+ weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards
378
+ weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
379
+ out = weight @ v
380
+
381
+ out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
382
+
383
+ return self.to_out(out)
384
+
385
+
386
+ class Resampler(nn.Module):
387
+ def __init__(
388
+ self,
389
+ dim=1024,
390
+ depth=8,
391
+ dim_head=64,
392
+ heads=16,
393
+ num_queries=8,
394
+ embedding_dim=768,
395
+ output_dim=1024,
396
+ ff_mult=4,
397
+ ):
398
+ super().__init__()
399
+ self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim ** 0.5)
400
+ self.proj_in = nn.Linear(embedding_dim, dim)
401
+ self.proj_out = nn.Linear(dim, output_dim)
402
+ self.norm_out = nn.LayerNorm(output_dim)
403
+
404
+ self.layers = nn.ModuleList([])
405
+ for _ in range(depth):
406
+ self.layers.append(
407
+ nn.ModuleList(
408
+ [
409
+ PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
410
+ nn.Sequential(
411
+ nn.LayerNorm(dim),
412
+ nn.Linear(dim, dim * ff_mult, bias=False),
413
+ nn.GELU(),
414
+ nn.Linear(dim * ff_mult, dim, bias=False),
415
+ )
416
+ ]
417
+ )
418
+ )
419
+
420
+ def forward(self, x):
421
+ latents = self.latents.repeat(x.size(0), 1, 1)
422
+ x = self.proj_in(x)
423
+ for attn, ff in self.layers:
424
+ latents = attn(x, latents) + latents
425
+ latents = ff(latents) + latents
426
+
427
+ latents = self.proj_out(latents)
428
+ return self.norm_out(latents)
429
+
430
+
431
+ class CondSequential(nn.Sequential):
432
+ """
433
+ A sequential module that passes timestep embeddings to the children that
434
+ support it as an extra input.
435
+ """
436
+
437
+ def forward(self, x, emb, context=None, num_frames=1):
438
+ for layer in self:
439
+ if isinstance(layer, ResBlock):
440
+ x = layer(x, emb)
441
+ elif isinstance(layer, SpatialTransformer3D):
442
+ x = layer(x, context, num_frames=num_frames)
443
+ else:
444
+ x = layer(x)
445
+ return x
446
+
447
+
448
+ class Upsample(nn.Module):
449
+ """
450
+ An upsampling layer with an optional convolution.
451
+ :param channels: channels in the inputs and outputs.
452
+ :param use_conv: a bool determining if a convolution is applied.
453
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
454
+ upsampling occurs in the inner-two dimensions.
455
+ """
456
+
457
+ def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
458
+ super().__init__()
459
+ self.channels = channels
460
+ self.out_channels = out_channels or channels
461
+ self.use_conv = use_conv
462
+ self.dims = dims
463
+ if use_conv:
464
+ self.conv = conv_nd(
465
+ dims, self.channels, self.out_channels, 3, padding=padding
466
+ )
467
+
468
+ def forward(self, x):
469
+ assert x.shape[1] == self.channels
470
+ if self.dims == 3:
471
+ x = F.interpolate(
472
+ x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
473
+ )
474
+ else:
475
+ x = F.interpolate(x, scale_factor=2, mode="nearest")
476
+ if self.use_conv:
477
+ x = self.conv(x)
478
+ return x
479
+
480
+
481
+ class Downsample(nn.Module):
482
+ """
483
+ A downsampling layer with an optional convolution.
484
+ :param channels: channels in the inputs and outputs.
485
+ :param use_conv: a bool determining if a convolution is applied.
486
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
487
+ downsampling occurs in the inner-two dimensions.
488
+ """
489
+
490
+ def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
491
+ super().__init__()
492
+ self.channels = channels
493
+ self.out_channels = out_channels or channels
494
+ self.use_conv = use_conv
495
+ self.dims = dims
496
+ stride = 2 if dims != 3 else (1, 2, 2)
497
+ if use_conv:
498
+ self.op = conv_nd(
499
+ dims,
500
+ self.channels,
501
+ self.out_channels,
502
+ 3,
503
+ stride=stride,
504
+ padding=padding,
505
+ )
506
+ else:
507
+ assert self.channels == self.out_channels
508
+ self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
509
+
510
+ def forward(self, x):
511
+ assert x.shape[1] == self.channels
512
+ return self.op(x)
513
+
514
+
515
+ class ResBlock(nn.Module):
516
+ """
517
+ A residual block that can optionally change the number of channels.
518
+ :param channels: the number of input channels.
519
+ :param emb_channels: the number of timestep embedding channels.
520
+ :param dropout: the rate of dropout.
521
+ :param out_channels: if specified, the number of out channels.
522
+ :param use_conv: if True and out_channels is specified, use a spatial
523
+ convolution instead of a smaller 1x1 convolution to change the
524
+ channels in the skip connection.
525
+ :param dims: determines if the signal is 1D, 2D, or 3D.
526
+ :param up: if True, use this block for upsampling.
527
+ :param down: if True, use this block for downsampling.
528
+ """
529
+
530
+ def __init__(
531
+ self,
532
+ channels,
533
+ emb_channels,
534
+ dropout,
535
+ out_channels=None,
536
+ use_conv=False,
537
+ use_scale_shift_norm=False,
538
+ dims=2,
539
+ up=False,
540
+ down=False,
541
+ ):
542
+ super().__init__()
543
+ self.channels = channels
544
+ self.emb_channels = emb_channels
545
+ self.dropout = dropout
546
+ self.out_channels = out_channels or channels
547
+ self.use_conv = use_conv
548
+ self.use_scale_shift_norm = use_scale_shift_norm
549
+
550
+ self.in_layers = nn.Sequential(
551
+ nn.GroupNorm(32, channels),
552
+ nn.SiLU(),
553
+ conv_nd(dims, channels, self.out_channels, 3, padding=1),
554
+ )
555
+
556
+ self.updown = up or down
557
+
558
+ if up:
559
+ self.h_upd = Upsample(channels, False, dims)
560
+ self.x_upd = Upsample(channels, False, dims)
561
+ elif down:
562
+ self.h_upd = Downsample(channels, False, dims)
563
+ self.x_upd = Downsample(channels, False, dims)
564
+ else:
565
+ self.h_upd = self.x_upd = nn.Identity()
566
+
567
+ self.emb_layers = nn.Sequential(
568
+ nn.SiLU(),
569
+ nn.Linear(
570
+ emb_channels,
571
+ 2 * self.out_channels if use_scale_shift_norm else self.out_channels,
572
+ ),
573
+ )
574
+ self.out_layers = nn.Sequential(
575
+ nn.GroupNorm(32, self.out_channels),
576
+ nn.SiLU(),
577
+ nn.Dropout(p=dropout),
578
+ zero_module(
579
+ conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
580
+ ),
581
+ )
582
+
583
+ if self.out_channels == channels:
584
+ self.skip_connection = nn.Identity()
585
+ elif use_conv:
586
+ self.skip_connection = conv_nd(
587
+ dims, channels, self.out_channels, 3, padding=1
588
+ )
589
+ else:
590
+ self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
591
+
592
+ def forward(self, x, emb):
593
+ if self.updown:
594
+ in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
595
+ h = in_rest(x)
596
+ h = self.h_upd(h)
597
+ x = self.x_upd(x)
598
+ h = in_conv(h)
599
+ else:
600
+ h = self.in_layers(x)
601
+ emb_out = self.emb_layers(emb).type(h.dtype)
602
+ while len(emb_out.shape) < len(h.shape):
603
+ emb_out = emb_out[..., None]
604
+ if self.use_scale_shift_norm:
605
+ out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
606
+ scale, shift = torch.chunk(emb_out, 2, dim=1)
607
+ h = out_norm(h) * (1 + scale) + shift
608
+ h = out_rest(h)
609
+ else:
610
+ h = h + emb_out
611
+ h = self.out_layers(h)
612
+ return self.skip_connection(x) + h
613
+
614
+
615
+ class MultiViewUNetModel(ModelMixin, ConfigMixin):
616
+ """
617
+ The full multi-view UNet model with attention, timestep embedding and camera embedding.
618
+ :param in_channels: channels in the input Tensor.
619
+ :param model_channels: base channel count for the model.
620
+ :param out_channels: channels in the output Tensor.
621
+ :param num_res_blocks: number of residual blocks per downsample.
622
+ :param attention_resolutions: a collection of downsample rates at which
623
+ attention will take place. May be a set, list, or tuple.
624
+ For example, if this contains 4, then at 4x downsampling, attention
625
+ will be used.
626
+ :param dropout: the dropout probability.
627
+ :param channel_mult: channel multiplier for each level of the UNet.
628
+ :param conv_resample: if True, use learned convolutions for upsampling and
629
+ downsampling.
630
+ :param dims: determines if the signal is 1D, 2D, or 3D.
631
+ :param num_classes: if specified (as an int), then this model will be
632
+ class-conditional with `num_classes` classes.
633
+ :param num_heads: the number of attention heads in each attention layer.
634
+ :param num_heads_channels: if specified, ignore num_heads and instead use
635
+ a fixed channel width per attention head.
636
+ :param num_heads_upsample: works with num_heads to set a different number
637
+ of heads for upsampling. Deprecated.
638
+ :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
639
+ :param resblock_updown: use residual blocks for up/downsampling.
640
+ :param use_new_attention_order: use a different attention pattern for potentially
641
+ increased efficiency.
642
+ :param camera_dim: dimensionality of camera input.
643
+ """
644
+
645
+ def __init__(
646
+ self,
647
+ image_size,
648
+ in_channels,
649
+ model_channels,
650
+ out_channels,
651
+ num_res_blocks,
652
+ attention_resolutions,
653
+ dropout=0,
654
+ channel_mult=(1, 2, 4, 8),
655
+ conv_resample=True,
656
+ dims=2,
657
+ num_classes=None,
658
+ num_heads=-1,
659
+ num_head_channels=-1,
660
+ num_heads_upsample=-1,
661
+ use_scale_shift_norm=False,
662
+ resblock_updown=False,
663
+ transformer_depth=1,
664
+ context_dim=None,
665
+ n_embed=None,
666
+ num_attention_blocks=None,
667
+ adm_in_channels=None,
668
+ camera_dim=None,
669
+ ip_dim=0, # imagedream uses ip_dim > 0
670
+ ip_weight=1.0,
671
+ **kwargs,
672
+ ):
673
+ super().__init__()
674
+ assert context_dim is not None
675
+
676
+ if num_heads_upsample == -1:
677
+ num_heads_upsample = num_heads
678
+
679
+ if num_heads == -1:
680
+ assert (
681
+ num_head_channels != -1
682
+ ), "Either num_heads or num_head_channels has to be set"
683
+
684
+ if num_head_channels == -1:
685
+ assert (
686
+ num_heads != -1
687
+ ), "Either num_heads or num_head_channels has to be set"
688
+
689
+ self.image_size = image_size
690
+ self.in_channels = in_channels
691
+ self.model_channels = model_channels
692
+ self.out_channels = out_channels
693
+ if isinstance(num_res_blocks, int):
694
+ self.num_res_blocks = len(channel_mult) * [num_res_blocks]
695
+ else:
696
+ if len(num_res_blocks) != len(channel_mult):
697
+ raise ValueError(
698
+ "provide num_res_blocks either as an int (globally constant) or "
699
+ "as a list/tuple (per-level) with the same length as channel_mult"
700
+ )
701
+ self.num_res_blocks = num_res_blocks
702
+
703
+ if num_attention_blocks is not None:
704
+ assert len(num_attention_blocks) == len(self.num_res_blocks)
705
+ assert all(
706
+ map(
707
+ lambda i: self.num_res_blocks[i] >= num_attention_blocks[i],
708
+ range(len(num_attention_blocks)),
709
+ )
710
+ )
711
+ print(
712
+ f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
713
+ f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
714
+ f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
715
+ f"attention will still not be set."
716
+ )
717
+
718
+ self.attention_resolutions = attention_resolutions
719
+ self.dropout = dropout
720
+ self.channel_mult = channel_mult
721
+ self.conv_resample = conv_resample
722
+ self.num_classes = num_classes
723
+ self.num_heads = num_heads
724
+ self.num_head_channels = num_head_channels
725
+ self.num_heads_upsample = num_heads_upsample
726
+ self.predict_codebook_ids = n_embed is not None
727
+
728
+ self.ip_dim = ip_dim
729
+ self.ip_weight = ip_weight
730
+
731
+ if self.ip_dim > 0:
732
+ self.image_embed = Resampler(
733
+ dim=context_dim,
734
+ depth=4,
735
+ dim_head=64,
736
+ heads=12,
737
+ num_queries=ip_dim, # num token
738
+ embedding_dim=1280,
739
+ output_dim=context_dim,
740
+ ff_mult=4,
741
+ )
742
+
743
+ time_embed_dim = model_channels * 4
744
+ self.time_embed = nn.Sequential(
745
+ nn.Linear(model_channels, time_embed_dim),
746
+ nn.SiLU(),
747
+ nn.Linear(time_embed_dim, time_embed_dim),
748
+ )
749
+
750
+ if camera_dim is not None:
751
+ time_embed_dim = model_channels * 4
752
+ self.camera_embed = nn.Sequential(
753
+ nn.Linear(camera_dim, time_embed_dim),
754
+ nn.SiLU(),
755
+ nn.Linear(time_embed_dim, time_embed_dim),
756
+ )
757
+
758
+ if self.num_classes is not None:
759
+ if isinstance(self.num_classes, int):
760
+ self.label_emb = nn.Embedding(self.num_classes, time_embed_dim)
761
+ elif self.num_classes == "continuous":
762
+ # print("setting up linear c_adm embedding layer")
763
+ self.label_emb = nn.Linear(1, time_embed_dim)
764
+ elif self.num_classes == "sequential":
765
+ assert adm_in_channels is not None
766
+ self.label_emb = nn.Sequential(
767
+ nn.Sequential(
768
+ nn.Linear(adm_in_channels, time_embed_dim),
769
+ nn.SiLU(),
770
+ nn.Linear(time_embed_dim, time_embed_dim),
771
+ )
772
+ )
773
+ else:
774
+ raise ValueError()
775
+
776
+ self.input_blocks = nn.ModuleList(
777
+ [
778
+ CondSequential(
779
+ conv_nd(dims, in_channels, model_channels, 3, padding=1)
780
+ )
781
+ ]
782
+ )
783
+ self._feature_size = model_channels
784
+ input_block_chans = [model_channels]
785
+ ch = model_channels
786
+ ds = 1
787
+ for level, mult in enumerate(channel_mult):
788
+ for nr in range(self.num_res_blocks[level]):
789
+ layers: List[Any] = [
790
+ ResBlock(
791
+ ch,
792
+ time_embed_dim,
793
+ dropout,
794
+ out_channels=mult * model_channels,
795
+ dims=dims,
796
+ use_scale_shift_norm=use_scale_shift_norm,
797
+ )
798
+ ]
799
+ ch = mult * model_channels
800
+ if ds in attention_resolutions:
801
+ if num_head_channels == -1:
802
+ dim_head = ch // num_heads
803
+ else:
804
+ num_heads = ch // num_head_channels
805
+ dim_head = num_head_channels
806
+
807
+ if num_attention_blocks is None or nr < num_attention_blocks[level]:
808
+ layers.append(
809
+ SpatialTransformer3D(
810
+ ch,
811
+ num_heads,
812
+ dim_head,
813
+ context_dim=context_dim,
814
+ depth=transformer_depth,
815
+ ip_dim=self.ip_dim,
816
+ ip_weight=self.ip_weight,
817
+ )
818
+ )
819
+ self.input_blocks.append(CondSequential(*layers))
820
+ self._feature_size += ch
821
+ input_block_chans.append(ch)
822
+ if level != len(channel_mult) - 1:
823
+ out_ch = ch
824
+ self.input_blocks.append(
825
+ CondSequential(
826
+ ResBlock(
827
+ ch,
828
+ time_embed_dim,
829
+ dropout,
830
+ out_channels=out_ch,
831
+ dims=dims,
832
+ use_scale_shift_norm=use_scale_shift_norm,
833
+ down=True,
834
+ )
835
+ if resblock_updown
836
+ else Downsample(
837
+ ch, conv_resample, dims=dims, out_channels=out_ch
838
+ )
839
+ )
840
+ )
841
+ ch = out_ch
842
+ input_block_chans.append(ch)
843
+ ds *= 2
844
+ self._feature_size += ch
845
+
846
+ if num_head_channels == -1:
847
+ dim_head = ch // num_heads
848
+ else:
849
+ num_heads = ch // num_head_channels
850
+ dim_head = num_head_channels
851
+
852
+ self.middle_block = CondSequential(
853
+ ResBlock(
854
+ ch,
855
+ time_embed_dim,
856
+ dropout,
857
+ dims=dims,
858
+ use_scale_shift_norm=use_scale_shift_norm,
859
+ ),
860
+ SpatialTransformer3D(
861
+ ch,
862
+ num_heads,
863
+ dim_head,
864
+ context_dim=context_dim,
865
+ depth=transformer_depth,
866
+ ip_dim=self.ip_dim,
867
+ ip_weight=self.ip_weight,
868
+ ),
869
+ ResBlock(
870
+ ch,
871
+ time_embed_dim,
872
+ dropout,
873
+ dims=dims,
874
+ use_scale_shift_norm=use_scale_shift_norm,
875
+ ),
876
+ )
877
+ self._feature_size += ch
878
+
879
+ self.output_blocks = nn.ModuleList([])
880
+ for level, mult in list(enumerate(channel_mult))[::-1]:
881
+ for i in range(self.num_res_blocks[level] + 1):
882
+ ich = input_block_chans.pop()
883
+ layers = [
884
+ ResBlock(
885
+ ch + ich,
886
+ time_embed_dim,
887
+ dropout,
888
+ out_channels=model_channels * mult,
889
+ dims=dims,
890
+ use_scale_shift_norm=use_scale_shift_norm,
891
+ )
892
+ ]
893
+ ch = model_channels * mult
894
+ if ds in attention_resolutions:
895
+ if num_head_channels == -1:
896
+ dim_head = ch // num_heads
897
+ else:
898
+ num_heads = ch // num_head_channels
899
+ dim_head = num_head_channels
900
+
901
+ if num_attention_blocks is None or i < num_attention_blocks[level]:
902
+ layers.append(
903
+ SpatialTransformer3D(
904
+ ch,
905
+ num_heads,
906
+ dim_head,
907
+ context_dim=context_dim,
908
+ depth=transformer_depth,
909
+ ip_dim=self.ip_dim,
910
+ ip_weight=self.ip_weight,
911
+ )
912
+ )
913
+ if level and i == self.num_res_blocks[level]:
914
+ out_ch = ch
915
+ layers.append(
916
+ ResBlock(
917
+ ch,
918
+ time_embed_dim,
919
+ dropout,
920
+ out_channels=out_ch,
921
+ dims=dims,
922
+ use_scale_shift_norm=use_scale_shift_norm,
923
+ up=True,
924
+ )
925
+ if resblock_updown
926
+ else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
927
+ )
928
+ ds //= 2
929
+ self.output_blocks.append(CondSequential(*layers))
930
+ self._feature_size += ch
931
+
932
+ self.out = nn.Sequential(
933
+ nn.GroupNorm(32, ch),
934
+ nn.SiLU(),
935
+ zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
936
+ )
937
+ if self.predict_codebook_ids:
938
+ self.id_predictor = nn.Sequential(
939
+ nn.GroupNorm(32, ch),
940
+ conv_nd(dims, model_channels, n_embed, 1),
941
+ # nn.LogSoftmax(dim=1) # change to cross_entropy and produce non-normalized logits
942
+ )
943
+
944
+ def forward(
945
+ self,
946
+ x,
947
+ timesteps=None,
948
+ context=None,
949
+ y=None,
950
+ camera=None,
951
+ num_frames=1,
952
+ ip=None,
953
+ ip_img=None,
954
+ **kwargs,
955
+ ):
956
+ """
957
+ Apply the model to an input batch.
958
+ :param x: an [(N x F) x C x ...] Tensor of inputs. F is the number of frames (views).
959
+ :param timesteps: a 1-D batch of timesteps.
960
+ :param context: conditioning plugged in via crossattn
961
+ :param y: an [N] Tensor of labels, if class-conditional.
962
+ :param num_frames: a integer indicating number of frames for tensor reshaping.
963
+ :return: an [(N x F) x C x ...] Tensor of outputs. F is the number of frames (views).
964
+ """
965
+ assert (
966
+ x.shape[0] % num_frames == 0
967
+ ), "input batch size must be dividable by num_frames!"
968
+ assert (y is not None) == (
969
+ self.num_classes is not None
970
+ ), "must specify y if and only if the model is class-conditional"
971
+
972
+ hs = []
973
+
974
+ t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
975
+
976
+ emb = self.time_embed(t_emb)
977
+
978
+ if self.num_classes is not None:
979
+ assert y is not None
980
+ assert y.shape[0] == x.shape[0]
981
+ emb = emb + self.label_emb(y)
982
+
983
+ # Add camera embeddings
984
+ if camera is not None:
985
+ emb = emb + self.camera_embed(camera)
986
+
987
+ # imagedream variant
988
+ if self.ip_dim > 0:
989
+ x[(num_frames - 1) :: num_frames, :, :, :] = ip_img # place at [4, 9]
990
+ ip_emb = self.image_embed(ip)
991
+ context = torch.cat((context, ip_emb), 1)
992
+
993
+ h = x
994
+ for module in self.input_blocks:
995
+ h = module(h, emb, context, num_frames=num_frames)
996
+ hs.append(h)
997
+ h = self.middle_block(h, emb, context, num_frames=num_frames)
998
+ for module in self.output_blocks:
999
+ h = torch.cat([h, hs.pop()], dim=1)
1000
+ h = module(h, emb, context, num_frames=num_frames)
1001
+ h = h.type(x.dtype)
1002
+ if self.predict_codebook_ids:
1003
+ return self.id_predictor(h)
1004
+ else:
1005
+ return self.out(h)
lgm/mvdream/pipeline_mvdream.py ADDED
@@ -0,0 +1,559 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn.functional as F
3
+ import inspect
4
+ import numpy as np
5
+ from typing import Callable, List, Optional, Union
6
+ from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModel, CLIPImageProcessor
7
+ from diffusers import AutoencoderKL, DiffusionPipeline
8
+ from diffusers.utils import (
9
+ deprecate,
10
+ is_accelerate_available,
11
+ is_accelerate_version,
12
+ logging,
13
+ )
14
+ from diffusers.configuration_utils import FrozenDict
15
+ from diffusers.schedulers import DDIMScheduler
16
+ from diffusers.utils.torch_utils import randn_tensor
17
+
18
+ from mvdream.mv_unet import MultiViewUNetModel, get_camera
19
+
20
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
21
+
22
+
23
+ class MVDreamPipeline(DiffusionPipeline):
24
+
25
+ _optional_components = ["feature_extractor", "image_encoder"]
26
+
27
+ def __init__(
28
+ self,
29
+ vae: AutoencoderKL,
30
+ unet: MultiViewUNetModel,
31
+ tokenizer: CLIPTokenizer,
32
+ text_encoder: CLIPTextModel,
33
+ scheduler: DDIMScheduler,
34
+ # imagedream variant
35
+ feature_extractor: CLIPImageProcessor,
36
+ image_encoder: CLIPVisionModel,
37
+ requires_safety_checker: bool = False,
38
+ ):
39
+ super().__init__()
40
+
41
+ if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1: # type: ignore
42
+ deprecation_message = (
43
+ f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
44
+ f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " # type: ignore
45
+ "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
46
+ " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
47
+ " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
48
+ " file"
49
+ )
50
+ deprecate(
51
+ "steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False
52
+ )
53
+ new_config = dict(scheduler.config)
54
+ new_config["steps_offset"] = 1
55
+ scheduler._internal_dict = FrozenDict(new_config)
56
+
57
+ if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True: # type: ignore
58
+ deprecation_message = (
59
+ f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
60
+ " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
61
+ " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
62
+ " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
63
+ " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
64
+ )
65
+ deprecate(
66
+ "clip_sample not set", "1.0.0", deprecation_message, standard_warn=False
67
+ )
68
+ new_config = dict(scheduler.config)
69
+ new_config["clip_sample"] = False
70
+ scheduler._internal_dict = FrozenDict(new_config)
71
+
72
+ self.register_modules(
73
+ vae=vae,
74
+ unet=unet,
75
+ scheduler=scheduler,
76
+ tokenizer=tokenizer,
77
+ text_encoder=text_encoder,
78
+ feature_extractor=feature_extractor,
79
+ image_encoder=image_encoder,
80
+ )
81
+ self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
82
+ self.register_to_config(requires_safety_checker=requires_safety_checker)
83
+
84
+ def enable_vae_slicing(self):
85
+ r"""
86
+ Enable sliced VAE decoding.
87
+
88
+ When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several
89
+ steps. This is useful to save some memory and allow larger batch sizes.
90
+ """
91
+ self.vae.enable_slicing()
92
+
93
+ def disable_vae_slicing(self):
94
+ r"""
95
+ Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to
96
+ computing decoding in one step.
97
+ """
98
+ self.vae.disable_slicing()
99
+
100
+ def enable_vae_tiling(self):
101
+ r"""
102
+ Enable tiled VAE decoding.
103
+
104
+ When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in
105
+ several steps. This is useful to save a large amount of memory and to allow the processing of larger images.
106
+ """
107
+ self.vae.enable_tiling()
108
+
109
+ def disable_vae_tiling(self):
110
+ r"""
111
+ Disable tiled VAE decoding. If `enable_vae_tiling` was previously invoked, this method will go back to
112
+ computing decoding in one step.
113
+ """
114
+ self.vae.disable_tiling()
115
+
116
+ def enable_sequential_cpu_offload(self, gpu_id=0):
117
+ r"""
118
+ Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
119
+ text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
120
+ `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
121
+ Note that offloading happens on a submodule basis. Memory savings are higher than with
122
+ `enable_model_cpu_offload`, but performance is lower.
123
+ """
124
+ if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):
125
+ from accelerate import cpu_offload
126
+ else:
127
+ raise ImportError(
128
+ "`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher"
129
+ )
130
+
131
+ device = torch.device(f"cuda:{gpu_id}")
132
+
133
+ if self.device.type != "cpu":
134
+ self.to("cpu", silence_dtype_warnings=True)
135
+ torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
136
+
137
+ for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
138
+ cpu_offload(cpu_offloaded_model, device)
139
+
140
+ def enable_model_cpu_offload(self, gpu_id=0):
141
+ r"""
142
+ Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
143
+ to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
144
+ method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
145
+ `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
146
+ """
147
+ if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
148
+ from accelerate import cpu_offload_with_hook
149
+ else:
150
+ raise ImportError(
151
+ "`enable_model_offload` requires `accelerate v0.17.0` or higher."
152
+ )
153
+
154
+ device = torch.device(f"cuda:{gpu_id}")
155
+
156
+ if self.device.type != "cpu":
157
+ self.to("cpu", silence_dtype_warnings=True)
158
+ torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
159
+
160
+ hook = None
161
+ for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:
162
+ _, hook = cpu_offload_with_hook(
163
+ cpu_offloaded_model, device, prev_module_hook=hook
164
+ )
165
+
166
+ # We'll offload the last model manually.
167
+ self.final_offload_hook = hook
168
+
169
+ @property
170
+ def _execution_device(self):
171
+ r"""
172
+ Returns the device on which the pipeline's models will be executed. After calling
173
+ `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
174
+ hooks.
175
+ """
176
+ if not hasattr(self.unet, "_hf_hook"):
177
+ return self.device
178
+ for module in self.unet.modules():
179
+ if (
180
+ hasattr(module, "_hf_hook")
181
+ and hasattr(module._hf_hook, "execution_device")
182
+ and module._hf_hook.execution_device is not None
183
+ ):
184
+ return torch.device(module._hf_hook.execution_device)
185
+ return self.device
186
+
187
+ def _encode_prompt(
188
+ self,
189
+ prompt,
190
+ device,
191
+ num_images_per_prompt,
192
+ do_classifier_free_guidance: bool,
193
+ negative_prompt=None,
194
+ ):
195
+ r"""
196
+ Encodes the prompt into text encoder hidden states.
197
+
198
+ Args:
199
+ prompt (`str` or `List[str]`, *optional*):
200
+ prompt to be encoded
201
+ device: (`torch.device`):
202
+ torch device
203
+ num_images_per_prompt (`int`):
204
+ number of images that should be generated per prompt
205
+ do_classifier_free_guidance (`bool`):
206
+ whether to use classifier free guidance or not
207
+ negative_prompt (`str` or `List[str]`, *optional*):
208
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
209
+ `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
210
+ Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
211
+ prompt_embeds (`torch.FloatTensor`, *optional*):
212
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
213
+ provided, text embeddings will be generated from `prompt` input argument.
214
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
215
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
216
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
217
+ argument.
218
+ """
219
+ if prompt is not None and isinstance(prompt, str):
220
+ batch_size = 1
221
+ elif prompt is not None and isinstance(prompt, list):
222
+ batch_size = len(prompt)
223
+ else:
224
+ raise ValueError(
225
+ f"`prompt` should be either a string or a list of strings, but got {type(prompt)}."
226
+ )
227
+
228
+ text_inputs = self.tokenizer(
229
+ prompt,
230
+ padding="max_length",
231
+ max_length=self.tokenizer.model_max_length,
232
+ truncation=True,
233
+ return_tensors="pt",
234
+ )
235
+ text_input_ids = text_inputs.input_ids
236
+ untruncated_ids = self.tokenizer(
237
+ prompt, padding="longest", return_tensors="pt"
238
+ ).input_ids
239
+
240
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
241
+ text_input_ids, untruncated_ids
242
+ ):
243
+ removed_text = self.tokenizer.batch_decode(
244
+ untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
245
+ )
246
+ logger.warning(
247
+ "The following part of your input was truncated because CLIP can only handle sequences up to"
248
+ f" {self.tokenizer.model_max_length} tokens: {removed_text}"
249
+ )
250
+
251
+ if (
252
+ hasattr(self.text_encoder.config, "use_attention_mask")
253
+ and self.text_encoder.config.use_attention_mask
254
+ ):
255
+ attention_mask = text_inputs.attention_mask.to(device)
256
+ else:
257
+ attention_mask = None
258
+
259
+ prompt_embeds = self.text_encoder(
260
+ text_input_ids.to(device),
261
+ attention_mask=attention_mask,
262
+ )
263
+ prompt_embeds = prompt_embeds[0]
264
+
265
+ prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
266
+
267
+ bs_embed, seq_len, _ = prompt_embeds.shape
268
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
269
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
270
+ prompt_embeds = prompt_embeds.view(
271
+ bs_embed * num_images_per_prompt, seq_len, -1
272
+ )
273
+
274
+ # get unconditional embeddings for classifier free guidance
275
+ if do_classifier_free_guidance:
276
+ uncond_tokens: List[str]
277
+ if negative_prompt is None:
278
+ uncond_tokens = [""] * batch_size
279
+ elif type(prompt) is not type(negative_prompt):
280
+ raise TypeError(
281
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
282
+ f" {type(prompt)}."
283
+ )
284
+ elif isinstance(negative_prompt, str):
285
+ uncond_tokens = [negative_prompt]
286
+ elif batch_size != len(negative_prompt):
287
+ raise ValueError(
288
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
289
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
290
+ " the batch size of `prompt`."
291
+ )
292
+ else:
293
+ uncond_tokens = negative_prompt
294
+
295
+ max_length = prompt_embeds.shape[1]
296
+ uncond_input = self.tokenizer(
297
+ uncond_tokens,
298
+ padding="max_length",
299
+ max_length=max_length,
300
+ truncation=True,
301
+ return_tensors="pt",
302
+ )
303
+
304
+ if (
305
+ hasattr(self.text_encoder.config, "use_attention_mask")
306
+ and self.text_encoder.config.use_attention_mask
307
+ ):
308
+ attention_mask = uncond_input.attention_mask.to(device)
309
+ else:
310
+ attention_mask = None
311
+
312
+ negative_prompt_embeds = self.text_encoder(
313
+ uncond_input.input_ids.to(device),
314
+ attention_mask=attention_mask,
315
+ )
316
+ negative_prompt_embeds = negative_prompt_embeds[0]
317
+
318
+ # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
319
+ seq_len = negative_prompt_embeds.shape[1]
320
+
321
+ negative_prompt_embeds = negative_prompt_embeds.to(
322
+ dtype=self.text_encoder.dtype, device=device
323
+ )
324
+
325
+ negative_prompt_embeds = negative_prompt_embeds.repeat(
326
+ 1, num_images_per_prompt, 1
327
+ )
328
+ negative_prompt_embeds = negative_prompt_embeds.view(
329
+ batch_size * num_images_per_prompt, seq_len, -1
330
+ )
331
+
332
+ # For classifier free guidance, we need to do two forward passes.
333
+ # Here we concatenate the unconditional and text embeddings into a single batch
334
+ # to avoid doing two forward passes
335
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
336
+
337
+ return prompt_embeds
338
+
339
+ def decode_latents(self, latents):
340
+ latents = 1 / self.vae.config.scaling_factor * latents
341
+ image = self.vae.decode(latents).sample
342
+ image = (image / 2 + 0.5).clamp(0, 1)
343
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
344
+ image = image.cpu().permute(0, 2, 3, 1).float().numpy()
345
+ return image
346
+
347
+ def prepare_extra_step_kwargs(self, generator, eta):
348
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
349
+ # eta (Ξ·) is only used with the DDIMScheduler, it will be ignored for other schedulers.
350
+ # eta corresponds to Ξ· in DDIM paper: https://arxiv.org/abs/2010.02502
351
+ # and should be between [0, 1]
352
+
353
+ accepts_eta = "eta" in set(
354
+ inspect.signature(self.scheduler.step).parameters.keys()
355
+ )
356
+ extra_step_kwargs = {}
357
+ if accepts_eta:
358
+ extra_step_kwargs["eta"] = eta
359
+
360
+ # check if the scheduler accepts generator
361
+ accepts_generator = "generator" in set(
362
+ inspect.signature(self.scheduler.step).parameters.keys()
363
+ )
364
+ if accepts_generator:
365
+ extra_step_kwargs["generator"] = generator
366
+ return extra_step_kwargs
367
+
368
+ def prepare_latents(
369
+ self,
370
+ batch_size,
371
+ num_channels_latents,
372
+ height,
373
+ width,
374
+ dtype,
375
+ device,
376
+ generator,
377
+ latents=None,
378
+ ):
379
+ shape = (
380
+ batch_size,
381
+ num_channels_latents,
382
+ height // self.vae_scale_factor,
383
+ width // self.vae_scale_factor,
384
+ )
385
+ if isinstance(generator, list) and len(generator) != batch_size:
386
+ raise ValueError(
387
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
388
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
389
+ )
390
+
391
+ if latents is None:
392
+ latents = randn_tensor(
393
+ shape, generator=generator, device=device, dtype=dtype
394
+ )
395
+ else:
396
+ latents = latents.to(device)
397
+
398
+ # scale the initial noise by the standard deviation required by the scheduler
399
+ latents = latents * self.scheduler.init_noise_sigma
400
+ return latents
401
+
402
+ def encode_image(self, image, device, num_images_per_prompt):
403
+ dtype = next(self.image_encoder.parameters()).dtype
404
+
405
+ if image.dtype == np.float32:
406
+ image = (image * 255).astype(np.uint8)
407
+
408
+ image = self.feature_extractor(image, return_tensors="pt").pixel_values
409
+ image = image.to(device=device, dtype=dtype)
410
+
411
+ image_embeds = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
412
+ image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
413
+
414
+ return torch.zeros_like(image_embeds), image_embeds
415
+
416
+ def encode_image_latents(self, image, device, num_images_per_prompt):
417
+
418
+ dtype = next(self.image_encoder.parameters()).dtype
419
+
420
+ image = torch.from_numpy(image).unsqueeze(0).permute(0, 3, 1, 2).to(device=device) # [1, 3, H, W]
421
+ image = 2 * image - 1
422
+ image = F.interpolate(image, (256, 256), mode='bilinear', align_corners=False)
423
+ image = image.to(dtype=dtype)
424
+
425
+ posterior = self.vae.encode(image).latent_dist
426
+ latents = posterior.sample() * self.vae.config.scaling_factor # [B, C, H, W]
427
+ latents = latents.repeat_interleave(num_images_per_prompt, dim=0)
428
+
429
+ return torch.zeros_like(latents), latents
430
+
431
+ @torch.no_grad()
432
+ def __call__(
433
+ self,
434
+ prompt: str = "",
435
+ image: Optional[np.ndarray] = None,
436
+ height: int = 256,
437
+ width: int = 256,
438
+ elevation: float = 0,
439
+ num_inference_steps: int = 50,
440
+ guidance_scale: float = 7.0,
441
+ negative_prompt: str = "",
442
+ num_images_per_prompt: int = 1,
443
+ eta: float = 0.0,
444
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
445
+ output_type: Optional[str] = "numpy", # pil, numpy, latents
446
+ callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
447
+ callback_steps: int = 1,
448
+ num_frames: int = 4,
449
+ device=torch.device("cuda:0"),
450
+ ):
451
+ self.unet = self.unet.to(device=device)
452
+ self.vae = self.vae.to(device=device)
453
+ self.text_encoder = self.text_encoder.to(device=device)
454
+
455
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
456
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
457
+ # corresponds to doing no classifier free guidance.
458
+ do_classifier_free_guidance = guidance_scale > 1.0
459
+
460
+ # Prepare timesteps
461
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
462
+ timesteps = self.scheduler.timesteps
463
+
464
+ # imagedream variant
465
+ if image is not None:
466
+ assert isinstance(image, np.ndarray) and image.dtype == np.float32
467
+ self.image_encoder = self.image_encoder.to(device=device)
468
+ image_embeds_neg, image_embeds_pos = self.encode_image(image, device, num_images_per_prompt)
469
+ image_latents_neg, image_latents_pos = self.encode_image_latents(image, device, num_images_per_prompt)
470
+
471
+ _prompt_embeds = self._encode_prompt(
472
+ prompt=prompt,
473
+ device=device,
474
+ num_images_per_prompt=num_images_per_prompt,
475
+ do_classifier_free_guidance=do_classifier_free_guidance,
476
+ negative_prompt=negative_prompt,
477
+ ) # type: ignore
478
+ prompt_embeds_neg, prompt_embeds_pos = _prompt_embeds.chunk(2)
479
+
480
+ # Prepare latent variables
481
+ actual_num_frames = num_frames if image is None else num_frames + 1
482
+ latents: torch.Tensor = self.prepare_latents(
483
+ actual_num_frames * num_images_per_prompt,
484
+ 4,
485
+ height,
486
+ width,
487
+ prompt_embeds_pos.dtype,
488
+ device,
489
+ generator,
490
+ None,
491
+ )
492
+
493
+ if image is not None:
494
+ camera = get_camera(num_frames, elevation=elevation, extra_view=True).to(dtype=latents.dtype, device=device)
495
+ else:
496
+ camera = get_camera(num_frames, elevation=elevation, extra_view=False).to(dtype=latents.dtype, device=device)
497
+ camera = camera.repeat_interleave(num_images_per_prompt, dim=0)
498
+
499
+ # Prepare extra step kwargs.
500
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
501
+
502
+ # Denoising loop
503
+ num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
504
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
505
+ for i, t in enumerate(timesteps):
506
+ # expand the latents if we are doing classifier free guidance
507
+ multiplier = 2 if do_classifier_free_guidance else 1
508
+ latent_model_input = torch.cat([latents] * multiplier)
509
+ latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
510
+
511
+ unet_inputs = {
512
+ 'x': latent_model_input,
513
+ 'timesteps': torch.tensor([t] * actual_num_frames * multiplier, dtype=latent_model_input.dtype, device=device),
514
+ 'context': torch.cat([prompt_embeds_neg] * actual_num_frames + [prompt_embeds_pos] * actual_num_frames),
515
+ 'num_frames': actual_num_frames,
516
+ 'camera': torch.cat([camera] * multiplier),
517
+ }
518
+
519
+ if image is not None:
520
+ unet_inputs['ip'] = torch.cat([image_embeds_neg] * actual_num_frames + [image_embeds_pos] * actual_num_frames)
521
+ unet_inputs['ip_img'] = torch.cat([image_latents_neg] + [image_latents_pos]) # no repeat
522
+
523
+ # predict the noise residual
524
+ noise_pred = self.unet.forward(**unet_inputs)
525
+
526
+ # perform guidance
527
+ if do_classifier_free_guidance:
528
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
529
+ noise_pred = noise_pred_uncond + guidance_scale * (
530
+ noise_pred_text - noise_pred_uncond
531
+ )
532
+
533
+ # compute the previous noisy sample x_t -> x_t-1
534
+ latents: torch.Tensor = self.scheduler.step(
535
+ noise_pred, t, latents, **extra_step_kwargs, return_dict=False
536
+ )[0]
537
+
538
+ # call the callback, if provided
539
+ if i == len(timesteps) - 1 or (
540
+ (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
541
+ ):
542
+ progress_bar.update()
543
+ if callback is not None and i % callback_steps == 0:
544
+ callback(i, t, latents) # type: ignore
545
+
546
+ # Post-processing
547
+ if output_type == "latent":
548
+ image = latents
549
+ elif output_type == "pil":
550
+ image = self.decode_latents(latents)
551
+ image = self.numpy_to_pil(image)
552
+ else: # numpy
553
+ image = self.decode_latents(latents)
554
+
555
+ # Offload last model to CPU
556
+ if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
557
+ self.final_offload_hook.offload()
558
+
559
+ return image
main_4d.py ADDED
@@ -0,0 +1,601 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import cv2
3
+ import time
4
+ import tqdm
5
+ import numpy as np
6
+
7
+ import torch
8
+ import torch.nn.functional as F
9
+
10
+ import rembg
11
+
12
+ from cam_utils import orbit_camera, OrbitCamera
13
+ from gs_renderer_4d import Renderer, MiniCam
14
+
15
+ from grid_put import mipmap_linear_grid_put_2d
16
+ import imageio
17
+
18
+ import copy
19
+
20
+
21
+ class GUI:
22
+ def __init__(self, opt):
23
+ self.opt = opt # shared with the trainer's opt to support in-place modification of rendering parameters.
24
+ self.gui = opt.gui # enable gui
25
+ self.W = opt.W
26
+ self.H = opt.H
27
+ self.cam = OrbitCamera(opt.W, opt.H, r=opt.radius, fovy=opt.fovy)
28
+
29
+ self.mode = "image"
30
+ # self.seed = "random"
31
+ self.seed = 888
32
+
33
+ self.buffer_image = np.ones((self.W, self.H, 3), dtype=np.float32)
34
+ self.need_update = True # update buffer_image
35
+
36
+ # models
37
+ self.device = torch.device("cuda")
38
+ self.bg_remover = None
39
+
40
+ self.guidance_sd = None
41
+ self.guidance_zero123 = None
42
+ self.guidance_svd = None
43
+
44
+
45
+ self.enable_sd = False
46
+ self.enable_zero123 = False
47
+ self.enable_svd = False
48
+
49
+
50
+ # renderer
51
+ self.renderer = Renderer(self.opt, sh_degree=self.opt.sh_degree)
52
+ self.gaussain_scale_factor = 1
53
+
54
+ # input image
55
+ self.input_img = None
56
+ self.input_mask = None
57
+ self.input_img_torch = None
58
+ self.input_mask_torch = None
59
+ self.overlay_input_img = False
60
+ self.overlay_input_img_ratio = 0.5
61
+
62
+ self.input_img_list = None
63
+ self.input_mask_list = None
64
+ self.input_img_torch_list = None
65
+ self.input_mask_torch_list = None
66
+
67
+ # input text
68
+ self.prompt = ""
69
+ self.negative_prompt = ""
70
+
71
+ # training stuff
72
+ self.training = False
73
+ self.optimizer = None
74
+ self.step = 0
75
+ self.train_steps = 1 # steps per rendering loop
76
+
77
+ # load input data from cmdline
78
+ if self.opt.input is not None: # True
79
+ self.load_input(self.opt.input) # load imgs, if has bg, then rm bg; or just load imgs
80
+
81
+ # override prompt from cmdline
82
+ if self.opt.prompt is not None: # None
83
+ self.prompt = self.opt.prompt
84
+
85
+ # override if provide a checkpoint
86
+ if self.opt.load is not None: # not None
87
+ self.renderer.initialize(self.opt.load)
88
+ # self.renderer.gaussians.load_model(opt.outdir, opt.save_path)
89
+ else:
90
+ # initialize gaussians to a blob
91
+ self.renderer.initialize(num_pts=self.opt.num_pts)
92
+
93
+ self.seed_everything()
94
+
95
+ def seed_everything(self):
96
+ try:
97
+ seed = int(self.seed)
98
+ except:
99
+ seed = np.random.randint(0, 1000000)
100
+
101
+ print(f'Seed: {seed:d}')
102
+ os.environ["PYTHONHASHSEED"] = str(seed)
103
+ np.random.seed(seed)
104
+ torch.manual_seed(seed)
105
+ torch.cuda.manual_seed(seed)
106
+ torch.backends.cudnn.deterministic = True
107
+ torch.backends.cudnn.benchmark = True
108
+
109
+ self.last_seed = seed
110
+
111
+ def prepare_train(self):
112
+
113
+ self.step = 0
114
+
115
+ # setup training
116
+ self.renderer.gaussians.training_setup(self.opt)
117
+
118
+ # # do not do progressive sh-level
119
+ self.renderer.gaussians.active_sh_degree = self.renderer.gaussians.max_sh_degree
120
+ self.optimizer = self.renderer.gaussians.optimizer
121
+
122
+ # default camera
123
+ if self.opt.mvdream or self.opt.imagedream:
124
+ # the second view is the front view for mvdream/imagedream.
125
+ pose = orbit_camera(self.opt.elevation, 90, self.opt.radius)
126
+ else:
127
+ pose = orbit_camera(self.opt.elevation, 0, self.opt.radius)
128
+ self.fixed_cam = MiniCam(
129
+ pose,
130
+ self.opt.ref_size,
131
+ self.opt.ref_size,
132
+ self.cam.fovy,
133
+ self.cam.fovx,
134
+ self.cam.near,
135
+ self.cam.far,
136
+ )
137
+
138
+ self.enable_sd = self.opt.lambda_sd > 0
139
+ self.enable_zero123 = self.opt.lambda_zero123 > 0
140
+ self.enable_svd = self.opt.lambda_svd > 0 and self.input_img is not None
141
+
142
+ # lazy load guidance model
143
+ if self.guidance_sd is None and self.enable_sd:
144
+ if self.opt.mvdream:
145
+ print(f"[INFO] loading MVDream...")
146
+ from guidance.mvdream_utils import MVDream
147
+ self.guidance_sd = MVDream(self.device)
148
+ print(f"[INFO] loaded MVDream!")
149
+ elif self.opt.imagedream:
150
+ print(f"[INFO] loading ImageDream...")
151
+ from guidance.imagedream_utils import ImageDream
152
+ self.guidance_sd = ImageDream(self.device)
153
+ print(f"[INFO] loaded ImageDream!")
154
+ else:
155
+ print(f"[INFO] loading SD...")
156
+ from guidance.sd_utils import StableDiffusion
157
+ self.guidance_sd = StableDiffusion(self.device)
158
+ print(f"[INFO] loaded SD!")
159
+
160
+ if self.guidance_zero123 is None and self.enable_zero123:
161
+ print(f"[INFO] loading zero123...")
162
+ from guidance.zero123_utils import Zero123
163
+ if self.opt.stable_zero123:
164
+ self.guidance_zero123 = Zero123(self.device, model_key='ashawkey/stable-zero123-diffusers')
165
+ else:
166
+ self.guidance_zero123 = Zero123(self.device, model_key='ashawkey/zero123-xl-diffusers')
167
+ print(f"[INFO] loaded zero123!")
168
+
169
+ if self.guidance_svd is None and self.enable_svd: # False
170
+ print(f"[INFO] loading SVD...")
171
+ from guidance.svd_utils import StableVideoDiffusion
172
+ self.guidance_svd = StableVideoDiffusion(self.device)
173
+ print(f"[INFO] loaded SVD!")
174
+
175
+ # input image
176
+ if self.input_img is not None:
177
+ self.input_img_torch = torch.from_numpy(self.input_img).permute(2, 0, 1).unsqueeze(0).to(self.device)
178
+ self.input_img_torch = F.interpolate(self.input_img_torch, (self.opt.ref_size, self.opt.ref_size), mode="bilinear", align_corners=False)
179
+
180
+ self.input_mask_torch = torch.from_numpy(self.input_mask).permute(2, 0, 1).unsqueeze(0).to(self.device)
181
+ self.input_mask_torch = F.interpolate(self.input_mask_torch, (self.opt.ref_size, self.opt.ref_size), mode="bilinear", align_corners=False)
182
+
183
+ if self.input_img_list is not None:
184
+ 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]
185
+ 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]
186
+
187
+ 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]
188
+ 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]
189
+ # prepare embeddings
190
+ with torch.no_grad():
191
+
192
+ if self.enable_sd:
193
+ if self.opt.imagedream:
194
+ img_pos_list, img_neg_list, ip_pos_list, ip_neg_list, emb_pos_list, emb_neg_list = [], [], [], [], [], []
195
+ for _ in range(self.opt.n_views):
196
+ for input_img_torch in self.input_img_torch_list:
197
+ 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])
198
+ img_pos_list.append(img_pos)
199
+ img_neg_list.append(img_neg)
200
+ ip_pos_list.append(ip_pos)
201
+ ip_neg_list.append(ip_neg)
202
+ emb_pos_list.append(emb_pos)
203
+ emb_neg_list.append(emb_neg)
204
+ self.guidance_sd.image_embeddings['pos'] = torch.cat(img_pos_list, 0)
205
+ self.guidance_sd.image_embeddings['neg'] = torch.cat(img_pos_list, 0)
206
+ self.guidance_sd.image_embeddings['ip_img'] = torch.cat(ip_pos_list, 0)
207
+ self.guidance_sd.image_embeddings['neg_ip_img'] = torch.cat(ip_neg_list, 0)
208
+ self.guidance_sd.embeddings['pos'] = torch.cat(emb_pos_list, 0)
209
+ self.guidance_sd.embeddings['neg'] = torch.cat(emb_neg_list, 0)
210
+ else:
211
+ self.guidance_sd.get_text_embeds([self.prompt], [self.negative_prompt])
212
+
213
+ if self.enable_zero123:
214
+ c_list, v_list = [], []
215
+ for _ in range(self.opt.n_views):
216
+ for input_img_torch in self.input_img_torch_list:
217
+ c, v = self.guidance_zero123.get_img_embeds(input_img_torch)
218
+ c_list.append(c)
219
+ v_list.append(v)
220
+ self.guidance_zero123.embeddings = [torch.cat(c_list, 0), torch.cat(v_list, 0)]
221
+
222
+ if self.enable_svd:
223
+ self.guidance_svd.get_img_embeds(self.input_img)
224
+
225
+ def train_step(self):
226
+ starter = torch.cuda.Event(enable_timing=True)
227
+ ender = torch.cuda.Event(enable_timing=True)
228
+ starter.record()
229
+
230
+ for _ in range(self.train_steps): # 1
231
+
232
+ self.step += 1 # self.step starts from 0
233
+ step_ratio = min(1, self.step / self.opt.iters) # 1, step / 500
234
+
235
+ # update lr
236
+ self.renderer.gaussians.update_learning_rate(self.step)
237
+
238
+ loss = 0
239
+
240
+ self.renderer.prepare_render()
241
+
242
+ ### known view
243
+ if not self.opt.imagedream:
244
+ for b_idx in range(self.opt.batch_size):
245
+ cur_cam = copy.deepcopy(self.fixed_cam)
246
+ cur_cam.time = b_idx
247
+ out = self.renderer.render(cur_cam)
248
+
249
+ # rgb loss
250
+ image = out["image"].unsqueeze(0) # [1, 3, H, W] in [0, 1]
251
+ loss = loss + 10000 * step_ratio * F.mse_loss(image, self.input_img_torch_list[b_idx]) / self.opt.batch_size
252
+
253
+ # mask loss
254
+ mask = out["alpha"].unsqueeze(0) # [1, 1, H, W] in [0, 1]
255
+ loss = loss + 1000 * step_ratio * F.mse_loss(mask, self.input_mask_torch_list[b_idx]) / self.opt.batch_size
256
+
257
+ ### novel view (manual batch)
258
+ render_resolution = 128 if step_ratio < 0.3 else (256 if step_ratio < 0.6 else 512)
259
+ # render_resolution = 512
260
+ images = []
261
+ poses = []
262
+ vers, hors, radii = [], [], []
263
+ # avoid too large elevation (> 80 or < -80), and make sure it always cover [-30, 30]
264
+ min_ver = max(min(self.opt.min_ver, self.opt.min_ver - self.opt.elevation), -80 - self.opt.elevation)
265
+ max_ver = min(max(self.opt.max_ver, self.opt.max_ver - self.opt.elevation), 80 - self.opt.elevation)
266
+
267
+ for _ in range(self.opt.n_views):
268
+ for b_idx in range(self.opt.batch_size):
269
+
270
+ # render random view
271
+ ver = np.random.randint(min_ver, max_ver)
272
+ hor = np.random.randint(-180, 180)
273
+ radius = 0
274
+
275
+ vers.append(ver)
276
+ hors.append(hor)
277
+ radii.append(radius)
278
+
279
+ pose = orbit_camera(self.opt.elevation + ver, hor, self.opt.radius + radius)
280
+ poses.append(pose)
281
+
282
+ cur_cam = MiniCam(pose, render_resolution, render_resolution, self.cam.fovy, self.cam.fovx, self.cam.near, self.cam.far, time=b_idx)
283
+
284
+ 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")
285
+ out = self.renderer.render(cur_cam, bg_color=bg_color)
286
+
287
+ image = out["image"].unsqueeze(0) # [1, 3, H, W] in [0, 1]
288
+ images.append(image)
289
+
290
+ # enable mvdream training
291
+ if self.opt.mvdream or self.opt.imagedream: # False
292
+ for view_i in range(1, 4):
293
+ pose_i = orbit_camera(self.opt.elevation + ver, hor + 90 * view_i, self.opt.radius + radius)
294
+ poses.append(pose_i)
295
+
296
+ cur_cam_i = MiniCam(pose_i, render_resolution, render_resolution, self.cam.fovy, self.cam.fovx, self.cam.near, self.cam.far)
297
+
298
+ # bg_color = torch.tensor([0.5, 0.5, 0.5], dtype=torch.float32, device="cuda")
299
+ out_i = self.renderer.render(cur_cam_i, bg_color=bg_color)
300
+
301
+ image = out_i["image"].unsqueeze(0) # [1, 3, H, W] in [0, 1]
302
+ images.append(image)
303
+
304
+
305
+
306
+ images = torch.cat(images, dim=0)
307
+ poses = torch.from_numpy(np.stack(poses, axis=0)).to(self.device)
308
+
309
+ # guidance loss
310
+ if self.enable_sd:
311
+ if self.opt.mvdream or self.opt.imagedream:
312
+ loss = loss + self.opt.lambda_sd * self.guidance_sd.train_step(images, poses, step_ratio)
313
+ else:
314
+ loss = loss + self.opt.lambda_sd * self.guidance_sd.train_step(images, step_ratio)
315
+
316
+ if self.enable_zero123:
317
+ 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)
318
+
319
+ if self.enable_svd:
320
+ loss = loss + self.opt.lambda_svd * self.guidance_svd.train_step(images, step_ratio)
321
+
322
+ # optimize step
323
+ loss.backward()
324
+ self.optimizer.step()
325
+ self.optimizer.zero_grad()
326
+
327
+ # densify and prune
328
+ if self.step >= self.opt.density_start_iter and self.step <= self.opt.density_end_iter:
329
+ viewspace_point_tensor, visibility_filter, radii = out["viewspace_points"], out["visibility_filter"], out["radii"]
330
+ self.renderer.gaussians.max_radii2D[visibility_filter] = torch.max(self.renderer.gaussians.max_radii2D[visibility_filter], radii[visibility_filter])
331
+ self.renderer.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter)
332
+
333
+ if self.step % self.opt.densification_interval == 0:
334
+ # size_threshold = 20 if self.step > self.opt.opacity_reset_interval else None
335
+ self.renderer.gaussians.densify_and_prune(self.opt.densify_grad_threshold, min_opacity=0.01, extent=0.5, max_screen_size=1)
336
+
337
+ if self.step % self.opt.opacity_reset_interval == 0:
338
+ self.renderer.gaussians.reset_opacity()
339
+
340
+ ender.record()
341
+ torch.cuda.synchronize()
342
+ t = starter.elapsed_time(ender)
343
+
344
+ self.need_update = True
345
+
346
+
347
+ def load_input(self, file):
348
+ if self.opt.data_mode == 'c4d':
349
+ file_list = [os.path.join(file, f'{x * self.opt.downsample_rate}.png') for x in range(self.opt.batch_size)]
350
+ elif self.opt.data_mode == 'svd':
351
+ # file_list = [file.replace('.png', f'_frames/{x* self.opt.downsample_rate:03d}_rgba.png') for x in range(self.opt.batch_size)]
352
+ # file_list = [x if os.path.exists(x) else (x.replace('_rgba.png', '.png')) for x in file_list]
353
+ file_list = [file.replace('.png', f'_frames/{x* self.opt.downsample_rate:03d}.png') for x in range(self.opt.batch_size)]
354
+ else:
355
+ raise NotImplementedError
356
+ self.input_img_list, self.input_mask_list = [], []
357
+ for file in file_list:
358
+ # load image
359
+ print(f'[INFO] load image from {file}...')
360
+ img = cv2.imread(file, cv2.IMREAD_UNCHANGED)
361
+ if img.shape[-1] == 3:
362
+ if self.bg_remover is None:
363
+ self.bg_remover = rembg.new_session()
364
+ img = rembg.remove(img, session=self.bg_remover)
365
+ # cv2.imwrite(file.replace('.png', '_rgba.png'), img)
366
+ img = cv2.resize(img, (self.W, self.H), interpolation=cv2.INTER_AREA)
367
+ img = img.astype(np.float32) / 255.0
368
+ input_mask = img[..., 3:]
369
+ # white bg
370
+ input_img = img[..., :3] * input_mask + (1 - input_mask)
371
+ # bgr to rgb
372
+ input_img = input_img[..., ::-1].copy()
373
+ self.input_img_list.append(input_img)
374
+ self.input_mask_list.append(input_mask)
375
+
376
+ @torch.no_grad()
377
+ def save_model(self, mode='geo', texture_size=1024, interp=1):
378
+ os.makedirs(self.opt.outdir, exist_ok=True)
379
+ if mode == 'geo':
380
+ path = f'logs/{opt.save_path}_mesh_{t:03d}.ply'
381
+ mesh = self.renderer.gaussians.extract_mesh_t(path, self.opt.density_thresh, t=t)
382
+ mesh.write_ply(path)
383
+
384
+ elif mode == 'geo+tex':
385
+ from mesh import Mesh, safe_normalize
386
+ os.makedirs(os.path.join(self.opt.outdir, self.opt.save_path+'_meshes'), exist_ok=True)
387
+ for t in range(self.opt.batch_size):
388
+ path = os.path.join(self.opt.outdir, self.opt.save_path+'_meshes', f'{t:03d}.obj')
389
+ mesh = self.renderer.gaussians.extract_mesh_t(path, self.opt.density_thresh, t=t)
390
+
391
+ # perform texture extraction
392
+ print(f"[INFO] unwrap uv...")
393
+ h = w = texture_size
394
+ mesh.auto_uv()
395
+ mesh.auto_normal()
396
+
397
+ albedo = torch.zeros((h, w, 3), device=self.device, dtype=torch.float32)
398
+ cnt = torch.zeros((h, w, 1), device=self.device, dtype=torch.float32)
399
+
400
+ vers = [0] * 8 + [-45] * 8 + [45] * 8 + [-89.9, 89.9]
401
+ hors = [0, 45, -45, 90, -90, 135, -135, 180] * 3 + [0, 0]
402
+
403
+ render_resolution = 512
404
+
405
+ import nvdiffrast.torch as dr
406
+
407
+ if not self.opt.force_cuda_rast and (not self.opt.gui or os.name == 'nt'):
408
+ glctx = dr.RasterizeGLContext()
409
+ else:
410
+ glctx = dr.RasterizeCudaContext()
411
+
412
+ for ver, hor in zip(vers, hors):
413
+ # render image
414
+ pose = orbit_camera(ver, hor, self.cam.radius)
415
+
416
+ cur_cam = MiniCam(
417
+ pose,
418
+ render_resolution,
419
+ render_resolution,
420
+ self.cam.fovy,
421
+ self.cam.fovx,
422
+ self.cam.near,
423
+ self.cam.far,
424
+ time=t
425
+ )
426
+
427
+ cur_out = self.renderer.render(cur_cam)
428
+
429
+ rgbs = cur_out["image"].unsqueeze(0) # [1, 3, H, W] in [0, 1]
430
+
431
+ # get coordinate in texture image
432
+ pose = torch.from_numpy(pose.astype(np.float32)).to(self.device)
433
+ proj = torch.from_numpy(self.cam.perspective.astype(np.float32)).to(self.device)
434
+
435
+ v_cam = torch.matmul(F.pad(mesh.v, pad=(0, 1), mode='constant', value=1.0), torch.inverse(pose).T).float().unsqueeze(0)
436
+ v_clip = v_cam @ proj.T
437
+ rast, rast_db = dr.rasterize(glctx, v_clip, mesh.f, (render_resolution, render_resolution))
438
+
439
+ depth, _ = dr.interpolate(-v_cam[..., [2]], rast, mesh.f) # [1, H, W, 1]
440
+ depth = depth.squeeze(0) # [H, W, 1]
441
+
442
+ alpha = (rast[0, ..., 3:] > 0).float()
443
+
444
+ uvs, _ = dr.interpolate(mesh.vt.unsqueeze(0), rast, mesh.ft) # [1, 512, 512, 2] in [0, 1]
445
+
446
+ # use normal to produce a back-project mask
447
+ normal, _ = dr.interpolate(mesh.vn.unsqueeze(0).contiguous(), rast, mesh.fn)
448
+ normal = safe_normalize(normal[0])
449
+
450
+ # rotated normal (where [0, 0, 1] always faces camera)
451
+ rot_normal = normal @ pose[:3, :3]
452
+ viewcos = rot_normal[..., [2]]
453
+
454
+ mask = (alpha > 0) & (viewcos > 0.5) # [H, W, 1]
455
+ mask = mask.view(-1)
456
+
457
+ uvs = uvs.view(-1, 2).clamp(0, 1)[mask]
458
+ rgbs = rgbs.view(3, -1).permute(1, 0)[mask].contiguous()
459
+
460
+ # update texture image
461
+ cur_albedo, cur_cnt = mipmap_linear_grid_put_2d(
462
+ h, w,
463
+ uvs[..., [1, 0]] * 2 - 1,
464
+ rgbs,
465
+ min_resolution=256,
466
+ return_count=True,
467
+ )
468
+
469
+ mask = cnt.squeeze(-1) < 0.1
470
+ albedo[mask] += cur_albedo[mask]
471
+ cnt[mask] += cur_cnt[mask]
472
+
473
+ mask = cnt.squeeze(-1) > 0
474
+ albedo[mask] = albedo[mask] / cnt[mask].repeat(1, 3)
475
+
476
+ mask = mask.view(h, w)
477
+
478
+ albedo = albedo.detach().cpu().numpy()
479
+ mask = mask.detach().cpu().numpy()
480
+
481
+ # dilate texture
482
+ from sklearn.neighbors import NearestNeighbors
483
+ from scipy.ndimage import binary_dilation, binary_erosion
484
+
485
+ inpaint_region = binary_dilation(mask, iterations=32)
486
+ inpaint_region[mask] = 0
487
+
488
+ search_region = mask.copy()
489
+ not_search_region = binary_erosion(search_region, iterations=3)
490
+ search_region[not_search_region] = 0
491
+
492
+ search_coords = np.stack(np.nonzero(search_region), axis=-1)
493
+ inpaint_coords = np.stack(np.nonzero(inpaint_region), axis=-1)
494
+
495
+ knn = NearestNeighbors(n_neighbors=1, algorithm="kd_tree").fit(
496
+ search_coords
497
+ )
498
+ _, indices = knn.kneighbors(inpaint_coords)
499
+
500
+ albedo[tuple(inpaint_coords.T)] = albedo[tuple(search_coords[indices[:, 0]].T)]
501
+
502
+ mesh.albedo = torch.from_numpy(albedo).to(self.device)
503
+ mesh.write(path)
504
+
505
+
506
+ elif mode == 'frames':
507
+ os.makedirs(os.path.join(self.opt.outdir, self.opt.save_path+'_frames'), exist_ok=True)
508
+ for t in range(self.opt.batch_size * interp):
509
+ tt = t / interp
510
+ path = os.path.join(self.opt.outdir, self.opt.save_path+'_frames', f'{t:03d}.ply')
511
+ self.renderer.gaussians.save_frame_ply(path, tt)
512
+ else:
513
+ path = os.path.join(self.opt.outdir, self.opt.save_path + '_4d_model.ply')
514
+ self.renderer.gaussians.save_ply(path)
515
+ self.renderer.gaussians.save_deformation(self.opt.outdir, self.opt.save_path)
516
+
517
+ print(f"[INFO] save model to {path}.")
518
+
519
+ # no gui mode
520
+ def train(self, iters=500, ui=False):
521
+ if self.gui:
522
+ from visualizer.visergui import ViserViewer
523
+ self.viser_gui = ViserViewer(device="cuda", viewer_port=8080)
524
+ if iters > 0:
525
+ self.prepare_train()
526
+ if self.gui:
527
+ self.viser_gui.set_renderer(self.renderer, self.fixed_cam)
528
+
529
+ for i in tqdm.trange(iters):
530
+ self.train_step()
531
+ if self.gui:
532
+ self.viser_gui.update()
533
+ if self.opt.mesh_format == 'frames':
534
+ self.save_model(mode='frames', interp=4)
535
+ elif self.opt.mesh_format == 'obj':
536
+ self.save_model(mode='geo+tex')
537
+
538
+ if self.opt.save_model:
539
+ self.save_model(mode='model')
540
+
541
+ # render eval
542
+ image_list =[]
543
+ nframes = self.opt.batch_size * 7 + 15 * 7
544
+ hor = 180
545
+ delta_hor = 45 / 15
546
+ delta_time = 1
547
+ for i in range(8):
548
+ time = 0
549
+ for j in range(self.opt.batch_size + 15):
550
+ pose = orbit_camera(self.opt.elevation, hor-180, self.opt.radius)
551
+ cur_cam = MiniCam(
552
+ pose,
553
+ 512,
554
+ 512,
555
+ self.cam.fovy,
556
+ self.cam.fovx,
557
+ self.cam.near,
558
+ self.cam.far,
559
+ time=time
560
+ )
561
+ with torch.no_grad():
562
+ outputs = self.renderer.render(cur_cam)
563
+
564
+ out = outputs["image"].cpu().detach().numpy().astype(np.float32)
565
+ out = np.transpose(out, (1, 2, 0))
566
+ out = np.uint8(out*255)
567
+ image_list.append(out)
568
+
569
+ time = (time + delta_time) % self.opt.batch_size
570
+ if j >= self.opt.batch_size:
571
+ hor = (hor+delta_hor) % 360
572
+
573
+
574
+ imageio.mimwrite(f'vis_data/{opt.save_path}.mp4', image_list, fps=7)
575
+
576
+ if self.gui:
577
+ while True:
578
+ self.viser_gui.update()
579
+
580
+ if __name__ == "__main__":
581
+ import argparse
582
+ from omegaconf import OmegaConf
583
+
584
+ parser = argparse.ArgumentParser()
585
+ parser.add_argument("--config", required=True, help="path to the yaml config file")
586
+ args, extras = parser.parse_known_args()
587
+
588
+ # override default config from cli
589
+ opt = OmegaConf.merge(OmegaConf.load(args.config), OmegaConf.from_cli(extras))
590
+ opt.save_path = os.path.splitext(os.path.basename(opt.input))[0] if opt.save_path == '' else opt.save_path
591
+
592
+
593
+ # auto find mesh from stage 1
594
+ opt.load = os.path.join(opt.outdir, opt.save_path + '_model.ply')
595
+
596
+ gui = GUI(opt)
597
+
598
+ gui.train(opt.iters)
599
+
600
+
601
+ # python main_4d.py --config configs/4d_low.yaml input=data/CONSISTENT4D_DATA/in-the-wild/blooming_rose
requirements.txt ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ tqdm
2
+ rich
3
+ ninja
4
+ numpy
5
+ pandas
6
+ scipy
7
+ scikit-learn
8
+ matplotlib
9
+ opencv-python
10
+ imageio
11
+ imageio-ffmpeg
12
+ omegaconf
13
+
14
+ torch==2.1.0 --index-url https://download.pytorch.org/whl/cu118
15
+ torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu118
16
+ torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118
17
+ xformer --index-url https://download.pytorch.org/whl/cu118 --no-deps
18
+ einops
19
+ plyfile
20
+ pygltflib
21
+ torchvision
22
+
23
+ # for stable-diffusion
24
+ huggingface_hub
25
+ diffusers
26
+ accelerate
27
+ transformers
28
+
29
+ rembg[gpu,cli]
30
+
31
+ # gradio demo
32
+ gradio
33
+ gradio-model4dgs
34
+
35
+ -e git+https://github.com/ashawkey/kiuikit.git@main#egg=kiui
scene/deformation.py ADDED
@@ -0,0 +1,241 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import functools
2
+ import math
3
+ import os
4
+ import time
5
+ from tkinter import W
6
+
7
+ import numpy as np
8
+ import torch
9
+ import torch.nn as nn
10
+ import torch.nn.functional as F
11
+ from torch.utils.cpp_extension import load
12
+ import torch.nn.init as init
13
+ from scene.hexplane import HexPlaneField
14
+
15
+
16
+ class Linear_Res(nn.Module):
17
+ def __init__(self, W):
18
+ super(Linear_Res, self).__init__()
19
+ self.main_stream = nn.Linear(W, W)
20
+
21
+ def forward(self, x):
22
+ x = F.relu(x)
23
+ return x + self.main_stream(x)
24
+
25
+
26
+ class Head_Res_Net(nn.Module):
27
+ def __init__(self, W, H):
28
+ super(Head_Res_Net, self).__init__()
29
+ self.W = W
30
+ self.H = H
31
+
32
+ self.feature_out = [Linear_Res(self.W)]
33
+ self.feature_out.append(nn.Linear(W, self.H))
34
+ self.feature_out = nn.Sequential(*self.feature_out)
35
+
36
+ def initialize_weights(self,):
37
+ for m in self.feature_out.modules():
38
+ if isinstance(m, nn.Linear):
39
+ init.constant_(m.weight, 0)
40
+ if m.bias is not None:
41
+ init.constant_(m.bias, 0)
42
+
43
+ def forward(self, x):
44
+ return self.feature_out(x)
45
+
46
+
47
+
48
+ class Deformation(nn.Module):
49
+ def __init__(self, D=8, W=256, input_ch=27, input_ch_time=9, skips=[], args=None, use_res=False):
50
+ super(Deformation, self).__init__()
51
+ self.D = D
52
+ self.W = W
53
+ self.input_ch = input_ch
54
+ self.input_ch_time = input_ch_time
55
+ self.skips = skips
56
+
57
+ self.no_grid = args.no_grid
58
+ self.grid = HexPlaneField(args.bounds, args.kplanes_config, args.multires)
59
+
60
+ self.use_res = use_res
61
+ if not self.use_res:
62
+ self.pos_deform, self.scales_deform, self.rotations_deform, self.opacity_deform = self.create_net()
63
+ else:
64
+ self.pos_deform, self.scales_deform, self.rotations_deform, self.opacity_deform = self.create_res_net()
65
+ self.args = args
66
+
67
+ def create_net(self):
68
+
69
+ mlp_out_dim = 0
70
+ if self.no_grid:
71
+ self.feature_out = [nn.Linear(4,self.W)]
72
+ else:
73
+ self.feature_out = [nn.Linear(mlp_out_dim + self.grid.feat_dim ,self.W)]
74
+
75
+ for i in range(self.D-1):
76
+ self.feature_out.append(nn.ReLU())
77
+ self.feature_out.append(nn.Linear(self.W,self.W))
78
+ self.feature_out = nn.Sequential(*self.feature_out)
79
+ output_dim = self.W
80
+ return \
81
+ nn.Sequential(nn.ReLU(),nn.Linear(self.W,self.W),nn.ReLU(),nn.Linear(self.W, 3)),\
82
+ nn.Sequential(nn.ReLU(),nn.Linear(self.W,self.W),nn.ReLU(),nn.Linear(self.W, 3)),\
83
+ nn.Sequential(nn.ReLU(),nn.Linear(self.W,self.W),nn.ReLU(),nn.Linear(self.W, 4)), \
84
+ nn.Sequential(nn.ReLU(),nn.Linear(self.W,self.W),nn.ReLU(),nn.Linear(self.W, 1))
85
+
86
+ def create_res_net(self,):
87
+
88
+ mlp_out_dim = 0
89
+
90
+ if self.no_grid:
91
+ self.feature_out = [nn.Linear(4,self.W)]
92
+ else:
93
+ self.feature_out = [nn.Linear(mlp_out_dim + self.grid.feat_dim ,self.W)]
94
+
95
+ for i in range(self.D-1):
96
+ self.feature_out.append(nn.ReLU())
97
+ self.feature_out.append(nn.Linear(self.W,self.W))
98
+ self.feature_out = nn.Sequential(*self.feature_out)
99
+
100
+ output_dim = self.W
101
+ return \
102
+ Head_Res_Net(self.W, 3), \
103
+ Head_Res_Net(self.W, 3), \
104
+ Head_Res_Net(self.W, 4), \
105
+ Head_Res_Net(self.W, 1)
106
+
107
+
108
+ def query_time(self, rays_pts_emb, scales_emb, rotations_emb, time_emb):
109
+ if self.args.no_mlp:
110
+ assert not self.no_grid
111
+ grid_feature = self.grid(rays_pts_emb[:,:3], time_emb[:,:1])
112
+ h = grid_feature
113
+ elif not self.use_res:
114
+ if self.no_grid:
115
+ h = torch.cat([rays_pts_emb[:,:3],time_emb[:,:1]],-1)
116
+ else:
117
+ grid_feature = self.grid(rays_pts_emb[:,:3], time_emb[:,:1])
118
+
119
+ h = grid_feature
120
+
121
+ h = self.feature_out(h)
122
+ else:
123
+ if self.no_grid:
124
+ h = torch.cat([rays_pts_emb[:,:3],time_emb[:,:1]],-1)
125
+ h = self.feature_out(h)
126
+ else:
127
+ grid_feature = self.grid(rays_pts_emb[:,:3], time_emb[:,:1])
128
+ h = self.feature_out(grid_feature)
129
+ return h
130
+
131
+ def forward(self, rays_pts_emb, scales_emb=None, rotations_emb=None, opacity = None, time_emb=None):
132
+ if time_emb is None:
133
+ return self.forward_static(rays_pts_emb[:,:3])
134
+ else:
135
+ return self.forward_dynamic(rays_pts_emb, scales_emb, rotations_emb, opacity, time_emb)
136
+
137
+ def forward_static(self, rays_pts_emb):
138
+ grid_feature = self.grid(rays_pts_emb[:,:3])
139
+ dx = self.static_mlp(grid_feature)
140
+ return rays_pts_emb[:, :3] + dx
141
+
142
+ def forward_dynamic(self,rays_pts_emb, scales_emb, rotations_emb, opacity_emb, time_emb):
143
+ hidden = self.query_time(rays_pts_emb, scales_emb, rotations_emb, time_emb).float()
144
+ if self.args.no_mlp:
145
+ return hidden[:, :3], hidden[:, 3:6], hidden[:, 6:10], hidden[:, 10:11]
146
+ dx = self.pos_deform(hidden)
147
+ pts = dx
148
+ if self.args.no_ds:
149
+ scales = scales_emb[:,:3]
150
+ else:
151
+ ds = self.scales_deform(hidden)
152
+ scales = ds
153
+ if self.args.no_dr:
154
+ rotations = rotations_emb[:,:4]
155
+ else:
156
+ dr = self.rotations_deform(hidden)
157
+ rotations = dr
158
+ if self.args.no_do:
159
+ opacity = opacity_emb[:,:1]
160
+ else:
161
+ do = self.opacity_deform(hidden)
162
+ opacity = do
163
+
164
+ return pts, scales, rotations, opacity
165
+ def get_mlp_parameters(self):
166
+ parameter_list = []
167
+ for name, param in self.named_parameters():
168
+ if "grid" not in name:
169
+ parameter_list.append(param)
170
+ return parameter_list
171
+ def get_grid_parameters(self):
172
+ return list(self.grid.parameters() )
173
+
174
+
175
+ class deform_network(nn.Module):
176
+ def __init__(self, args) :
177
+ super(deform_network, self).__init__()
178
+ net_width = args.net_width
179
+ timebase_pe = args.timebase_pe
180
+ defor_depth= args.defor_depth
181
+ posbase_pe= args.posebase_pe
182
+ scale_rotation_pe = args.scale_rotation_pe
183
+ opacity_pe = args.opacity_pe
184
+ timenet_width = args.timenet_width
185
+ timenet_output = args.timenet_output
186
+ times_ch = 2*timebase_pe+1
187
+ self.timenet = nn.Sequential(
188
+ nn.Linear(times_ch, timenet_width), nn.ReLU(),
189
+ nn.Linear(timenet_width, timenet_output))
190
+
191
+ self.use_res = args.use_res
192
+ if self.use_res:
193
+ print("Using zero-init and residual")
194
+ 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)
195
+ self.register_buffer('time_poc', torch.FloatTensor([(2**i) for i in range(timebase_pe)]))
196
+ self.register_buffer('pos_poc', torch.FloatTensor([(2**i) for i in range(posbase_pe)]))
197
+ self.register_buffer('rotation_scaling_poc', torch.FloatTensor([(2**i) for i in range(scale_rotation_pe)]))
198
+ self.register_buffer('opacity_poc', torch.FloatTensor([(2**i) for i in range(opacity_pe)]))
199
+ self.apply(initialize_weights)
200
+
201
+ if self.use_res:
202
+ self.deformation_net.pos_deform.initialize_weights()
203
+ self.deformation_net.scales_deform.initialize_weights()
204
+ self.deformation_net.rotations_deform.initialize_weights()
205
+ self.deformation_net.opacity_deform.initialize_weights()
206
+
207
+
208
+ def forward(self, point, scales=None, rotations=None, opacity=None, times_sel=None):
209
+ if times_sel is not None:
210
+ return self.forward_dynamic(point, scales, rotations, opacity, times_sel)
211
+ else:
212
+ return self.forward_static(point)
213
+
214
+
215
+ def forward_static(self, points):
216
+ points = self.deformation_net(points)
217
+ return points
218
+ def forward_dynamic(self, point, scales=None, rotations=None, opacity=None, times_sel=None):
219
+ means3D, scales, rotations, opacity = self.deformation_net( point,
220
+ scales,
221
+ rotations,
222
+ opacity,
223
+ times_sel)
224
+ return means3D, scales, rotations, opacity
225
+ def get_mlp_parameters(self):
226
+ return self.deformation_net.get_mlp_parameters() + list(self.timenet.parameters())
227
+ def get_grid_parameters(self):
228
+ return self.deformation_net.get_grid_parameters()
229
+
230
+
231
+ def initialize_weights(m):
232
+ if isinstance(m, nn.Linear):
233
+ init.xavier_uniform_(m.weight,gain=1)
234
+ if m.bias is not None:
235
+ init.xavier_uniform_(m.weight,gain=1)
236
+
237
+ def initialize_zeros_weights(m):
238
+ if isinstance(m, nn.Linear):
239
+ init.constant_(m.weight, 0)
240
+ if m.bias is not None:
241
+ init.constant_(m.bias, 0)
scene/hexplane.py ADDED
@@ -0,0 +1,182 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import itertools
2
+ import logging as log
3
+ from typing import Optional, Union, List, Dict, Sequence, Iterable, Collection, Callable
4
+
5
+ import torch
6
+ import torch.nn as nn
7
+ import torch.nn.functional as F
8
+
9
+
10
+ def get_normalized_directions(directions):
11
+ """SH encoding must be in the range [0, 1]
12
+
13
+ Args:
14
+ directions: batch of directions
15
+ """
16
+ return (directions + 1.0) / 2.0
17
+
18
+
19
+ def normalize_aabb(pts, aabb):
20
+ return (pts - aabb[0]) * (2.0 / (aabb[1] - aabb[0])) - 1.0
21
+ def grid_sample_wrapper(grid: torch.Tensor, coords: torch.Tensor, align_corners: bool = True) -> torch.Tensor:
22
+ grid_dim = coords.shape[-1]
23
+
24
+ if grid.dim() == grid_dim + 1:
25
+ # no batch dimension present, need to add it
26
+ grid = grid.unsqueeze(0)
27
+ if coords.dim() == 2:
28
+ coords = coords.unsqueeze(0)
29
+
30
+ if grid_dim == 2 or grid_dim == 3:
31
+ grid_sampler = F.grid_sample
32
+ else:
33
+ raise NotImplementedError(f"Grid-sample was called with {grid_dim}D data but is only "
34
+ f"implemented for 2 and 3D data.")
35
+
36
+ coords = coords.view([coords.shape[0]] + [1] * (grid_dim - 1) + list(coords.shape[1:]))
37
+ B, feature_dim = grid.shape[:2]
38
+ n = coords.shape[-2]
39
+ interp = grid_sampler(
40
+ grid, # [B, feature_dim, reso, ...]
41
+ coords, # [B, 1, ..., n, grid_dim]
42
+ align_corners=align_corners,
43
+ mode='bilinear', padding_mode='border')
44
+ interp = interp.view(B, feature_dim, n).transpose(-1, -2) # [B, n, feature_dim]
45
+ interp = interp.squeeze() # [B?, n, feature_dim?]
46
+ return interp
47
+
48
+ def init_grid_param(
49
+ grid_nd: int,
50
+ in_dim: int,
51
+ out_dim: int,
52
+ reso: Sequence[int],
53
+ a: float = 0.1,
54
+ b: float = 0.5):
55
+ assert in_dim == len(reso), "Resolution must have same number of elements as input-dimension"
56
+ has_time_planes = in_dim == 4
57
+ assert grid_nd <= in_dim
58
+ coo_combs = list(itertools.combinations(range(in_dim), grid_nd))
59
+ grid_coefs = nn.ParameterList()
60
+ for ci, coo_comb in enumerate(coo_combs):
61
+ new_grid_coef = nn.Parameter(torch.empty(
62
+ [1, out_dim] + [reso[cc] for cc in coo_comb[::-1]]
63
+ ))
64
+ if has_time_planes and 3 in coo_comb: # Initialize time planes to 1
65
+ nn.init.ones_(new_grid_coef)
66
+ else:
67
+ nn.init.uniform_(new_grid_coef, a=a, b=b)
68
+ grid_coefs.append(new_grid_coef)
69
+
70
+ return grid_coefs
71
+
72
+
73
+ def interpolate_ms_features(pts: torch.Tensor,
74
+ ms_grids: Collection[Iterable[nn.Module]],
75
+ grid_dimensions: int,
76
+ concat_features: bool,
77
+ num_levels: Optional[int],
78
+ ) -> torch.Tensor:
79
+ coo_combs = list(itertools.combinations(
80
+ range(pts.shape[-1]), grid_dimensions)
81
+ )
82
+ if num_levels is None:
83
+ num_levels = len(ms_grids)
84
+ multi_scale_interp = [] if concat_features else 0.
85
+ grid: nn.ParameterList
86
+ for scale_id, grid in enumerate(ms_grids[:num_levels]):
87
+ interp_space = 1.
88
+ for ci, coo_comb in enumerate(coo_combs):
89
+ # interpolate in plane
90
+ feature_dim = grid[ci].shape[1] # shape of grid[ci]: 1, out_dim, *reso
91
+ interp_out_plane = (
92
+ grid_sample_wrapper(grid[ci], pts[..., coo_comb])
93
+ .view(-1, feature_dim)
94
+ )
95
+ # compute product over planes
96
+ interp_space = interp_space * interp_out_plane
97
+
98
+ # combine over scales
99
+ if concat_features:
100
+ multi_scale_interp.append(interp_space)
101
+ else:
102
+ multi_scale_interp = multi_scale_interp + interp_space
103
+
104
+ if concat_features:
105
+ multi_scale_interp = torch.cat(multi_scale_interp, dim=-1)
106
+ return multi_scale_interp
107
+
108
+
109
+ class HexPlaneField(nn.Module):
110
+ def __init__(
111
+ self,
112
+
113
+ bounds,
114
+ planeconfig,
115
+ multires
116
+ ) -> None:
117
+ super().__init__()
118
+ aabb = torch.tensor([[bounds,bounds,bounds],
119
+ [-bounds,-bounds,-bounds]])
120
+ self.aabb = nn.Parameter(aabb, requires_grad=False)
121
+ self.grid_config = [planeconfig]
122
+ self.multiscale_res_multipliers = multires
123
+ self.concat_features = True
124
+
125
+ # 1. Init planes
126
+ self.grids = nn.ModuleList()
127
+ self.feat_dim = 0
128
+ for res in self.multiscale_res_multipliers:
129
+ # initialize coordinate grid
130
+ config = self.grid_config[0].copy()
131
+ # Resolution fix: multi-res only on spatial planes
132
+ config["resolution"] = [
133
+ r * res for r in config["resolution"][:3]
134
+ ] + config["resolution"][3:]
135
+ gp = init_grid_param(
136
+ grid_nd=config["grid_dimensions"],
137
+ in_dim=config["input_coordinate_dim"],
138
+ out_dim=config["output_coordinate_dim"],
139
+ reso=config["resolution"],
140
+ )
141
+ # shape[1] is out-dim - Concatenate over feature len for each scale
142
+ if self.concat_features:
143
+ self.feat_dim += gp[-1].shape[1]
144
+ else:
145
+ self.feat_dim = gp[-1].shape[1]
146
+ self.grids.append(gp)
147
+ # print(f"Initialized model grids: {self.grids}")
148
+ print("feature_dim:",self.feat_dim)
149
+
150
+
151
+ def set_aabb(self,xyz_max, xyz_min):
152
+ aabb = torch.tensor([
153
+ xyz_max,
154
+ xyz_min
155
+ ])
156
+ self.aabb = nn.Parameter(aabb,requires_grad=True)
157
+ print("Voxel Plane: set aabb=",self.aabb)
158
+
159
+ def get_density(self, pts: torch.Tensor, timestamps: Optional[torch.Tensor] = None):
160
+ """Computes and returns the densities."""
161
+
162
+ pts = normalize_aabb(pts, self.aabb)
163
+ pts = torch.cat((pts, timestamps), dim=-1) # [n_rays, n_samples, 4]
164
+
165
+ pts = pts.reshape(-1, pts.shape[-1])
166
+ features = interpolate_ms_features(
167
+ pts, ms_grids=self.grids, # noqa
168
+ grid_dimensions=self.grid_config[0]["grid_dimensions"],
169
+ concat_features=self.concat_features, num_levels=None)
170
+ if len(features) < 1:
171
+ features = torch.zeros((0, 1)).to(features.device)
172
+
173
+
174
+ return features
175
+
176
+ def forward(self,
177
+ pts: torch.Tensor,
178
+ timestamps: Optional[torch.Tensor] = None):
179
+
180
+ features = self.get_density(pts, timestamps)
181
+
182
+ return features
scene/regulation.py ADDED
@@ -0,0 +1,176 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import abc
2
+ import os
3
+ from typing import Sequence
4
+
5
+ import matplotlib.pyplot as plt
6
+ import numpy as np
7
+ import torch
8
+ import torch.optim.lr_scheduler
9
+ from torch import nn
10
+
11
+
12
+
13
+ def compute_plane_tv(t):
14
+ batch_size, c, h, w = t.shape
15
+ count_h = batch_size * c * (h - 1) * w
16
+ count_w = batch_size * c * h * (w - 1)
17
+ h_tv = torch.square(t[..., 1:, :] - t[..., :h-1, :]).sum()
18
+ w_tv = torch.square(t[..., :, 1:] - t[..., :, :w-1]).sum()
19
+ return 2 * (h_tv / count_h + w_tv / count_w) # This is summing over batch and c instead of avg
20
+
21
+
22
+ def compute_plane_smoothness(t):
23
+ batch_size, c, h, w = t.shape
24
+ # Convolve with a second derivative filter, in the time dimension which is dimension 2
25
+ first_difference = t[..., 1:, :] - t[..., :h-1, :] # [batch, c, h-1, w]
26
+ second_difference = first_difference[..., 1:, :] - first_difference[..., :h-2, :] # [batch, c, h-2, w]
27
+ # Take the L2 norm of the result
28
+ return torch.square(second_difference).mean()
29
+
30
+
31
+ class Regularizer():
32
+ def __init__(self, reg_type, initialization):
33
+ self.reg_type = reg_type
34
+ self.initialization = initialization
35
+ self.weight = float(self.initialization)
36
+ self.last_reg = None
37
+
38
+ def step(self, global_step):
39
+ pass
40
+
41
+ def report(self, d):
42
+ if self.last_reg is not None:
43
+ d[self.reg_type].update(self.last_reg.item())
44
+
45
+ def regularize(self, *args, **kwargs) -> torch.Tensor:
46
+ out = self._regularize(*args, **kwargs) * self.weight
47
+ self.last_reg = out.detach()
48
+ return out
49
+
50
+ @abc.abstractmethod
51
+ def _regularize(self, *args, **kwargs) -> torch.Tensor:
52
+ raise NotImplementedError()
53
+
54
+ def __str__(self):
55
+ return f"Regularizer({self.reg_type}, weight={self.weight})"
56
+
57
+
58
+ class PlaneTV(Regularizer):
59
+ def __init__(self, initial_value, what: str = 'field'):
60
+ if what not in {'field', 'proposal_network'}:
61
+ raise ValueError(f'what must be one of "field" or "proposal_network" '
62
+ f'but {what} was passed.')
63
+ name = f'planeTV-{what[:2]}'
64
+ super().__init__(name, initial_value)
65
+ self.what = what
66
+
67
+ def step(self, global_step):
68
+ pass
69
+
70
+ def _regularize(self, model, **kwargs):
71
+ multi_res_grids: Sequence[nn.ParameterList]
72
+ if self.what == 'field':
73
+ multi_res_grids = model.field.grids
74
+ elif self.what == 'proposal_network':
75
+ multi_res_grids = [p.grids for p in model.proposal_networks]
76
+ else:
77
+ raise NotImplementedError(self.what)
78
+ total = 0
79
+ # Note: input to compute_plane_tv should be of shape [batch_size, c, h, w]
80
+ for grids in multi_res_grids:
81
+ if len(grids) == 3:
82
+ spatial_grids = [0, 1, 2]
83
+ else:
84
+ spatial_grids = [0, 1, 3] # These are the spatial grids; the others are spatiotemporal
85
+ for grid_id in spatial_grids:
86
+ total += compute_plane_tv(grids[grid_id])
87
+ for grid in grids:
88
+ # grid: [1, c, h, w]
89
+ total += compute_plane_tv(grid)
90
+ return total
91
+
92
+
93
+ class TimeSmoothness(Regularizer):
94
+ def __init__(self, initial_value, what: str = 'field'):
95
+ if what not in {'field', 'proposal_network'}:
96
+ raise ValueError(f'what must be one of "field" or "proposal_network" '
97
+ f'but {what} was passed.')
98
+ name = f'time-smooth-{what[:2]}'
99
+ super().__init__(name, initial_value)
100
+ self.what = what
101
+
102
+ def _regularize(self, model, **kwargs) -> torch.Tensor:
103
+ multi_res_grids: Sequence[nn.ParameterList]
104
+ if self.what == 'field':
105
+ multi_res_grids = model.field.grids
106
+ elif self.what == 'proposal_network':
107
+ multi_res_grids = [p.grids for p in model.proposal_networks]
108
+ else:
109
+ raise NotImplementedError(self.what)
110
+ total = 0
111
+ # model.grids is 6 x [1, rank * F_dim, reso, reso]
112
+ for grids in multi_res_grids:
113
+ if len(grids) == 3:
114
+ time_grids = []
115
+ else:
116
+ time_grids = [2, 4, 5]
117
+ for grid_id in time_grids:
118
+ total += compute_plane_smoothness(grids[grid_id])
119
+ return torch.as_tensor(total)
120
+
121
+
122
+
123
+ class L1ProposalNetwork(Regularizer):
124
+ def __init__(self, initial_value):
125
+ super().__init__('l1-proposal-network', initial_value)
126
+
127
+ def _regularize(self, model, **kwargs) -> torch.Tensor:
128
+ grids = [p.grids for p in model.proposal_networks]
129
+ total = 0.0
130
+ for pn_grids in grids:
131
+ for grid in pn_grids:
132
+ total += torch.abs(grid).mean()
133
+ return torch.as_tensor(total)
134
+
135
+
136
+ class DepthTV(Regularizer):
137
+ def __init__(self, initial_value):
138
+ super().__init__('tv-depth', initial_value)
139
+
140
+ def _regularize(self, model, model_out, **kwargs) -> torch.Tensor:
141
+ depth = model_out['depth']
142
+ tv = compute_plane_tv(
143
+ depth.reshape(64, 64)[None, None, :, :]
144
+ )
145
+ return tv
146
+
147
+
148
+ class L1TimePlanes(Regularizer):
149
+ def __init__(self, initial_value, what='field'):
150
+ if what not in {'field', 'proposal_network'}:
151
+ raise ValueError(f'what must be one of "field" or "proposal_network" '
152
+ f'but {what} was passed.')
153
+ super().__init__(f'l1-time-{what[:2]}', initial_value)
154
+ self.what = what
155
+
156
+ def _regularize(self, model, **kwargs) -> torch.Tensor:
157
+ # model.grids is 6 x [1, rank * F_dim, reso, reso]
158
+ multi_res_grids: Sequence[nn.ParameterList]
159
+ if self.what == 'field':
160
+ multi_res_grids = model.field.grids
161
+ elif self.what == 'proposal_network':
162
+ multi_res_grids = [p.grids for p in model.proposal_networks]
163
+ else:
164
+ raise NotImplementedError(self.what)
165
+
166
+ total = 0.0
167
+ for grids in multi_res_grids:
168
+ if len(grids) == 3:
169
+ continue
170
+ else:
171
+ # These are the spatiotemporal grids
172
+ spatiotemporal_grids = [2, 4, 5]
173
+ for grid_id in spatiotemporal_grids:
174
+ total += torch.abs(1 - grids[grid_id]).mean()
175
+ return torch.as_tensor(total)
176
+
scene/utils.py ADDED
@@ -0,0 +1,429 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import copy
2
+ import json
3
+ import math
4
+ import os
5
+ import pathlib
6
+ from typing import Any, Callable, List, Optional, Text, Tuple, Union
7
+
8
+ import numpy as np
9
+ import scipy.signal
10
+ import torch
11
+ import torch.nn as nn
12
+ import torch.nn.functional as F
13
+ from torch import Tensor
14
+
15
+
16
+ PRNGKey = Any
17
+ Shape = Tuple[int]
18
+ Dtype = Any # this could be a real type?
19
+ Array = Any
20
+ Activation = Callable[[Array], Array]
21
+ Initializer = Callable[[PRNGKey, Shape, Dtype], Array]
22
+ Normalizer = Callable[[], Callable[[Array], Array]]
23
+ PathType = Union[Text, pathlib.PurePosixPath]
24
+
25
+ from pathlib import PurePosixPath as GPath
26
+
27
+
28
+ def _compute_residual_and_jacobian(
29
+ x: np.ndarray,
30
+ y: np.ndarray,
31
+ xd: np.ndarray,
32
+ yd: np.ndarray,
33
+ k1: float = 0.0,
34
+ k2: float = 0.0,
35
+ k3: float = 0.0,
36
+ p1: float = 0.0,
37
+ p2: float = 0.0,
38
+ ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray,
39
+ np.ndarray]:
40
+ """Auxiliary function of radial_and_tangential_undistort()."""
41
+
42
+ r = x * x + y * y
43
+ d = 1.0 + r * (k1 + r * (k2 + k3 * r))
44
+
45
+ fx = d * x + 2 * p1 * x * y + p2 * (r + 2 * x * x) - xd
46
+ fy = d * y + 2 * p2 * x * y + p1 * (r + 2 * y * y) - yd
47
+
48
+ # Compute derivative of d over [x, y]
49
+ d_r = (k1 + r * (2.0 * k2 + 3.0 * k3 * r))
50
+ d_x = 2.0 * x * d_r
51
+ d_y = 2.0 * y * d_r
52
+
53
+ # Compute derivative of fx over x and y.
54
+ fx_x = d + d_x * x + 2.0 * p1 * y + 6.0 * p2 * x
55
+ fx_y = d_y * x + 2.0 * p1 * x + 2.0 * p2 * y
56
+
57
+ # Compute derivative of fy over x and y.
58
+ fy_x = d_x * y + 2.0 * p2 * y + 2.0 * p1 * x
59
+ fy_y = d + d_y * y + 2.0 * p2 * x + 6.0 * p1 * y
60
+
61
+ return fx, fy, fx_x, fx_y, fy_x, fy_y
62
+
63
+
64
+ def _radial_and_tangential_undistort(
65
+ xd: np.ndarray,
66
+ yd: np.ndarray,
67
+ k1: float = 0,
68
+ k2: float = 0,
69
+ k3: float = 0,
70
+ p1: float = 0,
71
+ p2: float = 0,
72
+ eps: float = 1e-9,
73
+ max_iterations=10) -> Tuple[np.ndarray, np.ndarray]:
74
+ """Computes undistorted (x, y) from (xd, yd)."""
75
+ # Initialize from the distorted point.
76
+ x = xd.copy()
77
+ y = yd.copy()
78
+
79
+ for _ in range(max_iterations):
80
+ fx, fy, fx_x, fx_y, fy_x, fy_y = _compute_residual_and_jacobian(
81
+ x=x, y=y, xd=xd, yd=yd, k1=k1, k2=k2, k3=k3, p1=p1, p2=p2)
82
+ denominator = fy_x * fx_y - fx_x * fy_y
83
+ x_numerator = fx * fy_y - fy * fx_y
84
+ y_numerator = fy * fx_x - fx * fy_x
85
+ step_x = np.where(
86
+ np.abs(denominator) > eps, x_numerator / denominator,
87
+ np.zeros_like(denominator))
88
+ step_y = np.where(
89
+ np.abs(denominator) > eps, y_numerator / denominator,
90
+ np.zeros_like(denominator))
91
+
92
+ x = x + step_x
93
+ y = y + step_y
94
+
95
+ return x, y
96
+
97
+
98
+ class Camera:
99
+ """Class to handle camera geometry."""
100
+
101
+ def __init__(self,
102
+ orientation: np.ndarray,
103
+ position: np.ndarray,
104
+ focal_length: Union[np.ndarray, float],
105
+ principal_point: np.ndarray,
106
+ image_size: np.ndarray,
107
+ skew: Union[np.ndarray, float] = 0.0,
108
+ pixel_aspect_ratio: Union[np.ndarray, float] = 1.0,
109
+ radial_distortion: Optional[np.ndarray] = None,
110
+ tangential_distortion: Optional[np.ndarray] = None,
111
+ dtype=np.float32):
112
+ """Constructor for camera class."""
113
+ if radial_distortion is None:
114
+ radial_distortion = np.array([0.0, 0.0, 0.0], dtype)
115
+ if tangential_distortion is None:
116
+ tangential_distortion = np.array([0.0, 0.0], dtype)
117
+
118
+ self.orientation = np.array(orientation, dtype)
119
+ self.position = np.array(position, dtype)
120
+ self.focal_length = np.array(focal_length, dtype)
121
+ self.principal_point = np.array(principal_point, dtype)
122
+ self.skew = np.array(skew, dtype)
123
+ self.pixel_aspect_ratio = np.array(pixel_aspect_ratio, dtype)
124
+ self.radial_distortion = np.array(radial_distortion, dtype)
125
+ self.tangential_distortion = np.array(tangential_distortion, dtype)
126
+ self.image_size = np.array(image_size, np.uint32)
127
+ self.dtype = dtype
128
+
129
+ @classmethod
130
+ def from_json(cls, path: PathType):
131
+ """Loads a JSON camera into memory."""
132
+ path = GPath(path)
133
+ # with path.open('r') as fp:
134
+ with open(path, 'r') as fp:
135
+ camera_json = json.load(fp)
136
+
137
+ # Fix old camera JSON.
138
+ if 'tangential' in camera_json:
139
+ camera_json['tangential_distortion'] = camera_json['tangential']
140
+
141
+ return cls(
142
+ orientation=np.asarray(camera_json['orientation']),
143
+ position=np.asarray(camera_json['position']),
144
+ focal_length=camera_json['focal_length'],
145
+ principal_point=np.asarray(camera_json['principal_point']),
146
+ skew=camera_json['skew'],
147
+ pixel_aspect_ratio=camera_json['pixel_aspect_ratio'],
148
+ radial_distortion=np.asarray(camera_json['radial_distortion']),
149
+ tangential_distortion=np.asarray(camera_json['tangential_distortion']),
150
+ image_size=np.asarray(camera_json['image_size']),
151
+ )
152
+
153
+ def to_json(self):
154
+ return {
155
+ k: (v.tolist() if hasattr(v, 'tolist') else v)
156
+ for k, v in self.get_parameters().items()
157
+ }
158
+
159
+ def get_parameters(self):
160
+ return {
161
+ 'orientation': self.orientation,
162
+ 'position': self.position,
163
+ 'focal_length': self.focal_length,
164
+ 'principal_point': self.principal_point,
165
+ 'skew': self.skew,
166
+ 'pixel_aspect_ratio': self.pixel_aspect_ratio,
167
+ 'radial_distortion': self.radial_distortion,
168
+ 'tangential_distortion': self.tangential_distortion,
169
+ 'image_size': self.image_size,
170
+ }
171
+
172
+ @property
173
+ def scale_factor_x(self):
174
+ return self.focal_length
175
+
176
+ @property
177
+ def scale_factor_y(self):
178
+ return self.focal_length * self.pixel_aspect_ratio
179
+
180
+ @property
181
+ def principal_point_x(self):
182
+ return self.principal_point[0]
183
+
184
+ @property
185
+ def principal_point_y(self):
186
+ return self.principal_point[1]
187
+
188
+ @property
189
+ def has_tangential_distortion(self):
190
+ return any(self.tangential_distortion != 0.0)
191
+
192
+ @property
193
+ def has_radial_distortion(self):
194
+ return any(self.radial_distortion != 0.0)
195
+
196
+ @property
197
+ def image_size_y(self):
198
+ return self.image_size[1]
199
+
200
+ @property
201
+ def image_size_x(self):
202
+ return self.image_size[0]
203
+
204
+ @property
205
+ def image_shape(self):
206
+ return self.image_size_y, self.image_size_x
207
+
208
+ @property
209
+ def optical_axis(self):
210
+ return self.orientation[2, :]
211
+
212
+ @property
213
+ def translation(self):
214
+ return -np.matmul(self.orientation, self.position)
215
+
216
+ def pixel_to_local_rays(self, pixels: np.ndarray):
217
+ """Returns the local ray directions for the provided pixels."""
218
+ y = ((pixels[..., 1] - self.principal_point_y) / self.scale_factor_y)
219
+ x = ((pixels[..., 0] - self.principal_point_x - y * self.skew) /
220
+ self.scale_factor_x)
221
+
222
+ if self.has_radial_distortion or self.has_tangential_distortion:
223
+ x, y = _radial_and_tangential_undistort(
224
+ x,
225
+ y,
226
+ k1=self.radial_distortion[0],
227
+ k2=self.radial_distortion[1],
228
+ k3=self.radial_distortion[2],
229
+ p1=self.tangential_distortion[0],
230
+ p2=self.tangential_distortion[1])
231
+
232
+ dirs = np.stack([x, y, np.ones_like(x)], axis=-1)
233
+ return dirs / np.linalg.norm(dirs, axis=-1, keepdims=True)
234
+
235
+ def pixels_to_rays(self, pixels: np.ndarray) -> np.ndarray:
236
+ """Returns the rays for the provided pixels.
237
+
238
+ Args:
239
+ pixels: [A1, ..., An, 2] tensor or np.array containing 2d pixel positions.
240
+
241
+ Returns:
242
+ An array containing the normalized ray directions in world coordinates.
243
+ """
244
+ if pixels.shape[-1] != 2:
245
+ raise ValueError('The last dimension of pixels must be 2.')
246
+ if pixels.dtype != self.dtype:
247
+ raise ValueError(f'pixels dtype ({pixels.dtype!r}) must match camera '
248
+ f'dtype ({self.dtype!r})')
249
+
250
+ batch_shape = pixels.shape[:-1]
251
+ pixels = np.reshape(pixels, (-1, 2))
252
+
253
+ local_rays_dir = self.pixel_to_local_rays(pixels)
254
+ rays_dir = np.matmul(self.orientation.T, local_rays_dir[..., np.newaxis])
255
+ rays_dir = np.squeeze(rays_dir, axis=-1)
256
+
257
+ # Normalize rays.
258
+ rays_dir /= np.linalg.norm(rays_dir, axis=-1, keepdims=True)
259
+ rays_dir = rays_dir.reshape((*batch_shape, 3))
260
+ return rays_dir
261
+
262
+ def pixels_to_points(self, pixels: np.ndarray, depth: np.ndarray):
263
+ rays_through_pixels = self.pixels_to_rays(pixels)
264
+ cosa = np.matmul(rays_through_pixels, self.optical_axis)
265
+ points = (
266
+ rays_through_pixels * depth[..., np.newaxis] / cosa[..., np.newaxis] +
267
+ self.position)
268
+ return points
269
+
270
+ def points_to_local_points(self, points: np.ndarray):
271
+ translated_points = points - self.position
272
+ local_points = (np.matmul(self.orientation, translated_points.T)).T
273
+ return local_points
274
+
275
+ def project(self, points: np.ndarray):
276
+ """Projects a 3D point (x,y,z) to a pixel position (x,y)."""
277
+ batch_shape = points.shape[:-1]
278
+ points = points.reshape((-1, 3))
279
+ local_points = self.points_to_local_points(points)
280
+
281
+ # Get normalized local pixel positions.
282
+ x = local_points[..., 0] / local_points[..., 2]
283
+ y = local_points[..., 1] / local_points[..., 2]
284
+ r2 = x**2 + y**2
285
+
286
+ # Apply radial distortion.
287
+ distortion = 1.0 + r2 * (
288
+ self.radial_distortion[0] + r2 *
289
+ (self.radial_distortion[1] + self.radial_distortion[2] * r2))
290
+
291
+ # Apply tangential distortion.
292
+ x_times_y = x * y
293
+ x = (
294
+ x * distortion + 2.0 * self.tangential_distortion[0] * x_times_y +
295
+ self.tangential_distortion[1] * (r2 + 2.0 * x**2))
296
+ y = (
297
+ y * distortion + 2.0 * self.tangential_distortion[1] * x_times_y +
298
+ self.tangential_distortion[0] * (r2 + 2.0 * y**2))
299
+
300
+ # Map the distorted ray to the image plane and return the depth.
301
+ pixel_x = self.focal_length * x + self.skew * y + self.principal_point_x
302
+ pixel_y = (self.focal_length * self.pixel_aspect_ratio * y
303
+ + self.principal_point_y)
304
+
305
+ pixels = np.stack([pixel_x, pixel_y], axis=-1)
306
+ return pixels.reshape((*batch_shape, 2))
307
+
308
+ def get_pixel_centers(self):
309
+ """Returns the pixel centers."""
310
+ xx, yy = np.meshgrid(np.arange(self.image_size_x, dtype=self.dtype),
311
+ np.arange(self.image_size_y, dtype=self.dtype))
312
+ return np.stack([xx, yy], axis=-1) + 0.5
313
+
314
+ def scale(self, scale: float):
315
+ """Scales the camera."""
316
+ if scale <= 0:
317
+ raise ValueError('scale needs to be positive.')
318
+
319
+ new_camera = Camera(
320
+ orientation=self.orientation.copy(),
321
+ position=self.position.copy(),
322
+ focal_length=self.focal_length * scale,
323
+ principal_point=self.principal_point.copy() * scale,
324
+ skew=self.skew,
325
+ pixel_aspect_ratio=self.pixel_aspect_ratio,
326
+ radial_distortion=self.radial_distortion.copy(),
327
+ tangential_distortion=self.tangential_distortion.copy(),
328
+ image_size=np.array((int(round(self.image_size[0] * scale)),
329
+ int(round(self.image_size[1] * scale)))),
330
+ )
331
+ return new_camera
332
+
333
+ def look_at(self, position, look_at, up, eps=1e-6):
334
+ """Creates a copy of the camera which looks at a given point.
335
+
336
+ Copies the provided vision_sfm camera and returns a new camera that is
337
+ positioned at `camera_position` while looking at `look_at_position`.
338
+ Camera intrinsics are copied by this method. A common value for the
339
+ up_vector is (0, 1, 0).
340
+
341
+ Args:
342
+ position: A (3,) numpy array representing the position of the camera.
343
+ look_at: A (3,) numpy array representing the location the camera
344
+ looks at.
345
+ up: A (3,) numpy array representing the up direction, whose
346
+ projection is parallel to the y-axis of the image plane.
347
+ eps: a small number to prevent divides by zero.
348
+
349
+ Returns:
350
+ A new camera that is copied from the original but is positioned and
351
+ looks at the provided coordinates.
352
+
353
+ Raises:
354
+ ValueError: If the camera position and look at position are very close
355
+ to each other or if the up-vector is parallel to the requested optical
356
+ axis.
357
+ """
358
+
359
+ look_at_camera = self.copy()
360
+ optical_axis = look_at - position
361
+ norm = np.linalg.norm(optical_axis)
362
+ if norm < eps:
363
+ raise ValueError('The camera center and look at position are too close.')
364
+ optical_axis /= norm
365
+
366
+ right_vector = np.cross(optical_axis, up)
367
+ norm = np.linalg.norm(right_vector)
368
+ if norm < eps:
369
+ raise ValueError('The up-vector is parallel to the optical axis.')
370
+ right_vector /= norm
371
+
372
+ # The three directions here are orthogonal to each other and form a right
373
+ # handed coordinate system.
374
+ camera_rotation = np.identity(3)
375
+ camera_rotation[0, :] = right_vector
376
+ camera_rotation[1, :] = np.cross(optical_axis, right_vector)
377
+ camera_rotation[2, :] = optical_axis
378
+
379
+ look_at_camera.position = position
380
+ look_at_camera.orientation = camera_rotation
381
+ return look_at_camera
382
+
383
+ def crop_image_domain(
384
+ self, left: int = 0, right: int = 0, top: int = 0, bottom: int = 0):
385
+ """Returns a copy of the camera with adjusted image bounds.
386
+
387
+ Args:
388
+ left: number of pixels by which to reduce (or augment, if negative) the
389
+ image domain at the associated boundary.
390
+ right: likewise.
391
+ top: likewise.
392
+ bottom: likewise.
393
+
394
+ The crop parameters may not cause the camera image domain dimensions to
395
+ become non-positive.
396
+
397
+ Returns:
398
+ A camera with adjusted image dimensions. The focal length is unchanged,
399
+ and the principal point is updated to preserve the original principal
400
+ axis.
401
+ """
402
+
403
+ crop_left_top = np.array([left, top])
404
+ crop_right_bottom = np.array([right, bottom])
405
+ new_resolution = self.image_size - crop_left_top - crop_right_bottom
406
+ new_principal_point = self.principal_point - crop_left_top
407
+ if np.any(new_resolution <= 0):
408
+ raise ValueError('Crop would result in non-positive image dimensions.')
409
+
410
+ new_camera = self.copy()
411
+ new_camera.image_size = np.array([int(new_resolution[0]),
412
+ int(new_resolution[1])])
413
+ new_camera.principal_point = np.array([new_principal_point[0],
414
+ new_principal_point[1]])
415
+ return new_camera
416
+
417
+ def copy(self):
418
+ return copy.deepcopy(self)
419
+
420
+
421
+ ''' Misc
422
+ '''
423
+ mse2psnr = lambda x : -10. * torch.log10(x)
424
+ to8b = lambda x : (255*np.clip(x,0,1)).astype(np.uint8)
425
+
426
+
427
+
428
+ ''' Checkpoint utils
429
+ '''
scripts/add_bg_to_gt.py ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import cv2
3
+
4
+ os.makedirs('data/CONSISTENT4D_DATA/test_dataset/eval_gt_rgb', exist_ok=True)
5
+ file_list = []
6
+ for img_name in ['aurorus', 'crocodile', 'guppie', 'monster', 'pistol', 'skull', 'trump']:
7
+ os.makedirs(f'data/CONSISTENT4D_DATA/test_dataset/eval_gt_rgb/{img_name}', exist_ok=True)
8
+ for view in range(4):
9
+ os.makedirs(f'datdata/CONSISTENT4D_DATAa/test_dataset/eval_gt_rgb/{img_name}/eval_{view}', exist_ok=True)
10
+ for t in range(32):
11
+ file_list.append(f'data/CONSISTENT4D_DATA/test_dataset/eval_gt/{img_name}/eval_{view}/{t}.png')
12
+ for file in file_list:
13
+ img = cv2.imread(file, cv2.IMREAD_UNCHANGED)
14
+ input_mask = img[..., 3:]
15
+ input_mask = input_mask / 255.
16
+ input_img = img[..., :3] * input_mask + (1 - input_mask) * 255
17
+ fpath = file.replace('eval_gt', 'eval_gt_rgb')
18
+ cv2.imwrite(fpath, input_img)
scripts/convert_obj_to_video.py ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import glob
3
+ import argparse
4
+
5
+ parser = argparse.ArgumentParser()
6
+ parser.add_argument('--dir', default='logs', type=str, help='Directory where obj files are stored')
7
+ parser.add_argument('--out', default='videos', type=str, help='Directory where videos will be saved')
8
+ args = parser.parse_args()
9
+
10
+ out = args.out
11
+ os.makedirs(out, exist_ok=True)
12
+
13
+ files = glob.glob(f'{args.dir}/*.obj')
14
+ for f in files:
15
+ name = os.path.basename(f)
16
+ # first stage model, ignore
17
+ if name.endswith('_mesh.obj'):
18
+ continue
19
+ print(f'[INFO] process {name}')
20
+ os.system(f"python -m kiui.render {f} --save_video {os.path.join(out, name.replace('.obj', '.mp4'))} ")
scripts/gen_vid.py ADDED
@@ -0,0 +1,117 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+ from diffusers import StableVideoDiffusionPipeline
4
+
5
+ from PIL import Image
6
+ import numpy as np
7
+
8
+ import cv2
9
+ import rembg
10
+
11
+ import argparse
12
+ import imageio
13
+ import os
14
+
15
+ def add_margin(pil_img, top, right, bottom, left, color):
16
+ width, height = pil_img.size
17
+ new_width = width + right + left
18
+ new_height = height + top + bottom
19
+ result = Image.new(pil_img.mode, (new_width, new_height), color)
20
+ result.paste(pil_img, (left, top))
21
+ return result
22
+
23
+ def resize_image(image, output_size=(1024, 576)):
24
+ image = image.resize((output_size[1],output_size[1]))
25
+ pad_size = (output_size[0]-output_size[1]) //2
26
+ image = add_margin(image, 0, pad_size, 0, pad_size, tuple(np.array(image)[0,0]))
27
+ return image
28
+
29
+
30
+ def load_image(file, W, H, bg='white'):
31
+ # load image
32
+ print(f'[INFO] load image from {file}...')
33
+ img = cv2.imread(file, cv2.IMREAD_UNCHANGED)
34
+ bg_remover = rembg.new_session()
35
+ img = rembg.remove(img, session=bg_remover)
36
+ img = cv2.resize(img, (W, H), interpolation=cv2.INTER_AREA)
37
+ img = img.astype(np.float32) / 255.0
38
+ input_mask = img[..., 3:]
39
+ # white bg
40
+ if bg == 'white':
41
+ input_img = img[..., :3] * input_mask + (1 - input_mask)
42
+ elif bg == 'black':
43
+ input_img = img[..., :3]
44
+ else:
45
+ raise NotImplementedError
46
+ # bgr to rgb
47
+ input_img = input_img[..., ::-1].copy()
48
+ input_img = Image.fromarray(np.uint8(input_img*255))
49
+ return input_img
50
+
51
+ def load_image_w_bg(file, W, H):
52
+ # load image
53
+ print(f'[INFO] load image from {file}...')
54
+ img = cv2.imread(file, cv2.IMREAD_UNCHANGED)
55
+ img = cv2.resize(img, (W, H), interpolation=cv2.INTER_AREA)
56
+ img = img.astype(np.float32) / 255.0
57
+ input_img = img[..., :3]
58
+ # bgr to rgb
59
+ input_img = input_img[..., ::-1].copy()
60
+ input_img = Image.fromarray(np.uint8(input_img*255))
61
+ return input_img
62
+
63
+ def gen_vid(input_path, seed, bg, is_pad):
64
+ name = input_path.split('/')[-1].split('.')[0]
65
+ input_dir = os.path.dirname(input_path)
66
+ pipe = StableVideoDiffusionPipeline.from_pretrained(
67
+ "stabilityai/stable-video-diffusion-img2vid", torch_dtype=torch.float16, variant="fp16"
68
+ )
69
+ # pipe.enable_model_cpu_offload()
70
+ pipe.to("cuda")
71
+
72
+ if is_pad:
73
+ height, width = 576, 1024
74
+ else:
75
+ height, width = 512, 512
76
+
77
+ if seed is None:
78
+ for bg in ['white', 'black', 'orig']:
79
+ if bg == 'orig':
80
+ if 'rgba' in name:
81
+ continue
82
+ image = load_image_w_bg(input_path, width, height)
83
+ else:
84
+ image = load_image(input_path, width, height, bg)
85
+ if is_pad:
86
+ image = resize_image(image, output_size=(width, height))
87
+ for seed in range(20):
88
+ generator = torch.manual_seed(seed)
89
+ frames = pipe(image, height, width, generator=generator).frames[0]
90
+ imageio.mimwrite(f"{input_dir}/videos/{name}_{bg}_{seed:03}.mp4", frames, fps=7)
91
+ else:
92
+ if bg == 'orig':
93
+ if 'rgba' in name:
94
+ raise ValueError
95
+ image = load_image_w_bg(input_path, width, height)
96
+ else:
97
+ image = load_image(input_path, width, height, bg)
98
+ if is_pad:
99
+ image = resize_image(image, output_size=(width, height))
100
+ generator = torch.manual_seed(seed)
101
+ frames = pipe(image, height, width, generator=generator).frames[0]
102
+
103
+ imageio.mimwrite(f"{input_dir}/{name}_generated.mp4", frames, fps=7)
104
+ os.makedirs(f"{input_dir}/{name}_frames", exist_ok=True)
105
+ for idx, img in enumerate(frames):
106
+ if is_pad:
107
+ img = img.crop(((width-height) //2, 0, width - (width-height) //2, height))
108
+ img.save(f"{input_dir}/{name}_frames/{idx:03}.png")
109
+
110
+ if __name__ == '__main__':
111
+ parser = argparse.ArgumentParser()
112
+ parser.add_argument("--path", type=str, required=True)
113
+ parser.add_argument("--seed", type=int, default=None)
114
+ parser.add_argument("--bg", type=str, default='white')
115
+ parser.add_argument("--is_pad", type=bool, default=False)
116
+ args, extras = parser.parse_known_args()
117
+ gen_vid(args.path, args.seed, args.bg, args.is_pad)
scripts/process.py ADDED
@@ -0,0 +1,92 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import glob
3
+ import sys
4
+ import cv2
5
+ import argparse
6
+ import numpy as np
7
+ import matplotlib.pyplot as plt
8
+
9
+ import torch
10
+ import torch.nn as nn
11
+ import torch.nn.functional as F
12
+ from torchvision import transforms
13
+ from PIL import Image
14
+ import rembg
15
+
16
+ class BLIP2():
17
+ def __init__(self, device='cuda'):
18
+ self.device = device
19
+ from transformers import AutoProcessor, Blip2ForConditionalGeneration
20
+ self.processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
21
+ self.model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16).to(device)
22
+
23
+ @torch.no_grad()
24
+ def __call__(self, image):
25
+ image = Image.fromarray(image)
26
+ inputs = self.processor(image, return_tensors="pt").to(self.device, torch.float16)
27
+
28
+ generated_ids = self.model.generate(**inputs, max_new_tokens=20)
29
+ generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
30
+
31
+ return generated_text
32
+
33
+
34
+ if __name__ == '__main__':
35
+
36
+ parser = argparse.ArgumentParser()
37
+ parser.add_argument('path', type=str, help="path to image (png, jpeg, etc.)")
38
+ parser.add_argument('--model', default='u2net', type=str, help="rembg model, see https://github.com/danielgatis/rembg#models")
39
+ parser.add_argument('--size', default=256, type=int, help="output resolution")
40
+ parser.add_argument('--border_ratio', default=0.2, type=float, help="output border ratio")
41
+ parser.add_argument('--recenter', type=bool, default=True, help="recenter, potentially not helpful for multiview zero123")
42
+ opt = parser.parse_args()
43
+
44
+ session = rembg.new_session(model_name=opt.model)
45
+
46
+ if os.path.isdir(opt.path):
47
+ print(f'[INFO] processing directory {opt.path}...')
48
+ files = glob.glob(f'{opt.path}/*')
49
+ out_dir = opt.path
50
+ else: # isfile
51
+ files = [opt.path]
52
+ out_dir = os.path.dirname(opt.path)
53
+
54
+ for file in files:
55
+
56
+ out_base = os.path.basename(file).split('.')[0]
57
+ out_rgba = os.path.join(out_dir, out_base + '_rgba.png')
58
+
59
+ # load image
60
+ print(f'[INFO] loading image {file}...')
61
+ image = cv2.imread(file, cv2.IMREAD_UNCHANGED)
62
+
63
+ # carve background
64
+ print(f'[INFO] background removal...')
65
+ carved_image = rembg.remove(image, session=session) # [H, W, 4]
66
+ mask = carved_image[..., -1] > 0
67
+
68
+ # recenter
69
+ if opt.recenter:
70
+ print(f'[INFO] recenter...')
71
+ final_rgba = np.zeros((opt.size, opt.size, 4), dtype=np.uint8)
72
+
73
+ coords = np.nonzero(mask)
74
+ x_min, x_max = coords[0].min(), coords[0].max()
75
+ y_min, y_max = coords[1].min(), coords[1].max()
76
+ h = x_max - x_min
77
+ w = y_max - y_min
78
+ desired_size = int(opt.size * (1 - opt.border_ratio))
79
+ scale = desired_size / max(h, w)
80
+ h2 = int(h * scale)
81
+ w2 = int(w * scale)
82
+ x2_min = (opt.size - h2) // 2
83
+ x2_max = x2_min + h2
84
+ y2_min = (opt.size - w2) // 2
85
+ y2_max = y2_min + w2
86
+ 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)
87
+
88
+ else:
89
+ final_rgba = carved_image
90
+
91
+ # write image
92
+ cv2.imwrite(out_rgba, final_rgba)
scripts/runall.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import glob
3
+ import argparse
4
+
5
+ parser = argparse.ArgumentParser()
6
+ parser.add_argument('--dir', default='data', type=str, help='Directory where processed images are stored')
7
+ parser.add_argument('--out', default='logs', type=str, help='Directory where obj files will be saved')
8
+ parser.add_argument('--video-out', default='videos', type=str, help='Directory where videos will be saved')
9
+ parser.add_argument('--gpu', default=0, type=int, help='ID of GPU to use')
10
+ parser.add_argument('--elevation', default=0, type=int, help='Elevation angle of view in degrees')
11
+ parser.add_argument('--config', default='configs', type=str, help='Path to config directory, which contains image.yaml')
12
+ args = parser.parse_args()
13
+
14
+ files = glob.glob(f'{args.dir}/*_rgba.png')
15
+ configs_dir = args.config
16
+
17
+ # check if image.yaml exists
18
+ if not os.path.exists(os.path.join(configs_dir, 'image.yaml')):
19
+ raise FileNotFoundError(
20
+ f'image.yaml not found in {configs_dir} directory. Please check if the directory is correct.'
21
+ )
22
+
23
+ # create output directories if not exists
24
+ out_dir = args.out
25
+ os.makedirs(out_dir, exist_ok=True)
26
+ video_dir = args.video_out
27
+ os.makedirs(video_dir, exist_ok=True)
28
+
29
+
30
+ for file in files:
31
+ name = os.path.basename(file).replace("_rgba.png", "")
32
+ print(f'======== processing {name} ========')
33
+ # first stage
34
+ os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main.py '
35
+ f'--config {configs_dir}/image.yaml '
36
+ f'input={file} '
37
+ f'save_path={name} elevation={args.elevation}')
38
+ # second stage
39
+ os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main2.py '
40
+ f'--config {configs_dir}/image.yaml '
41
+ f'input={file} '
42
+ f'save_path={name} elevation={args.elevation}')
43
+ # export video
44
+ mesh_path = os.path.join(out_dir, f'{name}.obj')
45
+ os.system(f'python -m kiui.render {mesh_path} '
46
+ f'--save_video {video_dir}/{name}.mp4 '
47
+ f'--wogui '
48
+ f'--elevation {args.elevation}')
scripts/runall_mvdream.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import glob
3
+ import argparse
4
+
5
+ parser = argparse.ArgumentParser()
6
+ parser.add_argument('--gpu', default=0, type=int)
7
+ args = parser.parse_args()
8
+
9
+ prompts = [
10
+ # ('butterfly', 'a beautiful, intricate butterfly'),
11
+ # ('boy', 'a nendoroid of a chibi cute boy'),
12
+ # ('axe', 'a viking axe, fantasy, blender'),
13
+ # ('dog_rocket', 'corgi riding a rocket'),
14
+ ('teapot', 'a chinese teapot'),
15
+ ('squirrel_guitar', 'a DSLR photo of a squirrel playing guitar'),
16
+ # ('house', 'fisherman house, cute, cartoon, blender, stylized'),
17
+ # ('ship', 'Higly detailed, majestic royal tall ship, realistic painting'),
18
+ ('einstein', 'Albert Einstein with grey suit is riding a bicycle'),
19
+ # ('angle', 'a statue of an angle'),
20
+ ('lion', 'A 3D model of Simba, the lion cub from The Lion King, standing majestically on Pride Rock, character'),
21
+ # ('paris', 'mini Paris, highly detailed 3d model'),
22
+ # ('pig_backpack', 'a pig wearing a backpack'),
23
+ ('pisa_tower', 'Picture of the Leaning Tower of Pisa, featuring its tilted structure and marble facade'),
24
+ # ('robot', 'a human-like full body robot'),
25
+ ('coin', 'a golden coin'),
26
+ # ('cake', 'a delicious and beautiful cake'),
27
+ # ('horse', 'a DSLR photo of a horse'),
28
+ # ('cat', 'a photo of a cat'),
29
+ ('cat_hat', 'a photo of a cat wearing a wizard hat'),
30
+ # ('cat_ball', 'a photo of a cat playing with a red ball'),
31
+ # ('nendoroid', 'a nendoroid of a chibi girl'),
32
+
33
+ ]
34
+
35
+ for name, prompt in prompts:
36
+ print(f'======== processing {name} ========')
37
+ # first stage
38
+ os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main.py --config configs/text_mv.yaml prompt="{prompt}" save_path={name}')
39
+ # second stage
40
+ os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main2.py --config configs/text_mv.yaml prompt="{prompt}" save_path={name}')
41
+ # export video
42
+ mesh_path = os.path.join('logs', f'{name}.obj')
43
+ os.makedirs('videos', exist_ok=True)
44
+ os.system(f'python -m kiui.render {mesh_path} --save_video videos/{name}.mp4 --wogui')
scripts/runall_sd.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import glob
3
+ import argparse
4
+
5
+ parser = argparse.ArgumentParser()
6
+ parser.add_argument('--gpu', default=0, type=int)
7
+ args = parser.parse_args()
8
+
9
+ prompts = [
10
+ ('strawberry', 'a ripe strawberry'),
11
+ ('cactus_pot', 'a small saguaro cactus planted in a clay pot'),
12
+ ('hamburger', 'a delicious hamburger'),
13
+ ('icecream', 'an icecream'),
14
+ ('tulip', 'a blue tulip'),
15
+ ('pineapple', 'a ripe pineapple'),
16
+ ('goblet', 'a golden goblet'),
17
+ # ('squitopus', 'a squirrel-octopus hybrid'),
18
+ # ('astronaut', 'Michelangelo style statue of an astronaut'),
19
+ # ('teddy_bear', 'a teddy bear'),
20
+ # ('corgi_nurse', 'a plush toy of a corgi nurse'),
21
+ # ('teapot', 'a blue and white porcelain teapot'),
22
+ # ('skull', "a human skull"),
23
+ # ('penguin', 'a penguin'),
24
+ # ('campfire', 'a campfire'),
25
+ # ('donut', 'a donut with pink icing'),
26
+ # ('cupcake', 'a birthday cupcake'),
27
+ # ('pie', 'shepherds pie'),
28
+ # ('cone', 'a traffic cone'),
29
+ # ('schoolbus', 'a schoolbus'),
30
+ # ('avocado_chair', 'a chair that looks like an avocado'),
31
+ # ('glasses', 'a pair of sunglasses')
32
+ # ('potion', 'a bottle of green potion'),
33
+ # ('chalice', 'a delicate chalice'),
34
+ ]
35
+
36
+ for name, prompt in prompts:
37
+ print(f'======== processing {name} ========')
38
+ # first stage
39
+ os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main.py --config configs/text.yaml prompt="{prompt}" save_path={name}')
40
+ # second stage
41
+ os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main2.py --config configs/text.yaml prompt="{prompt}" save_path={name}')
42
+ # export video
43
+ mesh_path = os.path.join('logs', f'{name}.obj')
44
+ os.makedirs('videos', exist_ok=True)
45
+ os.system(f'python -m kiui.render {mesh_path} --save_video videos/{name}.mp4 --wogui')
sh_utils.py ADDED
@@ -0,0 +1,118 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2021 The PlenOctree Authors.
2
+ # Redistribution and use in source and binary forms, with or without
3
+ # modification, are permitted provided that the following conditions are met:
4
+ #
5
+ # 1. Redistributions of source code must retain the above copyright notice,
6
+ # this list of conditions and the following disclaimer.
7
+ #
8
+ # 2. Redistributions in binary form must reproduce the above copyright notice,
9
+ # this list of conditions and the following disclaimer in the documentation
10
+ # and/or other materials provided with the distribution.
11
+ #
12
+ # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
13
+ # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
14
+ # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
15
+ # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
16
+ # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
17
+ # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
18
+ # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
19
+ # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
20
+ # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
21
+ # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
22
+ # POSSIBILITY OF SUCH DAMAGE.
23
+
24
+ import torch
25
+
26
+ C0 = 0.28209479177387814
27
+ C1 = 0.4886025119029199
28
+ C2 = [
29
+ 1.0925484305920792,
30
+ -1.0925484305920792,
31
+ 0.31539156525252005,
32
+ -1.0925484305920792,
33
+ 0.5462742152960396
34
+ ]
35
+ C3 = [
36
+ -0.5900435899266435,
37
+ 2.890611442640554,
38
+ -0.4570457994644658,
39
+ 0.3731763325901154,
40
+ -0.4570457994644658,
41
+ 1.445305721320277,
42
+ -0.5900435899266435
43
+ ]
44
+ C4 = [
45
+ 2.5033429417967046,
46
+ -1.7701307697799304,
47
+ 0.9461746957575601,
48
+ -0.6690465435572892,
49
+ 0.10578554691520431,
50
+ -0.6690465435572892,
51
+ 0.47308734787878004,
52
+ -1.7701307697799304,
53
+ 0.6258357354491761,
54
+ ]
55
+
56
+
57
+ def eval_sh(deg, sh, dirs):
58
+ """
59
+ Evaluate spherical harmonics at unit directions
60
+ using hardcoded SH polynomials.
61
+ Works with torch/np/jnp.
62
+ ... Can be 0 or more batch dimensions.
63
+ Args:
64
+ deg: int SH deg. Currently, 0-3 supported
65
+ sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2]
66
+ dirs: jnp.ndarray unit directions [..., 3]
67
+ Returns:
68
+ [..., C]
69
+ """
70
+ assert deg <= 4 and deg >= 0
71
+ coeff = (deg + 1) ** 2
72
+ assert sh.shape[-1] >= coeff
73
+
74
+ result = C0 * sh[..., 0]
75
+ if deg > 0:
76
+ x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3]
77
+ result = (result -
78
+ C1 * y * sh[..., 1] +
79
+ C1 * z * sh[..., 2] -
80
+ C1 * x * sh[..., 3])
81
+
82
+ if deg > 1:
83
+ xx, yy, zz = x * x, y * y, z * z
84
+ xy, yz, xz = x * y, y * z, x * z
85
+ result = (result +
86
+ C2[0] * xy * sh[..., 4] +
87
+ C2[1] * yz * sh[..., 5] +
88
+ C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] +
89
+ C2[3] * xz * sh[..., 7] +
90
+ C2[4] * (xx - yy) * sh[..., 8])
91
+
92
+ if deg > 2:
93
+ result = (result +
94
+ C3[0] * y * (3 * xx - yy) * sh[..., 9] +
95
+ C3[1] * xy * z * sh[..., 10] +
96
+ C3[2] * y * (4 * zz - xx - yy)* sh[..., 11] +
97
+ C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] +
98
+ C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] +
99
+ C3[5] * z * (xx - yy) * sh[..., 14] +
100
+ C3[6] * x * (xx - 3 * yy) * sh[..., 15])
101
+
102
+ if deg > 3:
103
+ result = (result + C4[0] * xy * (xx - yy) * sh[..., 16] +
104
+ C4[1] * yz * (3 * xx - yy) * sh[..., 17] +
105
+ C4[2] * xy * (7 * zz - 1) * sh[..., 18] +
106
+ C4[3] * yz * (7 * zz - 3) * sh[..., 19] +
107
+ C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] +
108
+ C4[5] * xz * (7 * zz - 3) * sh[..., 21] +
109
+ C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] +
110
+ C4[7] * xz * (xx - 3 * yy) * sh[..., 23] +
111
+ C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) * sh[..., 24])
112
+ return result
113
+
114
+ def RGB2SH(rgb):
115
+ return (rgb - 0.5) / C0
116
+
117
+ def SH2RGB(sh):
118
+ return sh * C0 + 0.5
utils/camera_utils.py ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #
2
+ # Copyright (C) 2023, Inria
3
+ # GRAPHDECO research group, https://team.inria.fr/graphdeco
4
+ # All rights reserved.
5
+ #
6
+ # This software is free for non-commercial, research and evaluation use
7
+ # under the terms of the LICENSE.md file.
8
+ #
9
+ # For inquiries contact george.drettakis@inria.fr
10
+ #
11
+
12
+ from scene.cameras import Camera
13
+ import numpy as np
14
+ from utils.general_utils import PILtoTorch
15
+ from utils.graphics_utils import fov2focal
16
+
17
+ WARNED = False
18
+
19
+ def loadCam(args, id, cam_info, resolution_scale):
20
+
21
+
22
+ # resized_image_rgb = PILtoTorch(cam_info.image, resolution)
23
+
24
+ # gt_image = resized_image_rgb[:3, ...]
25
+ # loaded_mask = None
26
+
27
+ # if resized_image_rgb.shape[1] == 4:
28
+ # loaded_mask = resized_image_rgb[3:4, ...]
29
+
30
+ return Camera(colmap_id=cam_info.uid, R=cam_info.R, T=cam_info.T,
31
+ FoVx=cam_info.FovX, FoVy=cam_info.FovY,
32
+ image=cam_info.image, gt_alpha_mask=None,
33
+ image_name=cam_info.image_name, uid=id, data_device=args.data_device,
34
+ time = cam_info.time,
35
+ )
36
+
37
+ def cameraList_from_camInfos(cam_infos, resolution_scale, args):
38
+ camera_list = []
39
+
40
+ for id, c in enumerate(cam_infos):
41
+ camera_list.append(loadCam(args, id, c, resolution_scale))
42
+
43
+ return camera_list
44
+
45
+ def camera_to_JSON(id, camera : Camera):
46
+ Rt = np.zeros((4, 4))
47
+ Rt[:3, :3] = camera.R.transpose()
48
+ Rt[:3, 3] = camera.T
49
+ Rt[3, 3] = 1.0
50
+
51
+ W2C = np.linalg.inv(Rt)
52
+ pos = W2C[:3, 3]
53
+ rot = W2C[:3, :3]
54
+ serializable_array_2d = [x.tolist() for x in rot]
55
+ camera_entry = {
56
+ 'id' : id,
57
+ 'img_name' : camera.image_name,
58
+ 'width' : camera.width,
59
+ 'height' : camera.height,
60
+ 'position': pos.tolist(),
61
+ 'rotation': serializable_array_2d,
62
+ 'fy' : fov2focal(camera.FovY, camera.height),
63
+ 'fx' : fov2focal(camera.FovX, camera.width)
64
+ }
65
+ return camera_entry
utils/general_utils.py ADDED
@@ -0,0 +1,136 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #
2
+ # Copyright (C) 2023, Inria
3
+ # GRAPHDECO research group, https://team.inria.fr/graphdeco
4
+ # All rights reserved.
5
+ #
6
+ # This software is free for non-commercial, research and evaluation use
7
+ # under the terms of the LICENSE.md file.
8
+ #
9
+ # For inquiries contact george.drettakis@inria.fr
10
+ #
11
+
12
+ import torch
13
+ import sys
14
+ from datetime import datetime
15
+ import numpy as np
16
+ import random
17
+
18
+ def inverse_sigmoid(x):
19
+ return torch.log(x/(1-x))
20
+
21
+ def PILtoTorch(pil_image, resolution):
22
+ if resolution is not None:
23
+ resized_image_PIL = pil_image.resize(resolution)
24
+ else:
25
+ resized_image_PIL = pil_image
26
+ resized_image = torch.from_numpy(np.array(resized_image_PIL)) / 255.0
27
+ if len(resized_image.shape) == 3:
28
+ return resized_image.permute(2, 0, 1)
29
+ else:
30
+ return resized_image.unsqueeze(dim=-1).permute(2, 0, 1)
31
+
32
+ def get_expon_lr_func(
33
+ lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000
34
+ ):
35
+ """
36
+ Copied from Plenoxels
37
+
38
+ Continuous learning rate decay function. Adapted from JaxNeRF
39
+ The returned rate is lr_init when step=0 and lr_final when step=max_steps, and
40
+ is log-linearly interpolated elsewhere (equivalent to exponential decay).
41
+ If lr_delay_steps>0 then the learning rate will be scaled by some smooth
42
+ function of lr_delay_mult, such that the initial learning rate is
43
+ lr_init*lr_delay_mult at the beginning of optimization but will be eased back
44
+ to the normal learning rate when steps>lr_delay_steps.
45
+ :param conf: config subtree 'lr' or similar
46
+ :param max_steps: int, the number of steps during optimization.
47
+ :return HoF which takes step as input
48
+ """
49
+
50
+ def helper(step):
51
+ if step < 0 or (lr_init == 0.0 and lr_final == 0.0):
52
+ # Disable this parameter
53
+ return 0.0
54
+ if lr_delay_steps > 0:
55
+ # A kind of reverse cosine decay.
56
+ delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin(
57
+ 0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1)
58
+ )
59
+ else:
60
+ delay_rate = 1.0
61
+ t = np.clip(step / max_steps, 0, 1)
62
+ log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t)
63
+ return delay_rate * log_lerp
64
+
65
+ return helper
66
+
67
+ def strip_lowerdiag(L):
68
+ uncertainty = torch.zeros((L.shape[0], 6), dtype=torch.float, device="cuda")
69
+
70
+ uncertainty[:, 0] = L[:, 0, 0]
71
+ uncertainty[:, 1] = L[:, 0, 1]
72
+ uncertainty[:, 2] = L[:, 0, 2]
73
+ uncertainty[:, 3] = L[:, 1, 1]
74
+ uncertainty[:, 4] = L[:, 1, 2]
75
+ uncertainty[:, 5] = L[:, 2, 2]
76
+ return uncertainty
77
+
78
+ def strip_symmetric(sym):
79
+ return strip_lowerdiag(sym)
80
+
81
+ def build_rotation(r):
82
+ norm = torch.sqrt(r[:,0]*r[:,0] + r[:,1]*r[:,1] + r[:,2]*r[:,2] + r[:,3]*r[:,3])
83
+
84
+ q = r / norm[:, None]
85
+
86
+ R = torch.zeros((q.size(0), 3, 3), device='cuda')
87
+
88
+ w = q[:, 0]
89
+ x = q[:, 1]
90
+ y = q[:, 2]
91
+ z = q[:, 3]
92
+
93
+ R[:, 0, 0] = 1 - 2 * (y*y + z*z)
94
+ R[:, 0, 1] = 2 * (x*y - w*z)
95
+ R[:, 0, 2] = 2 * (x*z + w*y)
96
+ R[:, 1, 0] = 2 * (x*y + w*z)
97
+ R[:, 1, 1] = 1 - 2 * (x*x + z*z)
98
+ R[:, 1, 2] = 2 * (y*z - w*x)
99
+ R[:, 2, 0] = 2 * (x*z - w*y)
100
+ R[:, 2, 1] = 2 * (y*z + w*x)
101
+ R[:, 2, 2] = 1 - 2 * (x*x + y*y)
102
+ return R
103
+
104
+ def build_scaling_rotation(s, r):
105
+ L = torch.zeros((s.shape[0], 3, 3), dtype=torch.float, device="cuda")
106
+ R = build_rotation(r)
107
+
108
+ L[:,0,0] = s[:,0]
109
+ L[:,1,1] = s[:,1]
110
+ L[:,2,2] = s[:,2]
111
+
112
+ L = R @ L
113
+ return L
114
+
115
+ def safe_state(silent):
116
+ old_f = sys.stdout
117
+ class F:
118
+ def __init__(self, silent):
119
+ self.silent = silent
120
+
121
+ def write(self, x):
122
+ if not self.silent:
123
+ if x.endswith("\n"):
124
+ old_f.write(x.replace("\n", " [{}]\n".format(str(datetime.now().strftime("%d/%m %H:%M:%S")))))
125
+ else:
126
+ old_f.write(x)
127
+
128
+ def flush(self):
129
+ old_f.flush()
130
+
131
+ sys.stdout = F(silent)
132
+
133
+ random.seed(0)
134
+ np.random.seed(0)
135
+ torch.manual_seed(0)
136
+ torch.cuda.set_device(torch.device("cuda:0"))