sammyview80 commited on
Commit
7306307
1 Parent(s): d8cdc8e

Upload 27 files

Browse files
.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
+ wheel/torchmcubes-0.1.0-cp310-cp310-linux_x86_64.whl filter=lfs diff=lfs merge=lfs -text
Dockerfile ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM python:3.10-slim-buster
2
+
3
+ ### Set up user with permissions
4
+ # Set up a new user named "user" with user ID 1000
5
+ RUN useradd -m -u 1000 user
6
+
7
+ # Switch to the "user" user
8
+ USER user
9
+
10
+ # Set home to the user's home directory
11
+ ENV HOME=/home/user \
12
+ PATH=/home/user/.local/bin:$PATH
13
+
14
+ # Set the working directory to the user's home directory
15
+ WORKDIR $HOME/app
16
+
17
+ # Copy the current directory contents into the container at $HOME/app setting the owner to the user
18
+ COPY --chown=user . $HOME/app
19
+
20
+ ### Set up app-specific content
21
+ COPY requirements.txt requirements.txt
22
+ RUN pip3 install -r requirements.txt
23
+
24
+ COPY . .
25
+
26
+ ### Update permissions for the app
27
+ USER root
28
+ RUN chmod 777 ~/app/*
29
+ USER user
30
+
31
+ CMD ["gunicorn", "-b", "0.0.0.0:7860", "app:app"]
README.md CHANGED
@@ -1,13 +1,14 @@
1
  ---
2
- title: Img 3d
3
- emoji: 👁
4
- colorFrom: pink
5
- colorTo: blue
6
- sdk: streamlit
7
- sdk_version: 1.32.2
 
8
  app_file: app.py
9
  pinned: false
10
- license: apache-2.0
11
  ---
12
 
13
  Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
1
  ---
2
+ title: TripoSR
3
+ emoji: 🐳
4
+ colorFrom: gray
5
+ colorTo: red
6
+ sdk: gradio
7
+ sdk_version: 4.20.1
8
+ python_version: 3.10.13
9
  app_file: app.py
10
  pinned: false
11
+ license: mit
12
  ---
13
 
14
  Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
app.py ADDED
@@ -0,0 +1,182 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import shlex
2
+ import subprocess
3
+ import tempfile
4
+ import time
5
+
6
+ import numpy as np
7
+ import rembg
8
+ import spaces
9
+ import torch
10
+ from PIL import Image
11
+
12
+ subprocess.run(shlex.split('pip install wheel/torchmcubes-0.1.0-cp310-cp310-linux_x86_64.whl'))
13
+
14
+ from tsr.system import TSR
15
+ from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation
16
+
17
+ from flask import Flask, flash, request
18
+ from flask_session import Session
19
+
20
+
21
+ app = Flask(__name__)
22
+ app.config["SESSION_PERMANENT"] = False
23
+ app.config["SESSION_TYPE"] = "filesystem"
24
+
25
+ Session(app)
26
+
27
+
28
+ if torch.cuda.is_available():
29
+ device = "cuda:0"
30
+ else:
31
+ device = "cpu"
32
+
33
+ model = TSR.from_pretrained(
34
+ "stabilityai/TripoSR",
35
+ config_name="config.yaml",
36
+ weight_name="model.ckpt",
37
+ )
38
+ model.renderer.set_chunk_size(131072)
39
+ model.to(device)
40
+
41
+ rembg_session = rembg.new_session()
42
+
43
+
44
+ def check_input_image(input_image):
45
+ if input_image is None:
46
+ raise ValueError('Please provide an input image.')
47
+
48
+
49
+ def preprocess(input_image, do_remove_background, foreground_ratio):
50
+ def fill_background(image):
51
+ image = np.array(image).astype(np.float32) / 255.0
52
+ image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
53
+ image = Image.fromarray((image * 255.0).astype(np.uint8))
54
+ return image
55
+
56
+ if do_remove_background:
57
+ image = input_image.convert("RGB")
58
+ image = remove_background(image, rembg_session)
59
+ image = resize_foreground(image, foreground_ratio)
60
+ image = fill_background(image)
61
+ else:
62
+ image = input_image
63
+ if image.mode == "RGBA":
64
+ image = fill_background(image)
65
+ return image
66
+
67
+
68
+ @spaces.GPU
69
+ def generate(image, mc_resolution, formats=["obj", "glb"]):
70
+ scene_codes = model(image, device=device)
71
+ mesh = model.extract_mesh(scene_codes, resolution=mc_resolution)[0]
72
+ mesh = to_gradio_3d_orientation(mesh)
73
+
74
+ mesh_path_glb = tempfile.NamedTemporaryFile(suffix=f".glb", delete=False)
75
+ mesh.export(mesh_path_glb.name)
76
+
77
+ mesh_path_obj = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False)
78
+ mesh.apply_scale([-1, 1, 1]) # Otherwise the visualized .obj will be flipped
79
+ mesh.export(mesh_path_obj.name)
80
+
81
+ return mesh_path_obj.name, mesh_path_glb.name
82
+
83
+ def run_example(image_pil):
84
+ preprocessed = preprocess(image_pil, False, 0.9)
85
+ mesh_name_obj, mesh_name_glb = generate(preprocessed, 256, ["obj", "glb"])
86
+ return preprocessed, mesh_name_obj, mesh_name_glb
87
+
88
+
89
+ @app.route("/", methods=['GET', 'POST'])
90
+ def hello():
91
+ if request.method == 'POST':
92
+ if 'file' not in request.files:
93
+ flash('No file part')
94
+ return {"status": "Failed", "message": "Please Provide file name(file)."}
95
+ file = request.files['file']
96
+ image = Image.open(file)
97
+ preprocess_image = run_example(image)
98
+ print(preprocess_image)
99
+ return {"status": "Success", "message": "You can download the 3D model.", "data": preprocess_image}
100
+
101
+ else:
102
+ return {
103
+ "status": "Success",
104
+ "message":"You can upload an image file to get the 3D model."
105
+ }
106
+
107
+
108
+
109
+ if __name__ == "__main__":
110
+ app.run()
111
+
112
+ # with gr.Blocks() as demo:
113
+ # gr.Markdown(HEADER)
114
+ # with gr.Row(variant="panel"):
115
+ # with gr.Column():
116
+ # with gr.Row():
117
+ # input_image = gr.Image(
118
+ # label="Input Image",
119
+ # image_mode="RGBA",
120
+ # sources="upload",
121
+ # type="pil",
122
+ # elem_id="content_image",
123
+ # )
124
+ # processed_image = gr.Image(label="Processed Image", interactive=False)
125
+ # with gr.Row():
126
+ # with gr.Group():
127
+ # do_remove_background = gr.Checkbox(
128
+ # label="Remove Background", value=True
129
+ # )
130
+ # foreground_ratio = gr.Slider(
131
+ # label="Foreground Ratio",
132
+ # minimum=0.5,
133
+ # maximum=1.0,
134
+ # value=0.85,
135
+ # step=0.05,
136
+ # )
137
+ # mc_resolution = gr.Slider(
138
+ # label="Marching Cubes Resolution",
139
+ # minimum=32,
140
+ # maximum=320,
141
+ # value=256,
142
+ # step=32
143
+ # )
144
+ # with gr.Row():
145
+ # submit = gr.Button("Generate", elem_id="generate", variant="primary")
146
+ # with gr.Column():
147
+ # with gr.Tab("OBJ"):
148
+ # output_model_obj = gr.Model3D(
149
+ # label="Output Model (OBJ Format)",
150
+ # interactive=False,
151
+ # )
152
+ # gr.Markdown("Note: Downloaded object will be flipped in case of .obj export. Export .glb instead or manually flip it before usage.")
153
+ # with gr.Tab("GLB"):
154
+ # output_model_glb = gr.Model3D(
155
+ # label="Output Model (GLB Format)",
156
+ # interactive=False,
157
+ # )
158
+ # gr.Markdown("Note: The model shown here has a darker appearance. Download to get correct results.")
159
+ # with gr.Row(variant="panel"):
160
+ # gr.Examples(
161
+ # examples=[
162
+ # os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))
163
+ # ],
164
+ # inputs=[input_image],
165
+ # outputs=[processed_image, output_model_obj, output_model_glb],
166
+ # cache_examples=True,
167
+ # fn=partial(run_example),
168
+ # label="Examples",
169
+ # examples_per_page=20
170
+ # )
171
+ # submit.click(fn=check_input_image, inputs=[input_image]).success(
172
+ # fn=preprocess,
173
+ # inputs=[input_image, do_remove_background, foreground_ratio],
174
+ # outputs=[processed_image],
175
+ # ).success(
176
+ # fn=generate,
177
+ # inputs=[processed_image, mc_resolution],
178
+ # outputs=[output_model_obj, output_model_glb],
179
+ # )
180
+
181
+ # demo.queue(max_size=10)
182
+ # demo.launch()
requirements.txt ADDED
@@ -0,0 +1,194 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ absl-py==2.1.0
2
+ accelerate==0.28.0
3
+ aiofiles==23.2.1
4
+ aiohttp==3.9.3
5
+ aiosignal==1.3.1
6
+ altair==5.2.0
7
+ annotated-types==0.6.0
8
+ antlr4-python3-runtime==4.9.3
9
+ anyio==4.3.0
10
+ async-timeout==4.0.3
11
+ attributedict==0.3.0
12
+ attrs==23.2.0
13
+ autoawq @ https://github.com/casper-hansen/AutoAWQ/releases/download/v0.1.6/autoawq-0.1.6+cu118-cp310-cp310-linux_x86_64.whl
14
+ blessings==1.7
15
+ blinker==1.7.0
16
+ cachelib==0.12.0
17
+ cachetools==5.3.3
18
+ certifi==2024.2.2
19
+ chardet==5.2.0
20
+ charset-normalizer==3.3.2
21
+ click==8.1.7
22
+ cmake==3.28.3
23
+ codecov==2.1.13
24
+ colorama==0.4.6
25
+ coloredlogs==15.0.1
26
+ colour-runner==0.1.1
27
+ contourpy==1.2.0
28
+ coverage==7.4.4
29
+ cycler==0.12.1
30
+ DataProperty==1.0.1
31
+ datasets==2.14.7
32
+ deepdiff==6.7.1
33
+ dill==0.3.7
34
+ distlib==0.3.8
35
+ einops==0.7.0
36
+ evaluate==0.4.1
37
+ exceptiongroup==1.2.0
38
+ fastapi==0.110.0
39
+ ffmpy==0.3.2
40
+ filelock==3.13.1
41
+ Flask==3.0.2
42
+ Flask-Session==0.7.0
43
+ flatbuffers==24.3.7
44
+ fonttools==4.50.0
45
+ frozenlist==1.4.1
46
+ fsspec==2023.10.0
47
+ gradio==4.8.0
48
+ gradio_client==0.7.1
49
+ h11==0.14.0
50
+ httpcore==1.0.4
51
+ httpx==0.27.0
52
+ huggingface-hub==0.17.3
53
+ humanfriendly==10.0
54
+ idna==3.6
55
+ imageio==2.34.0
56
+ importlib_resources==6.3.1
57
+ inspecta==0.1.3
58
+ itsdangerous==2.1.2
59
+ Jinja2==3.1.3
60
+ joblib==1.3.2
61
+ jsonlines==4.0.0
62
+ jsonschema==4.21.1
63
+ jsonschema-specifications==2023.12.1
64
+ kiwisolver==1.4.5
65
+ lazy_loader==0.3
66
+ lit==18.1.1
67
+ llvmlite==0.42.0
68
+ lm_eval==0.4.1
69
+ lxml==5.1.0
70
+ markdown-it-py==3.0.0
71
+ MarkupSafe==2.1.5
72
+ matplotlib==3.8.3
73
+ mbstrdecoder==1.1.3
74
+ mdurl==0.1.2
75
+ mpmath==1.3.0
76
+ msgspec==0.18.6
77
+ multidict==6.0.5
78
+ multiprocess==0.70.15
79
+ networkx==3.2.1
80
+ nltk==3.8.1
81
+ numba==0.59.0
82
+ numexpr==2.9.0
83
+ numpy==1.26.4
84
+ nvidia-cublas-cu11==11.10.3.66
85
+ nvidia-cublas-cu12==12.1.3.1
86
+ nvidia-cuda-cupti-cu11==11.7.101
87
+ nvidia-cuda-cupti-cu12==12.1.105
88
+ nvidia-cuda-nvrtc-cu11==11.7.99
89
+ nvidia-cuda-nvrtc-cu12==12.1.105
90
+ nvidia-cuda-runtime-cu11==11.7.99
91
+ nvidia-cuda-runtime-cu12==12.1.105
92
+ nvidia-cudnn-cu11==8.5.0.96
93
+ nvidia-cudnn-cu12==8.9.2.26
94
+ nvidia-cufft-cu11==10.9.0.58
95
+ nvidia-cufft-cu12==11.0.2.54
96
+ nvidia-curand-cu11==10.2.10.91
97
+ nvidia-curand-cu12==10.3.2.106
98
+ nvidia-cusolver-cu11==11.4.0.1
99
+ nvidia-cusolver-cu12==11.4.5.107
100
+ nvidia-cusparse-cu11==11.7.4.91
101
+ nvidia-cusparse-cu12==12.1.0.106
102
+ nvidia-nccl-cu11==2.14.3
103
+ nvidia-nccl-cu12==2.19.3
104
+ nvidia-nvjitlink-cu12==12.4.99
105
+ nvidia-nvtx-cu11==11.7.91
106
+ nvidia-nvtx-cu12==12.1.105
107
+ omegaconf==2.3.0
108
+ onnxruntime==1.17.1
109
+ opencv-python-headless==4.9.0.80
110
+ ordered-set==4.1.0
111
+ orjson==3.9.15
112
+ packaging==24.0
113
+ pandas==2.2.1
114
+ pathvalidate==3.2.0
115
+ peft==0.9.0
116
+ Pillow==10.1.0
117
+ platformdirs==4.2.0
118
+ pluggy==1.4.0
119
+ pooch==1.8.1
120
+ portalocker==2.8.2
121
+ protobuf==5.26.0
122
+ psutil==5.9.8
123
+ pyarrow==15.0.2
124
+ pyarrow-hotfix==0.6
125
+ pybind11==2.11.1
126
+ pydantic==2.6.4
127
+ pydantic_core==2.16.3
128
+ pydub==0.25.1
129
+ Pygments==2.17.2
130
+ PyMatting==1.1.12
131
+ pyparsing==3.1.2
132
+ pyproject-api==1.6.1
133
+ pytablewriter==1.2.0
134
+ python-dateutil==2.9.0.post0
135
+ python-multipart==0.0.9
136
+ pytz==2024.1
137
+ PyYAML==6.0.1
138
+ referencing==0.34.0
139
+ regex==2023.12.25
140
+ rembg==2.0.56
141
+ requests==2.31.0
142
+ responses==0.18.0
143
+ rich==13.7.1
144
+ rootpath==0.1.1
145
+ rouge-score==0.1.2
146
+ rpds-py==0.18.0
147
+ sacrebleu==2.4.1
148
+ safetensors==0.4.2
149
+ scikit-image==0.22.0
150
+ scikit-learn==1.4.1.post1
151
+ scipy==1.12.0
152
+ semantic-version==2.10.0
153
+ sentencepiece==0.2.0
154
+ shellingham==1.5.4
155
+ six==1.16.0
156
+ sniffio==1.3.1
157
+ spaces==0.24.2
158
+ sqlitedict==2.1.0
159
+ starlette==0.36.3
160
+ sympy==1.12
161
+ tabledata==1.3.3
162
+ tabulate==0.9.0
163
+ tcolorpy==0.1.4
164
+ termcolor==2.4.0
165
+ texttable==1.7.0
166
+ threadpoolctl==3.3.0
167
+ tifffile==2024.2.12
168
+ tokenizers==0.14.1
169
+ toml==0.10.2
170
+ tomli==2.0.1
171
+ tomlkit==0.12.0
172
+ toolz==0.12.1
173
+ torch==2.0.1
174
+ torchmcubes @ file:///home/codehimalaya/Documents/huggingface/TripoSR/wheel/torchmcubes-0.1.0-cp310-cp310-linux_x86_64.whl
175
+ torchvision==0.17.1
176
+ tox==4.14.1
177
+ tqdm==4.66.2
178
+ tqdm-multiprocess==0.0.11
179
+ transformers==4.35.0
180
+ trimesh==4.0.5
181
+ triton==2.0.0
182
+ typepy==1.3.2
183
+ typer==0.9.0
184
+ typing_extensions==4.10.0
185
+ tzdata==2024.1
186
+ urllib3==2.2.1
187
+ uvicorn==0.28.0
188
+ virtualenv==20.25.1
189
+ websockets==11.0.3
190
+ Werkzeug==3.0.1
191
+ xformers==0.0.22
192
+ xxhash==3.4.1
193
+ yarl==1.9.4
194
+ zstandard==0.22.0
tsr/__pycache__/system.cpython-310.pyc ADDED
Binary file (5.19 kB). View file
 
tsr/__pycache__/utils.cpython-310.pyc ADDED
Binary file (13.7 kB). View file
 
tsr/models/__pycache__/camera.cpython-310.pyc ADDED
Binary file (1.48 kB). View file
 
tsr/models/__pycache__/isosurface.cpython-310.pyc ADDED
Binary file (2.05 kB). View file
 
tsr/models/__pycache__/nerf_renderer.cpython-310.pyc ADDED
Binary file (5.31 kB). View file
 
tsr/models/__pycache__/network_utils.cpython-310.pyc ADDED
Binary file (3.44 kB). View file
 
tsr/models/isosurface.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Callable, Optional, Tuple
2
+
3
+ import numpy as np
4
+ import torch
5
+ import torch.nn as nn
6
+ from torchmcubes import marching_cubes
7
+
8
+
9
+ class IsosurfaceHelper(nn.Module):
10
+ points_range: Tuple[float, float] = (0, 1)
11
+
12
+ @property
13
+ def grid_vertices(self) -> torch.FloatTensor:
14
+ raise NotImplementedError
15
+
16
+
17
+ class MarchingCubeHelper(IsosurfaceHelper):
18
+ def __init__(self, resolution: int) -> None:
19
+ super().__init__()
20
+ self.resolution = resolution
21
+ self.mc_func: Callable = marching_cubes
22
+ self._grid_vertices: Optional[torch.FloatTensor] = None
23
+
24
+ @property
25
+ def grid_vertices(self) -> torch.FloatTensor:
26
+ if self._grid_vertices is None:
27
+ # keep the vertices on CPU so that we can support very large resolution
28
+ x, y, z = (
29
+ torch.linspace(*self.points_range, self.resolution),
30
+ torch.linspace(*self.points_range, self.resolution),
31
+ torch.linspace(*self.points_range, self.resolution),
32
+ )
33
+ x, y, z = torch.meshgrid(x, y, z, indexing="ij")
34
+ verts = torch.cat(
35
+ [x.reshape(-1, 1), y.reshape(-1, 1), z.reshape(-1, 1)], dim=-1
36
+ ).reshape(-1, 3)
37
+ self._grid_vertices = verts
38
+ return self._grid_vertices
39
+
40
+ def forward(
41
+ self,
42
+ level: torch.FloatTensor,
43
+ ) -> Tuple[torch.FloatTensor, torch.LongTensor]:
44
+ level = -level.view(self.resolution, self.resolution, self.resolution)
45
+ v_pos, t_pos_idx = self.mc_func(level.detach(), 0.0)
46
+ v_pos = v_pos[..., [2, 1, 0]]
47
+ v_pos = v_pos / (self.resolution - 1.0)
48
+ return v_pos, t_pos_idx
tsr/models/nerf_renderer.py ADDED
@@ -0,0 +1,180 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from dataclasses import dataclass, field
2
+ from typing import Dict
3
+
4
+ import torch
5
+ import torch.nn.functional as F
6
+ from einops import rearrange, reduce
7
+
8
+ from ..utils import (
9
+ BaseModule,
10
+ chunk_batch,
11
+ get_activation,
12
+ rays_intersect_bbox,
13
+ scale_tensor,
14
+ )
15
+
16
+
17
+ class TriplaneNeRFRenderer(BaseModule):
18
+ @dataclass
19
+ class Config(BaseModule.Config):
20
+ radius: float
21
+
22
+ feature_reduction: str = "concat"
23
+ density_activation: str = "trunc_exp"
24
+ density_bias: float = -1.0
25
+ color_activation: str = "sigmoid"
26
+ num_samples_per_ray: int = 128
27
+ randomized: bool = False
28
+
29
+ cfg: Config
30
+
31
+ def configure(self) -> None:
32
+ assert self.cfg.feature_reduction in ["concat", "mean"]
33
+ self.chunk_size = 0
34
+
35
+ def set_chunk_size(self, chunk_size: int):
36
+ assert (
37
+ chunk_size >= 0
38
+ ), "chunk_size must be a non-negative integer (0 for no chunking)."
39
+ self.chunk_size = chunk_size
40
+
41
+ def query_triplane(
42
+ self,
43
+ decoder: torch.nn.Module,
44
+ positions: torch.Tensor,
45
+ triplane: torch.Tensor,
46
+ ) -> Dict[str, torch.Tensor]:
47
+ input_shape = positions.shape[:-1]
48
+ positions = positions.view(-1, 3)
49
+
50
+ # positions in (-radius, radius)
51
+ # normalized to (-1, 1) for grid sample
52
+ positions = scale_tensor(
53
+ positions, (-self.cfg.radius, self.cfg.radius), (-1, 1)
54
+ )
55
+
56
+ def _query_chunk(x):
57
+ indices2D: torch.Tensor = torch.stack(
58
+ (x[..., [0, 1]], x[..., [0, 2]], x[..., [1, 2]]),
59
+ dim=-3,
60
+ )
61
+ out: torch.Tensor = F.grid_sample(
62
+ rearrange(triplane, "Np Cp Hp Wp -> Np Cp Hp Wp", Np=3),
63
+ rearrange(indices2D, "Np N Nd -> Np () N Nd", Np=3),
64
+ align_corners=False,
65
+ mode="bilinear",
66
+ )
67
+ if self.cfg.feature_reduction == "concat":
68
+ out = rearrange(out, "Np Cp () N -> N (Np Cp)", Np=3)
69
+ elif self.cfg.feature_reduction == "mean":
70
+ out = reduce(out, "Np Cp () N -> N Cp", Np=3, reduction="mean")
71
+ else:
72
+ raise NotImplementedError
73
+
74
+ net_out: Dict[str, torch.Tensor] = decoder(out)
75
+ return net_out
76
+
77
+ if self.chunk_size > 0:
78
+ net_out = chunk_batch(_query_chunk, self.chunk_size, positions)
79
+ else:
80
+ net_out = _query_chunk(positions)
81
+
82
+ net_out["density_act"] = get_activation(self.cfg.density_activation)(
83
+ net_out["density"] + self.cfg.density_bias
84
+ )
85
+ net_out["color"] = get_activation(self.cfg.color_activation)(
86
+ net_out["features"]
87
+ )
88
+
89
+ net_out = {k: v.view(*input_shape, -1) for k, v in net_out.items()}
90
+
91
+ return net_out
92
+
93
+ def _forward(
94
+ self,
95
+ decoder: torch.nn.Module,
96
+ triplane: torch.Tensor,
97
+ rays_o: torch.Tensor,
98
+ rays_d: torch.Tensor,
99
+ **kwargs,
100
+ ):
101
+ rays_shape = rays_o.shape[:-1]
102
+ rays_o = rays_o.view(-1, 3)
103
+ rays_d = rays_d.view(-1, 3)
104
+ n_rays = rays_o.shape[0]
105
+
106
+ t_near, t_far, rays_valid = rays_intersect_bbox(rays_o, rays_d, self.cfg.radius)
107
+ t_near, t_far = t_near[rays_valid], t_far[rays_valid]
108
+
109
+ t_vals = torch.linspace(
110
+ 0, 1, self.cfg.num_samples_per_ray + 1, device=triplane.device
111
+ )
112
+ t_mid = (t_vals[:-1] + t_vals[1:]) / 2.0
113
+ z_vals = t_near * (1 - t_mid[None]) + t_far * t_mid[None] # (N_rays, N_samples)
114
+
115
+ xyz = (
116
+ rays_o[:, None, :] + z_vals[..., None] * rays_d[..., None, :]
117
+ ) # (N_rays, N_sample, 3)
118
+
119
+ mlp_out = self.query_triplane(
120
+ decoder=decoder,
121
+ positions=xyz,
122
+ triplane=triplane,
123
+ )
124
+
125
+ eps = 1e-10
126
+ # deltas = z_vals[:, 1:] - z_vals[:, :-1] # (N_rays, N_samples)
127
+ deltas = t_vals[1:] - t_vals[:-1] # (N_rays, N_samples)
128
+ alpha = 1 - torch.exp(
129
+ -deltas * mlp_out["density_act"][..., 0]
130
+ ) # (N_rays, N_samples)
131
+ accum_prod = torch.cat(
132
+ [
133
+ torch.ones_like(alpha[:, :1]),
134
+ torch.cumprod(1 - alpha[:, :-1] + eps, dim=-1),
135
+ ],
136
+ dim=-1,
137
+ )
138
+ weights = alpha * accum_prod # (N_rays, N_samples)
139
+ comp_rgb_ = (weights[..., None] * mlp_out["color"]).sum(dim=-2) # (N_rays, 3)
140
+ opacity_ = weights.sum(dim=-1) # (N_rays)
141
+
142
+ comp_rgb = torch.zeros(
143
+ n_rays, 3, dtype=comp_rgb_.dtype, device=comp_rgb_.device
144
+ )
145
+ opacity = torch.zeros(n_rays, dtype=opacity_.dtype, device=opacity_.device)
146
+ comp_rgb[rays_valid] = comp_rgb_
147
+ opacity[rays_valid] = opacity_
148
+
149
+ comp_rgb += 1 - opacity[..., None]
150
+ comp_rgb = comp_rgb.view(*rays_shape, 3)
151
+
152
+ return comp_rgb
153
+
154
+ def forward(
155
+ self,
156
+ decoder: torch.nn.Module,
157
+ triplane: torch.Tensor,
158
+ rays_o: torch.Tensor,
159
+ rays_d: torch.Tensor,
160
+ ) -> Dict[str, torch.Tensor]:
161
+ if triplane.ndim == 4:
162
+ comp_rgb = self._forward(decoder, triplane, rays_o, rays_d)
163
+ else:
164
+ comp_rgb = torch.stack(
165
+ [
166
+ self._forward(decoder, triplane[i], rays_o[i], rays_d[i])
167
+ for i in range(triplane.shape[0])
168
+ ],
169
+ dim=0,
170
+ )
171
+
172
+ return comp_rgb
173
+
174
+ def train(self, mode=True):
175
+ self.randomized = mode and self.cfg.randomized
176
+ return super().train(mode=mode)
177
+
178
+ def eval(self):
179
+ self.randomized = False
180
+ return super().eval()
tsr/models/network_utils.py ADDED
@@ -0,0 +1,124 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from dataclasses import dataclass, field
2
+ from typing import Optional
3
+
4
+ import torch
5
+ import torch.nn as nn
6
+ from einops import rearrange
7
+
8
+ from ..utils import BaseModule
9
+
10
+
11
+ class TriplaneUpsampleNetwork(BaseModule):
12
+ @dataclass
13
+ class Config(BaseModule.Config):
14
+ in_channels: int
15
+ out_channels: int
16
+
17
+ cfg: Config
18
+
19
+ def configure(self) -> None:
20
+ self.upsample = nn.ConvTranspose2d(
21
+ self.cfg.in_channels, self.cfg.out_channels, kernel_size=2, stride=2
22
+ )
23
+
24
+ def forward(self, triplanes: torch.Tensor) -> torch.Tensor:
25
+ triplanes_up = rearrange(
26
+ self.upsample(
27
+ rearrange(triplanes, "B Np Ci Hp Wp -> (B Np) Ci Hp Wp", Np=3)
28
+ ),
29
+ "(B Np) Co Hp Wp -> B Np Co Hp Wp",
30
+ Np=3,
31
+ )
32
+ return triplanes_up
33
+
34
+
35
+ class NeRFMLP(BaseModule):
36
+ @dataclass
37
+ class Config(BaseModule.Config):
38
+ in_channels: int
39
+ n_neurons: int
40
+ n_hidden_layers: int
41
+ activation: str = "relu"
42
+ bias: bool = True
43
+ weight_init: Optional[str] = "kaiming_uniform"
44
+ bias_init: Optional[str] = None
45
+
46
+ cfg: Config
47
+
48
+ def configure(self) -> None:
49
+ layers = [
50
+ self.make_linear(
51
+ self.cfg.in_channels,
52
+ self.cfg.n_neurons,
53
+ bias=self.cfg.bias,
54
+ weight_init=self.cfg.weight_init,
55
+ bias_init=self.cfg.bias_init,
56
+ ),
57
+ self.make_activation(self.cfg.activation),
58
+ ]
59
+ for i in range(self.cfg.n_hidden_layers - 1):
60
+ layers += [
61
+ self.make_linear(
62
+ self.cfg.n_neurons,
63
+ self.cfg.n_neurons,
64
+ bias=self.cfg.bias,
65
+ weight_init=self.cfg.weight_init,
66
+ bias_init=self.cfg.bias_init,
67
+ ),
68
+ self.make_activation(self.cfg.activation),
69
+ ]
70
+ layers += [
71
+ self.make_linear(
72
+ self.cfg.n_neurons,
73
+ 4, # density 1 + features 3
74
+ bias=self.cfg.bias,
75
+ weight_init=self.cfg.weight_init,
76
+ bias_init=self.cfg.bias_init,
77
+ )
78
+ ]
79
+ self.layers = nn.Sequential(*layers)
80
+
81
+ def make_linear(
82
+ self,
83
+ dim_in,
84
+ dim_out,
85
+ bias=True,
86
+ weight_init=None,
87
+ bias_init=None,
88
+ ):
89
+ layer = nn.Linear(dim_in, dim_out, bias=bias)
90
+
91
+ if weight_init is None:
92
+ pass
93
+ elif weight_init == "kaiming_uniform":
94
+ torch.nn.init.kaiming_uniform_(layer.weight, nonlinearity="relu")
95
+ else:
96
+ raise NotImplementedError
97
+
98
+ if bias:
99
+ if bias_init is None:
100
+ pass
101
+ elif bias_init == "zero":
102
+ torch.nn.init.zeros_(layer.bias)
103
+ else:
104
+ raise NotImplementedError
105
+
106
+ return layer
107
+
108
+ def make_activation(self, activation):
109
+ if activation == "relu":
110
+ return nn.ReLU(inplace=True)
111
+ elif activation == "silu":
112
+ return nn.SiLU(inplace=True)
113
+ else:
114
+ raise NotImplementedError
115
+
116
+ def forward(self, x):
117
+ inp_shape = x.shape[:-1]
118
+ x = x.reshape(-1, x.shape[-1])
119
+
120
+ features = self.layers(x)
121
+ features = features.reshape(*inp_shape, -1)
122
+ out = {"density": features[..., 0:1], "features": features[..., 1:4]}
123
+
124
+ return out
tsr/models/tokenizers/__pycache__/dino.cpython-310.pyc ADDED
Binary file (18.6 kB). View file
 
tsr/models/tokenizers/__pycache__/image.cpython-310.pyc ADDED
Binary file (2.43 kB). View file
 
tsr/models/tokenizers/__pycache__/triplane.cpython-310.pyc ADDED
Binary file (1.77 kB). View file
 
tsr/models/tokenizers/image.py ADDED
@@ -0,0 +1,67 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from dataclasses import dataclass
2
+ from typing import Optional
3
+
4
+ import torch
5
+ import torch.nn as nn
6
+ from einops import rearrange
7
+ from huggingface_hub import hf_hub_download
8
+ from transformers.models.vit.modeling_vit import ViTModel
9
+
10
+ from ...utils import BaseModule
11
+
12
+
13
+ class DINOSingleImageTokenizer(BaseModule):
14
+ @dataclass
15
+ class Config(BaseModule.Config):
16
+ pretrained_model_name_or_path: str = "facebook/dino-vitb16"
17
+ enable_gradient_checkpointing: bool = False
18
+
19
+ cfg: Config
20
+
21
+ def configure(self) -> None:
22
+ self.model: ViTModel = ViTModel(
23
+ ViTModel.config_class.from_pretrained(
24
+ hf_hub_download(
25
+ repo_id=self.cfg.pretrained_model_name_or_path,
26
+ filename="config.json",
27
+ )
28
+ )
29
+ )
30
+
31
+ if self.cfg.enable_gradient_checkpointing:
32
+ self.model.encoder.gradient_checkpointing = True
33
+
34
+ self.register_buffer(
35
+ "image_mean",
36
+ torch.as_tensor([0.485, 0.456, 0.406]).reshape(1, 1, 3, 1, 1),
37
+ persistent=False,
38
+ )
39
+ self.register_buffer(
40
+ "image_std",
41
+ torch.as_tensor([0.229, 0.224, 0.225]).reshape(1, 1, 3, 1, 1),
42
+ persistent=False,
43
+ )
44
+
45
+ def forward(self, images: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
46
+ packed = False
47
+ if images.ndim == 4:
48
+ packed = True
49
+ images = images.unsqueeze(1)
50
+
51
+ batch_size, n_input_views = images.shape[:2]
52
+ images = (images - self.image_mean) / self.image_std
53
+ out = self.model(
54
+ rearrange(images, "B N C H W -> (B N) C H W"), interpolate_pos_encoding=True
55
+ )
56
+ local_features, global_features = out.last_hidden_state, out.pooler_output
57
+ local_features = local_features.permute(0, 2, 1)
58
+ local_features = rearrange(
59
+ local_features, "(B N) Ct Nt -> B N Ct Nt", B=batch_size
60
+ )
61
+ if packed:
62
+ local_features = local_features.squeeze(1)
63
+
64
+ return local_features
65
+
66
+ def detokenize(self, *args, **kwargs):
67
+ raise NotImplementedError
tsr/models/tokenizers/triplane.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ from dataclasses import dataclass
3
+
4
+ import torch
5
+ import torch.nn as nn
6
+ from einops import rearrange, repeat
7
+
8
+ from ...utils import BaseModule
9
+
10
+
11
+ class Triplane1DTokenizer(BaseModule):
12
+ @dataclass
13
+ class Config(BaseModule.Config):
14
+ plane_size: int
15
+ num_channels: int
16
+
17
+ cfg: Config
18
+
19
+ def configure(self) -> None:
20
+ self.embeddings = nn.Parameter(
21
+ torch.randn(
22
+ (3, self.cfg.num_channels, self.cfg.plane_size, self.cfg.plane_size),
23
+ dtype=torch.float32,
24
+ )
25
+ * 1
26
+ / math.sqrt(self.cfg.num_channels)
27
+ )
28
+
29
+ def forward(self, batch_size: int) -> torch.Tensor:
30
+ return rearrange(
31
+ repeat(self.embeddings, "Np Ct Hp Wp -> B Np Ct Hp Wp", B=batch_size),
32
+ "B Np Ct Hp Wp -> B Ct (Np Hp Wp)",
33
+ )
34
+
35
+ def detokenize(self, tokens: torch.Tensor) -> torch.Tensor:
36
+ batch_size, Ct, Nt = tokens.shape
37
+ assert Nt == self.cfg.plane_size**2 * 3
38
+ assert Ct == self.cfg.num_channels
39
+ return rearrange(
40
+ tokens,
41
+ "B Ct (Np Hp Wp) -> B Np Ct Hp Wp",
42
+ Np=3,
43
+ Hp=self.cfg.plane_size,
44
+ Wp=self.cfg.plane_size,
45
+ )
tsr/models/transformer/__pycache__/attention.cpython-310.pyc ADDED
Binary file (15.3 kB). View file
 
tsr/models/transformer/__pycache__/basic_transformer_block.cpython-310.pyc ADDED
Binary file (9.97 kB). View file
 
tsr/models/transformer/__pycache__/transformer_1d.cpython-310.pyc ADDED
Binary file (7.49 kB). View file
 
tsr/models/transformer/attention.py ADDED
@@ -0,0 +1,628 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2023 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ from typing import Optional
15
+
16
+ import torch
17
+ import torch.nn.functional as F
18
+ from torch import nn
19
+
20
+
21
+ class Attention(nn.Module):
22
+ r"""
23
+ A cross attention layer.
24
+
25
+ Parameters:
26
+ query_dim (`int`):
27
+ The number of channels in the query.
28
+ cross_attention_dim (`int`, *optional*):
29
+ The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`.
30
+ heads (`int`, *optional*, defaults to 8):
31
+ The number of heads to use for multi-head attention.
32
+ dim_head (`int`, *optional*, defaults to 64):
33
+ The number of channels in each head.
34
+ dropout (`float`, *optional*, defaults to 0.0):
35
+ The dropout probability to use.
36
+ bias (`bool`, *optional*, defaults to False):
37
+ Set to `True` for the query, key, and value linear layers to contain a bias parameter.
38
+ upcast_attention (`bool`, *optional*, defaults to False):
39
+ Set to `True` to upcast the attention computation to `float32`.
40
+ upcast_softmax (`bool`, *optional*, defaults to False):
41
+ Set to `True` to upcast the softmax computation to `float32`.
42
+ cross_attention_norm (`str`, *optional*, defaults to `None`):
43
+ The type of normalization to use for the cross attention. Can be `None`, `layer_norm`, or `group_norm`.
44
+ cross_attention_norm_num_groups (`int`, *optional*, defaults to 32):
45
+ The number of groups to use for the group norm in the cross attention.
46
+ added_kv_proj_dim (`int`, *optional*, defaults to `None`):
47
+ The number of channels to use for the added key and value projections. If `None`, no projection is used.
48
+ norm_num_groups (`int`, *optional*, defaults to `None`):
49
+ The number of groups to use for the group norm in the attention.
50
+ spatial_norm_dim (`int`, *optional*, defaults to `None`):
51
+ The number of channels to use for the spatial normalization.
52
+ out_bias (`bool`, *optional*, defaults to `True`):
53
+ Set to `True` to use a bias in the output linear layer.
54
+ scale_qk (`bool`, *optional*, defaults to `True`):
55
+ Set to `True` to scale the query and key by `1 / sqrt(dim_head)`.
56
+ only_cross_attention (`bool`, *optional*, defaults to `False`):
57
+ Set to `True` to only use cross attention and not added_kv_proj_dim. Can only be set to `True` if
58
+ `added_kv_proj_dim` is not `None`.
59
+ eps (`float`, *optional*, defaults to 1e-5):
60
+ An additional value added to the denominator in group normalization that is used for numerical stability.
61
+ rescale_output_factor (`float`, *optional*, defaults to 1.0):
62
+ A factor to rescale the output by dividing it with this value.
63
+ residual_connection (`bool`, *optional*, defaults to `False`):
64
+ Set to `True` to add the residual connection to the output.
65
+ _from_deprecated_attn_block (`bool`, *optional*, defaults to `False`):
66
+ Set to `True` if the attention block is loaded from a deprecated state dict.
67
+ processor (`AttnProcessor`, *optional*, defaults to `None`):
68
+ The attention processor to use. If `None`, defaults to `AttnProcessor2_0` if `torch 2.x` is used and
69
+ `AttnProcessor` otherwise.
70
+ """
71
+
72
+ def __init__(
73
+ self,
74
+ query_dim: int,
75
+ cross_attention_dim: Optional[int] = None,
76
+ heads: int = 8,
77
+ dim_head: int = 64,
78
+ dropout: float = 0.0,
79
+ bias: bool = False,
80
+ upcast_attention: bool = False,
81
+ upcast_softmax: bool = False,
82
+ cross_attention_norm: Optional[str] = None,
83
+ cross_attention_norm_num_groups: int = 32,
84
+ added_kv_proj_dim: Optional[int] = None,
85
+ norm_num_groups: Optional[int] = None,
86
+ out_bias: bool = True,
87
+ scale_qk: bool = True,
88
+ only_cross_attention: bool = False,
89
+ eps: float = 1e-5,
90
+ rescale_output_factor: float = 1.0,
91
+ residual_connection: bool = False,
92
+ _from_deprecated_attn_block: bool = False,
93
+ processor: Optional["AttnProcessor"] = None,
94
+ out_dim: int = None,
95
+ ):
96
+ super().__init__()
97
+ self.inner_dim = out_dim if out_dim is not None else dim_head * heads
98
+ self.query_dim = query_dim
99
+ self.cross_attention_dim = (
100
+ cross_attention_dim if cross_attention_dim is not None else query_dim
101
+ )
102
+ self.upcast_attention = upcast_attention
103
+ self.upcast_softmax = upcast_softmax
104
+ self.rescale_output_factor = rescale_output_factor
105
+ self.residual_connection = residual_connection
106
+ self.dropout = dropout
107
+ self.fused_projections = False
108
+ self.out_dim = out_dim if out_dim is not None else query_dim
109
+
110
+ # we make use of this private variable to know whether this class is loaded
111
+ # with an deprecated state dict so that we can convert it on the fly
112
+ self._from_deprecated_attn_block = _from_deprecated_attn_block
113
+
114
+ self.scale_qk = scale_qk
115
+ self.scale = dim_head**-0.5 if self.scale_qk else 1.0
116
+
117
+ self.heads = out_dim // dim_head if out_dim is not None else heads
118
+ # for slice_size > 0 the attention score computation
119
+ # is split across the batch axis to save memory
120
+ # You can set slice_size with `set_attention_slice`
121
+ self.sliceable_head_dim = heads
122
+
123
+ self.added_kv_proj_dim = added_kv_proj_dim
124
+ self.only_cross_attention = only_cross_attention
125
+
126
+ if self.added_kv_proj_dim is None and self.only_cross_attention:
127
+ raise ValueError(
128
+ "`only_cross_attention` can only be set to True if `added_kv_proj_dim` is not None. Make sure to set either `only_cross_attention=False` or define `added_kv_proj_dim`."
129
+ )
130
+
131
+ if norm_num_groups is not None:
132
+ self.group_norm = nn.GroupNorm(
133
+ num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True
134
+ )
135
+ else:
136
+ self.group_norm = None
137
+
138
+ self.spatial_norm = None
139
+
140
+ if cross_attention_norm is None:
141
+ self.norm_cross = None
142
+ elif cross_attention_norm == "layer_norm":
143
+ self.norm_cross = nn.LayerNorm(self.cross_attention_dim)
144
+ elif cross_attention_norm == "group_norm":
145
+ if self.added_kv_proj_dim is not None:
146
+ # The given `encoder_hidden_states` are initially of shape
147
+ # (batch_size, seq_len, added_kv_proj_dim) before being projected
148
+ # to (batch_size, seq_len, cross_attention_dim). The norm is applied
149
+ # before the projection, so we need to use `added_kv_proj_dim` as
150
+ # the number of channels for the group norm.
151
+ norm_cross_num_channels = added_kv_proj_dim
152
+ else:
153
+ norm_cross_num_channels = self.cross_attention_dim
154
+
155
+ self.norm_cross = nn.GroupNorm(
156
+ num_channels=norm_cross_num_channels,
157
+ num_groups=cross_attention_norm_num_groups,
158
+ eps=1e-5,
159
+ affine=True,
160
+ )
161
+ else:
162
+ raise ValueError(
163
+ f"unknown cross_attention_norm: {cross_attention_norm}. Should be None, 'layer_norm' or 'group_norm'"
164
+ )
165
+
166
+ linear_cls = nn.Linear
167
+
168
+ self.linear_cls = linear_cls
169
+ self.to_q = linear_cls(query_dim, self.inner_dim, bias=bias)
170
+
171
+ if not self.only_cross_attention:
172
+ # only relevant for the `AddedKVProcessor` classes
173
+ self.to_k = linear_cls(self.cross_attention_dim, self.inner_dim, bias=bias)
174
+ self.to_v = linear_cls(self.cross_attention_dim, self.inner_dim, bias=bias)
175
+ else:
176
+ self.to_k = None
177
+ self.to_v = None
178
+
179
+ if self.added_kv_proj_dim is not None:
180
+ self.add_k_proj = linear_cls(added_kv_proj_dim, self.inner_dim)
181
+ self.add_v_proj = linear_cls(added_kv_proj_dim, self.inner_dim)
182
+
183
+ self.to_out = nn.ModuleList([])
184
+ self.to_out.append(linear_cls(self.inner_dim, self.out_dim, bias=out_bias))
185
+ self.to_out.append(nn.Dropout(dropout))
186
+
187
+ # set attention processor
188
+ # We use the AttnProcessor2_0 by default when torch 2.x is used which uses
189
+ # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention
190
+ # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
191
+ if processor is None:
192
+ processor = (
193
+ AttnProcessor2_0()
194
+ if hasattr(F, "scaled_dot_product_attention") and self.scale_qk
195
+ else AttnProcessor()
196
+ )
197
+ self.set_processor(processor)
198
+
199
+ def set_processor(self, processor: "AttnProcessor") -> None:
200
+ self.processor = processor
201
+
202
+ def forward(
203
+ self,
204
+ hidden_states: torch.FloatTensor,
205
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
206
+ attention_mask: Optional[torch.FloatTensor] = None,
207
+ **cross_attention_kwargs,
208
+ ) -> torch.Tensor:
209
+ r"""
210
+ The forward method of the `Attention` class.
211
+
212
+ Args:
213
+ hidden_states (`torch.Tensor`):
214
+ The hidden states of the query.
215
+ encoder_hidden_states (`torch.Tensor`, *optional*):
216
+ The hidden states of the encoder.
217
+ attention_mask (`torch.Tensor`, *optional*):
218
+ The attention mask to use. If `None`, no mask is applied.
219
+ **cross_attention_kwargs:
220
+ Additional keyword arguments to pass along to the cross attention.
221
+
222
+ Returns:
223
+ `torch.Tensor`: The output of the attention layer.
224
+ """
225
+ # The `Attention` class can call different attention processors / attention functions
226
+ # here we simply pass along all tensors to the selected processor class
227
+ # For standard processors that are defined here, `**cross_attention_kwargs` is empty
228
+ return self.processor(
229
+ self,
230
+ hidden_states,
231
+ encoder_hidden_states=encoder_hidden_states,
232
+ attention_mask=attention_mask,
233
+ **cross_attention_kwargs,
234
+ )
235
+
236
+ def batch_to_head_dim(self, tensor: torch.Tensor) -> torch.Tensor:
237
+ r"""
238
+ Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size // heads, seq_len, dim * heads]`. `heads`
239
+ is the number of heads initialized while constructing the `Attention` class.
240
+
241
+ Args:
242
+ tensor (`torch.Tensor`): The tensor to reshape.
243
+
244
+ Returns:
245
+ `torch.Tensor`: The reshaped tensor.
246
+ """
247
+ head_size = self.heads
248
+ batch_size, seq_len, dim = tensor.shape
249
+ tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)
250
+ tensor = tensor.permute(0, 2, 1, 3).reshape(
251
+ batch_size // head_size, seq_len, dim * head_size
252
+ )
253
+ return tensor
254
+
255
+ def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Tensor:
256
+ r"""
257
+ Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size, seq_len, heads, dim // heads]` `heads` is
258
+ the number of heads initialized while constructing the `Attention` class.
259
+
260
+ Args:
261
+ tensor (`torch.Tensor`): The tensor to reshape.
262
+ out_dim (`int`, *optional*, defaults to `3`): The output dimension of the tensor. If `3`, the tensor is
263
+ reshaped to `[batch_size * heads, seq_len, dim // heads]`.
264
+
265
+ Returns:
266
+ `torch.Tensor`: The reshaped tensor.
267
+ """
268
+ head_size = self.heads
269
+ batch_size, seq_len, dim = tensor.shape
270
+ tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)
271
+ tensor = tensor.permute(0, 2, 1, 3)
272
+
273
+ if out_dim == 3:
274
+ tensor = tensor.reshape(batch_size * head_size, seq_len, dim // head_size)
275
+
276
+ return tensor
277
+
278
+ def get_attention_scores(
279
+ self,
280
+ query: torch.Tensor,
281
+ key: torch.Tensor,
282
+ attention_mask: torch.Tensor = None,
283
+ ) -> torch.Tensor:
284
+ r"""
285
+ Compute the attention scores.
286
+
287
+ Args:
288
+ query (`torch.Tensor`): The query tensor.
289
+ key (`torch.Tensor`): The key tensor.
290
+ attention_mask (`torch.Tensor`, *optional*): The attention mask to use. If `None`, no mask is applied.
291
+
292
+ Returns:
293
+ `torch.Tensor`: The attention probabilities/scores.
294
+ """
295
+ dtype = query.dtype
296
+ if self.upcast_attention:
297
+ query = query.float()
298
+ key = key.float()
299
+
300
+ if attention_mask is None:
301
+ baddbmm_input = torch.empty(
302
+ query.shape[0],
303
+ query.shape[1],
304
+ key.shape[1],
305
+ dtype=query.dtype,
306
+ device=query.device,
307
+ )
308
+ beta = 0
309
+ else:
310
+ baddbmm_input = attention_mask
311
+ beta = 1
312
+
313
+ attention_scores = torch.baddbmm(
314
+ baddbmm_input,
315
+ query,
316
+ key.transpose(-1, -2),
317
+ beta=beta,
318
+ alpha=self.scale,
319
+ )
320
+ del baddbmm_input
321
+
322
+ if self.upcast_softmax:
323
+ attention_scores = attention_scores.float()
324
+
325
+ attention_probs = attention_scores.softmax(dim=-1)
326
+ del attention_scores
327
+
328
+ attention_probs = attention_probs.to(dtype)
329
+
330
+ return attention_probs
331
+
332
+ def prepare_attention_mask(
333
+ self,
334
+ attention_mask: torch.Tensor,
335
+ target_length: int,
336
+ batch_size: int,
337
+ out_dim: int = 3,
338
+ ) -> torch.Tensor:
339
+ r"""
340
+ Prepare the attention mask for the attention computation.
341
+
342
+ Args:
343
+ attention_mask (`torch.Tensor`):
344
+ The attention mask to prepare.
345
+ target_length (`int`):
346
+ The target length of the attention mask. This is the length of the attention mask after padding.
347
+ batch_size (`int`):
348
+ The batch size, which is used to repeat the attention mask.
349
+ out_dim (`int`, *optional*, defaults to `3`):
350
+ The output dimension of the attention mask. Can be either `3` or `4`.
351
+
352
+ Returns:
353
+ `torch.Tensor`: The prepared attention mask.
354
+ """
355
+ head_size = self.heads
356
+ if attention_mask is None:
357
+ return attention_mask
358
+
359
+ current_length: int = attention_mask.shape[-1]
360
+ if current_length != target_length:
361
+ if attention_mask.device.type == "mps":
362
+ # HACK: MPS: Does not support padding by greater than dimension of input tensor.
363
+ # Instead, we can manually construct the padding tensor.
364
+ padding_shape = (
365
+ attention_mask.shape[0],
366
+ attention_mask.shape[1],
367
+ target_length,
368
+ )
369
+ padding = torch.zeros(
370
+ padding_shape,
371
+ dtype=attention_mask.dtype,
372
+ device=attention_mask.device,
373
+ )
374
+ attention_mask = torch.cat([attention_mask, padding], dim=2)
375
+ else:
376
+ # TODO: for pipelines such as stable-diffusion, padding cross-attn mask:
377
+ # we want to instead pad by (0, remaining_length), where remaining_length is:
378
+ # remaining_length: int = target_length - current_length
379
+ # TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding
380
+ attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
381
+
382
+ if out_dim == 3:
383
+ if attention_mask.shape[0] < batch_size * head_size:
384
+ attention_mask = attention_mask.repeat_interleave(head_size, dim=0)
385
+ elif out_dim == 4:
386
+ attention_mask = attention_mask.unsqueeze(1)
387
+ attention_mask = attention_mask.repeat_interleave(head_size, dim=1)
388
+
389
+ return attention_mask
390
+
391
+ def norm_encoder_hidden_states(
392
+ self, encoder_hidden_states: torch.Tensor
393
+ ) -> torch.Tensor:
394
+ r"""
395
+ Normalize the encoder hidden states. Requires `self.norm_cross` to be specified when constructing the
396
+ `Attention` class.
397
+
398
+ Args:
399
+ encoder_hidden_states (`torch.Tensor`): Hidden states of the encoder.
400
+
401
+ Returns:
402
+ `torch.Tensor`: The normalized encoder hidden states.
403
+ """
404
+ assert (
405
+ self.norm_cross is not None
406
+ ), "self.norm_cross must be defined to call self.norm_encoder_hidden_states"
407
+
408
+ if isinstance(self.norm_cross, nn.LayerNorm):
409
+ encoder_hidden_states = self.norm_cross(encoder_hidden_states)
410
+ elif isinstance(self.norm_cross, nn.GroupNorm):
411
+ # Group norm norms along the channels dimension and expects
412
+ # input to be in the shape of (N, C, *). In this case, we want
413
+ # to norm along the hidden dimension, so we need to move
414
+ # (batch_size, sequence_length, hidden_size) ->
415
+ # (batch_size, hidden_size, sequence_length)
416
+ encoder_hidden_states = encoder_hidden_states.transpose(1, 2)
417
+ encoder_hidden_states = self.norm_cross(encoder_hidden_states)
418
+ encoder_hidden_states = encoder_hidden_states.transpose(1, 2)
419
+ else:
420
+ assert False
421
+
422
+ return encoder_hidden_states
423
+
424
+ @torch.no_grad()
425
+ def fuse_projections(self, fuse=True):
426
+ is_cross_attention = self.cross_attention_dim != self.query_dim
427
+ device = self.to_q.weight.data.device
428
+ dtype = self.to_q.weight.data.dtype
429
+
430
+ if not is_cross_attention:
431
+ # fetch weight matrices.
432
+ concatenated_weights = torch.cat(
433
+ [self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data]
434
+ )
435
+ in_features = concatenated_weights.shape[1]
436
+ out_features = concatenated_weights.shape[0]
437
+
438
+ # create a new single projection layer and copy over the weights.
439
+ self.to_qkv = self.linear_cls(
440
+ in_features, out_features, bias=False, device=device, dtype=dtype
441
+ )
442
+ self.to_qkv.weight.copy_(concatenated_weights)
443
+
444
+ else:
445
+ concatenated_weights = torch.cat(
446
+ [self.to_k.weight.data, self.to_v.weight.data]
447
+ )
448
+ in_features = concatenated_weights.shape[1]
449
+ out_features = concatenated_weights.shape[0]
450
+
451
+ self.to_kv = self.linear_cls(
452
+ in_features, out_features, bias=False, device=device, dtype=dtype
453
+ )
454
+ self.to_kv.weight.copy_(concatenated_weights)
455
+
456
+ self.fused_projections = fuse
457
+
458
+
459
+ class AttnProcessor:
460
+ r"""
461
+ Default processor for performing attention-related computations.
462
+ """
463
+
464
+ def __call__(
465
+ self,
466
+ attn: Attention,
467
+ hidden_states: torch.FloatTensor,
468
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
469
+ attention_mask: Optional[torch.FloatTensor] = None,
470
+ ) -> torch.Tensor:
471
+ residual = hidden_states
472
+
473
+ input_ndim = hidden_states.ndim
474
+
475
+ if input_ndim == 4:
476
+ batch_size, channel, height, width = hidden_states.shape
477
+ hidden_states = hidden_states.view(
478
+ batch_size, channel, height * width
479
+ ).transpose(1, 2)
480
+
481
+ batch_size, sequence_length, _ = (
482
+ hidden_states.shape
483
+ if encoder_hidden_states is None
484
+ else encoder_hidden_states.shape
485
+ )
486
+ attention_mask = attn.prepare_attention_mask(
487
+ attention_mask, sequence_length, batch_size
488
+ )
489
+
490
+ if attn.group_norm is not None:
491
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
492
+ 1, 2
493
+ )
494
+
495
+ query = attn.to_q(hidden_states)
496
+
497
+ if encoder_hidden_states is None:
498
+ encoder_hidden_states = hidden_states
499
+ elif attn.norm_cross:
500
+ encoder_hidden_states = attn.norm_encoder_hidden_states(
501
+ encoder_hidden_states
502
+ )
503
+
504
+ key = attn.to_k(encoder_hidden_states)
505
+ value = attn.to_v(encoder_hidden_states)
506
+
507
+ query = attn.head_to_batch_dim(query)
508
+ key = attn.head_to_batch_dim(key)
509
+ value = attn.head_to_batch_dim(value)
510
+
511
+ attention_probs = attn.get_attention_scores(query, key, attention_mask)
512
+ hidden_states = torch.bmm(attention_probs, value)
513
+ hidden_states = attn.batch_to_head_dim(hidden_states)
514
+
515
+ # linear proj
516
+ hidden_states = attn.to_out[0](hidden_states)
517
+ # dropout
518
+ hidden_states = attn.to_out[1](hidden_states)
519
+
520
+ if input_ndim == 4:
521
+ hidden_states = hidden_states.transpose(-1, -2).reshape(
522
+ batch_size, channel, height, width
523
+ )
524
+
525
+ if attn.residual_connection:
526
+ hidden_states = hidden_states + residual
527
+
528
+ hidden_states = hidden_states / attn.rescale_output_factor
529
+
530
+ return hidden_states
531
+
532
+
533
+ class AttnProcessor2_0:
534
+ r"""
535
+ Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
536
+ """
537
+
538
+ def __init__(self):
539
+ if not hasattr(F, "scaled_dot_product_attention"):
540
+ raise ImportError(
541
+ "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
542
+ )
543
+
544
+ def __call__(
545
+ self,
546
+ attn: Attention,
547
+ hidden_states: torch.FloatTensor,
548
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
549
+ attention_mask: Optional[torch.FloatTensor] = None,
550
+ ) -> torch.FloatTensor:
551
+ residual = hidden_states
552
+
553
+ input_ndim = hidden_states.ndim
554
+
555
+ if input_ndim == 4:
556
+ batch_size, channel, height, width = hidden_states.shape
557
+ hidden_states = hidden_states.view(
558
+ batch_size, channel, height * width
559
+ ).transpose(1, 2)
560
+
561
+ batch_size, sequence_length, _ = (
562
+ hidden_states.shape
563
+ if encoder_hidden_states is None
564
+ else encoder_hidden_states.shape
565
+ )
566
+
567
+ if attention_mask is not None:
568
+ attention_mask = attn.prepare_attention_mask(
569
+ attention_mask, sequence_length, batch_size
570
+ )
571
+ # scaled_dot_product_attention expects attention_mask shape to be
572
+ # (batch, heads, source_length, target_length)
573
+ attention_mask = attention_mask.view(
574
+ batch_size, attn.heads, -1, attention_mask.shape[-1]
575
+ )
576
+
577
+ if attn.group_norm is not None:
578
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
579
+ 1, 2
580
+ )
581
+
582
+ query = attn.to_q(hidden_states)
583
+
584
+ if encoder_hidden_states is None:
585
+ encoder_hidden_states = hidden_states
586
+ elif attn.norm_cross:
587
+ encoder_hidden_states = attn.norm_encoder_hidden_states(
588
+ encoder_hidden_states
589
+ )
590
+
591
+ key = attn.to_k(encoder_hidden_states)
592
+ value = attn.to_v(encoder_hidden_states)
593
+
594
+ inner_dim = key.shape[-1]
595
+ head_dim = inner_dim // attn.heads
596
+
597
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
598
+
599
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
600
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
601
+
602
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
603
+ # TODO: add support for attn.scale when we move to Torch 2.1
604
+ hidden_states = F.scaled_dot_product_attention(
605
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
606
+ )
607
+
608
+ hidden_states = hidden_states.transpose(1, 2).reshape(
609
+ batch_size, -1, attn.heads * head_dim
610
+ )
611
+ hidden_states = hidden_states.to(query.dtype)
612
+
613
+ # linear proj
614
+ hidden_states = attn.to_out[0](hidden_states)
615
+ # dropout
616
+ hidden_states = attn.to_out[1](hidden_states)
617
+
618
+ if input_ndim == 4:
619
+ hidden_states = hidden_states.transpose(-1, -2).reshape(
620
+ batch_size, channel, height, width
621
+ )
622
+
623
+ if attn.residual_connection:
624
+ hidden_states = hidden_states + residual
625
+
626
+ hidden_states = hidden_states / attn.rescale_output_factor
627
+
628
+ return hidden_states
tsr/models/transformer/basic_transformer_block.py ADDED
@@ -0,0 +1,314 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2023 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from typing import Any, Dict, Optional
16
+
17
+ import torch
18
+ import torch.nn.functional as F
19
+ from torch import nn
20
+
21
+ from .attention import Attention
22
+
23
+
24
+ class BasicTransformerBlock(nn.Module):
25
+ r"""
26
+ A basic Transformer block.
27
+
28
+ Parameters:
29
+ dim (`int`): The number of channels in the input and output.
30
+ num_attention_heads (`int`): The number of heads to use for multi-head attention.
31
+ attention_head_dim (`int`): The number of channels in each head.
32
+ dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
33
+ cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
34
+ activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
35
+ num_embeds_ada_norm (:
36
+ obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
37
+ attention_bias (:
38
+ obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
39
+ only_cross_attention (`bool`, *optional*):
40
+ Whether to use only cross-attention layers. In this case two cross attention layers are used.
41
+ double_self_attention (`bool`, *optional*):
42
+ Whether to use two self-attention layers. In this case no cross attention layers are used.
43
+ upcast_attention (`bool`, *optional*):
44
+ Whether to upcast the attention computation to float32. This is useful for mixed precision training.
45
+ norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
46
+ Whether to use learnable elementwise affine parameters for normalization.
47
+ norm_type (`str`, *optional*, defaults to `"layer_norm"`):
48
+ The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
49
+ final_dropout (`bool` *optional*, defaults to False):
50
+ Whether to apply a final dropout after the last feed-forward layer.
51
+ attention_type (`str`, *optional*, defaults to `"default"`):
52
+ The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`.
53
+ """
54
+
55
+ def __init__(
56
+ self,
57
+ dim: int,
58
+ num_attention_heads: int,
59
+ attention_head_dim: int,
60
+ dropout=0.0,
61
+ cross_attention_dim: Optional[int] = None,
62
+ activation_fn: str = "geglu",
63
+ attention_bias: bool = False,
64
+ only_cross_attention: bool = False,
65
+ double_self_attention: bool = False,
66
+ upcast_attention: bool = False,
67
+ norm_elementwise_affine: bool = True,
68
+ norm_type: str = "layer_norm",
69
+ final_dropout: bool = False,
70
+ ):
71
+ super().__init__()
72
+ self.only_cross_attention = only_cross_attention
73
+
74
+ assert norm_type == "layer_norm"
75
+
76
+ # Define 3 blocks. Each block has its own normalization layer.
77
+ # 1. Self-Attn
78
+ self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
79
+ self.attn1 = Attention(
80
+ query_dim=dim,
81
+ heads=num_attention_heads,
82
+ dim_head=attention_head_dim,
83
+ dropout=dropout,
84
+ bias=attention_bias,
85
+ cross_attention_dim=cross_attention_dim if only_cross_attention else None,
86
+ upcast_attention=upcast_attention,
87
+ )
88
+
89
+ # 2. Cross-Attn
90
+ if cross_attention_dim is not None or double_self_attention:
91
+ # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
92
+ # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
93
+ # the second cross attention block.
94
+ self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
95
+
96
+ self.attn2 = Attention(
97
+ query_dim=dim,
98
+ cross_attention_dim=cross_attention_dim
99
+ if not double_self_attention
100
+ else None,
101
+ heads=num_attention_heads,
102
+ dim_head=attention_head_dim,
103
+ dropout=dropout,
104
+ bias=attention_bias,
105
+ upcast_attention=upcast_attention,
106
+ ) # is self-attn if encoder_hidden_states is none
107
+ else:
108
+ self.norm2 = None
109
+ self.attn2 = None
110
+
111
+ # 3. Feed-forward
112
+ self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
113
+ self.ff = FeedForward(
114
+ dim,
115
+ dropout=dropout,
116
+ activation_fn=activation_fn,
117
+ final_dropout=final_dropout,
118
+ )
119
+
120
+ # let chunk size default to None
121
+ self._chunk_size = None
122
+ self._chunk_dim = 0
123
+
124
+ def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int):
125
+ # Sets chunk feed-forward
126
+ self._chunk_size = chunk_size
127
+ self._chunk_dim = dim
128
+
129
+ def forward(
130
+ self,
131
+ hidden_states: torch.FloatTensor,
132
+ attention_mask: Optional[torch.FloatTensor] = None,
133
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
134
+ encoder_attention_mask: Optional[torch.FloatTensor] = None,
135
+ ) -> torch.FloatTensor:
136
+ # Notice that normalization is always applied before the real computation in the following blocks.
137
+ # 0. Self-Attention
138
+ norm_hidden_states = self.norm1(hidden_states)
139
+
140
+ attn_output = self.attn1(
141
+ norm_hidden_states,
142
+ encoder_hidden_states=encoder_hidden_states
143
+ if self.only_cross_attention
144
+ else None,
145
+ attention_mask=attention_mask,
146
+ )
147
+
148
+ hidden_states = attn_output + hidden_states
149
+
150
+ # 3. Cross-Attention
151
+ if self.attn2 is not None:
152
+ norm_hidden_states = self.norm2(hidden_states)
153
+
154
+ attn_output = self.attn2(
155
+ norm_hidden_states,
156
+ encoder_hidden_states=encoder_hidden_states,
157
+ attention_mask=encoder_attention_mask,
158
+ )
159
+ hidden_states = attn_output + hidden_states
160
+
161
+ # 4. Feed-forward
162
+ norm_hidden_states = self.norm3(hidden_states)
163
+
164
+ if self._chunk_size is not None:
165
+ # "feed_forward_chunk_size" can be used to save memory
166
+ if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
167
+ raise ValueError(
168
+ f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
169
+ )
170
+
171
+ num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
172
+ ff_output = torch.cat(
173
+ [
174
+ self.ff(hid_slice)
175
+ for hid_slice in norm_hidden_states.chunk(
176
+ num_chunks, dim=self._chunk_dim
177
+ )
178
+ ],
179
+ dim=self._chunk_dim,
180
+ )
181
+ else:
182
+ ff_output = self.ff(norm_hidden_states)
183
+
184
+ hidden_states = ff_output + hidden_states
185
+
186
+ return hidden_states
187
+
188
+
189
+ class FeedForward(nn.Module):
190
+ r"""
191
+ A feed-forward layer.
192
+
193
+ Parameters:
194
+ dim (`int`): The number of channels in the input.
195
+ dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.
196
+ mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.
197
+ dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
198
+ activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
199
+ final_dropout (`bool` *optional*, defaults to False): Apply a final dropout.
200
+ """
201
+
202
+ def __init__(
203
+ self,
204
+ dim: int,
205
+ dim_out: Optional[int] = None,
206
+ mult: int = 4,
207
+ dropout: float = 0.0,
208
+ activation_fn: str = "geglu",
209
+ final_dropout: bool = False,
210
+ ):
211
+ super().__init__()
212
+ inner_dim = int(dim * mult)
213
+ dim_out = dim_out if dim_out is not None else dim
214
+ linear_cls = nn.Linear
215
+
216
+ if activation_fn == "gelu":
217
+ act_fn = GELU(dim, inner_dim)
218
+ if activation_fn == "gelu-approximate":
219
+ act_fn = GELU(dim, inner_dim, approximate="tanh")
220
+ elif activation_fn == "geglu":
221
+ act_fn = GEGLU(dim, inner_dim)
222
+ elif activation_fn == "geglu-approximate":
223
+ act_fn = ApproximateGELU(dim, inner_dim)
224
+
225
+ self.net = nn.ModuleList([])
226
+ # project in
227
+ self.net.append(act_fn)
228
+ # project dropout
229
+ self.net.append(nn.Dropout(dropout))
230
+ # project out
231
+ self.net.append(linear_cls(inner_dim, dim_out))
232
+ # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
233
+ if final_dropout:
234
+ self.net.append(nn.Dropout(dropout))
235
+
236
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
237
+ for module in self.net:
238
+ hidden_states = module(hidden_states)
239
+ return hidden_states
240
+
241
+
242
+ class GELU(nn.Module):
243
+ r"""
244
+ GELU activation function with tanh approximation support with `approximate="tanh"`.
245
+
246
+ Parameters:
247
+ dim_in (`int`): The number of channels in the input.
248
+ dim_out (`int`): The number of channels in the output.
249
+ approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation.
250
+ """
251
+
252
+ def __init__(self, dim_in: int, dim_out: int, approximate: str = "none"):
253
+ super().__init__()
254
+ self.proj = nn.Linear(dim_in, dim_out)
255
+ self.approximate = approximate
256
+
257
+ def gelu(self, gate: torch.Tensor) -> torch.Tensor:
258
+ if gate.device.type != "mps":
259
+ return F.gelu(gate, approximate=self.approximate)
260
+ # mps: gelu is not implemented for float16
261
+ return F.gelu(gate.to(dtype=torch.float32), approximate=self.approximate).to(
262
+ dtype=gate.dtype
263
+ )
264
+
265
+ def forward(self, hidden_states):
266
+ hidden_states = self.proj(hidden_states)
267
+ hidden_states = self.gelu(hidden_states)
268
+ return hidden_states
269
+
270
+
271
+ class GEGLU(nn.Module):
272
+ r"""
273
+ A variant of the gated linear unit activation function from https://arxiv.org/abs/2002.05202.
274
+
275
+ Parameters:
276
+ dim_in (`int`): The number of channels in the input.
277
+ dim_out (`int`): The number of channels in the output.
278
+ """
279
+
280
+ def __init__(self, dim_in: int, dim_out: int):
281
+ super().__init__()
282
+ linear_cls = nn.Linear
283
+
284
+ self.proj = linear_cls(dim_in, dim_out * 2)
285
+
286
+ def gelu(self, gate: torch.Tensor) -> torch.Tensor:
287
+ if gate.device.type != "mps":
288
+ return F.gelu(gate)
289
+ # mps: gelu is not implemented for float16
290
+ return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype)
291
+
292
+ def forward(self, hidden_states, scale: float = 1.0):
293
+ args = ()
294
+ hidden_states, gate = self.proj(hidden_states, *args).chunk(2, dim=-1)
295
+ return hidden_states * self.gelu(gate)
296
+
297
+
298
+ class ApproximateGELU(nn.Module):
299
+ r"""
300
+ The approximate form of Gaussian Error Linear Unit (GELU). For more details, see section 2:
301
+ https://arxiv.org/abs/1606.08415.
302
+
303
+ Parameters:
304
+ dim_in (`int`): The number of channels in the input.
305
+ dim_out (`int`): The number of channels in the output.
306
+ """
307
+
308
+ def __init__(self, dim_in: int, dim_out: int):
309
+ super().__init__()
310
+ self.proj = nn.Linear(dim_in, dim_out)
311
+
312
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
313
+ x = self.proj(x)
314
+ return x * torch.sigmoid(1.702 * x)
tsr/models/transformer/transformer_1d.py ADDED
@@ -0,0 +1,216 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from dataclasses import dataclass, field
2
+ from typing import Any, Dict, Optional
3
+
4
+ import torch
5
+ import torch.nn.functional as F
6
+ from torch import nn
7
+
8
+ from ...utils import BaseModule
9
+ from .basic_transformer_block import BasicTransformerBlock
10
+
11
+
12
+ class Transformer1D(BaseModule):
13
+ """
14
+ A 1D Transformer model for sequence data.
15
+
16
+ Parameters:
17
+ num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.
18
+ attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.
19
+ in_channels (`int`, *optional*):
20
+ The number of channels in the input and output (specify if the input is **continuous**).
21
+ num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.
22
+ dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
23
+ cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
24
+ activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward.
25
+ num_embeds_ada_norm ( `int`, *optional*):
26
+ The number of diffusion steps used during training. Pass if at least one of the norm_layers is
27
+ `AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are
28
+ added to the hidden states.
29
+
30
+ During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`.
31
+ attention_bias (`bool`, *optional*):
32
+ Configure if the `TransformerBlocks` attention should contain a bias parameter.
33
+ """
34
+
35
+ @dataclass
36
+ class Config(BaseModule.Config):
37
+ num_attention_heads: int = 16
38
+ attention_head_dim: int = 88
39
+ in_channels: Optional[int] = None
40
+ out_channels: Optional[int] = None
41
+ num_layers: int = 1
42
+ dropout: float = 0.0
43
+ norm_num_groups: int = 32
44
+ cross_attention_dim: Optional[int] = None
45
+ attention_bias: bool = False
46
+ activation_fn: str = "geglu"
47
+ only_cross_attention: bool = False
48
+ double_self_attention: bool = False
49
+ upcast_attention: bool = False
50
+ norm_type: str = "layer_norm"
51
+ norm_elementwise_affine: bool = True
52
+ gradient_checkpointing: bool = False
53
+
54
+ cfg: Config
55
+
56
+ def configure(self) -> None:
57
+ self.num_attention_heads = self.cfg.num_attention_heads
58
+ self.attention_head_dim = self.cfg.attention_head_dim
59
+ inner_dim = self.num_attention_heads * self.attention_head_dim
60
+
61
+ linear_cls = nn.Linear
62
+
63
+ # 2. Define input layers
64
+ self.in_channels = self.cfg.in_channels
65
+
66
+ self.norm = torch.nn.GroupNorm(
67
+ num_groups=self.cfg.norm_num_groups,
68
+ num_channels=self.cfg.in_channels,
69
+ eps=1e-6,
70
+ affine=True,
71
+ )
72
+ self.proj_in = linear_cls(self.cfg.in_channels, inner_dim)
73
+
74
+ # 3. Define transformers blocks
75
+ self.transformer_blocks = nn.ModuleList(
76
+ [
77
+ BasicTransformerBlock(
78
+ inner_dim,
79
+ self.num_attention_heads,
80
+ self.attention_head_dim,
81
+ dropout=self.cfg.dropout,
82
+ cross_attention_dim=self.cfg.cross_attention_dim,
83
+ activation_fn=self.cfg.activation_fn,
84
+ attention_bias=self.cfg.attention_bias,
85
+ only_cross_attention=self.cfg.only_cross_attention,
86
+ double_self_attention=self.cfg.double_self_attention,
87
+ upcast_attention=self.cfg.upcast_attention,
88
+ norm_type=self.cfg.norm_type,
89
+ norm_elementwise_affine=self.cfg.norm_elementwise_affine,
90
+ )
91
+ for d in range(self.cfg.num_layers)
92
+ ]
93
+ )
94
+
95
+ # 4. Define output layers
96
+ self.out_channels = (
97
+ self.cfg.in_channels
98
+ if self.cfg.out_channels is None
99
+ else self.cfg.out_channels
100
+ )
101
+
102
+ self.proj_out = linear_cls(inner_dim, self.cfg.in_channels)
103
+
104
+ self.gradient_checkpointing = self.cfg.gradient_checkpointing
105
+
106
+ def forward(
107
+ self,
108
+ hidden_states: torch.Tensor,
109
+ encoder_hidden_states: Optional[torch.Tensor] = None,
110
+ attention_mask: Optional[torch.Tensor] = None,
111
+ encoder_attention_mask: Optional[torch.Tensor] = None,
112
+ ):
113
+ """
114
+ The [`Transformer1DModel`] forward method.
115
+
116
+ Args:
117
+ hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
118
+ Input `hidden_states`.
119
+ encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
120
+ Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
121
+ self-attention.
122
+ timestep ( `torch.LongTensor`, *optional*):
123
+ Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
124
+ class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
125
+ Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
126
+ `AdaLayerZeroNorm`.
127
+ cross_attention_kwargs ( `Dict[str, Any]`, *optional*):
128
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
129
+ `self.processor` in
130
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
131
+ attention_mask ( `torch.Tensor`, *optional*):
132
+ An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
133
+ is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
134
+ negative values to the attention scores corresponding to "discard" tokens.
135
+ encoder_attention_mask ( `torch.Tensor`, *optional*):
136
+ Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:
137
+
138
+ * Mask `(batch, sequence_length)` True = keep, False = discard.
139
+ * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.
140
+
141
+ If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
142
+ above. This bias will be added to the cross-attention scores.
143
+ return_dict (`bool`, *optional*, defaults to `True`):
144
+ Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
145
+ tuple.
146
+
147
+ Returns:
148
+ If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
149
+ `tuple` where the first element is the sample tensor.
150
+ """
151
+ # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
152
+ # we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
153
+ # we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
154
+ # expects mask of shape:
155
+ # [batch, key_tokens]
156
+ # adds singleton query_tokens dimension:
157
+ # [batch, 1, key_tokens]
158
+ # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
159
+ # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)
160
+ # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
161
+ if attention_mask is not None and attention_mask.ndim == 2:
162
+ # assume that mask is expressed as:
163
+ # (1 = keep, 0 = discard)
164
+ # convert mask into a bias that can be added to attention scores:
165
+ # (keep = +0, discard = -10000.0)
166
+ attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0
167
+ attention_mask = attention_mask.unsqueeze(1)
168
+
169
+ # convert encoder_attention_mask to a bias the same way we do for attention_mask
170
+ if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
171
+ encoder_attention_mask = (
172
+ 1 - encoder_attention_mask.to(hidden_states.dtype)
173
+ ) * -10000.0
174
+ encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
175
+
176
+ # 1. Input
177
+ batch, _, seq_len = hidden_states.shape
178
+ residual = hidden_states
179
+
180
+ hidden_states = self.norm(hidden_states)
181
+ inner_dim = hidden_states.shape[1]
182
+ hidden_states = hidden_states.permute(0, 2, 1).reshape(
183
+ batch, seq_len, inner_dim
184
+ )
185
+ hidden_states = self.proj_in(hidden_states)
186
+
187
+ # 2. Blocks
188
+ for block in self.transformer_blocks:
189
+ if self.training and self.gradient_checkpointing:
190
+ hidden_states = torch.utils.checkpoint.checkpoint(
191
+ block,
192
+ hidden_states,
193
+ attention_mask,
194
+ encoder_hidden_states,
195
+ encoder_attention_mask,
196
+ use_reentrant=False,
197
+ )
198
+ else:
199
+ hidden_states = block(
200
+ hidden_states,
201
+ attention_mask=attention_mask,
202
+ encoder_hidden_states=encoder_hidden_states,
203
+ encoder_attention_mask=encoder_attention_mask,
204
+ )
205
+
206
+ # 3. Output
207
+ hidden_states = self.proj_out(hidden_states)
208
+ hidden_states = (
209
+ hidden_states.reshape(batch, seq_len, inner_dim)
210
+ .permute(0, 2, 1)
211
+ .contiguous()
212
+ )
213
+
214
+ output = hidden_states + residual
215
+
216
+ return output
tsr/system.py ADDED
@@ -0,0 +1,203 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import os
3
+ from dataclasses import dataclass, field
4
+ from typing import List, Union
5
+
6
+ import numpy as np
7
+ import PIL.Image
8
+ import torch
9
+ import torch.nn.functional as F
10
+ import trimesh
11
+ from einops import rearrange
12
+ from huggingface_hub import hf_hub_download
13
+ from omegaconf import OmegaConf
14
+ from PIL import Image
15
+
16
+ from .models.isosurface import MarchingCubeHelper
17
+ from .utils import (
18
+ BaseModule,
19
+ ImagePreprocessor,
20
+ find_class,
21
+ get_spherical_cameras,
22
+ scale_tensor,
23
+ )
24
+
25
+
26
+ class TSR(BaseModule):
27
+ @dataclass
28
+ class Config(BaseModule.Config):
29
+ cond_image_size: int
30
+
31
+ image_tokenizer_cls: str
32
+ image_tokenizer: dict
33
+
34
+ tokenizer_cls: str
35
+ tokenizer: dict
36
+
37
+ backbone_cls: str
38
+ backbone: dict
39
+
40
+ post_processor_cls: str
41
+ post_processor: dict
42
+
43
+ decoder_cls: str
44
+ decoder: dict
45
+
46
+ renderer_cls: str
47
+ renderer: dict
48
+
49
+ cfg: Config
50
+
51
+ @classmethod
52
+ def from_pretrained(
53
+ cls, pretrained_model_name_or_path: str, config_name: str, weight_name: str, token=None
54
+ ):
55
+ if os.path.isdir(pretrained_model_name_or_path):
56
+ config_path = os.path.join(pretrained_model_name_or_path, config_name)
57
+ weight_path = os.path.join(pretrained_model_name_or_path, weight_name)
58
+ else:
59
+ config_path = hf_hub_download(
60
+ repo_id=pretrained_model_name_or_path, filename=config_name, token=token
61
+ )
62
+ weight_path = hf_hub_download(
63
+ repo_id=pretrained_model_name_or_path, filename=weight_name, token=token
64
+ )
65
+
66
+ cfg = OmegaConf.load(config_path)
67
+ OmegaConf.resolve(cfg)
68
+ model = cls(cfg)
69
+ ckpt = torch.load(weight_path, map_location="cpu")
70
+ model.load_state_dict(ckpt)
71
+ return model
72
+
73
+ def configure(self):
74
+ self.image_tokenizer = find_class(self.cfg.image_tokenizer_cls)(
75
+ self.cfg.image_tokenizer
76
+ )
77
+ self.tokenizer = find_class(self.cfg.tokenizer_cls)(self.cfg.tokenizer)
78
+ self.backbone = find_class(self.cfg.backbone_cls)(self.cfg.backbone)
79
+ self.post_processor = find_class(self.cfg.post_processor_cls)(
80
+ self.cfg.post_processor
81
+ )
82
+ self.decoder = find_class(self.cfg.decoder_cls)(self.cfg.decoder)
83
+ self.renderer = find_class(self.cfg.renderer_cls)(self.cfg.renderer)
84
+ self.image_processor = ImagePreprocessor()
85
+ self.isosurface_helper = None
86
+
87
+ def forward(
88
+ self,
89
+ image: Union[
90
+ PIL.Image.Image,
91
+ np.ndarray,
92
+ torch.FloatTensor,
93
+ List[PIL.Image.Image],
94
+ List[np.ndarray],
95
+ List[torch.FloatTensor],
96
+ ],
97
+ device: str,
98
+ ) -> torch.FloatTensor:
99
+ rgb_cond = self.image_processor(image, self.cfg.cond_image_size)[:, None].to(
100
+ device
101
+ )
102
+ batch_size = rgb_cond.shape[0]
103
+
104
+ input_image_tokens: torch.Tensor = self.image_tokenizer(
105
+ rearrange(rgb_cond, "B Nv H W C -> B Nv C H W", Nv=1),
106
+ )
107
+
108
+ input_image_tokens = rearrange(
109
+ input_image_tokens, "B Nv C Nt -> B (Nv Nt) C", Nv=1
110
+ )
111
+
112
+ tokens: torch.Tensor = self.tokenizer(batch_size)
113
+
114
+ tokens = self.backbone(
115
+ tokens,
116
+ encoder_hidden_states=input_image_tokens,
117
+ )
118
+
119
+ scene_codes = self.post_processor(self.tokenizer.detokenize(tokens))
120
+ return scene_codes
121
+
122
+ def render(
123
+ self,
124
+ scene_codes,
125
+ n_views: int,
126
+ elevation_deg: float = 0.0,
127
+ camera_distance: float = 1.9,
128
+ fovy_deg: float = 40.0,
129
+ height: int = 256,
130
+ width: int = 256,
131
+ return_type: str = "pil",
132
+ ):
133
+ rays_o, rays_d = get_spherical_cameras(
134
+ n_views, elevation_deg, camera_distance, fovy_deg, height, width
135
+ )
136
+ rays_o, rays_d = rays_o.to(scene_codes.device), rays_d.to(scene_codes.device)
137
+
138
+ def process_output(image: torch.FloatTensor):
139
+ if return_type == "pt":
140
+ return image
141
+ elif return_type == "np":
142
+ return image.detach().cpu().numpy()
143
+ elif return_type == "pil":
144
+ return Image.fromarray(
145
+ (image.detach().cpu().numpy() * 255.0).astype(np.uint8)
146
+ )
147
+ else:
148
+ raise NotImplementedError
149
+
150
+ images = []
151
+ for scene_code in scene_codes:
152
+ images_ = []
153
+ for i in range(n_views):
154
+ with torch.no_grad():
155
+ image = self.renderer(
156
+ self.decoder, scene_code, rays_o[i], rays_d[i]
157
+ )
158
+ images_.append(process_output(image))
159
+ images.append(images_)
160
+
161
+ return images
162
+
163
+ def set_marching_cubes_resolution(self, resolution: int):
164
+ if (
165
+ self.isosurface_helper is not None
166
+ and self.isosurface_helper.resolution == resolution
167
+ ):
168
+ return
169
+ self.isosurface_helper = MarchingCubeHelper(resolution)
170
+
171
+ def extract_mesh(self, scene_codes, resolution: int = 256, threshold: float = 25.0):
172
+ self.set_marching_cubes_resolution(resolution)
173
+ meshes = []
174
+ for scene_code in scene_codes:
175
+ with torch.no_grad():
176
+ density = self.renderer.query_triplane(
177
+ self.decoder,
178
+ scale_tensor(
179
+ self.isosurface_helper.grid_vertices.to(scene_codes.device),
180
+ self.isosurface_helper.points_range,
181
+ (-self.renderer.cfg.radius, self.renderer.cfg.radius),
182
+ ),
183
+ scene_code,
184
+ )["density_act"]
185
+ v_pos, t_pos_idx = self.isosurface_helper(-(density - threshold))
186
+ v_pos = scale_tensor(
187
+ v_pos,
188
+ self.isosurface_helper.points_range,
189
+ (-self.renderer.cfg.radius, self.renderer.cfg.radius),
190
+ )
191
+ with torch.no_grad():
192
+ color = self.renderer.query_triplane(
193
+ self.decoder,
194
+ v_pos,
195
+ scene_code,
196
+ )["color"]
197
+ mesh = trimesh.Trimesh(
198
+ vertices=v_pos.cpu().numpy(),
199
+ faces=t_pos_idx.cpu().numpy(),
200
+ vertex_colors=color.cpu().numpy(),
201
+ )
202
+ meshes.append(mesh)
203
+ return meshes
tsr/utils.py ADDED
@@ -0,0 +1,482 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import importlib
2
+ import math
3
+ from collections import defaultdict
4
+ from dataclasses import dataclass
5
+ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
6
+
7
+ import imageio
8
+ import numpy as np
9
+ import PIL.Image
10
+ import rembg
11
+ import torch
12
+ import torch.nn as nn
13
+ import torch.nn.functional as F
14
+ import trimesh
15
+ from omegaconf import DictConfig, OmegaConf
16
+ from PIL import Image
17
+
18
+
19
+ def parse_structured(fields: Any, cfg: Optional[Union[dict, DictConfig]] = None) -> Any:
20
+ scfg = OmegaConf.merge(OmegaConf.structured(fields), cfg)
21
+ return scfg
22
+
23
+
24
+ def find_class(cls_string):
25
+ module_string = ".".join(cls_string.split(".")[:-1])
26
+ cls_name = cls_string.split(".")[-1]
27
+ module = importlib.import_module(module_string, package=None)
28
+ cls = getattr(module, cls_name)
29
+ return cls
30
+
31
+
32
+ def get_intrinsic_from_fov(fov, H, W, bs=-1):
33
+ focal_length = 0.5 * H / np.tan(0.5 * fov)
34
+ intrinsic = np.identity(3, dtype=np.float32)
35
+ intrinsic[0, 0] = focal_length
36
+ intrinsic[1, 1] = focal_length
37
+ intrinsic[0, 2] = W / 2.0
38
+ intrinsic[1, 2] = H / 2.0
39
+
40
+ if bs > 0:
41
+ intrinsic = intrinsic[None].repeat(bs, axis=0)
42
+
43
+ return torch.from_numpy(intrinsic)
44
+
45
+
46
+ class BaseModule(nn.Module):
47
+ @dataclass
48
+ class Config:
49
+ pass
50
+
51
+ cfg: Config # add this to every subclass of BaseModule to enable static type checking
52
+
53
+ def __init__(
54
+ self, cfg: Optional[Union[dict, DictConfig]] = None, *args, **kwargs
55
+ ) -> None:
56
+ super().__init__()
57
+ self.cfg = parse_structured(self.Config, cfg)
58
+ self.configure(*args, **kwargs)
59
+
60
+ def configure(self, *args, **kwargs) -> None:
61
+ raise NotImplementedError
62
+
63
+
64
+ class ImagePreprocessor:
65
+ def convert_and_resize(
66
+ self,
67
+ image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
68
+ size: int,
69
+ ):
70
+ if isinstance(image, PIL.Image.Image):
71
+ image = torch.from_numpy(np.array(image).astype(np.float32) / 255.0)
72
+ elif isinstance(image, np.ndarray):
73
+ if image.dtype == np.uint8:
74
+ image = torch.from_numpy(image.astype(np.float32) / 255.0)
75
+ else:
76
+ image = torch.from_numpy(image)
77
+ elif isinstance(image, torch.Tensor):
78
+ pass
79
+
80
+ batched = image.ndim == 4
81
+
82
+ if not batched:
83
+ image = image[None, ...]
84
+ image = F.interpolate(
85
+ image.permute(0, 3, 1, 2),
86
+ (size, size),
87
+ mode="bilinear",
88
+ align_corners=False,
89
+ antialias=True,
90
+ ).permute(0, 2, 3, 1)
91
+ if not batched:
92
+ image = image[0]
93
+ return image
94
+
95
+ def __call__(
96
+ self,
97
+ image: Union[
98
+ PIL.Image.Image,
99
+ np.ndarray,
100
+ torch.FloatTensor,
101
+ List[PIL.Image.Image],
102
+ List[np.ndarray],
103
+ List[torch.FloatTensor],
104
+ ],
105
+ size: int,
106
+ ) -> Any:
107
+ if isinstance(image, (np.ndarray, torch.FloatTensor)) and image.ndim == 4:
108
+ image = self.convert_and_resize(image, size)
109
+ else:
110
+ if not isinstance(image, list):
111
+ image = [image]
112
+ image = [self.convert_and_resize(im, size) for im in image]
113
+ image = torch.stack(image, dim=0)
114
+ return image
115
+
116
+
117
+ def rays_intersect_bbox(
118
+ rays_o: torch.Tensor,
119
+ rays_d: torch.Tensor,
120
+ radius: float,
121
+ near: float = 0.0,
122
+ valid_thresh: float = 0.01,
123
+ ):
124
+ input_shape = rays_o.shape[:-1]
125
+ rays_o, rays_d = rays_o.view(-1, 3), rays_d.view(-1, 3)
126
+ rays_d_valid = torch.where(
127
+ rays_d.abs() < 1e-6, torch.full_like(rays_d, 1e-6), rays_d
128
+ )
129
+ if type(radius) in [int, float]:
130
+ radius = torch.FloatTensor(
131
+ [[-radius, radius], [-radius, radius], [-radius, radius]]
132
+ ).to(rays_o.device)
133
+ radius = (
134
+ 1.0 - 1.0e-3
135
+ ) * radius # tighten the radius to make sure the intersection point lies in the bounding box
136
+ interx0 = (radius[..., 1] - rays_o) / rays_d_valid
137
+ interx1 = (radius[..., 0] - rays_o) / rays_d_valid
138
+ t_near = torch.minimum(interx0, interx1).amax(dim=-1).clamp_min(near)
139
+ t_far = torch.maximum(interx0, interx1).amin(dim=-1)
140
+
141
+ # check wheter a ray intersects the bbox or not
142
+ rays_valid = t_far - t_near > valid_thresh
143
+
144
+ t_near[torch.where(~rays_valid)] = 0.0
145
+ t_far[torch.where(~rays_valid)] = 0.0
146
+
147
+ t_near = t_near.view(*input_shape, 1)
148
+ t_far = t_far.view(*input_shape, 1)
149
+ rays_valid = rays_valid.view(*input_shape)
150
+
151
+ return t_near, t_far, rays_valid
152
+
153
+
154
+ def chunk_batch(func: Callable, chunk_size: int, *args, **kwargs) -> Any:
155
+ if chunk_size <= 0:
156
+ return func(*args, **kwargs)
157
+ B = None
158
+ for arg in list(args) + list(kwargs.values()):
159
+ if isinstance(arg, torch.Tensor):
160
+ B = arg.shape[0]
161
+ break
162
+ assert (
163
+ B is not None
164
+ ), "No tensor found in args or kwargs, cannot determine batch size."
165
+ out = defaultdict(list)
166
+ out_type = None
167
+ # max(1, B) to support B == 0
168
+ for i in range(0, max(1, B), chunk_size):
169
+ out_chunk = func(
170
+ *[
171
+ arg[i : i + chunk_size] if isinstance(arg, torch.Tensor) else arg
172
+ for arg in args
173
+ ],
174
+ **{
175
+ k: arg[i : i + chunk_size] if isinstance(arg, torch.Tensor) else arg
176
+ for k, arg in kwargs.items()
177
+ },
178
+ )
179
+ if out_chunk is None:
180
+ continue
181
+ out_type = type(out_chunk)
182
+ if isinstance(out_chunk, torch.Tensor):
183
+ out_chunk = {0: out_chunk}
184
+ elif isinstance(out_chunk, tuple) or isinstance(out_chunk, list):
185
+ chunk_length = len(out_chunk)
186
+ out_chunk = {i: chunk for i, chunk in enumerate(out_chunk)}
187
+ elif isinstance(out_chunk, dict):
188
+ pass
189
+ else:
190
+ print(
191
+ f"Return value of func must be in type [torch.Tensor, list, tuple, dict], get {type(out_chunk)}."
192
+ )
193
+ exit(1)
194
+ for k, v in out_chunk.items():
195
+ v = v if torch.is_grad_enabled() else v.detach()
196
+ out[k].append(v)
197
+
198
+ if out_type is None:
199
+ return None
200
+
201
+ out_merged: Dict[Any, Optional[torch.Tensor]] = {}
202
+ for k, v in out.items():
203
+ if all([vv is None for vv in v]):
204
+ # allow None in return value
205
+ out_merged[k] = None
206
+ elif all([isinstance(vv, torch.Tensor) for vv in v]):
207
+ out_merged[k] = torch.cat(v, dim=0)
208
+ else:
209
+ raise TypeError(
210
+ f"Unsupported types in return value of func: {[type(vv) for vv in v if not isinstance(vv, torch.Tensor)]}"
211
+ )
212
+
213
+ if out_type is torch.Tensor:
214
+ return out_merged[0]
215
+ elif out_type in [tuple, list]:
216
+ return out_type([out_merged[i] for i in range(chunk_length)])
217
+ elif out_type is dict:
218
+ return out_merged
219
+
220
+
221
+ ValidScale = Union[Tuple[float, float], torch.FloatTensor]
222
+
223
+
224
+ def scale_tensor(dat: torch.FloatTensor, inp_scale: ValidScale, tgt_scale: ValidScale):
225
+ if inp_scale is None:
226
+ inp_scale = (0, 1)
227
+ if tgt_scale is None:
228
+ tgt_scale = (0, 1)
229
+ if isinstance(tgt_scale, torch.FloatTensor):
230
+ assert dat.shape[-1] == tgt_scale.shape[-1]
231
+ dat = (dat - inp_scale[0]) / (inp_scale[1] - inp_scale[0])
232
+ dat = dat * (tgt_scale[1] - tgt_scale[0]) + tgt_scale[0]
233
+ return dat
234
+
235
+
236
+ def get_activation(name) -> Callable:
237
+ if name is None:
238
+ return lambda x: x
239
+ name = name.lower()
240
+ if name == "none":
241
+ return lambda x: x
242
+ elif name == "exp":
243
+ return lambda x: torch.exp(x)
244
+ elif name == "sigmoid":
245
+ return lambda x: torch.sigmoid(x)
246
+ elif name == "tanh":
247
+ return lambda x: torch.tanh(x)
248
+ elif name == "softplus":
249
+ return lambda x: F.softplus(x)
250
+ else:
251
+ try:
252
+ return getattr(F, name)
253
+ except AttributeError:
254
+ raise ValueError(f"Unknown activation function: {name}")
255
+
256
+
257
+ def get_ray_directions(
258
+ H: int,
259
+ W: int,
260
+ focal: Union[float, Tuple[float, float]],
261
+ principal: Optional[Tuple[float, float]] = None,
262
+ use_pixel_centers: bool = True,
263
+ normalize: bool = True,
264
+ ) -> torch.FloatTensor:
265
+ """
266
+ Get ray directions for all pixels in camera coordinate.
267
+ Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
268
+ ray-tracing-generating-camera-rays/standard-coordinate-systems
269
+
270
+ Inputs:
271
+ H, W, focal, principal, use_pixel_centers: image height, width, focal length, principal point and whether use pixel centers
272
+ Outputs:
273
+ directions: (H, W, 3), the direction of the rays in camera coordinate
274
+ """
275
+ pixel_center = 0.5 if use_pixel_centers else 0
276
+
277
+ if isinstance(focal, float):
278
+ fx, fy = focal, focal
279
+ cx, cy = W / 2, H / 2
280
+ else:
281
+ fx, fy = focal
282
+ assert principal is not None
283
+ cx, cy = principal
284
+
285
+ i, j = torch.meshgrid(
286
+ torch.arange(W, dtype=torch.float32) + pixel_center,
287
+ torch.arange(H, dtype=torch.float32) + pixel_center,
288
+ indexing="xy",
289
+ )
290
+
291
+ directions = torch.stack([(i - cx) / fx, -(j - cy) / fy, -torch.ones_like(i)], -1)
292
+
293
+ if normalize:
294
+ directions = F.normalize(directions, dim=-1)
295
+
296
+ return directions
297
+
298
+
299
+ def get_rays(
300
+ directions,
301
+ c2w,
302
+ keepdim=False,
303
+ noise_scale=0.0,
304
+ normalize=False,
305
+ ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
306
+ # Rotate ray directions from camera coordinate to the world coordinate
307
+ assert directions.shape[-1] == 3
308
+
309
+ if directions.ndim == 2: # (N_rays, 3)
310
+ if c2w.ndim == 2: # (4, 4)
311
+ c2w = c2w[None, :, :]
312
+ assert c2w.ndim == 3 # (N_rays, 4, 4) or (1, 4, 4)
313
+ rays_d = (directions[:, None, :] * c2w[:, :3, :3]).sum(-1) # (N_rays, 3)
314
+ rays_o = c2w[:, :3, 3].expand(rays_d.shape)
315
+ elif directions.ndim == 3: # (H, W, 3)
316
+ assert c2w.ndim in [2, 3]
317
+ if c2w.ndim == 2: # (4, 4)
318
+ rays_d = (directions[:, :, None, :] * c2w[None, None, :3, :3]).sum(
319
+ -1
320
+ ) # (H, W, 3)
321
+ rays_o = c2w[None, None, :3, 3].expand(rays_d.shape)
322
+ elif c2w.ndim == 3: # (B, 4, 4)
323
+ rays_d = (directions[None, :, :, None, :] * c2w[:, None, None, :3, :3]).sum(
324
+ -1
325
+ ) # (B, H, W, 3)
326
+ rays_o = c2w[:, None, None, :3, 3].expand(rays_d.shape)
327
+ elif directions.ndim == 4: # (B, H, W, 3)
328
+ assert c2w.ndim == 3 # (B, 4, 4)
329
+ rays_d = (directions[:, :, :, None, :] * c2w[:, None, None, :3, :3]).sum(
330
+ -1
331
+ ) # (B, H, W, 3)
332
+ rays_o = c2w[:, None, None, :3, 3].expand(rays_d.shape)
333
+
334
+ # add camera noise to avoid grid-like artifect
335
+ # https://github.com/ashawkey/stable-dreamfusion/blob/49c3d4fa01d68a4f027755acf94e1ff6020458cc/nerf/utils.py#L373
336
+ if noise_scale > 0:
337
+ rays_o = rays_o + torch.randn(3, device=rays_o.device) * noise_scale
338
+ rays_d = rays_d + torch.randn(3, device=rays_d.device) * noise_scale
339
+
340
+ if normalize:
341
+ rays_d = F.normalize(rays_d, dim=-1)
342
+ if not keepdim:
343
+ rays_o, rays_d = rays_o.reshape(-1, 3), rays_d.reshape(-1, 3)
344
+
345
+ return rays_o, rays_d
346
+
347
+
348
+ def get_spherical_cameras(
349
+ n_views: int,
350
+ elevation_deg: float,
351
+ camera_distance: float,
352
+ fovy_deg: float,
353
+ height: int,
354
+ width: int,
355
+ ):
356
+ azimuth_deg = torch.linspace(0, 360.0, n_views + 1)[:n_views]
357
+ elevation_deg = torch.full_like(azimuth_deg, elevation_deg)
358
+ camera_distances = torch.full_like(elevation_deg, camera_distance)
359
+
360
+ elevation = elevation_deg * math.pi / 180
361
+ azimuth = azimuth_deg * math.pi / 180
362
+
363
+ # convert spherical coordinates to cartesian coordinates
364
+ # right hand coordinate system, x back, y right, z up
365
+ # elevation in (-90, 90), azimuth from +x to +y in (-180, 180)
366
+ camera_positions = torch.stack(
367
+ [
368
+ camera_distances * torch.cos(elevation) * torch.cos(azimuth),
369
+ camera_distances * torch.cos(elevation) * torch.sin(azimuth),
370
+ camera_distances * torch.sin(elevation),
371
+ ],
372
+ dim=-1,
373
+ )
374
+
375
+ # default scene center at origin
376
+ center = torch.zeros_like(camera_positions)
377
+ # default camera up direction as +z
378
+ up = torch.as_tensor([0, 0, 1], dtype=torch.float32)[None, :].repeat(n_views, 1)
379
+
380
+ fovy = torch.full_like(elevation_deg, fovy_deg) * math.pi / 180
381
+
382
+ lookat = F.normalize(center - camera_positions, dim=-1)
383
+ right = F.normalize(torch.cross(lookat, up), dim=-1)
384
+ up = F.normalize(torch.cross(right, lookat), dim=-1)
385
+ c2w3x4 = torch.cat(
386
+ [torch.stack([right, up, -lookat], dim=-1), camera_positions[:, :, None]],
387
+ dim=-1,
388
+ )
389
+ c2w = torch.cat([c2w3x4, torch.zeros_like(c2w3x4[:, :1])], dim=1)
390
+ c2w[:, 3, 3] = 1.0
391
+
392
+ # get directions by dividing directions_unit_focal by focal length
393
+ focal_length = 0.5 * height / torch.tan(0.5 * fovy)
394
+ directions_unit_focal = get_ray_directions(
395
+ H=height,
396
+ W=width,
397
+ focal=1.0,
398
+ )
399
+ directions = directions_unit_focal[None, :, :, :].repeat(n_views, 1, 1, 1)
400
+ directions[:, :, :, :2] = (
401
+ directions[:, :, :, :2] / focal_length[:, None, None, None]
402
+ )
403
+ # must use normalize=True to normalize directions here
404
+ rays_o, rays_d = get_rays(directions, c2w, keepdim=True, normalize=True)
405
+
406
+ return rays_o, rays_d
407
+
408
+
409
+ def remove_background(
410
+ image: PIL.Image.Image,
411
+ rembg_session: Any = None,
412
+ force: bool = False,
413
+ **rembg_kwargs,
414
+ ) -> PIL.Image.Image:
415
+ do_remove = True
416
+ if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
417
+ do_remove = False
418
+ do_remove = do_remove or force
419
+ if do_remove:
420
+ image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
421
+ return image
422
+
423
+
424
+ def resize_foreground(
425
+ image: PIL.Image.Image,
426
+ ratio: float,
427
+ ) -> PIL.Image.Image:
428
+ image = np.array(image)
429
+ assert image.shape[-1] == 4
430
+ alpha = np.where(image[..., 3] > 0)
431
+ y1, y2, x1, x2 = (
432
+ alpha[0].min(),
433
+ alpha[0].max(),
434
+ alpha[1].min(),
435
+ alpha[1].max(),
436
+ )
437
+ # crop the foreground
438
+ fg = image[y1:y2, x1:x2]
439
+ # pad to square
440
+ size = max(fg.shape[0], fg.shape[1])
441
+ ph0, pw0 = (size - fg.shape[0]) // 2, (size - fg.shape[1]) // 2
442
+ ph1, pw1 = size - fg.shape[0] - ph0, size - fg.shape[1] - pw0
443
+ new_image = np.pad(
444
+ fg,
445
+ ((ph0, ph1), (pw0, pw1), (0, 0)),
446
+ mode="constant",
447
+ constant_values=((0, 0), (0, 0), (0, 0)),
448
+ )
449
+
450
+ # compute padding according to the ratio
451
+ new_size = int(new_image.shape[0] / ratio)
452
+ # pad to size, double side
453
+ ph0, pw0 = (new_size - size) // 2, (new_size - size) // 2
454
+ ph1, pw1 = new_size - size - ph0, new_size - size - pw0
455
+ new_image = np.pad(
456
+ new_image,
457
+ ((ph0, ph1), (pw0, pw1), (0, 0)),
458
+ mode="constant",
459
+ constant_values=((0, 0), (0, 0), (0, 0)),
460
+ )
461
+ new_image = PIL.Image.fromarray(new_image)
462
+ return new_image
463
+
464
+
465
+ def save_video(
466
+ frames: List[PIL.Image.Image],
467
+ output_path: str,
468
+ fps: int = 30,
469
+ ):
470
+ # use imageio to save video
471
+ frames = [np.array(frame) for frame in frames]
472
+ writer = imageio.get_writer(output_path, fps=fps)
473
+ for frame in frames:
474
+ writer.append_data(frame)
475
+ writer.close()
476
+
477
+
478
+ def to_gradio_3d_orientation(mesh):
479
+ mesh.apply_transform(trimesh.transformations.rotation_matrix(-np.pi/2, [1, 0, 0]))
480
+ # mesh.apply_scale([1, 1, -1])
481
+ mesh.apply_transform(trimesh.transformations.rotation_matrix(np.pi/2, [0, 1, 0]))
482
+ return mesh
wheel/torchmcubes-0.1.0-cp310-cp310-linux_x86_64.whl ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4af160ba1274e2205d3529a7b82efdb6946c2158a78e19631ed840301055b8d6
3
+ size 5824388