GonzaloMG commited on
Commit
9558d50
1 Parent(s): 5c08b12

Upload 22 files

Browse files
GeoWizard/README.md ADDED
@@ -0,0 +1 @@
 
 
1
+ Code is copied from https://github.com/fuxiao0719/GeoWizard. Modifications are indicated within the code.
GeoWizard/geowizard/models/attention.py ADDED
@@ -0,0 +1,777 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+
16
+ # Some modifications are reimplemented in public environments by Xiao Fu and Mu Hu
17
+
18
+
19
+ from typing import Any, Dict, Optional
20
+
21
+ import torch
22
+ import torch.nn.functional as F
23
+ from torch import nn
24
+ import xformers
25
+
26
+ from diffusers.utils import USE_PEFT_BACKEND
27
+ from diffusers.utils.torch_utils import maybe_allow_in_graph
28
+ from diffusers.models.activations import GEGLU, GELU, ApproximateGELU
29
+ from diffusers.models.attention_processor import Attention
30
+ from diffusers.models.embeddings import SinusoidalPositionalEmbedding
31
+ from diffusers.models.lora import LoRACompatibleLinear
32
+ from diffusers.models.normalization import AdaLayerNorm, AdaLayerNormContinuous, AdaLayerNormZero, RMSNorm
33
+
34
+
35
+ def _chunked_feed_forward(
36
+ ff: nn.Module, hidden_states: torch.Tensor, chunk_dim: int, chunk_size: int, lora_scale: Optional[float] = None
37
+ ):
38
+ # "feed_forward_chunk_size" can be used to save memory
39
+ if hidden_states.shape[chunk_dim] % chunk_size != 0:
40
+ raise ValueError(
41
+ f"`hidden_states` dimension to be chunked: {hidden_states.shape[chunk_dim]} has to be divisible by chunk size: {chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
42
+ )
43
+
44
+ num_chunks = hidden_states.shape[chunk_dim] // chunk_size
45
+ if lora_scale is None:
46
+ ff_output = torch.cat(
47
+ [ff(hid_slice) for hid_slice in hidden_states.chunk(num_chunks, dim=chunk_dim)],
48
+ dim=chunk_dim,
49
+ )
50
+ else:
51
+ # TOOD(Patrick): LoRA scale can be removed once PEFT refactor is complete
52
+ ff_output = torch.cat(
53
+ [ff(hid_slice, scale=lora_scale) for hid_slice in hidden_states.chunk(num_chunks, dim=chunk_dim)],
54
+ dim=chunk_dim,
55
+ )
56
+
57
+ return ff_output
58
+
59
+
60
+ @maybe_allow_in_graph
61
+ class GatedSelfAttentionDense(nn.Module):
62
+ r"""
63
+ A gated self-attention dense layer that combines visual features and object features.
64
+
65
+ Parameters:
66
+ query_dim (`int`): The number of channels in the query.
67
+ context_dim (`int`): The number of channels in the context.
68
+ n_heads (`int`): The number of heads to use for attention.
69
+ d_head (`int`): The number of channels in each head.
70
+ """
71
+
72
+ def __init__(self, query_dim: int, context_dim: int, n_heads: int, d_head: int):
73
+ super().__init__()
74
+
75
+ # we need a linear projection since we need cat visual feature and obj feature
76
+ self.linear = nn.Linear(context_dim, query_dim)
77
+
78
+ self.attn = Attention(query_dim=query_dim, heads=n_heads, dim_head=d_head)
79
+ self.ff = FeedForward(query_dim, activation_fn="geglu")
80
+
81
+ self.norm1 = nn.LayerNorm(query_dim)
82
+ self.norm2 = nn.LayerNorm(query_dim)
83
+
84
+ self.register_parameter("alpha_attn", nn.Parameter(torch.tensor(0.0)))
85
+ self.register_parameter("alpha_dense", nn.Parameter(torch.tensor(0.0)))
86
+
87
+ self.enabled = True
88
+
89
+ def forward(self, x: torch.Tensor, objs: torch.Tensor) -> torch.Tensor:
90
+ if not self.enabled:
91
+ return x
92
+
93
+ n_visual = x.shape[1]
94
+ objs = self.linear(objs)
95
+
96
+ x = x + self.alpha_attn.tanh() * self.attn(self.norm1(torch.cat([x, objs], dim=1)))[:, :n_visual, :]
97
+ x = x + self.alpha_dense.tanh() * self.ff(self.norm2(x))
98
+
99
+ return x
100
+
101
+
102
+ @maybe_allow_in_graph
103
+ class BasicTransformerBlock(nn.Module):
104
+ r"""
105
+ A basic Transformer block.
106
+
107
+ Parameters:
108
+ dim (`int`): The number of channels in the input and output.
109
+ num_attention_heads (`int`): The number of heads to use for multi-head attention.
110
+ attention_head_dim (`int`): The number of channels in each head.
111
+ dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
112
+ cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
113
+ activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
114
+ num_embeds_ada_norm (:
115
+ obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
116
+ attention_bias (:
117
+ obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
118
+ only_cross_attention (`bool`, *optional*):
119
+ Whether to use only cross-attention layers. In this case two cross attention layers are used.
120
+ double_self_attention (`bool`, *optional*):
121
+ Whether to use two self-attention layers. In this case no cross attention layers are used.
122
+ upcast_attention (`bool`, *optional*):
123
+ Whether to upcast the attention computation to float32. This is useful for mixed precision training.
124
+ norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
125
+ Whether to use learnable elementwise affine parameters for normalization.
126
+ norm_type (`str`, *optional*, defaults to `"layer_norm"`):
127
+ The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
128
+ final_dropout (`bool` *optional*, defaults to False):
129
+ Whether to apply a final dropout after the last feed-forward layer.
130
+ attention_type (`str`, *optional*, defaults to `"default"`):
131
+ The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`.
132
+ positional_embeddings (`str`, *optional*, defaults to `None`):
133
+ The type of positional embeddings to apply to.
134
+ num_positional_embeddings (`int`, *optional*, defaults to `None`):
135
+ The maximum number of positional embeddings to apply.
136
+ """
137
+
138
+ def __init__(
139
+ self,
140
+ dim: int,
141
+ num_attention_heads: int,
142
+ attention_head_dim: int,
143
+ dropout=0.0,
144
+ cross_attention_dim: Optional[int] = None,
145
+ activation_fn: str = "geglu",
146
+ num_embeds_ada_norm: Optional[int] = None,
147
+ attention_bias: bool = False,
148
+ only_cross_attention: bool = False,
149
+ double_self_attention: bool = False,
150
+ upcast_attention: bool = False,
151
+ norm_elementwise_affine: bool = True,
152
+ norm_type: str = "layer_norm", # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single'
153
+ norm_eps: float = 1e-5,
154
+ final_dropout: bool = False,
155
+ attention_type: str = "default",
156
+ positional_embeddings: Optional[str] = None,
157
+ num_positional_embeddings: Optional[int] = None,
158
+ ada_norm_continous_conditioning_embedding_dim: Optional[int] = None,
159
+ ada_norm_bias: Optional[int] = None,
160
+ ff_inner_dim: Optional[int] = None,
161
+ ff_bias: bool = True,
162
+ attention_out_bias: bool = True,
163
+ ):
164
+ super().__init__()
165
+ self.only_cross_attention = only_cross_attention
166
+
167
+ self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
168
+ self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm"
169
+ self.use_ada_layer_norm_single = norm_type == "ada_norm_single"
170
+ self.use_layer_norm = norm_type == "layer_norm"
171
+ self.use_ada_layer_norm_continuous = norm_type == "ada_norm_continuous"
172
+
173
+ if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
174
+ raise ValueError(
175
+ f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"
176
+ f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."
177
+ )
178
+
179
+ if positional_embeddings and (num_positional_embeddings is None):
180
+ raise ValueError(
181
+ "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined."
182
+ )
183
+
184
+ if positional_embeddings == "sinusoidal":
185
+ self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings)
186
+ else:
187
+ self.pos_embed = None
188
+
189
+ # Define 3 blocks. Each block has its own normalization layer.
190
+ # 1. Self-Attn
191
+ if self.use_ada_layer_norm:
192
+ self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm)
193
+ elif self.use_ada_layer_norm_zero:
194
+ self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)
195
+ elif self.use_ada_layer_norm_continuous:
196
+ self.norm1 = AdaLayerNormContinuous(
197
+ dim,
198
+ ada_norm_continous_conditioning_embedding_dim,
199
+ norm_elementwise_affine,
200
+ norm_eps,
201
+ ada_norm_bias,
202
+ "rms_norm",
203
+ )
204
+ else:
205
+ self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
206
+
207
+
208
+ self.attn1 = CustomJointAttention(
209
+ query_dim=dim,
210
+ heads=num_attention_heads,
211
+ dim_head=attention_head_dim,
212
+ dropout=dropout,
213
+ bias=attention_bias,
214
+ cross_attention_dim=cross_attention_dim if only_cross_attention else None,
215
+ upcast_attention=upcast_attention,
216
+ out_bias=attention_out_bias
217
+ )
218
+
219
+ # 2. Cross-Attn
220
+ if cross_attention_dim is not None or double_self_attention:
221
+ # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
222
+ # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
223
+ # the second cross attention block.
224
+
225
+ if self.use_ada_layer_norm:
226
+ self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm)
227
+ elif self.use_ada_layer_norm_continuous:
228
+ self.norm2 = AdaLayerNormContinuous(
229
+ dim,
230
+ ada_norm_continous_conditioning_embedding_dim,
231
+ norm_elementwise_affine,
232
+ norm_eps,
233
+ ada_norm_bias,
234
+ "rms_norm",
235
+ )
236
+ else:
237
+ self.norm2 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
238
+
239
+ self.attn2 = Attention(
240
+ query_dim=dim,
241
+ cross_attention_dim=cross_attention_dim if not double_self_attention else None,
242
+ heads=num_attention_heads,
243
+ dim_head=attention_head_dim,
244
+ dropout=dropout,
245
+ bias=attention_bias,
246
+ upcast_attention=upcast_attention,
247
+ out_bias=attention_out_bias,
248
+ ) # is self-attn if encoder_hidden_states is none
249
+ else:
250
+ self.norm2 = None
251
+ self.attn2 = None
252
+
253
+ # 3. Feed-forward
254
+ if self.use_ada_layer_norm_continuous:
255
+ self.norm3 = AdaLayerNormContinuous(
256
+ dim,
257
+ ada_norm_continous_conditioning_embedding_dim,
258
+ norm_elementwise_affine,
259
+ norm_eps,
260
+ ada_norm_bias,
261
+ "layer_norm",
262
+ )
263
+ elif not self.use_ada_layer_norm_single:
264
+ self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
265
+
266
+ self.ff = FeedForward(
267
+ dim,
268
+ dropout=dropout,
269
+ activation_fn=activation_fn,
270
+ final_dropout=final_dropout,
271
+ inner_dim=ff_inner_dim,
272
+ bias=ff_bias,
273
+ )
274
+
275
+ # 4. Fuser
276
+ if attention_type == "gated" or attention_type == "gated-text-image":
277
+ self.fuser = GatedSelfAttentionDense(dim, cross_attention_dim, num_attention_heads, attention_head_dim)
278
+
279
+ # 5. Scale-shift for PixArt-Alpha.
280
+ if self.use_ada_layer_norm_single:
281
+ self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
282
+
283
+ # let chunk size default to None
284
+ self._chunk_size = None
285
+ self._chunk_dim = 0
286
+
287
+ def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
288
+ # Sets chunk feed-forward
289
+ self._chunk_size = chunk_size
290
+ self._chunk_dim = dim
291
+
292
+ def forward(
293
+ self,
294
+ hidden_states: torch.FloatTensor,
295
+ attention_mask: Optional[torch.FloatTensor] = None,
296
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
297
+ encoder_attention_mask: Optional[torch.FloatTensor] = None,
298
+ timestep: Optional[torch.LongTensor] = None,
299
+ cross_attention_kwargs: Dict[str, Any] = None,
300
+ class_labels: Optional[torch.LongTensor] = None,
301
+ added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
302
+ ) -> torch.FloatTensor:
303
+ # Notice that normalization is always applied before the real computation in the following blocks.
304
+
305
+ # 0. Self-Attention
306
+ batch_size = hidden_states.shape[0]
307
+
308
+ if self.use_ada_layer_norm:
309
+ norm_hidden_states = self.norm1(hidden_states, timestep)
310
+ elif self.use_ada_layer_norm_zero:
311
+ norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
312
+ hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
313
+ )
314
+ elif self.use_layer_norm:
315
+ norm_hidden_states = self.norm1(hidden_states)
316
+ elif self.use_ada_layer_norm_continuous:
317
+ norm_hidden_states = self.norm1(hidden_states, added_cond_kwargs["pooled_text_emb"])
318
+ elif self.use_ada_layer_norm_single:
319
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
320
+ self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1)
321
+ ).chunk(6, dim=1)
322
+ norm_hidden_states = self.norm1(hidden_states)
323
+ norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
324
+ norm_hidden_states = norm_hidden_states.squeeze(1)
325
+ else:
326
+ raise ValueError("Incorrect norm used")
327
+
328
+ if self.pos_embed is not None:
329
+ norm_hidden_states = self.pos_embed(norm_hidden_states)
330
+
331
+ # 1. Retrieve lora scale.
332
+ lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
333
+
334
+ # 2. Prepare GLIGEN inputs
335
+ cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}
336
+ gligen_kwargs = cross_attention_kwargs.pop("gligen", None)
337
+
338
+ attn_output = self.attn1(
339
+ norm_hidden_states,
340
+ encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
341
+ attention_mask=attention_mask,
342
+ **cross_attention_kwargs,
343
+ )
344
+ if self.use_ada_layer_norm_zero:
345
+ attn_output = gate_msa.unsqueeze(1) * attn_output
346
+ elif self.use_ada_layer_norm_single:
347
+ attn_output = gate_msa * attn_output
348
+
349
+ hidden_states = attn_output + hidden_states
350
+ if hidden_states.ndim == 4:
351
+ hidden_states = hidden_states.squeeze(1)
352
+
353
+ # 2.5 GLIGEN Control
354
+ if gligen_kwargs is not None:
355
+ hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"])
356
+
357
+ # 3. Cross-Attention
358
+ if self.attn2 is not None:
359
+ if self.use_ada_layer_norm:
360
+ norm_hidden_states = self.norm2(hidden_states, timestep)
361
+ elif self.use_ada_layer_norm_zero or self.use_layer_norm:
362
+ norm_hidden_states = self.norm2(hidden_states)
363
+ elif self.use_ada_layer_norm_single:
364
+ # For PixArt norm2 isn't applied here:
365
+ # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103
366
+ norm_hidden_states = hidden_states
367
+ elif self.use_ada_layer_norm_continuous:
368
+ norm_hidden_states = self.norm2(hidden_states, added_cond_kwargs["pooled_text_emb"])
369
+ else:
370
+ raise ValueError("Incorrect norm")
371
+
372
+ if self.pos_embed is not None and self.use_ada_layer_norm_single is False:
373
+ norm_hidden_states = self.pos_embed(norm_hidden_states)
374
+
375
+ attn_output = self.attn2(
376
+ norm_hidden_states,
377
+ encoder_hidden_states=encoder_hidden_states,
378
+ attention_mask=encoder_attention_mask,
379
+ **cross_attention_kwargs,
380
+ )
381
+ hidden_states = attn_output + hidden_states
382
+
383
+ # 4. Feed-forward
384
+ if self.use_ada_layer_norm_continuous:
385
+ norm_hidden_states = self.norm3(hidden_states, added_cond_kwargs["pooled_text_emb"])
386
+ elif not self.use_ada_layer_norm_single:
387
+ norm_hidden_states = self.norm3(hidden_states)
388
+
389
+ if self.use_ada_layer_norm_zero:
390
+ norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
391
+
392
+ if self.use_ada_layer_norm_single:
393
+ norm_hidden_states = self.norm2(hidden_states)
394
+ norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
395
+
396
+ if self._chunk_size is not None:
397
+ # "feed_forward_chunk_size" can be used to save memory
398
+ ff_output = _chunked_feed_forward(
399
+ self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size, lora_scale=lora_scale
400
+ )
401
+ else:
402
+ ff_output = self.ff(norm_hidden_states, scale=lora_scale)
403
+
404
+ if self.use_ada_layer_norm_zero:
405
+ ff_output = gate_mlp.unsqueeze(1) * ff_output
406
+ elif self.use_ada_layer_norm_single:
407
+ ff_output = gate_mlp * ff_output
408
+
409
+ hidden_states = ff_output + hidden_states
410
+ if hidden_states.ndim == 4:
411
+ hidden_states = hidden_states.squeeze(1)
412
+
413
+ return hidden_states
414
+
415
+
416
+ class CustomJointAttention(Attention):
417
+ def set_use_memory_efficient_attention_xformers(
418
+ self, use_memory_efficient_attention_xformers: bool, *args, **kwargs
419
+ ):
420
+ processor = XFormersJointAttnProcessor()
421
+ self.set_processor(processor)
422
+ # print("using xformers attention processor")
423
+
424
+
425
+ class XFormersJointAttnProcessor:
426
+ r"""
427
+ Default processor for performing attention-related computations.
428
+ """
429
+
430
+ def __call__(
431
+ self,
432
+ attn: Attention,
433
+ hidden_states,
434
+ encoder_hidden_states=None,
435
+ attention_mask=None,
436
+ temb=None,
437
+ num_tasks=2
438
+ ):
439
+
440
+ residual = hidden_states
441
+
442
+ if attn.spatial_norm is not None:
443
+ hidden_states = attn.spatial_norm(hidden_states, temb)
444
+
445
+ input_ndim = hidden_states.ndim
446
+
447
+ if input_ndim == 4:
448
+ batch_size, channel, height, width = hidden_states.shape
449
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
450
+
451
+ batch_size, sequence_length, _ = (
452
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
453
+ )
454
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
455
+
456
+ # from yuancheng; here attention_mask is None
457
+ if attention_mask is not None:
458
+ # expand our mask's singleton query_tokens dimension:
459
+ # [batch*heads, 1, key_tokens] ->
460
+ # [batch*heads, query_tokens, key_tokens]
461
+ # so that it can be added as a bias onto the attention scores that xformers computes:
462
+ # [batch*heads, query_tokens, key_tokens]
463
+ # we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
464
+ _, query_tokens, _ = hidden_states.shape
465
+ attention_mask = attention_mask.expand(-1, query_tokens, -1)
466
+
467
+ if attn.group_norm is not None:
468
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
469
+
470
+ query = attn.to_q(hidden_states)
471
+
472
+ if encoder_hidden_states is None:
473
+ encoder_hidden_states = hidden_states
474
+ elif attn.norm_cross:
475
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
476
+
477
+ key = attn.to_k(encoder_hidden_states)
478
+ value = attn.to_v(encoder_hidden_states)
479
+
480
+ assert num_tasks == 2 # only support two tasks now
481
+
482
+ key_0, key_1 = torch.chunk(key, dim=0, chunks=2) # keys shape (b t) d c
483
+ value_0, value_1 = torch.chunk(value, dim=0, chunks=2)
484
+
485
+ # key = torch.cat([key_1, key_0], dim=0)
486
+ # value = torch.cat([value_1, value_0], dim=0)
487
+
488
+ key = torch.cat([key_0, key_1], dim=1) # (b t) 2d c
489
+ value = torch.cat([value_0, value_1], dim=1) # (b t) 2d c
490
+ key = torch.cat([key]*2, dim=0) # (2 b t) 2d c
491
+ value = torch.cat([value]*2, dim=0) # (2 b t) 2d c
492
+
493
+ query = attn.head_to_batch_dim(query).contiguous()
494
+ key = attn.head_to_batch_dim(key).contiguous()
495
+ value = attn.head_to_batch_dim(value).contiguous()
496
+
497
+ hidden_states = xformers.ops.memory_efficient_attention(query, key, value, attn_bias=attention_mask)
498
+ hidden_states = attn.batch_to_head_dim(hidden_states)
499
+
500
+ # linear proj
501
+ hidden_states = attn.to_out[0](hidden_states)
502
+ # dropout
503
+ hidden_states = attn.to_out[1](hidden_states)
504
+
505
+ if input_ndim == 4:
506
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
507
+
508
+ if attn.residual_connection:
509
+ hidden_states = hidden_states + residual
510
+
511
+ hidden_states = hidden_states / attn.rescale_output_factor
512
+
513
+ return hidden_states
514
+
515
+
516
+ @maybe_allow_in_graph
517
+ class TemporalBasicTransformerBlock(nn.Module):
518
+ r"""
519
+ A basic Transformer block for video like data.
520
+
521
+ Parameters:
522
+ dim (`int`): The number of channels in the input and output.
523
+ time_mix_inner_dim (`int`): The number of channels for temporal attention.
524
+ num_attention_heads (`int`): The number of heads to use for multi-head attention.
525
+ attention_head_dim (`int`): The number of channels in each head.
526
+ cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
527
+ """
528
+
529
+ def __init__(
530
+ self,
531
+ dim: int,
532
+ time_mix_inner_dim: int,
533
+ num_attention_heads: int,
534
+ attention_head_dim: int,
535
+ cross_attention_dim: Optional[int] = None,
536
+ ):
537
+ super().__init__()
538
+ self.is_res = dim == time_mix_inner_dim
539
+
540
+ self.norm_in = nn.LayerNorm(dim)
541
+
542
+ # Define 3 blocks. Each block has its own normalization layer.
543
+ # 1. Self-Attn
544
+ self.norm_in = nn.LayerNorm(dim)
545
+ self.ff_in = FeedForward(
546
+ dim,
547
+ dim_out=time_mix_inner_dim,
548
+ activation_fn="geglu",
549
+ )
550
+
551
+ self.norm1 = nn.LayerNorm(time_mix_inner_dim)
552
+ self.attn1 = Attention(
553
+ query_dim=time_mix_inner_dim,
554
+ heads=num_attention_heads,
555
+ dim_head=attention_head_dim,
556
+ cross_attention_dim=None,
557
+ )
558
+
559
+ # 2. Cross-Attn
560
+ if cross_attention_dim is not None:
561
+ # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
562
+ # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
563
+ # the second cross attention block.
564
+ self.norm2 = nn.LayerNorm(time_mix_inner_dim)
565
+ self.attn2 = Attention(
566
+ query_dim=time_mix_inner_dim,
567
+ cross_attention_dim=cross_attention_dim,
568
+ heads=num_attention_heads,
569
+ dim_head=attention_head_dim,
570
+ ) # is self-attn if encoder_hidden_states is none
571
+ else:
572
+ self.norm2 = None
573
+ self.attn2 = None
574
+
575
+ # 3. Feed-forward
576
+ self.norm3 = nn.LayerNorm(time_mix_inner_dim)
577
+ self.ff = FeedForward(time_mix_inner_dim, activation_fn="geglu")
578
+
579
+ # let chunk size default to None
580
+ self._chunk_size = None
581
+ self._chunk_dim = None
582
+
583
+ def set_chunk_feed_forward(self, chunk_size: Optional[int], **kwargs):
584
+ # Sets chunk feed-forward
585
+ self._chunk_size = chunk_size
586
+ # chunk dim should be hardcoded to 1 to have better speed vs. memory trade-off
587
+ self._chunk_dim = 1
588
+
589
+ def forward(
590
+ self,
591
+ hidden_states: torch.FloatTensor,
592
+ num_frames: int,
593
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
594
+ ) -> torch.FloatTensor:
595
+ # Notice that normalization is always applied before the real computation in the following blocks.
596
+ # 0. Self-Attention
597
+ batch_size = hidden_states.shape[0]
598
+
599
+ batch_frames, seq_length, channels = hidden_states.shape
600
+ batch_size = batch_frames // num_frames
601
+
602
+ hidden_states = hidden_states[None, :].reshape(batch_size, num_frames, seq_length, channels)
603
+ hidden_states = hidden_states.permute(0, 2, 1, 3)
604
+ hidden_states = hidden_states.reshape(batch_size * seq_length, num_frames, channels)
605
+
606
+ residual = hidden_states
607
+ hidden_states = self.norm_in(hidden_states)
608
+
609
+ if self._chunk_size is not None:
610
+ hidden_states = _chunked_feed_forward(self.ff_in, hidden_states, self._chunk_dim, self._chunk_size)
611
+ else:
612
+ hidden_states = self.ff_in(hidden_states)
613
+
614
+ if self.is_res:
615
+ hidden_states = hidden_states + residual
616
+
617
+ norm_hidden_states = self.norm1(hidden_states)
618
+ attn_output = self.attn1(norm_hidden_states, encoder_hidden_states=None)
619
+ hidden_states = attn_output + hidden_states
620
+
621
+ # 3. Cross-Attention
622
+ if self.attn2 is not None:
623
+ norm_hidden_states = self.norm2(hidden_states)
624
+ attn_output = self.attn2(norm_hidden_states, encoder_hidden_states=encoder_hidden_states)
625
+ hidden_states = attn_output + hidden_states
626
+
627
+ # 4. Feed-forward
628
+ norm_hidden_states = self.norm3(hidden_states)
629
+
630
+ if self._chunk_size is not None:
631
+ ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)
632
+ else:
633
+ ff_output = self.ff(norm_hidden_states)
634
+
635
+ if self.is_res:
636
+ hidden_states = ff_output + hidden_states
637
+ else:
638
+ hidden_states = ff_output
639
+
640
+ hidden_states = hidden_states[None, :].reshape(batch_size, seq_length, num_frames, channels)
641
+ hidden_states = hidden_states.permute(0, 2, 1, 3)
642
+ hidden_states = hidden_states.reshape(batch_size * num_frames, seq_length, channels)
643
+
644
+ return hidden_states
645
+
646
+
647
+ class SkipFFTransformerBlock(nn.Module):
648
+ def __init__(
649
+ self,
650
+ dim: int,
651
+ num_attention_heads: int,
652
+ attention_head_dim: int,
653
+ kv_input_dim: int,
654
+ kv_input_dim_proj_use_bias: bool,
655
+ dropout=0.0,
656
+ cross_attention_dim: Optional[int] = None,
657
+ attention_bias: bool = False,
658
+ attention_out_bias: bool = True,
659
+ ):
660
+ super().__init__()
661
+ if kv_input_dim != dim:
662
+ self.kv_mapper = nn.Linear(kv_input_dim, dim, kv_input_dim_proj_use_bias)
663
+ else:
664
+ self.kv_mapper = None
665
+
666
+ self.norm1 = RMSNorm(dim, 1e-06)
667
+
668
+ self.attn1 = Attention(
669
+ query_dim=dim,
670
+ heads=num_attention_heads,
671
+ dim_head=attention_head_dim,
672
+ dropout=dropout,
673
+ bias=attention_bias,
674
+ cross_attention_dim=cross_attention_dim,
675
+ out_bias=attention_out_bias,
676
+ )
677
+
678
+ self.norm2 = RMSNorm(dim, 1e-06)
679
+
680
+ self.attn2 = Attention(
681
+ query_dim=dim,
682
+ cross_attention_dim=cross_attention_dim,
683
+ heads=num_attention_heads,
684
+ dim_head=attention_head_dim,
685
+ dropout=dropout,
686
+ bias=attention_bias,
687
+ out_bias=attention_out_bias,
688
+ )
689
+
690
+ def forward(self, hidden_states, encoder_hidden_states, cross_attention_kwargs):
691
+ cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}
692
+
693
+ if self.kv_mapper is not None:
694
+ encoder_hidden_states = self.kv_mapper(F.silu(encoder_hidden_states))
695
+
696
+ norm_hidden_states = self.norm1(hidden_states)
697
+
698
+ attn_output = self.attn1(
699
+ norm_hidden_states,
700
+ encoder_hidden_states=encoder_hidden_states,
701
+ **cross_attention_kwargs,
702
+ )
703
+
704
+ hidden_states = attn_output + hidden_states
705
+
706
+ norm_hidden_states = self.norm2(hidden_states)
707
+
708
+ attn_output = self.attn2(
709
+ norm_hidden_states,
710
+ encoder_hidden_states=encoder_hidden_states,
711
+ **cross_attention_kwargs,
712
+ )
713
+
714
+ hidden_states = attn_output + hidden_states
715
+
716
+ return hidden_states
717
+
718
+
719
+ class FeedForward(nn.Module):
720
+ r"""
721
+ A feed-forward layer.
722
+
723
+ Parameters:
724
+ dim (`int`): The number of channels in the input.
725
+ dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.
726
+ mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.
727
+ dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
728
+ activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
729
+ final_dropout (`bool` *optional*, defaults to False): Apply a final dropout.
730
+ bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
731
+ """
732
+
733
+ def __init__(
734
+ self,
735
+ dim: int,
736
+ dim_out: Optional[int] = None,
737
+ mult: int = 4,
738
+ dropout: float = 0.0,
739
+ activation_fn: str = "geglu",
740
+ final_dropout: bool = False,
741
+ inner_dim=None,
742
+ bias: bool = True,
743
+ ):
744
+ super().__init__()
745
+ if inner_dim is None:
746
+ inner_dim = int(dim * mult)
747
+ dim_out = dim_out if dim_out is not None else dim
748
+ linear_cls = LoRACompatibleLinear if not USE_PEFT_BACKEND else nn.Linear
749
+
750
+ if activation_fn == "gelu":
751
+ act_fn = GELU(dim, inner_dim, bias=bias)
752
+ if activation_fn == "gelu-approximate":
753
+ act_fn = GELU(dim, inner_dim, approximate="tanh", bias=bias)
754
+ elif activation_fn == "geglu":
755
+ act_fn = GEGLU(dim, inner_dim, bias=bias)
756
+ elif activation_fn == "geglu-approximate":
757
+ act_fn = ApproximateGELU(dim, inner_dim, bias=bias)
758
+
759
+ self.net = nn.ModuleList([])
760
+ # project in
761
+ self.net.append(act_fn)
762
+ # project dropout
763
+ self.net.append(nn.Dropout(dropout))
764
+ # project out
765
+ self.net.append(linear_cls(inner_dim, dim_out, bias=bias))
766
+ # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
767
+ if final_dropout:
768
+ self.net.append(nn.Dropout(dropout))
769
+
770
+ def forward(self, hidden_states: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
771
+ compatible_cls = (GEGLU,) if USE_PEFT_BACKEND else (GEGLU, LoRACompatibleLinear)
772
+ for module in self.net:
773
+ if isinstance(module, compatible_cls):
774
+ hidden_states = module(hidden_states, scale)
775
+ else:
776
+ hidden_states = module(hidden_states)
777
+ return hidden_states
GeoWizard/geowizard/models/geowizard_pipeline.py ADDED
@@ -0,0 +1,383 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Adapted from Marigold :https://github.com/prs-eth/Marigold
2
+
3
+ # @GonzaloMartinGarcia, all new additions to the GeoWizard code have been marked with # add.
4
+
5
+ from typing import Any, Dict, Union
6
+
7
+ import torch.nn as nn
8
+ import torch
9
+ from torch.utils.data import DataLoader, TensorDataset
10
+ import numpy as np
11
+ from tqdm.auto import tqdm
12
+ from PIL import Image
13
+ from diffusers import (
14
+ DiffusionPipeline,
15
+ DDIMScheduler,
16
+ AutoencoderKL,
17
+ )
18
+ from GeoWizard.geowizard.models.unet_2d_condition import UNet2DConditionModel
19
+ from diffusers.utils import BaseOutput
20
+ from transformers import CLIPTextModel, CLIPTokenizer
21
+ from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
22
+ import torchvision.transforms.functional as TF
23
+ from torchvision.transforms import InterpolationMode
24
+
25
+ from GeoWizard.geowizard.utils.image_util import resize_max_res,chw2hwc,colorize_depth_maps
26
+ from GeoWizard.geowizard.utils.depth_ensemble import ensemble_depths
27
+ from GeoWizard.geowizard.utils.normal_ensemble import ensemble_normals
28
+ import cv2
29
+
30
+ # add
31
+ from GeoWizard.geowizard.utils.noise import pyramid_noise_like
32
+
33
+
34
+ class DepthNormalPipelineOutput(BaseOutput):
35
+ """
36
+ Output class for GeoWizard monocular depth & normal prediction pipeline.
37
+ Args:
38
+ depth_np (`np.ndarray`):
39
+ Predicted depth map, with depth values in the range of [0, 1].
40
+ depth_colored (`PIL.Image.Image`):
41
+ Colorized depth map, with the shape of [3, H, W] and values in [0, 1].
42
+ normal_np (`np.ndarray`):
43
+ Predicted normal map, with depth values in the range of [0, 1].
44
+ normal_colored (`PIL.Image.Image`):
45
+ Colorized normal map, with the shape of [3, H, W] and values in [0, 1].
46
+ uncertainty (`None` or `np.ndarray`):
47
+ Uncalibrated uncertainty(MAD, median absolute deviation) coming from ensembling.
48
+ """
49
+ depth_np: np.ndarray
50
+ depth_colored: Image.Image
51
+ normal_np: np.ndarray
52
+ normal_colored: Image.Image
53
+ uncertainty: Union[None, np.ndarray]
54
+
55
+ class DepthNormalEstimationPipeline(DiffusionPipeline):
56
+ # two hyper-parameters
57
+ latent_scale_factor = 0.18215
58
+
59
+ def __init__(self,
60
+ unet:UNet2DConditionModel,
61
+ vae:AutoencoderKL,
62
+ scheduler:DDIMScheduler,
63
+ image_encoder:CLIPVisionModelWithProjection,
64
+ feature_extractor:CLIPImageProcessor,
65
+ ):
66
+ super().__init__()
67
+
68
+ self.register_modules(
69
+ unet=unet,
70
+ vae=vae,
71
+ scheduler=scheduler,
72
+ image_encoder=image_encoder,
73
+ feature_extractor=feature_extractor,
74
+ )
75
+ self.img_embed = None
76
+
77
+ @torch.no_grad()
78
+ def __call__(self,
79
+ input_image:Image,
80
+ denoising_steps: int = 10,
81
+ ensemble_size: int = 10,
82
+ processing_res: int = 768,
83
+ match_input_res:bool =True,
84
+ batch_size:int = 0,
85
+ domain: str = "indoor",
86
+ color_map: str="Spectral",
87
+ show_progress_bar:bool = True,
88
+ ensemble_kwargs: Dict = None,
89
+ # add
90
+ noise="gaussian",
91
+ ) -> DepthNormalPipelineOutput:
92
+
93
+ # inherit from thea Diffusion Pipeline
94
+ device = self.device
95
+ input_size = input_image.size
96
+
97
+ # adjust the input resolution.
98
+ if not match_input_res:
99
+ assert (
100
+ processing_res is not None
101
+ )," Value Error: `resize_output_back` is only valid with "
102
+
103
+ assert processing_res >=0
104
+ assert denoising_steps >=1
105
+ assert ensemble_size >=1
106
+
107
+ # --------------- Image Processing ------------------------
108
+ # Resize image
109
+ if processing_res >0:
110
+ input_image = resize_max_res(
111
+ input_image, max_edge_resolution=processing_res
112
+ )
113
+
114
+ # Convert the image to RGB, to 1. reomve the alpha channel.
115
+ input_image = input_image.convert("RGB")
116
+ image = np.array(input_image)
117
+
118
+ # Normalize RGB Values.
119
+ rgb = np.transpose(image,(2,0,1))
120
+ rgb_norm = rgb / 255.0 * 2.0 - 1.0 # [0, 255] -> [-1, 1]
121
+ rgb_norm = torch.from_numpy(rgb_norm).to(self.dtype)
122
+ rgb_norm = rgb_norm.to(device)
123
+
124
+ assert rgb_norm.min() >= -1.0 and rgb_norm.max() <= 1.0
125
+
126
+ # ----------------- predicting depth -----------------
127
+ duplicated_rgb = torch.stack([rgb_norm] * ensemble_size)
128
+ single_rgb_dataset = TensorDataset(duplicated_rgb)
129
+
130
+ # find the batch size
131
+ if batch_size>0:
132
+ _bs = batch_size
133
+ else:
134
+ _bs = 1
135
+
136
+ single_rgb_loader = DataLoader(single_rgb_dataset, batch_size=_bs, shuffle=False)
137
+
138
+ # predicted the depth
139
+ depth_pred_ls = []
140
+ normal_pred_ls = []
141
+
142
+ if show_progress_bar:
143
+ iterable_bar = tqdm(
144
+ single_rgb_loader, desc=" " * 2 + "Inference batches", leave=False
145
+ )
146
+ else:
147
+ iterable_bar = single_rgb_loader
148
+
149
+ for batch in iterable_bar:
150
+ (batched_image, )= batch # here the image is still around 0-1
151
+
152
+ depth_pred_raw, normal_pred_raw = self.single_infer(
153
+ input_rgb=batched_image,
154
+ num_inference_steps=denoising_steps,
155
+ domain=domain,
156
+ show_pbar=show_progress_bar,
157
+ # add
158
+ noise=noise,
159
+ )
160
+ depth_pred_ls.append(depth_pred_raw.detach().clone())
161
+ normal_pred_ls.append(normal_pred_raw.detach().clone())
162
+
163
+ depth_preds = torch.concat(depth_pred_ls, axis=0).squeeze() #(10,224,768)
164
+ normal_preds = torch.concat(normal_pred_ls, axis=0).squeeze()
165
+ torch.cuda.empty_cache() # clear vram cache for ensembling
166
+
167
+ # ----------------- Test-time ensembling -----------------
168
+ if ensemble_size > 1:
169
+ depth_pred, pred_uncert = ensemble_depths(
170
+ depth_preds, **(ensemble_kwargs or {})
171
+ )
172
+ normal_pred = ensemble_normals(normal_preds)
173
+ else:
174
+ depth_pred = depth_preds
175
+ normal_pred = normal_preds
176
+ pred_uncert = None
177
+
178
+ # ----------------- Post processing -----------------
179
+ # Scale prediction to [0, 1]
180
+ min_d = torch.min(depth_pred)
181
+ max_d = torch.max(depth_pred)
182
+ depth_pred = (depth_pred - min_d) / (max_d - min_d)
183
+
184
+ # Convert to numpy
185
+ depth_pred = depth_pred.cpu().numpy().astype(np.float32)
186
+ normal_pred = normal_pred.cpu().numpy().astype(np.float32)
187
+
188
+ # Resize back to original resolution
189
+ if match_input_res:
190
+ pred_img = Image.fromarray(depth_pred)
191
+ pred_img = pred_img.resize(input_size)
192
+ depth_pred = np.asarray(pred_img)
193
+ normal_pred = cv2.resize(chw2hwc(normal_pred), input_size, interpolation = cv2.INTER_NEAREST)
194
+
195
+ # Clip output range: current size is the original size
196
+ depth_pred = depth_pred.clip(0, 1)
197
+ normal_pred = normal_pred.clip(-1, 1)
198
+
199
+ # Colorize
200
+ depth_colored = colorize_depth_maps(
201
+ depth_pred, 0, 1, cmap=color_map
202
+ ).squeeze() # [3, H, W], value in (0, 1)
203
+ depth_colored = (depth_colored * 255).astype(np.uint8)
204
+ depth_colored_hwc = chw2hwc(depth_colored)
205
+ depth_colored_img = Image.fromarray(depth_colored_hwc)
206
+
207
+ normal_colored = ((normal_pred + 1)/2 * 255).astype(np.uint8)
208
+ normal_colored_img = Image.fromarray(normal_colored)
209
+
210
+ self.img_embed = None
211
+
212
+ return DepthNormalPipelineOutput(
213
+ depth_np = depth_pred,
214
+ depth_colored = depth_colored_img,
215
+ normal_np = normal_pred,
216
+ normal_colored = normal_colored_img,
217
+ uncertainty=pred_uncert,
218
+ )
219
+
220
+ def __encode_img_embed(self, rgb):
221
+ """
222
+ Encode clip embeddings for img
223
+ """
224
+ clip_image_mean = torch.as_tensor(self.feature_extractor.image_mean)[:,None,None].to(device=self.device, dtype=self.dtype)
225
+ clip_image_std = torch.as_tensor(self.feature_extractor.image_std)[:,None,None].to(device=self.device, dtype=self.dtype)
226
+
227
+ img_in_proc = TF.resize((rgb +1)/2,
228
+ (self.feature_extractor.crop_size['height'], self.feature_extractor.crop_size['width']),
229
+ interpolation=InterpolationMode.BICUBIC,
230
+ antialias=True
231
+ )
232
+ # do the normalization in float32 to preserve precision
233
+ img_in_proc = ((img_in_proc.float() - clip_image_mean) / clip_image_std).to(self.dtype)
234
+ img_embed = self.image_encoder(img_in_proc).image_embeds.unsqueeze(1).to(self.dtype)
235
+
236
+ self.img_embed = img_embed
237
+
238
+
239
+ @torch.no_grad()
240
+ def single_infer(self,input_rgb:torch.Tensor,
241
+ num_inference_steps:int,
242
+ domain:str,
243
+ show_pbar:bool,
244
+ # add
245
+ noise="gaussian",
246
+ ):
247
+
248
+ device = input_rgb.device
249
+
250
+ # Set timesteps: inherit from the diffuison pipeline
251
+ self.scheduler.set_timesteps(num_inference_steps, device=device) # here the numbers of the steps is only 10.
252
+ timesteps = self.scheduler.timesteps # [T]
253
+
254
+ # encode image
255
+ rgb_latent = self.encode_RGB(input_rgb)
256
+
257
+ # add
258
+ # Initial geometric maps
259
+ if noise == "gaussian":
260
+ geo_latent = torch.randn(rgb_latent.shape, device=device, dtype=self.dtype).repeat(2,1,1,1)
261
+ elif noise == "zeros":
262
+ geo_latent = torch.zeros(rgb_latent.shape, device=device, dtype=self.dtype).repeat(2,1,1,1)
263
+ elif noise == "pyramid":
264
+ geo_latent = pyramid_noise_like(rgb_latent, timesteps).repeat(2,1,1,1)
265
+ else:
266
+ raise ValueError(f"Invalid noise type: {noise}")
267
+
268
+ rgb_latent = rgb_latent.repeat(2,1,1,1)
269
+
270
+ # Batched img embedding
271
+ if self.img_embed is None:
272
+ self.__encode_img_embed(input_rgb)
273
+
274
+ batch_img_embed = self.img_embed.repeat(
275
+ (rgb_latent.shape[0], 1, 1)
276
+ ) # [B, 1, 768]
277
+
278
+ # hybrid switcher
279
+ geo_class = torch.tensor([[0., 1.], [1, 0]], device=device, dtype=self.dtype)
280
+ geo_embedding = torch.cat([torch.sin(geo_class), torch.cos(geo_class)], dim=-1)
281
+
282
+ if domain == "indoor":
283
+ domain_class = torch.tensor([[1., 0., 0]], device=device, dtype=self.dtype).repeat(2,1)
284
+ elif domain == "outdoor":
285
+ domain_class = torch.tensor([[0., 1., 0]], device=device, dtype=self.dtype).repeat(2,1)
286
+ elif domain == "object":
287
+ domain_class = torch.tensor([[0., 0., 1]], device=device, dtype=self.dtype).repeat(2,1)
288
+ domain_embedding = torch.cat([torch.sin(domain_class), torch.cos(domain_class)], dim=-1)
289
+
290
+ class_embedding = torch.cat((geo_embedding, domain_embedding), dim=-1)
291
+
292
+ # Denoising loop
293
+ if show_pbar:
294
+ iterable = tqdm(
295
+ enumerate(timesteps),
296
+ total=len(timesteps),
297
+ leave=False,
298
+ desc=" " * 4 + "Diffusion denoising",
299
+ )
300
+ else:
301
+ iterable = enumerate(timesteps)
302
+
303
+ for i, t in iterable:
304
+ unet_input = torch.cat([rgb_latent, geo_latent], dim=1)
305
+
306
+ # predict the noise residual
307
+ noise_pred = self.unet(
308
+ unet_input, t.repeat(2), encoder_hidden_states=batch_img_embed, class_labels=class_embedding
309
+ ).sample # [B, 4, h, w]
310
+
311
+ # compute the previous noisy sample x_t -> x_t-1
312
+ geo_latent = self.scheduler.step(noise_pred, t, geo_latent).prev_sample
313
+
314
+ geo_latent = geo_latent
315
+ torch.cuda.empty_cache()
316
+
317
+ depth = self.decode_depth(geo_latent[0][None])
318
+ depth = torch.clip(depth, -1.0, 1.0)
319
+ depth = (depth + 1.0) / 2.0
320
+
321
+ normal = self.decode_normal(geo_latent[1][None])
322
+ normal /= (torch.norm(normal, p=2, dim=1, keepdim=True)+1e-5)
323
+ normal *= -1.
324
+
325
+ return depth, normal
326
+
327
+
328
+ def encode_RGB(self, rgb_in: torch.Tensor) -> torch.Tensor:
329
+ """
330
+ Encode RGB image into latent.
331
+ Args:
332
+ rgb_in (`torch.Tensor`):
333
+ Input RGB image to be encoded.
334
+ Returns:
335
+ `torch.Tensor`: Image latent.
336
+ """
337
+
338
+ # encode
339
+ h = self.vae.encoder(rgb_in)
340
+
341
+ moments = self.vae.quant_conv(h)
342
+ mean, logvar = torch.chunk(moments, 2, dim=1)
343
+ # scale latent
344
+ rgb_latent = mean * self.latent_scale_factor
345
+
346
+ return rgb_latent
347
+
348
+ def decode_depth(self, depth_latent: torch.Tensor) -> torch.Tensor:
349
+ """
350
+ Decode depth latent into depth map.
351
+ Args:
352
+ depth_latent (`torch.Tensor`):
353
+ Depth latent to be decoded.
354
+ Returns:
355
+ `torch.Tensor`: Decoded depth map.
356
+ """
357
+
358
+ # scale latent
359
+ depth_latent = depth_latent / self.latent_scale_factor
360
+ # decode
361
+ z = self.vae.post_quant_conv(depth_latent)
362
+ stacked = self.vae.decoder(z)
363
+ # mean of output channels
364
+ depth_mean = stacked.mean(dim=1, keepdim=True)
365
+ return depth_mean
366
+
367
+ def decode_normal(self, normal_latent: torch.Tensor) -> torch.Tensor:
368
+ """
369
+ Decode normal latent into normal map.
370
+ Args:
371
+ normal_latent (`torch.Tensor`):
372
+ Depth latent to be decoded.
373
+ Returns:
374
+ `torch.Tensor`: Decoded normal map.
375
+ """
376
+
377
+ # scale latent
378
+ normal_latent = normal_latent / self.latent_scale_factor
379
+ # decode
380
+ z = self.vae.post_quant_conv(normal_latent)
381
+ normal = self.vae.decoder(z)
382
+ return normal
383
+
GeoWizard/geowizard/models/transformer_2d.py ADDED
@@ -0,0 +1,463 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ # Some modifications are reimplemented in public environments by Xiao Fu and Mu Hu
16
+
17
+ from dataclasses import dataclass
18
+ from typing import Any, Dict, Optional
19
+
20
+ import torch
21
+ import torch.nn.functional as F
22
+ from torch import nn
23
+
24
+ from diffusers.configuration_utils import ConfigMixin, register_to_config
25
+ from diffusers.models.embeddings import ImagePositionalEmbeddings
26
+ from diffusers.utils import USE_PEFT_BACKEND, BaseOutput, deprecate, is_torch_version
27
+ from GeoWizard.geowizard.models.attention import BasicTransformerBlock
28
+ from diffusers.models.embeddings import PatchEmbed, PixArtAlphaTextProjection
29
+ from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear
30
+ from diffusers.models.modeling_utils import ModelMixin
31
+ from diffusers.models.normalization import AdaLayerNormSingle
32
+
33
+
34
+ @dataclass
35
+ class Transformer2DModelOutput(BaseOutput):
36
+ """
37
+ The output of [`Transformer2DModel`].
38
+
39
+ Args:
40
+ sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete):
41
+ The hidden states output conditioned on the `encoder_hidden_states` input. If discrete, returns probability
42
+ distributions for the unnoised latent pixels.
43
+ """
44
+
45
+ sample: torch.FloatTensor
46
+
47
+
48
+ class Transformer2DModel(ModelMixin, ConfigMixin):
49
+ """
50
+ A 2D Transformer model for image-like data.
51
+
52
+ Parameters:
53
+ num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.
54
+ attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.
55
+ in_channels (`int`, *optional*):
56
+ The number of channels in the input and output (specify if the input is **continuous**).
57
+ num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.
58
+ dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
59
+ cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
60
+ sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**).
61
+ This is fixed during training since it is used to learn a number of position embeddings.
62
+ num_vector_embeds (`int`, *optional*):
63
+ The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**).
64
+ Includes the class for the masked latent pixel.
65
+ activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward.
66
+ num_embeds_ada_norm ( `int`, *optional*):
67
+ The number of diffusion steps used during training. Pass if at least one of the norm_layers is
68
+ `AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are
69
+ added to the hidden states.
70
+
71
+ During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`.
72
+ attention_bias (`bool`, *optional*):
73
+ Configure if the `TransformerBlocks` attention should contain a bias parameter.
74
+ """
75
+
76
+ _supports_gradient_checkpointing = True
77
+
78
+ @register_to_config
79
+ def __init__(
80
+ self,
81
+ num_attention_heads: int = 16,
82
+ attention_head_dim: int = 88,
83
+ in_channels: Optional[int] = None,
84
+ out_channels: Optional[int] = None,
85
+ num_layers: int = 1,
86
+ dropout: float = 0.0,
87
+ norm_num_groups: int = 32,
88
+ cross_attention_dim: Optional[int] = None,
89
+ attention_bias: bool = False,
90
+ sample_size: Optional[int] = None,
91
+ num_vector_embeds: Optional[int] = None,
92
+ patch_size: Optional[int] = None,
93
+ activation_fn: str = "geglu",
94
+ num_embeds_ada_norm: Optional[int] = None,
95
+ use_linear_projection: bool = False,
96
+ only_cross_attention: bool = False,
97
+ double_self_attention: bool = False,
98
+ upcast_attention: bool = False,
99
+ norm_type: str = "layer_norm",
100
+ norm_elementwise_affine: bool = True,
101
+ norm_eps: float = 1e-5,
102
+ attention_type: str = "default",
103
+ caption_channels: int = None,
104
+ ):
105
+ super().__init__()
106
+ self.use_linear_projection = use_linear_projection
107
+ self.num_attention_heads = num_attention_heads
108
+ self.attention_head_dim = attention_head_dim
109
+ inner_dim = num_attention_heads * attention_head_dim
110
+
111
+ conv_cls = nn.Conv2d if USE_PEFT_BACKEND else LoRACompatibleConv
112
+ linear_cls = nn.Linear if USE_PEFT_BACKEND else LoRACompatibleLinear
113
+
114
+ # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
115
+ # Define whether input is continuous or discrete depending on configuration
116
+ self.is_input_continuous = (in_channels is not None) and (patch_size is None)
117
+ self.is_input_vectorized = num_vector_embeds is not None
118
+ self.is_input_patches = in_channels is not None and patch_size is not None
119
+
120
+ if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
121
+ deprecation_message = (
122
+ f"The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"
123
+ " incorrectly set to `'layer_norm'`.Make sure to set `norm_type` to `'ada_norm'` in the config."
124
+ " Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect"
125
+ " results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it"
126
+ " would be very nice if you could open a Pull request for the `transformer/config.json` file"
127
+ )
128
+ deprecate("norm_type!=num_embeds_ada_norm", "1.0.0", deprecation_message, standard_warn=False)
129
+ norm_type = "ada_norm"
130
+
131
+ if self.is_input_continuous and self.is_input_vectorized:
132
+ raise ValueError(
133
+ f"Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"
134
+ " sure that either `in_channels` or `num_vector_embeds` is None."
135
+ )
136
+ elif self.is_input_vectorized and self.is_input_patches:
137
+ raise ValueError(
138
+ f"Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make"
139
+ " sure that either `num_vector_embeds` or `num_patches` is None."
140
+ )
141
+ elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
142
+ raise ValueError(
143
+ f"Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:"
144
+ f" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None."
145
+ )
146
+
147
+ # 2. Define input layers
148
+ if self.is_input_continuous:
149
+ self.in_channels = in_channels
150
+
151
+ self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)
152
+ if use_linear_projection:
153
+ self.proj_in = linear_cls(in_channels, inner_dim)
154
+ else:
155
+ self.proj_in = conv_cls(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
156
+ elif self.is_input_vectorized:
157
+ assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
158
+ assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"
159
+
160
+ self.height = sample_size
161
+ self.width = sample_size
162
+ self.num_vector_embeds = num_vector_embeds
163
+ self.num_latent_pixels = self.height * self.width
164
+
165
+ self.latent_image_embedding = ImagePositionalEmbeddings(
166
+ num_embed=num_vector_embeds, embed_dim=inner_dim, height=self.height, width=self.width
167
+ )
168
+ elif self.is_input_patches:
169
+ assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"
170
+
171
+ self.height = sample_size
172
+ self.width = sample_size
173
+
174
+ self.patch_size = patch_size
175
+ interpolation_scale = self.config.sample_size // 64 # => 64 (= 512 pixart) has interpolation scale 1
176
+ interpolation_scale = max(interpolation_scale, 1)
177
+ self.pos_embed = PatchEmbed(
178
+ height=sample_size,
179
+ width=sample_size,
180
+ patch_size=patch_size,
181
+ in_channels=in_channels,
182
+ embed_dim=inner_dim,
183
+ interpolation_scale=interpolation_scale,
184
+ )
185
+
186
+ # 3. Define transformers blocks
187
+ self.transformer_blocks = nn.ModuleList(
188
+ [
189
+ BasicTransformerBlock(
190
+ inner_dim,
191
+ num_attention_heads,
192
+ attention_head_dim,
193
+ dropout=dropout,
194
+ cross_attention_dim=cross_attention_dim,
195
+ activation_fn=activation_fn,
196
+ num_embeds_ada_norm=num_embeds_ada_norm,
197
+ attention_bias=attention_bias,
198
+ only_cross_attention=only_cross_attention,
199
+ double_self_attention=double_self_attention,
200
+ upcast_attention=upcast_attention,
201
+ norm_type=norm_type,
202
+ norm_elementwise_affine=norm_elementwise_affine,
203
+ norm_eps=norm_eps,
204
+ attention_type=attention_type,
205
+ )
206
+ for d in range(num_layers)
207
+ ]
208
+ )
209
+
210
+ # 4. Define output layers
211
+ self.out_channels = in_channels if out_channels is None else out_channels
212
+ if self.is_input_continuous:
213
+ # TODO: should use out_channels for continuous projections
214
+ if use_linear_projection:
215
+ self.proj_out = linear_cls(inner_dim, in_channels)
216
+ else:
217
+ self.proj_out = conv_cls(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
218
+ elif self.is_input_vectorized:
219
+ self.norm_out = nn.LayerNorm(inner_dim)
220
+ self.out = nn.Linear(inner_dim, self.num_vector_embeds - 1)
221
+ elif self.is_input_patches and norm_type != "ada_norm_single":
222
+ self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6)
223
+ self.proj_out_1 = nn.Linear(inner_dim, 2 * inner_dim)
224
+ self.proj_out_2 = nn.Linear(inner_dim, patch_size * patch_size * self.out_channels)
225
+ elif self.is_input_patches and norm_type == "ada_norm_single":
226
+ self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6)
227
+ self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim**0.5)
228
+ self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * self.out_channels)
229
+
230
+ # 5. PixArt-Alpha blocks.
231
+ self.adaln_single = None
232
+ self.use_additional_conditions = False
233
+ if norm_type == "ada_norm_single":
234
+ self.use_additional_conditions = self.config.sample_size == 128
235
+ # TODO(Sayak, PVP) clean this, for now we use sample size to determine whether to use
236
+ # additional conditions until we find better name
237
+ self.adaln_single = AdaLayerNormSingle(inner_dim, use_additional_conditions=self.use_additional_conditions)
238
+
239
+ self.caption_projection = None
240
+ if caption_channels is not None:
241
+ self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim)
242
+
243
+ self.gradient_checkpointing = False
244
+
245
+ def _set_gradient_checkpointing(self, module, value=False):
246
+ if hasattr(module, "gradient_checkpointing"):
247
+ module.gradient_checkpointing = value
248
+
249
+ def forward(
250
+ self,
251
+ hidden_states: torch.Tensor,
252
+ encoder_hidden_states: Optional[torch.Tensor] = None,
253
+ timestep: Optional[torch.LongTensor] = None,
254
+ added_cond_kwargs: Dict[str, torch.Tensor] = None,
255
+ class_labels: Optional[torch.LongTensor] = None,
256
+ cross_attention_kwargs: Dict[str, Any] = None,
257
+ attention_mask: Optional[torch.Tensor] = None,
258
+ encoder_attention_mask: Optional[torch.Tensor] = None,
259
+ return_dict: bool = True,
260
+ ):
261
+ """
262
+ The [`Transformer2DModel`] forward method.
263
+
264
+ Args:
265
+ hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
266
+ Input `hidden_states`.
267
+ encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
268
+ Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
269
+ self-attention.
270
+ timestep ( `torch.LongTensor`, *optional*):
271
+ Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
272
+ class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
273
+ Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
274
+ `AdaLayerZeroNorm`.
275
+ cross_attention_kwargs ( `Dict[str, Any]`, *optional*):
276
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
277
+ `self.processor` in
278
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
279
+ attention_mask ( `torch.Tensor`, *optional*):
280
+ An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
281
+ is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
282
+ negative values to the attention scores corresponding to "discard" tokens.
283
+ encoder_attention_mask ( `torch.Tensor`, *optional*):
284
+ Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:
285
+
286
+ * Mask `(batch, sequence_length)` True = keep, False = discard.
287
+ * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.
288
+
289
+ If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
290
+ above. This bias will be added to the cross-attention scores.
291
+ return_dict (`bool`, *optional*, defaults to `True`):
292
+ Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
293
+ tuple.
294
+
295
+ Returns:
296
+ If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
297
+ `tuple` where the first element is the sample tensor.
298
+ """
299
+ # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
300
+ # we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
301
+ # we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
302
+ # expects mask of shape:
303
+ # [batch, key_tokens]
304
+ # adds singleton query_tokens dimension:
305
+ # [batch, 1, key_tokens]
306
+ # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
307
+ # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)
308
+ # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
309
+
310
+ if attention_mask is not None and attention_mask.ndim == 2:
311
+ # assume that mask is expressed as:
312
+ # (1 = keep, 0 = discard)
313
+ # convert mask into a bias that can be added to attention scores:
314
+ # (keep = +0, discard = -10000.0)
315
+ attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0
316
+ attention_mask = attention_mask.unsqueeze(1)
317
+
318
+ # convert encoder_attention_mask to a bias the same way we do for attention_mask
319
+ if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
320
+ encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
321
+ encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
322
+
323
+ # Retrieve lora scale.
324
+ lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
325
+
326
+ # 1. Input
327
+ if self.is_input_continuous:
328
+ batch, _, height, width = hidden_states.shape
329
+ residual = hidden_states
330
+
331
+ hidden_states = self.norm(hidden_states)
332
+ if not self.use_linear_projection:
333
+ hidden_states = (
334
+ self.proj_in(hidden_states, scale=lora_scale)
335
+ if not USE_PEFT_BACKEND
336
+ else self.proj_in(hidden_states)
337
+ )
338
+ inner_dim = hidden_states.shape[1]
339
+ hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim)
340
+ else:
341
+ inner_dim = hidden_states.shape[1]
342
+ hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim)
343
+ hidden_states = (
344
+ self.proj_in(hidden_states, scale=lora_scale)
345
+ if not USE_PEFT_BACKEND
346
+ else self.proj_in(hidden_states)
347
+ )
348
+
349
+ elif self.is_input_vectorized:
350
+ hidden_states = self.latent_image_embedding(hidden_states)
351
+ elif self.is_input_patches:
352
+ height, width = hidden_states.shape[-2] // self.patch_size, hidden_states.shape[-1] // self.patch_size
353
+ hidden_states = self.pos_embed(hidden_states)
354
+
355
+ if self.adaln_single is not None:
356
+ if self.use_additional_conditions and added_cond_kwargs is None:
357
+ raise ValueError(
358
+ "`added_cond_kwargs` cannot be None when using additional conditions for `adaln_single`."
359
+ )
360
+ batch_size = hidden_states.shape[0]
361
+ timestep, embedded_timestep = self.adaln_single(
362
+ timestep, added_cond_kwargs, batch_size=batch_size, hidden_dtype=hidden_states.dtype
363
+ )
364
+
365
+ # 2. Blocks
366
+ if self.caption_projection is not None:
367
+ batch_size = hidden_states.shape[0]
368
+ encoder_hidden_states = self.caption_projection(encoder_hidden_states)
369
+ encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1])
370
+
371
+ for block in self.transformer_blocks:
372
+ if self.training and self.gradient_checkpointing:
373
+
374
+ def create_custom_forward(module, return_dict=None):
375
+ def custom_forward(*inputs):
376
+ if return_dict is not None:
377
+ return module(*inputs, return_dict=return_dict)
378
+ else:
379
+ return module(*inputs)
380
+
381
+ return custom_forward
382
+
383
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
384
+ hidden_states = torch.utils.checkpoint.checkpoint(
385
+ create_custom_forward(block),
386
+ hidden_states,
387
+ attention_mask,
388
+ encoder_hidden_states,
389
+ encoder_attention_mask,
390
+ timestep,
391
+ cross_attention_kwargs,
392
+ class_labels,
393
+ **ckpt_kwargs,
394
+ )
395
+ else:
396
+ hidden_states = block(
397
+ hidden_states,
398
+ attention_mask=attention_mask,
399
+ encoder_hidden_states=encoder_hidden_states,
400
+ encoder_attention_mask=encoder_attention_mask,
401
+ timestep=timestep,
402
+ cross_attention_kwargs=cross_attention_kwargs,
403
+ class_labels=class_labels,
404
+ )
405
+
406
+ # 3. Output
407
+ if self.is_input_continuous:
408
+ if not self.use_linear_projection:
409
+ hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()
410
+ hidden_states = (
411
+ self.proj_out(hidden_states, scale=lora_scale)
412
+ if not USE_PEFT_BACKEND
413
+ else self.proj_out(hidden_states)
414
+ )
415
+ else:
416
+ hidden_states = (
417
+ self.proj_out(hidden_states, scale=lora_scale)
418
+ if not USE_PEFT_BACKEND
419
+ else self.proj_out(hidden_states)
420
+ )
421
+ hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()
422
+
423
+ output = hidden_states + residual
424
+ elif self.is_input_vectorized:
425
+ hidden_states = self.norm_out(hidden_states)
426
+ logits = self.out(hidden_states)
427
+ # (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
428
+ logits = logits.permute(0, 2, 1)
429
+
430
+ # log(p(x_0))
431
+ output = F.log_softmax(logits.double(), dim=1).float()
432
+
433
+ if self.is_input_patches:
434
+ if self.config.norm_type != "ada_norm_single":
435
+ conditioning = self.transformer_blocks[0].norm1.emb(
436
+ timestep, class_labels, hidden_dtype=hidden_states.dtype
437
+ )
438
+ shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)
439
+ hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
440
+ hidden_states = self.proj_out_2(hidden_states)
441
+ elif self.config.norm_type == "ada_norm_single":
442
+ shift, scale = (self.scale_shift_table[None] + embedded_timestep[:, None]).chunk(2, dim=1)
443
+ hidden_states = self.norm_out(hidden_states)
444
+ # Modulation
445
+ hidden_states = hidden_states * (1 + scale) + shift
446
+ hidden_states = self.proj_out(hidden_states)
447
+ hidden_states = hidden_states.squeeze(1)
448
+
449
+ # unpatchify
450
+ if self.adaln_single is None:
451
+ height = width = int(hidden_states.shape[1] ** 0.5)
452
+ hidden_states = hidden_states.reshape(
453
+ shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)
454
+ )
455
+ hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
456
+ output = hidden_states.reshape(
457
+ shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)
458
+ )
459
+
460
+ if not return_dict:
461
+ return (output,)
462
+
463
+ return Transformer2DModelOutput(sample=output)
GeoWizard/geowizard/models/unet_2d_blocks.py ADDED
The diff for this file is too large to render. See raw diff
 
GeoWizard/geowizard/models/unet_2d_condition.py ADDED
@@ -0,0 +1,1222 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ # Some modifications are reimplemented in public environments by Xiao Fu and Mu Hu
16
+
17
+ from dataclasses import dataclass
18
+ from typing import Any, Dict, List, Optional, Tuple, Union
19
+
20
+ import torch
21
+ import torch.nn as nn
22
+ import torch.utils.checkpoint
23
+
24
+ from diffusers.configuration_utils import ConfigMixin, register_to_config
25
+ from diffusers.loaders import UNet2DConditionLoadersMixin
26
+ from diffusers.utils import USE_PEFT_BACKEND, BaseOutput, deprecate, logging, scale_lora_layers, unscale_lora_layers
27
+ from diffusers.models.activations import get_activation
28
+ from diffusers.models.attention_processor import (
29
+ ADDED_KV_ATTENTION_PROCESSORS,
30
+ CROSS_ATTENTION_PROCESSORS,
31
+ Attention,
32
+ AttentionProcessor,
33
+ AttnAddedKVProcessor,
34
+ AttnProcessor,
35
+ )
36
+ from diffusers.models.embeddings import (
37
+ GaussianFourierProjection,
38
+ ImageHintTimeEmbedding,
39
+ ImageProjection,
40
+ ImageTimeEmbedding,
41
+ #PositionNet,
42
+ TextImageProjection,
43
+ TextImageTimeEmbedding,
44
+ TextTimeEmbedding,
45
+ TimestepEmbedding,
46
+ Timesteps,
47
+ )
48
+
49
+ # add
50
+ from diffusers.models.modeling_utils import ModelMixin
51
+ import diffusers
52
+ if diffusers.__version__ >'0.25':
53
+ from diffusers.models.embeddings import GLIGENTextBoundingboxProjection as PositionNet
54
+ else:
55
+ from diffusers.models.embeddings import PositionNet
56
+
57
+ #from models.unet_2d_blocks import (
58
+ from GeoWizard.geowizard.models.unet_2d_blocks import (
59
+ UNetMidBlock2D,
60
+ UNetMidBlock2DCrossAttn,
61
+ UNetMidBlock2DSimpleCrossAttn,
62
+ get_down_block,
63
+ get_up_block,
64
+ )
65
+
66
+
67
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
68
+
69
+
70
+ @dataclass
71
+ class UNet2DConditionOutput(BaseOutput):
72
+ """
73
+ The output of [`UNet2DConditionModel`].
74
+
75
+ Args:
76
+ sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
77
+ The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model.
78
+ """
79
+
80
+ sample: torch.FloatTensor = None
81
+
82
+
83
+ class UNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin):
84
+ r"""
85
+ A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample
86
+ shaped output.
87
+
88
+ This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
89
+ for all models (such as downloading or saving).
90
+
91
+ Parameters:
92
+ sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
93
+ Height and width of input/output sample.
94
+ in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample.
95
+ out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
96
+ center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample.
97
+ flip_sin_to_cos (`bool`, *optional*, defaults to `False`):
98
+ Whether to flip the sin to cos in the time embedding.
99
+ freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding.
100
+ down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
101
+ The tuple of downsample blocks to use.
102
+ mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`):
103
+ Block type for middle of UNet, it can be one of `UNetMidBlock2DCrossAttn`, `UNetMidBlock2D`, or
104
+ `UNetMidBlock2DSimpleCrossAttn`. If `None`, the mid block layer is skipped.
105
+ up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`):
106
+ The tuple of upsample blocks to use.
107
+ only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`):
108
+ Whether to include self-attention in the basic transformer blocks, see
109
+ [`~models.attention.BasicTransformerBlock`].
110
+ block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
111
+ The tuple of output channels for each block.
112
+ layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
113
+ downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.
114
+ mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block.
115
+ dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
116
+ act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
117
+ norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.
118
+ If `None`, normalization and activation layers is skipped in post-processing.
119
+ norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.
120
+ cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
121
+ The dimension of the cross attention features.
122
+ transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1):
123
+ The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
124
+ [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],
125
+ [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
126
+ reverse_transformer_layers_per_block : (`Tuple[Tuple]`, *optional*, defaults to None):
127
+ The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`], in the upsampling
128
+ blocks of the U-Net. Only relevant if `transformer_layers_per_block` is of type `Tuple[Tuple]` and for
129
+ [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],
130
+ [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
131
+ encoder_hid_dim (`int`, *optional*, defaults to None):
132
+ If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim`
133
+ dimension to `cross_attention_dim`.
134
+ encoder_hid_dim_type (`str`, *optional*, defaults to `None`):
135
+ If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text
136
+ embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`.
137
+ attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.
138
+ num_attention_heads (`int`, *optional*):
139
+ The number of attention heads. If not defined, defaults to `attention_head_dim`
140
+ resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config
141
+ for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`.
142
+ class_embed_type (`str`, *optional*, defaults to `None`):
143
+ The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`,
144
+ `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`.
145
+ addition_embed_type (`str`, *optional*, defaults to `None`):
146
+ Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or
147
+ "text". "text" will use the `TextTimeEmbedding` layer.
148
+ addition_time_embed_dim: (`int`, *optional*, defaults to `None`):
149
+ Dimension for the timestep embeddings.
150
+ num_class_embeds (`int`, *optional*, defaults to `None`):
151
+ Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing
152
+ class conditioning with `class_embed_type` equal to `None`.
153
+ time_embedding_type (`str`, *optional*, defaults to `positional`):
154
+ The type of position embedding to use for timesteps. Choose from `positional` or `fourier`.
155
+ time_embedding_dim (`int`, *optional*, defaults to `None`):
156
+ An optional override for the dimension of the projected time embedding.
157
+ time_embedding_act_fn (`str`, *optional*, defaults to `None`):
158
+ Optional activation function to use only once on the time embeddings before they are passed to the rest of
159
+ the UNet. Choose from `silu`, `mish`, `gelu`, and `swish`.
160
+ timestep_post_act (`str`, *optional*, defaults to `None`):
161
+ The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`.
162
+ time_cond_proj_dim (`int`, *optional*, defaults to `None`):
163
+ The dimension of `cond_proj` layer in the timestep embedding.
164
+ conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer. conv_out_kernel (`int`,
165
+ *optional*, default to `3`): The kernel size of `conv_out` layer. projection_class_embeddings_input_dim (`int`,
166
+ *optional*): The dimension of the `class_labels` input when
167
+ `class_embed_type="projection"`. Required when `class_embed_type="projection"`.
168
+ class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time
169
+ embeddings with the class embeddings.
170
+ mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`):
171
+ Whether to use cross attention with the mid block when using the `UNetMidBlock2DSimpleCrossAttn`. If
172
+ `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is `None`, the
173
+ `only_cross_attention` value is used as the value for `mid_block_only_cross_attention`. Default to `False`
174
+ otherwise.
175
+ """
176
+
177
+ _supports_gradient_checkpointing = True
178
+
179
+ @register_to_config
180
+ def __init__(
181
+ self,
182
+ sample_size: Optional[int] = None,
183
+ in_channels: int = 4,
184
+ out_channels: int = 4,
185
+ center_input_sample: bool = False,
186
+ flip_sin_to_cos: bool = True,
187
+ freq_shift: int = 0,
188
+ down_block_types: Tuple[str] = (
189
+ "CrossAttnDownBlock2D",
190
+ "CrossAttnDownBlock2D",
191
+ "CrossAttnDownBlock2D",
192
+ "DownBlock2D",
193
+ ),
194
+ mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
195
+ up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
196
+ only_cross_attention: Union[bool, Tuple[bool]] = False,
197
+ block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
198
+ layers_per_block: Union[int, Tuple[int]] = 2,
199
+ downsample_padding: int = 1,
200
+ mid_block_scale_factor: float = 1,
201
+ dropout: float = 0.0,
202
+ act_fn: str = "silu",
203
+ norm_num_groups: Optional[int] = 32,
204
+ norm_eps: float = 1e-5,
205
+ cross_attention_dim: Union[int, Tuple[int]] = 1280,
206
+ transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
207
+ reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None,
208
+ encoder_hid_dim: Optional[int] = None,
209
+ encoder_hid_dim_type: Optional[str] = None,
210
+ attention_head_dim: Union[int, Tuple[int]] = 8,
211
+ num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
212
+ dual_cross_attention: bool = False,
213
+ use_linear_projection: bool = False,
214
+ class_embed_type: Optional[str] = None,
215
+ addition_embed_type: Optional[str] = None,
216
+ addition_time_embed_dim: Optional[int] = None,
217
+ num_class_embeds: Optional[int] = None,
218
+ upcast_attention: bool = False,
219
+ resnet_time_scale_shift: str = "default",
220
+ resnet_skip_time_act: bool = False,
221
+ resnet_out_scale_factor: int = 1.0,
222
+ time_embedding_type: str = "positional",
223
+ time_embedding_dim: Optional[int] = None,
224
+ time_embedding_act_fn: Optional[str] = None,
225
+ timestep_post_act: Optional[str] = None,
226
+ time_cond_proj_dim: Optional[int] = None,
227
+ conv_in_kernel: int = 3,
228
+ conv_out_kernel: int = 3,
229
+ projection_class_embeddings_input_dim: Optional[int] = None,
230
+ attention_type: str = "default",
231
+ class_embeddings_concat: bool = False,
232
+ mid_block_only_cross_attention: Optional[bool] = None,
233
+ cross_attention_norm: Optional[str] = None,
234
+ addition_embed_type_num_heads=64,
235
+ ):
236
+ super().__init__()
237
+
238
+ self.sample_size = sample_size
239
+
240
+ if num_attention_heads is not None:
241
+ raise ValueError(
242
+ "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19."
243
+ )
244
+
245
+ # If `num_attention_heads` is not defined (which is the case for most models)
246
+ # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
247
+ # The reason for this behavior is to correct for incorrectly named variables that were introduced
248
+ # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
249
+ # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
250
+ # which is why we correct for the naming here.
251
+ num_attention_heads = num_attention_heads or attention_head_dim
252
+
253
+ # Check inputs
254
+ if len(down_block_types) != len(up_block_types):
255
+ raise ValueError(
256
+ f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
257
+ )
258
+
259
+ if len(block_out_channels) != len(down_block_types):
260
+ raise ValueError(
261
+ f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
262
+ )
263
+
264
+ if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):
265
+ raise ValueError(
266
+ f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."
267
+ )
268
+
269
+ if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
270
+ raise ValueError(
271
+ f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
272
+ )
273
+
274
+ if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):
275
+ raise ValueError(
276
+ f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}."
277
+ )
278
+
279
+ if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types):
280
+ raise ValueError(
281
+ f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}."
282
+ )
283
+
284
+ if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types):
285
+ raise ValueError(
286
+ f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}."
287
+ )
288
+ if isinstance(transformer_layers_per_block, list) and reverse_transformer_layers_per_block is None:
289
+ for layer_number_per_block in transformer_layers_per_block:
290
+ if isinstance(layer_number_per_block, list):
291
+ raise ValueError("Must provide 'reverse_transformer_layers_per_block` if using asymmetrical UNet.")
292
+
293
+ # input
294
+ conv_in_padding = (conv_in_kernel - 1) // 2
295
+ self.conv_in = nn.Conv2d(
296
+ in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding
297
+ )
298
+
299
+ # time
300
+ if time_embedding_type == "fourier":
301
+ time_embed_dim = time_embedding_dim or block_out_channels[0] * 2
302
+ if time_embed_dim % 2 != 0:
303
+ raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.")
304
+ self.time_proj = GaussianFourierProjection(
305
+ time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos
306
+ )
307
+ timestep_input_dim = time_embed_dim
308
+ elif time_embedding_type == "positional":
309
+ time_embed_dim = time_embedding_dim or block_out_channels[0] * 4
310
+
311
+ self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
312
+ timestep_input_dim = block_out_channels[0]
313
+ else:
314
+ raise ValueError(
315
+ f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`."
316
+ )
317
+
318
+ self.time_embedding = TimestepEmbedding(
319
+ timestep_input_dim,
320
+ time_embed_dim,
321
+ act_fn=act_fn,
322
+ post_act_fn=timestep_post_act,
323
+ cond_proj_dim=time_cond_proj_dim,
324
+ )
325
+
326
+ if encoder_hid_dim_type is None and encoder_hid_dim is not None:
327
+ encoder_hid_dim_type = "text_proj"
328
+ self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type)
329
+ logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.")
330
+
331
+ if encoder_hid_dim is None and encoder_hid_dim_type is not None:
332
+ raise ValueError(
333
+ f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}."
334
+ )
335
+
336
+ if encoder_hid_dim_type == "text_proj":
337
+ self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim)
338
+ elif encoder_hid_dim_type == "text_image_proj":
339
+ # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much
340
+ # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
341
+ # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)`
342
+ self.encoder_hid_proj = TextImageProjection(
343
+ text_embed_dim=encoder_hid_dim,
344
+ image_embed_dim=cross_attention_dim,
345
+ cross_attention_dim=cross_attention_dim,
346
+ )
347
+ elif encoder_hid_dim_type == "image_proj":
348
+ # Kandinsky 2.2
349
+ self.encoder_hid_proj = ImageProjection(
350
+ image_embed_dim=encoder_hid_dim,
351
+ cross_attention_dim=cross_attention_dim,
352
+ )
353
+ elif encoder_hid_dim_type is not None:
354
+ raise ValueError(
355
+ f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'."
356
+ )
357
+ else:
358
+ self.encoder_hid_proj = None
359
+
360
+ # class embedding
361
+ if class_embed_type is None and num_class_embeds is not None:
362
+ self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
363
+ elif class_embed_type == "timestep":
364
+ self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim, act_fn=act_fn)
365
+ elif class_embed_type == "identity":
366
+ self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
367
+ elif class_embed_type == "projection":
368
+ if projection_class_embeddings_input_dim is None:
369
+ raise ValueError(
370
+ "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"
371
+ )
372
+ # The projection `class_embed_type` is the same as the timestep `class_embed_type` except
373
+ # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings
374
+ # 2. it projects from an arbitrary input dimension.
375
+ #
376
+ # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.
377
+ # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.
378
+ # As a result, `TimestepEmbedding` can be passed arbitrary vectors.
379
+ self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
380
+ elif class_embed_type == "simple_projection":
381
+ if projection_class_embeddings_input_dim is None:
382
+ raise ValueError(
383
+ "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set"
384
+ )
385
+ self.class_embedding = nn.Linear(projection_class_embeddings_input_dim, time_embed_dim)
386
+ else:
387
+ self.class_embedding = None
388
+
389
+ if addition_embed_type == "text":
390
+ if encoder_hid_dim is not None:
391
+ text_time_embedding_from_dim = encoder_hid_dim
392
+ else:
393
+ text_time_embedding_from_dim = cross_attention_dim
394
+
395
+ self.add_embedding = TextTimeEmbedding(
396
+ text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads
397
+ )
398
+ elif addition_embed_type == "text_image":
399
+ # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much
400
+ # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
401
+ # case when `addition_embed_type == "text_image"` (Kadinsky 2.1)`
402
+ self.add_embedding = TextImageTimeEmbedding(
403
+ text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim
404
+ )
405
+ elif addition_embed_type == "text_time":
406
+ self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift)
407
+ self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
408
+ elif addition_embed_type == "image":
409
+ # Kandinsky 2.2
410
+ self.add_embedding = ImageTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim)
411
+ elif addition_embed_type == "image_hint":
412
+ # Kandinsky 2.2 ControlNet
413
+ self.add_embedding = ImageHintTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim)
414
+ elif addition_embed_type is not None:
415
+ raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.")
416
+
417
+ if time_embedding_act_fn is None:
418
+ self.time_embed_act = None
419
+ else:
420
+ self.time_embed_act = get_activation(time_embedding_act_fn)
421
+
422
+ self.down_blocks = nn.ModuleList([])
423
+ self.up_blocks = nn.ModuleList([])
424
+
425
+ if isinstance(only_cross_attention, bool):
426
+ if mid_block_only_cross_attention is None:
427
+ mid_block_only_cross_attention = only_cross_attention
428
+
429
+ only_cross_attention = [only_cross_attention] * len(down_block_types)
430
+
431
+ if mid_block_only_cross_attention is None:
432
+ mid_block_only_cross_attention = False
433
+
434
+ if isinstance(num_attention_heads, int):
435
+ num_attention_heads = (num_attention_heads,) * len(down_block_types)
436
+
437
+ if isinstance(attention_head_dim, int):
438
+ attention_head_dim = (attention_head_dim,) * len(down_block_types)
439
+
440
+ if isinstance(cross_attention_dim, int):
441
+ cross_attention_dim = (cross_attention_dim,) * len(down_block_types)
442
+
443
+ if isinstance(layers_per_block, int):
444
+ layers_per_block = [layers_per_block] * len(down_block_types)
445
+
446
+ if isinstance(transformer_layers_per_block, int):
447
+ transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
448
+
449
+ if class_embeddings_concat:
450
+ # The time embeddings are concatenated with the class embeddings. The dimension of the
451
+ # time embeddings passed to the down, middle, and up blocks is twice the dimension of the
452
+ # regular time embeddings
453
+ blocks_time_embed_dim = time_embed_dim * 2
454
+ else:
455
+ blocks_time_embed_dim = time_embed_dim
456
+
457
+ # down
458
+ output_channel = block_out_channels[0]
459
+ for i, down_block_type in enumerate(down_block_types):
460
+ input_channel = output_channel
461
+ output_channel = block_out_channels[i]
462
+ is_final_block = i == len(block_out_channels) - 1
463
+
464
+ down_block = get_down_block(
465
+ down_block_type,
466
+ num_layers=layers_per_block[i],
467
+ transformer_layers_per_block=transformer_layers_per_block[i],
468
+ in_channels=input_channel,
469
+ out_channels=output_channel,
470
+ temb_channels=blocks_time_embed_dim,
471
+ add_downsample=not is_final_block,
472
+ resnet_eps=norm_eps,
473
+ resnet_act_fn=act_fn,
474
+ resnet_groups=norm_num_groups,
475
+ cross_attention_dim=cross_attention_dim[i],
476
+ num_attention_heads=num_attention_heads[i],
477
+ downsample_padding=downsample_padding,
478
+ dual_cross_attention=dual_cross_attention,
479
+ use_linear_projection=use_linear_projection,
480
+ only_cross_attention=only_cross_attention[i],
481
+ upcast_attention=upcast_attention,
482
+ resnet_time_scale_shift=resnet_time_scale_shift,
483
+ attention_type=attention_type,
484
+ resnet_skip_time_act=resnet_skip_time_act,
485
+ resnet_out_scale_factor=resnet_out_scale_factor,
486
+ cross_attention_norm=cross_attention_norm,
487
+ attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
488
+ dropout=dropout,
489
+ )
490
+ self.down_blocks.append(down_block)
491
+
492
+ # mid
493
+ if mid_block_type == "UNetMidBlock2DCrossAttn":
494
+ self.mid_block = UNetMidBlock2DCrossAttn(
495
+ transformer_layers_per_block=transformer_layers_per_block[-1],
496
+ in_channels=block_out_channels[-1],
497
+ temb_channels=blocks_time_embed_dim,
498
+ dropout=dropout,
499
+ resnet_eps=norm_eps,
500
+ resnet_act_fn=act_fn,
501
+ output_scale_factor=mid_block_scale_factor,
502
+ resnet_time_scale_shift=resnet_time_scale_shift,
503
+ cross_attention_dim=cross_attention_dim[-1],
504
+ num_attention_heads=num_attention_heads[-1],
505
+ resnet_groups=norm_num_groups,
506
+ dual_cross_attention=dual_cross_attention,
507
+ use_linear_projection=use_linear_projection,
508
+ upcast_attention=upcast_attention,
509
+ attention_type=attention_type,
510
+ )
511
+ elif mid_block_type == "UNetMidBlock2DSimpleCrossAttn":
512
+ self.mid_block = UNetMidBlock2DSimpleCrossAttn(
513
+ in_channels=block_out_channels[-1],
514
+ temb_channels=blocks_time_embed_dim,
515
+ dropout=dropout,
516
+ resnet_eps=norm_eps,
517
+ resnet_act_fn=act_fn,
518
+ output_scale_factor=mid_block_scale_factor,
519
+ cross_attention_dim=cross_attention_dim[-1],
520
+ attention_head_dim=attention_head_dim[-1],
521
+ resnet_groups=norm_num_groups,
522
+ resnet_time_scale_shift=resnet_time_scale_shift,
523
+ skip_time_act=resnet_skip_time_act,
524
+ only_cross_attention=mid_block_only_cross_attention,
525
+ cross_attention_norm=cross_attention_norm,
526
+ )
527
+ elif mid_block_type == "UNetMidBlock2D":
528
+ self.mid_block = UNetMidBlock2D(
529
+ in_channels=block_out_channels[-1],
530
+ temb_channels=blocks_time_embed_dim,
531
+ dropout=dropout,
532
+ num_layers=0,
533
+ resnet_eps=norm_eps,
534
+ resnet_act_fn=act_fn,
535
+ output_scale_factor=mid_block_scale_factor,
536
+ resnet_groups=norm_num_groups,
537
+ resnet_time_scale_shift=resnet_time_scale_shift,
538
+ add_attention=False,
539
+ )
540
+ elif mid_block_type is None:
541
+ self.mid_block = None
542
+ else:
543
+ raise ValueError(f"unknown mid_block_type : {mid_block_type}")
544
+
545
+ # count how many layers upsample the images
546
+ self.num_upsamplers = 0
547
+
548
+ # up
549
+ reversed_block_out_channels = list(reversed(block_out_channels))
550
+ reversed_num_attention_heads = list(reversed(num_attention_heads))
551
+ reversed_layers_per_block = list(reversed(layers_per_block))
552
+ reversed_cross_attention_dim = list(reversed(cross_attention_dim))
553
+ reversed_transformer_layers_per_block = (
554
+ list(reversed(transformer_layers_per_block))
555
+ if reverse_transformer_layers_per_block is None
556
+ else reverse_transformer_layers_per_block
557
+ )
558
+ only_cross_attention = list(reversed(only_cross_attention))
559
+
560
+ output_channel = reversed_block_out_channels[0]
561
+ for i, up_block_type in enumerate(up_block_types):
562
+ is_final_block = i == len(block_out_channels) - 1
563
+
564
+ prev_output_channel = output_channel
565
+ output_channel = reversed_block_out_channels[i]
566
+ input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
567
+
568
+ # add upsample block for all BUT final layer
569
+ if not is_final_block:
570
+ add_upsample = True
571
+ self.num_upsamplers += 1
572
+ else:
573
+ add_upsample = False
574
+
575
+ up_block = get_up_block(
576
+ up_block_type,
577
+ num_layers=reversed_layers_per_block[i] + 1,
578
+ transformer_layers_per_block=reversed_transformer_layers_per_block[i],
579
+ in_channels=input_channel,
580
+ out_channels=output_channel,
581
+ prev_output_channel=prev_output_channel,
582
+ temb_channels=blocks_time_embed_dim,
583
+ add_upsample=add_upsample,
584
+ resnet_eps=norm_eps,
585
+ resnet_act_fn=act_fn,
586
+ resolution_idx=i,
587
+ resnet_groups=norm_num_groups,
588
+ cross_attention_dim=reversed_cross_attention_dim[i],
589
+ num_attention_heads=reversed_num_attention_heads[i],
590
+ dual_cross_attention=dual_cross_attention,
591
+ use_linear_projection=use_linear_projection,
592
+ only_cross_attention=only_cross_attention[i],
593
+ upcast_attention=upcast_attention,
594
+ resnet_time_scale_shift=resnet_time_scale_shift,
595
+ attention_type=attention_type,
596
+ resnet_skip_time_act=resnet_skip_time_act,
597
+ resnet_out_scale_factor=resnet_out_scale_factor,
598
+ cross_attention_norm=cross_attention_norm,
599
+ attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
600
+ dropout=dropout,
601
+ )
602
+ self.up_blocks.append(up_block)
603
+ prev_output_channel = output_channel
604
+
605
+ # out
606
+ if norm_num_groups is not None:
607
+ self.conv_norm_out = nn.GroupNorm(
608
+ num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps
609
+ )
610
+
611
+ self.conv_act = get_activation(act_fn)
612
+
613
+ else:
614
+ self.conv_norm_out = None
615
+ self.conv_act = None
616
+
617
+ conv_out_padding = (conv_out_kernel - 1) // 2
618
+ self.conv_out = nn.Conv2d(
619
+ block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding
620
+ )
621
+
622
+ if attention_type in ["gated", "gated-text-image"]:
623
+ positive_len = 768
624
+ if isinstance(cross_attention_dim, int):
625
+ positive_len = cross_attention_dim
626
+ elif isinstance(cross_attention_dim, tuple) or isinstance(cross_attention_dim, list):
627
+ positive_len = cross_attention_dim[0]
628
+
629
+ feature_type = "text-only" if attention_type == "gated" else "text-image"
630
+ self.position_net = PositionNet(
631
+ positive_len=positive_len, out_dim=cross_attention_dim, feature_type=feature_type
632
+ )
633
+
634
+ @property
635
+ def attn_processors(self) -> Dict[str, AttentionProcessor]:
636
+ r"""
637
+ Returns:
638
+ `dict` of attention processors: A dictionary containing all attention processors used in the model with
639
+ indexed by its weight name.
640
+ """
641
+ # set recursively
642
+ processors = {}
643
+
644
+ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
645
+ if hasattr(module, "get_processor"):
646
+ processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
647
+
648
+ for sub_name, child in module.named_children():
649
+ fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
650
+
651
+ return processors
652
+
653
+ for name, module in self.named_children():
654
+ fn_recursive_add_processors(name, module, processors)
655
+
656
+ return processors
657
+
658
+ def set_attn_processor(
659
+ self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False
660
+ ):
661
+ r"""
662
+ Sets the attention processor to use to compute attention.
663
+
664
+ Parameters:
665
+ processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
666
+ The instantiated processor class or a dictionary of processor classes that will be set as the processor
667
+ for **all** `Attention` layers.
668
+
669
+ If `processor` is a dict, the key needs to define the path to the corresponding cross attention
670
+ processor. This is strongly recommended when setting trainable attention processors.
671
+
672
+ """
673
+ count = len(self.attn_processors.keys())
674
+
675
+ if isinstance(processor, dict) and len(processor) != count:
676
+ raise ValueError(
677
+ f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
678
+ f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
679
+ )
680
+
681
+ def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
682
+ if hasattr(module, "set_processor"):
683
+ if not isinstance(processor, dict):
684
+ module.set_processor(processor, _remove_lora=_remove_lora)
685
+ else:
686
+ module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora)
687
+
688
+ for sub_name, child in module.named_children():
689
+ fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
690
+
691
+ for name, module in self.named_children():
692
+ fn_recursive_attn_processor(name, module, processor)
693
+
694
+ def set_default_attn_processor(self):
695
+ """
696
+ Disables custom attention processors and sets the default attention implementation.
697
+ """
698
+ if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
699
+ processor = AttnAddedKVProcessor()
700
+ elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
701
+ processor = AttnProcessor()
702
+ else:
703
+ raise ValueError(
704
+ f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
705
+ )
706
+
707
+ self.set_attn_processor(processor, _remove_lora=True)
708
+
709
+ def set_attention_slice(self, slice_size):
710
+ r"""
711
+ Enable sliced attention computation.
712
+
713
+ When this option is enabled, the attention module splits the input tensor in slices to compute attention in
714
+ several steps. This is useful for saving some memory in exchange for a small decrease in speed.
715
+
716
+ Args:
717
+ slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
718
+ When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If
719
+ `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is
720
+ provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
721
+ must be a multiple of `slice_size`.
722
+ """
723
+ sliceable_head_dims = []
724
+
725
+ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
726
+ if hasattr(module, "set_attention_slice"):
727
+ sliceable_head_dims.append(module.sliceable_head_dim)
728
+
729
+ for child in module.children():
730
+ fn_recursive_retrieve_sliceable_dims(child)
731
+
732
+ # retrieve number of attention layers
733
+ for module in self.children():
734
+ fn_recursive_retrieve_sliceable_dims(module)
735
+
736
+ num_sliceable_layers = len(sliceable_head_dims)
737
+
738
+ if slice_size == "auto":
739
+ # half the attention head size is usually a good trade-off between
740
+ # speed and memory
741
+ slice_size = [dim // 2 for dim in sliceable_head_dims]
742
+ elif slice_size == "max":
743
+ # make smallest slice possible
744
+ slice_size = num_sliceable_layers * [1]
745
+
746
+ slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size
747
+
748
+ if len(slice_size) != len(sliceable_head_dims):
749
+ raise ValueError(
750
+ f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
751
+ f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
752
+ )
753
+
754
+ for i in range(len(slice_size)):
755
+ size = slice_size[i]
756
+ dim = sliceable_head_dims[i]
757
+ if size is not None and size > dim:
758
+ raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
759
+
760
+ # Recursively walk through all the children.
761
+ # Any children which exposes the set_attention_slice method
762
+ # gets the message
763
+ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
764
+ if hasattr(module, "set_attention_slice"):
765
+ module.set_attention_slice(slice_size.pop())
766
+
767
+ for child in module.children():
768
+ fn_recursive_set_attention_slice(child, slice_size)
769
+
770
+ reversed_slice_size = list(reversed(slice_size))
771
+ for module in self.children():
772
+ fn_recursive_set_attention_slice(module, reversed_slice_size)
773
+
774
+ def _set_gradient_checkpointing(self, module, value=False):
775
+ if hasattr(module, "gradient_checkpointing"):
776
+ module.gradient_checkpointing = value
777
+
778
+ def enable_freeu(self, s1, s2, b1, b2):
779
+ r"""Enables the FreeU mechanism from https://arxiv.org/abs/2309.11497.
780
+
781
+ The suffixes after the scaling factors represent the stage blocks where they are being applied.
782
+
783
+ Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of values that
784
+ are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL.
785
+
786
+ Args:
787
+ s1 (`float`):
788
+ Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to
789
+ mitigate the "oversmoothing effect" in the enhanced denoising process.
790
+ s2 (`float`):
791
+ Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to
792
+ mitigate the "oversmoothing effect" in the enhanced denoising process.
793
+ b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features.
794
+ b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features.
795
+ """
796
+ for i, upsample_block in enumerate(self.up_blocks):
797
+ setattr(upsample_block, "s1", s1)
798
+ setattr(upsample_block, "s2", s2)
799
+ setattr(upsample_block, "b1", b1)
800
+ setattr(upsample_block, "b2", b2)
801
+
802
+ def disable_freeu(self):
803
+ """Disables the FreeU mechanism."""
804
+ freeu_keys = {"s1", "s2", "b1", "b2"}
805
+ for i, upsample_block in enumerate(self.up_blocks):
806
+ for k in freeu_keys:
807
+ if hasattr(upsample_block, k) or getattr(upsample_block, k, None) is not None:
808
+ setattr(upsample_block, k, None)
809
+
810
+ def fuse_qkv_projections(self):
811
+ """
812
+ Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query,
813
+ key, value) are fused. For cross-attention modules, key and value projection matrices are fused.
814
+
815
+ <Tip warning={true}>
816
+
817
+ This API is 🧪 experimental.
818
+
819
+ </Tip>
820
+ """
821
+ self.original_attn_processors = None
822
+
823
+ for _, attn_processor in self.attn_processors.items():
824
+ if "Added" in str(attn_processor.__class__.__name__):
825
+ raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
826
+
827
+ self.original_attn_processors = self.attn_processors
828
+
829
+ for module in self.modules():
830
+ if isinstance(module, Attention):
831
+ module.fuse_projections(fuse=True)
832
+
833
+ def unfuse_qkv_projections(self):
834
+ """Disables the fused QKV projection if enabled.
835
+
836
+ <Tip warning={true}>
837
+
838
+ This API is 🧪 experimental.
839
+
840
+ </Tip>
841
+
842
+ """
843
+ if self.original_attn_processors is not None:
844
+ self.set_attn_processor(self.original_attn_processors)
845
+
846
+ def forward(
847
+ self,
848
+ sample: torch.FloatTensor,
849
+ timestep: Union[torch.Tensor, float, int],
850
+ encoder_hidden_states: torch.Tensor,
851
+ class_labels: Optional[torch.Tensor] = None,
852
+ timestep_cond: Optional[torch.Tensor] = None,
853
+ attention_mask: Optional[torch.Tensor] = None,
854
+ cross_attention_kwargs: Optional[Dict[str, Any]] = None,
855
+ added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
856
+ down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
857
+ mid_block_additional_residual: Optional[torch.Tensor] = None,
858
+ down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
859
+ encoder_attention_mask: Optional[torch.Tensor] = None,
860
+ return_dict: bool = True,
861
+ ) -> Union[UNet2DConditionOutput, Tuple]:
862
+ r"""
863
+ The [`UNet2DConditionModel`] forward method.
864
+
865
+ Args:
866
+ sample (`torch.FloatTensor`):
867
+ The noisy input tensor with the following shape `(batch, channel, height, width)`.
868
+ timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input.
869
+ encoder_hidden_states (`torch.FloatTensor`):
870
+ The encoder hidden states with shape `(batch, sequence_length, feature_dim)`.
871
+ class_labels (`torch.Tensor`, *optional*, defaults to `None`):
872
+ Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
873
+ timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`):
874
+ Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed
875
+ through the `self.time_embedding` layer to obtain the timestep embeddings.
876
+ attention_mask (`torch.Tensor`, *optional*, defaults to `None`):
877
+ An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
878
+ is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
879
+ negative values to the attention scores corresponding to "discard" tokens.
880
+ cross_attention_kwargs (`dict`, *optional*):
881
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
882
+ `self.processor` in
883
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
884
+ added_cond_kwargs: (`dict`, *optional*):
885
+ A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that
886
+ are passed along to the UNet blocks.
887
+ down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*):
888
+ A tuple of tensors that if specified are added to the residuals of down unet blocks.
889
+ mid_block_additional_residual: (`torch.Tensor`, *optional*):
890
+ A tensor that if specified is added to the residual of the middle unet block.
891
+ encoder_attention_mask (`torch.Tensor`):
892
+ A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If
893
+ `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias,
894
+ which adds large negative values to the attention scores corresponding to "discard" tokens.
895
+ return_dict (`bool`, *optional*, defaults to `True`):
896
+ Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
897
+ tuple.
898
+ cross_attention_kwargs (`dict`, *optional*):
899
+ A kwargs dictionary that if specified is passed along to the [`AttnProcessor`].
900
+ added_cond_kwargs: (`dict`, *optional*):
901
+ A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that
902
+ are passed along to the UNet blocks.
903
+ down_block_additional_residuals (`tuple` of `torch.Tensor`, *optional*):
904
+ additional residuals to be added to UNet long skip connections from down blocks to up blocks for
905
+ example from ControlNet side model(s)
906
+ mid_block_additional_residual (`torch.Tensor`, *optional*):
907
+ additional residual to be added to UNet mid block output, for example from ControlNet side model
908
+ down_intrablock_additional_residuals (`tuple` of `torch.Tensor`, *optional*):
909
+ additional residuals to be added within UNet down blocks, for example from T2I-Adapter side model(s)
910
+
911
+ Returns:
912
+ [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
913
+ If `return_dict` is True, an [`~models.unet_2d_condition.UNet2DConditionOutput`] is returned, otherwise
914
+ a `tuple` is returned where the first element is the sample tensor.
915
+ """
916
+
917
+ # By default samples have to be AT least a multiple of the overall upsampling factor.
918
+ # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).
919
+ # However, the upsampling interpolation output size can be forced to fit any upsampling size
920
+ # on the fly if necessary.
921
+ default_overall_up_factor = 2**self.num_upsamplers
922
+
923
+ # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
924
+ forward_upsample_size = False
925
+ upsample_size = None
926
+
927
+ for dim in sample.shape[-2:]:
928
+ if dim % default_overall_up_factor != 0:
929
+ # Forward upsample size to force interpolation output size.
930
+ forward_upsample_size = True
931
+ break
932
+
933
+ # ensure attention_mask is a bias, and give it a singleton query_tokens dimension
934
+ # expects mask of shape:
935
+ # [batch, key_tokens]
936
+ # adds singleton query_tokens dimension:
937
+ # [batch, 1, key_tokens]
938
+ # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
939
+ # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)
940
+ # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
941
+ if attention_mask is not None:
942
+ # assume that mask is expressed as:
943
+ # (1 = keep, 0 = discard)
944
+ # convert mask into a bias that can be added to attention scores:
945
+ # (keep = +0, discard = -10000.0)
946
+ attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
947
+ attention_mask = attention_mask.unsqueeze(1)
948
+
949
+ # convert encoder_attention_mask to a bias the same way we do for attention_mask
950
+ if encoder_attention_mask is not None:
951
+ encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
952
+ encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
953
+
954
+ # 0. center input if necessary
955
+ if self.config.center_input_sample:
956
+ sample = 2 * sample - 1.0
957
+
958
+ # 1. time
959
+ timesteps = timestep
960
+ if not torch.is_tensor(timesteps):
961
+ # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
962
+ # This would be a good case for the `match` statement (Python 3.10+)
963
+ is_mps = sample.device.type == "mps"
964
+ if isinstance(timestep, float):
965
+ dtype = torch.float32 if is_mps else torch.float64
966
+ else:
967
+ dtype = torch.int32 if is_mps else torch.int64
968
+ timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
969
+ elif len(timesteps.shape) == 0:
970
+ timesteps = timesteps[None].to(sample.device)
971
+
972
+ # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
973
+ timesteps = timesteps.expand(sample.shape[0])
974
+
975
+ t_emb = self.time_proj(timesteps)
976
+
977
+ # `Timesteps` does not contain any weights and will always return f32 tensors
978
+ # but time_embedding might actually be running in fp16. so we need to cast here.
979
+ # there might be better ways to encapsulate this.
980
+ t_emb = t_emb.to(dtype=sample.dtype)
981
+
982
+ emb = self.time_embedding(t_emb, timestep_cond)
983
+ aug_emb = None
984
+
985
+ if self.class_embedding is not None:
986
+ if class_labels is None:
987
+ raise ValueError("class_labels should be provided when num_class_embeds > 0")
988
+
989
+ if self.config.class_embed_type == "timestep":
990
+ class_labels = self.time_proj(class_labels)
991
+
992
+ # `Timesteps` does not contain any weights and will always return f32 tensors
993
+ # there might be better ways to encapsulate this.
994
+ class_labels = class_labels.to(dtype=sample.dtype)
995
+
996
+ class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype)
997
+
998
+ if self.config.class_embeddings_concat:
999
+ emb = torch.cat([emb, class_emb], dim=-1)
1000
+ else:
1001
+ emb = emb + class_emb
1002
+
1003
+ if self.config.addition_embed_type == "text":
1004
+ aug_emb = self.add_embedding(encoder_hidden_states)
1005
+ elif self.config.addition_embed_type == "text_image":
1006
+ # Kandinsky 2.1 - style
1007
+ if "image_embeds" not in added_cond_kwargs:
1008
+ raise ValueError(
1009
+ f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
1010
+ )
1011
+
1012
+ image_embs = added_cond_kwargs.get("image_embeds")
1013
+ text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states)
1014
+ aug_emb = self.add_embedding(text_embs, image_embs)
1015
+ elif self.config.addition_embed_type == "text_time":
1016
+ # SDXL - style
1017
+ if "text_embeds" not in added_cond_kwargs:
1018
+ raise ValueError(
1019
+ f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
1020
+ )
1021
+ text_embeds = added_cond_kwargs.get("text_embeds")
1022
+ if "time_ids" not in added_cond_kwargs:
1023
+ raise ValueError(
1024
+ f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
1025
+ )
1026
+ time_ids = added_cond_kwargs.get("time_ids")
1027
+ time_embeds = self.add_time_proj(time_ids.flatten())
1028
+ time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
1029
+ add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
1030
+ add_embeds = add_embeds.to(emb.dtype)
1031
+ aug_emb = self.add_embedding(add_embeds)
1032
+ elif self.config.addition_embed_type == "image":
1033
+ # Kandinsky 2.2 - style
1034
+ if "image_embeds" not in added_cond_kwargs:
1035
+ raise ValueError(
1036
+ f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
1037
+ )
1038
+ image_embs = added_cond_kwargs.get("image_embeds")
1039
+ aug_emb = self.add_embedding(image_embs)
1040
+ elif self.config.addition_embed_type == "image_hint":
1041
+ # Kandinsky 2.2 - style
1042
+ if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs:
1043
+ raise ValueError(
1044
+ f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`"
1045
+ )
1046
+ image_embs = added_cond_kwargs.get("image_embeds")
1047
+ hint = added_cond_kwargs.get("hint")
1048
+ aug_emb, hint = self.add_embedding(image_embs, hint)
1049
+ sample = torch.cat([sample, hint], dim=1)
1050
+
1051
+ emb = emb + aug_emb if aug_emb is not None else emb
1052
+
1053
+ if self.time_embed_act is not None:
1054
+ emb = self.time_embed_act(emb)
1055
+
1056
+ if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj":
1057
+ encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states)
1058
+ elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj":
1059
+ # Kadinsky 2.1 - style
1060
+ if "image_embeds" not in added_cond_kwargs:
1061
+ raise ValueError(
1062
+ f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`"
1063
+ )
1064
+
1065
+ image_embeds = added_cond_kwargs.get("image_embeds")
1066
+ encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds)
1067
+ elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj":
1068
+ # Kandinsky 2.2 - style
1069
+ if "image_embeds" not in added_cond_kwargs:
1070
+ raise ValueError(
1071
+ f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`"
1072
+ )
1073
+ image_embeds = added_cond_kwargs.get("image_embeds")
1074
+ encoder_hidden_states = self.encoder_hid_proj(image_embeds)
1075
+ elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "ip_image_proj":
1076
+ if "image_embeds" not in added_cond_kwargs:
1077
+ raise ValueError(
1078
+ f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'ip_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`"
1079
+ )
1080
+ image_embeds = added_cond_kwargs.get("image_embeds")
1081
+ image_embeds = self.encoder_hid_proj(image_embeds).to(encoder_hidden_states.dtype)
1082
+ encoder_hidden_states = torch.cat([encoder_hidden_states, image_embeds], dim=1)
1083
+
1084
+ # 2. pre-process
1085
+ sample = self.conv_in(sample)
1086
+
1087
+ # 2.5 GLIGEN position net
1088
+ if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None:
1089
+ cross_attention_kwargs = cross_attention_kwargs.copy()
1090
+ gligen_args = cross_attention_kwargs.pop("gligen")
1091
+ cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)}
1092
+
1093
+ # 3. down
1094
+ lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
1095
+ if USE_PEFT_BACKEND:
1096
+ # weight the lora layers by setting `lora_scale` for each PEFT layer
1097
+ scale_lora_layers(self, lora_scale)
1098
+
1099
+ is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
1100
+ # using new arg down_intrablock_additional_residuals for T2I-Adapters, to distinguish from controlnets
1101
+ is_adapter = down_intrablock_additional_residuals is not None
1102
+ # maintain backward compatibility for legacy usage, where
1103
+ # T2I-Adapter and ControlNet both use down_block_additional_residuals arg
1104
+ # but can only use one or the other
1105
+ if not is_adapter and mid_block_additional_residual is None and down_block_additional_residuals is not None:
1106
+ deprecate(
1107
+ "T2I should not use down_block_additional_residuals",
1108
+ "1.3.0",
1109
+ "Passing intrablock residual connections with `down_block_additional_residuals` is deprecated \
1110
+ and will be removed in diffusers 1.3.0. `down_block_additional_residuals` should only be used \
1111
+ for ControlNet. Please make sure use `down_intrablock_additional_residuals` instead. ",
1112
+ standard_warn=False,
1113
+ )
1114
+ down_intrablock_additional_residuals = down_block_additional_residuals
1115
+ is_adapter = True
1116
+
1117
+ down_block_res_samples = (sample,)
1118
+ for downsample_block in self.down_blocks:
1119
+ if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
1120
+ # For t2i-adapter CrossAttnDownBlock2D
1121
+ additional_residuals = {}
1122
+ if is_adapter and len(down_intrablock_additional_residuals) > 0:
1123
+ additional_residuals["additional_residuals"] = down_intrablock_additional_residuals.pop(0)
1124
+
1125
+ sample, res_samples = downsample_block(
1126
+ hidden_states=sample,
1127
+ temb=emb,
1128
+ encoder_hidden_states=encoder_hidden_states,
1129
+ attention_mask=attention_mask,
1130
+ cross_attention_kwargs=cross_attention_kwargs,
1131
+ encoder_attention_mask=encoder_attention_mask,
1132
+ **additional_residuals,
1133
+ )
1134
+ else:
1135
+ sample, res_samples = downsample_block(hidden_states=sample, temb=emb, scale=lora_scale)
1136
+ if is_adapter and len(down_intrablock_additional_residuals) > 0:
1137
+ sample += down_intrablock_additional_residuals.pop(0)
1138
+
1139
+ down_block_res_samples += res_samples
1140
+
1141
+ if is_controlnet:
1142
+ new_down_block_res_samples = ()
1143
+
1144
+ for down_block_res_sample, down_block_additional_residual in zip(
1145
+ down_block_res_samples, down_block_additional_residuals
1146
+ ):
1147
+ down_block_res_sample = down_block_res_sample + down_block_additional_residual
1148
+ new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
1149
+
1150
+ down_block_res_samples = new_down_block_res_samples
1151
+
1152
+ # 4. mid
1153
+ if self.mid_block is not None:
1154
+ if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention:
1155
+ sample = self.mid_block(
1156
+ sample,
1157
+ emb,
1158
+ encoder_hidden_states=encoder_hidden_states,
1159
+ attention_mask=attention_mask,
1160
+ cross_attention_kwargs=cross_attention_kwargs,
1161
+ encoder_attention_mask=encoder_attention_mask,
1162
+ )
1163
+ else:
1164
+ sample = self.mid_block(sample, emb)
1165
+
1166
+ # To support T2I-Adapter-XL
1167
+ if (
1168
+ is_adapter
1169
+ and len(down_intrablock_additional_residuals) > 0
1170
+ and sample.shape == down_intrablock_additional_residuals[0].shape
1171
+ ):
1172
+ sample += down_intrablock_additional_residuals.pop(0)
1173
+
1174
+ if is_controlnet:
1175
+ sample = sample + mid_block_additional_residual
1176
+
1177
+ # 5. up
1178
+ for i, upsample_block in enumerate(self.up_blocks):
1179
+ is_final_block = i == len(self.up_blocks) - 1
1180
+
1181
+ res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
1182
+ down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
1183
+
1184
+ # if we have not reached the final block and need to forward the
1185
+ # upsample size, we do it here
1186
+ if not is_final_block and forward_upsample_size:
1187
+ upsample_size = down_block_res_samples[-1].shape[2:]
1188
+
1189
+ if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
1190
+ sample = upsample_block(
1191
+ hidden_states=sample,
1192
+ temb=emb,
1193
+ res_hidden_states_tuple=res_samples,
1194
+ encoder_hidden_states=encoder_hidden_states,
1195
+ cross_attention_kwargs=cross_attention_kwargs,
1196
+ upsample_size=upsample_size,
1197
+ attention_mask=attention_mask,
1198
+ encoder_attention_mask=encoder_attention_mask,
1199
+ )
1200
+ else:
1201
+ sample = upsample_block(
1202
+ hidden_states=sample,
1203
+ temb=emb,
1204
+ res_hidden_states_tuple=res_samples,
1205
+ upsample_size=upsample_size,
1206
+ scale=lora_scale,
1207
+ )
1208
+
1209
+ # 6. post-process
1210
+ if self.conv_norm_out:
1211
+ sample = self.conv_norm_out(sample)
1212
+ sample = self.conv_act(sample)
1213
+ sample = self.conv_out(sample)
1214
+
1215
+ if USE_PEFT_BACKEND:
1216
+ # remove `lora_scale` from each PEFT layer
1217
+ unscale_lora_layers(self, lora_scale)
1218
+
1219
+ if not return_dict:
1220
+ return (sample,)
1221
+
1222
+ return UNet2DConditionOutput(sample=sample)
GeoWizard/geowizard/utils/README.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ Some files are adapted from Marigold :https://github.com/prs-eth/Marigold,
2
+ Thanks for their great work!
GeoWizard/geowizard/utils/alignment.py ADDED
@@ -0,0 +1,73 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Author: Bingxin Ke
2
+ # Last modified: 2024-01-11
3
+
4
+ import numpy as np
5
+ import torch
6
+
7
+
8
+ def align_depth_least_square(
9
+ gt_arr: np.ndarray,
10
+ pred_arr: np.ndarray,
11
+ valid_mask_arr: np.ndarray,
12
+ return_scale_shift=True,
13
+ max_resolution=None,
14
+ ):
15
+ ori_shape = pred_arr.shape # input shape
16
+
17
+ gt = gt_arr.squeeze() # [H, W]
18
+ pred = pred_arr.squeeze()
19
+ valid_mask = valid_mask_arr.squeeze()
20
+
21
+ # Downsample
22
+ if max_resolution is not None:
23
+ scale_factor = np.min(max_resolution / np.array(ori_shape[-2:]))
24
+ if scale_factor < 1:
25
+ downscaler = torch.nn.Upsample(scale_factor=scale_factor, mode="nearest")
26
+ gt = downscaler(torch.as_tensor(gt).unsqueeze(0)).numpy()
27
+ pred = downscaler(torch.as_tensor(pred).unsqueeze(0)).numpy()
28
+ valid_mask = (
29
+ downscaler(torch.as_tensor(valid_mask).unsqueeze(0).float())
30
+ .bool()
31
+ .numpy()
32
+ )
33
+
34
+ assert (
35
+ gt.shape == pred.shape == valid_mask.shape
36
+ ), f"{gt.shape}, {pred.shape}, {valid_mask.shape}"
37
+
38
+ gt_masked = gt[valid_mask].reshape((-1, 1))
39
+ pred_masked = pred[valid_mask].reshape((-1, 1))
40
+
41
+ # numpy solver
42
+ _ones = np.ones_like(pred_masked)
43
+ A = np.concatenate([pred_masked, _ones], axis=-1)
44
+ X = np.linalg.lstsq(A, gt_masked, rcond=None)[0]
45
+ scale, shift = X
46
+
47
+ aligned_pred = pred_arr * scale + shift
48
+
49
+ # restore dimensions
50
+ aligned_pred = aligned_pred.reshape(ori_shape)
51
+
52
+ if return_scale_shift:
53
+ return aligned_pred, scale, shift
54
+ else:
55
+ return aligned_pred
56
+
57
+
58
+ # ******************** disparity space ********************
59
+ def depth2disparity(depth, return_mask=False):
60
+ if isinstance(depth, torch.Tensor):
61
+ disparity = torch.zeros_like(depth)
62
+ elif isinstance(depth, np.ndarray):
63
+ disparity = np.zeros_like(depth)
64
+ non_negtive_mask = depth > 0
65
+ disparity[non_negtive_mask] = 1.0 / depth[non_negtive_mask]
66
+ if return_mask:
67
+ return disparity, non_negtive_mask
68
+ else:
69
+ return disparity
70
+
71
+
72
+ def disparity2depth(disparity, **kwargs):
73
+ return depth2disparity(disparity, **kwargs)
GeoWizard/geowizard/utils/batch_size.py ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # A reimplemented version in public environments by Xiao Fu and Mu Hu
2
+
3
+ import torch
4
+ import math
5
+
6
+
7
+ # Search table for suggested max. inference batch size
8
+ bs_search_table = [
9
+ # tested on A100-PCIE-80GB
10
+ {"res": 768, "total_vram": 79, "bs": 35, "dtype": torch.float32},
11
+ {"res": 1024, "total_vram": 79, "bs": 20, "dtype": torch.float32},
12
+ # tested on A100-PCIE-40GB
13
+ {"res": 768, "total_vram": 39, "bs": 15, "dtype": torch.float32},
14
+ {"res": 1024, "total_vram": 39, "bs": 8, "dtype": torch.float32},
15
+ {"res": 768, "total_vram": 39, "bs": 30, "dtype": torch.float16},
16
+ {"res": 1024, "total_vram": 39, "bs": 15, "dtype": torch.float16},
17
+ # tested on RTX3090, RTX4090
18
+ {"res": 512, "total_vram": 23, "bs": 20, "dtype": torch.float32},
19
+ {"res": 768, "total_vram": 23, "bs": 7, "dtype": torch.float32},
20
+ {"res": 1024, "total_vram": 23, "bs": 3, "dtype": torch.float32},
21
+ {"res": 512, "total_vram": 23, "bs": 40, "dtype": torch.float16},
22
+ {"res": 768, "total_vram": 23, "bs": 18, "dtype": torch.float16},
23
+ {"res": 1024, "total_vram": 23, "bs": 10, "dtype": torch.float16},
24
+ # tested on GTX1080Ti
25
+ {"res": 512, "total_vram": 10, "bs": 5, "dtype": torch.float32},
26
+ {"res": 768, "total_vram": 10, "bs": 2, "dtype": torch.float32},
27
+ {"res": 512, "total_vram": 10, "bs": 10, "dtype": torch.float16},
28
+ {"res": 768, "total_vram": 10, "bs": 5, "dtype": torch.float16},
29
+ {"res": 1024, "total_vram": 10, "bs": 3, "dtype": torch.float16},
30
+ ]
31
+
32
+
33
+ def find_batch_size(ensemble_size: int, input_res: int, dtype: torch.dtype) -> int:
34
+ """
35
+ Automatically search for suitable operating batch size.
36
+
37
+ Args:
38
+ ensemble_size (`int`):
39
+ Number of predictions to be ensembled.
40
+ input_res (`int`):
41
+ Operating resolution of the input image.
42
+
43
+ Returns:
44
+ `int`: Operating batch size.
45
+ """
46
+ if not torch.cuda.is_available():
47
+ return 1
48
+
49
+ total_vram = torch.cuda.mem_get_info()[1] / 1024.0**3
50
+ filtered_bs_search_table = [s for s in bs_search_table if s["dtype"] == dtype]
51
+ for settings in sorted(
52
+ filtered_bs_search_table,
53
+ key=lambda k: (k["res"], -k["total_vram"]),
54
+ ):
55
+ if input_res <= settings["res"] and total_vram >= settings["total_vram"]:
56
+ bs = settings["bs"]
57
+ if bs > ensemble_size:
58
+ bs = ensemble_size
59
+ elif bs > math.ceil(ensemble_size / 2) and bs < ensemble_size:
60
+ bs = math.ceil(ensemble_size / 2)
61
+ return bs
62
+
63
+ return 1
GeoWizard/geowizard/utils/colormap.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # A reimplemented version in public environments by Xiao Fu and Mu Hu
2
+
3
+ import numpy as np
4
+ import cv2
5
+
6
+ def kitti_colormap(disparity, maxval=-1):
7
+ """
8
+ A utility function to reproduce KITTI fake colormap
9
+ Arguments:
10
+ - disparity: numpy float32 array of dimension HxW
11
+ - maxval: maximum disparity value for normalization (if equal to -1, the maximum value in disparity will be used)
12
+
13
+ Returns a numpy uint8 array of shape HxWx3.
14
+ """
15
+ if maxval < 0:
16
+ maxval = np.max(disparity)
17
+
18
+ colormap = np.asarray([[0,0,0,114],[0,0,1,185],[1,0,0,114],[1,0,1,174],[0,1,0,114],[0,1,1,185],[1,1,0,114],[1,1,1,0]])
19
+ weights = np.asarray([8.771929824561404,5.405405405405405,8.771929824561404,5.747126436781609,8.771929824561404,5.405405405405405,8.771929824561404,0])
20
+ cumsum = np.asarray([0,0.114,0.299,0.413,0.587,0.701,0.8859999999999999,0.9999999999999999])
21
+
22
+ colored_disp = np.zeros([disparity.shape[0], disparity.shape[1], 3])
23
+ values = np.expand_dims(np.minimum(np.maximum(disparity/maxval, 0.), 1.), -1)
24
+ bins = np.repeat(np.repeat(np.expand_dims(np.expand_dims(cumsum,axis=0),axis=0), disparity.shape[1], axis=1), disparity.shape[0], axis=0)
25
+ diffs = np.where((np.repeat(values, 8, axis=-1) - bins) > 0, -1000, (np.repeat(values, 8, axis=-1) - bins))
26
+ index = np.argmax(diffs, axis=-1)-1
27
+
28
+ w = 1-(values[:,:,0]-cumsum[index])*np.asarray(weights)[index]
29
+
30
+
31
+ colored_disp[:,:,2] = (w*colormap[index][:,:,0] + (1.-w)*colormap[index+1][:,:,0])
32
+ colored_disp[:,:,1] = (w*colormap[index][:,:,1] + (1.-w)*colormap[index+1][:,:,1])
33
+ colored_disp[:,:,0] = (w*colormap[index][:,:,2] + (1.-w)*colormap[index+1][:,:,2])
34
+
35
+ return (colored_disp*np.expand_dims((disparity>0),-1)*255).astype(np.uint8)
36
+
37
+ def read_16bit_gt(path):
38
+ """
39
+ A utility function to read KITTI 16bit gt
40
+ Arguments:
41
+ - path: filepath
42
+ Returns a numpy float32 array of shape HxW.
43
+ """
44
+ gt = cv2.imread(path,-1).astype(np.float32)/256.
45
+ return gt
GeoWizard/geowizard/utils/common.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # A reimplemented version in public environments by Xiao Fu and Mu Hu
2
+
3
+ import json
4
+ import yaml
5
+ import logging
6
+ import os
7
+ import numpy as np
8
+ import sys
9
+
10
+ def load_loss_scheme(loss_config):
11
+ with open(loss_config, 'r') as f:
12
+ loss_json = yaml.safe_load(f)
13
+ return loss_json
14
+
15
+
16
+ DEBUG =0
17
+ logger = logging.getLogger()
18
+
19
+
20
+ if DEBUG:
21
+ #coloredlogs.install(level='DEBUG')
22
+ logger.setLevel(logging.DEBUG)
23
+ else:
24
+ #coloredlogs.install(level='INFO')
25
+ logger.setLevel(logging.INFO)
26
+
27
+
28
+ strhdlr = logging.StreamHandler()
29
+ logger.addHandler(strhdlr)
30
+ formatter = logging.Formatter('%(asctime)s [%(filename)s:%(lineno)d] %(levelname)s %(message)s')
31
+ strhdlr.setFormatter(formatter)
32
+
33
+
34
+
35
+ def count_parameters(model):
36
+ return sum(p.numel() for p in model.parameters() if p.requires_grad)
37
+
38
+ def check_path(path):
39
+ if not os.path.exists(path):
40
+ os.makedirs(path, exist_ok=True)
41
+
42
+
GeoWizard/geowizard/utils/dataset_configuration.py ADDED
@@ -0,0 +1,81 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # A reimplemented version in public environments by Xiao Fu and Mu Hu
2
+
3
+ import torch
4
+ import torch.nn as nn
5
+ import torch.nn.functional as F
6
+ import numpy as np
7
+ import sys
8
+ sys.path.append("..")
9
+
10
+ from dataloader.mix_loader import MixDataset
11
+ from torch.utils.data import DataLoader
12
+ from dataloader import transforms
13
+ import os
14
+
15
+
16
+ # Get Dataset Here
17
+ def prepare_dataset(data_dir=None,
18
+ batch_size=1,
19
+ test_batch=1,
20
+ datathread=4,
21
+ logger=None):
22
+
23
+ # set the config parameters
24
+ dataset_config_dict = dict()
25
+
26
+ train_dataset = MixDataset(data_dir=data_dir)
27
+
28
+ img_height, img_width = train_dataset.get_img_size()
29
+
30
+ datathread = datathread
31
+ if os.environ.get('datathread') is not None:
32
+ datathread = int(os.environ.get('datathread'))
33
+
34
+ if logger is not None:
35
+ logger.info("Use %d processes to load data..." % datathread)
36
+
37
+ train_loader = DataLoader(train_dataset, batch_size = batch_size, \
38
+ shuffle = True, num_workers = datathread, \
39
+ pin_memory = True)
40
+
41
+ num_batches_per_epoch = len(train_loader)
42
+
43
+ dataset_config_dict['num_batches_per_epoch'] = num_batches_per_epoch
44
+ dataset_config_dict['img_size'] = (img_height,img_width)
45
+
46
+ return train_loader, dataset_config_dict
47
+
48
+ def depth_scale_shift_normalization(depth):
49
+
50
+ bsz = depth.shape[0]
51
+
52
+ depth_ = depth[:,0,:,:].reshape(bsz,-1).cpu().numpy()
53
+ min_value = torch.from_numpy(np.percentile(a=depth_,q=2,axis=1)).to(depth)[...,None,None,None]
54
+ max_value = torch.from_numpy(np.percentile(a=depth_,q=98,axis=1)).to(depth)[...,None,None,None]
55
+
56
+ normalized_depth = ((depth - min_value)/(max_value-min_value+1e-5) - 0.5) * 2
57
+ normalized_depth = torch.clip(normalized_depth, -1., 1.)
58
+
59
+ return normalized_depth
60
+
61
+
62
+
63
+ def resize_max_res_tensor(input_tensor, mode, recom_resolution=768):
64
+ assert input_tensor.shape[1]==3
65
+ original_H, original_W = input_tensor.shape[2:]
66
+ downscale_factor = min(recom_resolution/original_H, recom_resolution/original_W)
67
+
68
+ if mode == 'normal':
69
+ resized_input_tensor = F.interpolate(input_tensor,
70
+ scale_factor=downscale_factor,
71
+ mode='nearest')
72
+ else:
73
+ resized_input_tensor = F.interpolate(input_tensor,
74
+ scale_factor=downscale_factor,
75
+ mode='bilinear',
76
+ align_corners=False)
77
+
78
+ if mode == 'depth':
79
+ return resized_input_tensor / downscale_factor
80
+ else:
81
+ return resized_input_tensor
GeoWizard/geowizard/utils/de_normalized.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ from scipy.optimize import least_squares
3
+ import torch
4
+
5
+ def align_scale_shift(pred, target, clip_max):
6
+ mask = (target > 0) & (target < clip_max)
7
+ if mask.sum() > 10:
8
+ target_mask = target[mask]
9
+ pred_mask = pred[mask]
10
+ scale, shift = np.polyfit(pred_mask, target_mask, deg=1)
11
+ return scale, shift
12
+ else:
13
+ return 1, 0
14
+
15
+ def align_scale(pred: torch.tensor, target: torch.tensor):
16
+ mask = target > 0
17
+ if torch.sum(mask) > 10:
18
+ scale = torch.median(target[mask]) / (torch.median(pred[mask]) + 1e-8)
19
+ else:
20
+ scale = 1
21
+ pred_scale = pred * scale
22
+ return pred_scale, scale
23
+
24
+ def align_shift(pred: torch.tensor, target: torch.tensor):
25
+ mask = target > 0
26
+ if torch.sum(mask) > 10:
27
+ shift = torch.median(target[mask]) - (torch.median(pred[mask]) + 1e-8)
28
+ else:
29
+ shift = 0
30
+ pred_shift = pred + shift
31
+ return pred_shift, shift
GeoWizard/geowizard/utils/depth2normal.py ADDED
@@ -0,0 +1,186 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # A reimplemented version in public environments by Xiao Fu and Mu Hu
2
+
3
+ import pickle
4
+ import os
5
+ import h5py
6
+ import numpy as np
7
+ import cv2
8
+ import torch
9
+ import torch.nn as nn
10
+ import glob
11
+
12
+
13
+ def init_image_coor(height, width):
14
+ x_row = np.arange(0, width)
15
+ x = np.tile(x_row, (height, 1))
16
+ x = x[np.newaxis, :, :]
17
+ x = x.astype(np.float32)
18
+ x = torch.from_numpy(x.copy()).cuda()
19
+ u_u0 = x - width/2.0
20
+
21
+ y_col = np.arange(0, height) # y_col = np.arange(0, height)
22
+ y = np.tile(y_col, (width, 1)).T
23
+ y = y[np.newaxis, :, :]
24
+ y = y.astype(np.float32)
25
+ y = torch.from_numpy(y.copy()).cuda()
26
+ v_v0 = y - height/2.0
27
+ return u_u0, v_v0
28
+
29
+
30
+ def depth_to_xyz(depth, focal_length):
31
+ b, c, h, w = depth.shape
32
+ u_u0, v_v0 = init_image_coor(h, w)
33
+ x = u_u0 * depth / focal_length[0]
34
+ y = v_v0 * depth / focal_length[1]
35
+ z = depth
36
+ pw = torch.cat([x, y, z], 1).permute(0, 2, 3, 1) # [b, h, w, c]
37
+ return pw
38
+
39
+
40
+ def get_surface_normal(xyz, patch_size=5):
41
+ # xyz: [1, h, w, 3]
42
+ x, y, z = torch.unbind(xyz, dim=3)
43
+ x = torch.unsqueeze(x, 0)
44
+ y = torch.unsqueeze(y, 0)
45
+ z = torch.unsqueeze(z, 0)
46
+
47
+ xx = x * x
48
+ yy = y * y
49
+ zz = z * z
50
+ xy = x * y
51
+ xz = x * z
52
+ yz = y * z
53
+ patch_weight = torch.ones((1, 1, patch_size, patch_size), requires_grad=False).cuda()
54
+ xx_patch = nn.functional.conv2d(xx, weight=patch_weight, padding=int(patch_size / 2))
55
+ yy_patch = nn.functional.conv2d(yy, weight=patch_weight, padding=int(patch_size / 2))
56
+ zz_patch = nn.functional.conv2d(zz, weight=patch_weight, padding=int(patch_size / 2))
57
+ xy_patch = nn.functional.conv2d(xy, weight=patch_weight, padding=int(patch_size / 2))
58
+ xz_patch = nn.functional.conv2d(xz, weight=patch_weight, padding=int(patch_size / 2))
59
+ yz_patch = nn.functional.conv2d(yz, weight=patch_weight, padding=int(patch_size / 2))
60
+ ATA = torch.stack([xx_patch, xy_patch, xz_patch, xy_patch, yy_patch, yz_patch, xz_patch, yz_patch, zz_patch],
61
+ dim=4)
62
+ ATA = torch.squeeze(ATA)
63
+ ATA = torch.reshape(ATA, (ATA.size(0), ATA.size(1), 3, 3))
64
+ eps_identity = 1e-6 * torch.eye(3, device=ATA.device, dtype=ATA.dtype)[None, None, :, :].repeat([ATA.size(0), ATA.size(1), 1, 1])
65
+ ATA = ATA + eps_identity
66
+ x_patch = nn.functional.conv2d(x, weight=patch_weight, padding=int(patch_size / 2))
67
+ y_patch = nn.functional.conv2d(y, weight=patch_weight, padding=int(patch_size / 2))
68
+ z_patch = nn.functional.conv2d(z, weight=patch_weight, padding=int(patch_size / 2))
69
+ AT1 = torch.stack([x_patch, y_patch, z_patch], dim=4)
70
+ AT1 = torch.squeeze(AT1)
71
+ AT1 = torch.unsqueeze(AT1, 3)
72
+
73
+ patch_num = 4
74
+ patch_x = int(AT1.size(1) / patch_num)
75
+ patch_y = int(AT1.size(0) / patch_num)
76
+ n_img = torch.randn(AT1.shape).cuda()
77
+ overlap = patch_size // 2 + 1
78
+ for x in range(int(patch_num)):
79
+ for y in range(int(patch_num)):
80
+ left_flg = 0 if x == 0 else 1
81
+ right_flg = 0 if x == patch_num -1 else 1
82
+ top_flg = 0 if y == 0 else 1
83
+ btm_flg = 0 if y == patch_num - 1 else 1
84
+ at1 = AT1[y * patch_y - top_flg * overlap:(y + 1) * patch_y + btm_flg * overlap,
85
+ x * patch_x - left_flg * overlap:(x + 1) * patch_x + right_flg * overlap]
86
+ ata = ATA[y * patch_y - top_flg * overlap:(y + 1) * patch_y + btm_flg * overlap,
87
+ x * patch_x - left_flg * overlap:(x + 1) * patch_x + right_flg * overlap]
88
+ # n_img_tmp, _ = torch.solve(at1, ata)
89
+ n_img_tmp = torch.linalg.solve(ata, at1)
90
+
91
+ n_img_tmp_select = n_img_tmp[top_flg * overlap:patch_y + top_flg * overlap, left_flg * overlap:patch_x + left_flg * overlap, :, :]
92
+ n_img[y * patch_y:y * patch_y + patch_y, x * patch_x:x * patch_x + patch_x, :, :] = n_img_tmp_select
93
+
94
+ n_img_L2 = torch.sqrt(torch.sum(n_img ** 2, dim=2, keepdim=True))
95
+ n_img_norm = n_img / n_img_L2
96
+
97
+ # re-orient normals consistently
98
+ orient_mask = torch.sum(torch.squeeze(n_img_norm) * torch.squeeze(xyz), dim=2) > 0
99
+ n_img_norm[orient_mask] *= -1
100
+ return n_img_norm
101
+
102
+ def get_surface_normalv2(xyz, patch_size=5):
103
+ """
104
+ xyz: xyz coordinates
105
+ patch: [p1, p2, p3,
106
+ p4, p5, p6,
107
+ p7, p8, p9]
108
+ surface_normal = [(p9-p1) x (p3-p7)] + [(p6-p4) - (p8-p2)]
109
+ return: normal [h, w, 3, b]
110
+ """
111
+ b, h, w, c = xyz.shape
112
+ half_patch = patch_size // 2
113
+ xyz_pad = torch.zeros((b, h + patch_size - 1, w + patch_size - 1, c), dtype=xyz.dtype, device=xyz.device)
114
+ xyz_pad[:, half_patch:-half_patch, half_patch:-half_patch, :] = xyz
115
+
116
+ # xyz_left_top = xyz_pad[:, :h, :w, :] # p1
117
+ # xyz_right_bottom = xyz_pad[:, -h:, -w:, :]# p9
118
+ # xyz_left_bottom = xyz_pad[:, -h:, :w, :] # p7
119
+ # xyz_right_top = xyz_pad[:, :h, -w:, :] # p3
120
+ # xyz_cross1 = xyz_left_top - xyz_right_bottom # p1p9
121
+ # xyz_cross2 = xyz_left_bottom - xyz_right_top # p7p3
122
+
123
+ xyz_left = xyz_pad[:, half_patch:half_patch + h, :w, :] # p4
124
+ xyz_right = xyz_pad[:, half_patch:half_patch + h, -w:, :] # p6
125
+ xyz_top = xyz_pad[:, :h, half_patch:half_patch + w, :] # p2
126
+ xyz_bottom = xyz_pad[:, -h:, half_patch:half_patch + w, :] # p8
127
+ xyz_horizon = xyz_left - xyz_right # p4p6
128
+ xyz_vertical = xyz_top - xyz_bottom # p2p8
129
+
130
+ xyz_left_in = xyz_pad[:, half_patch:half_patch + h, 1:w+1, :] # p4
131
+ xyz_right_in = xyz_pad[:, half_patch:half_patch + h, patch_size-1:patch_size-1+w, :] # p6
132
+ xyz_top_in = xyz_pad[:, 1:h+1, half_patch:half_patch + w, :] # p2
133
+ xyz_bottom_in = xyz_pad[:, patch_size-1:patch_size-1+h, half_patch:half_patch + w, :] # p8
134
+ xyz_horizon_in = xyz_left_in - xyz_right_in # p4p6
135
+ xyz_vertical_in = xyz_top_in - xyz_bottom_in # p2p8
136
+
137
+ n_img_1 = torch.cross(xyz_horizon_in, xyz_vertical_in, dim=3)
138
+ n_img_2 = torch.cross(xyz_horizon, xyz_vertical, dim=3)
139
+
140
+ # re-orient normals consistently
141
+ orient_mask = torch.sum(n_img_1 * xyz, dim=3) > 0
142
+ n_img_1[orient_mask] *= -1
143
+ orient_mask = torch.sum(n_img_2 * xyz, dim=3) > 0
144
+ n_img_2[orient_mask] *= -1
145
+
146
+ n_img1_L2 = torch.sqrt(torch.sum(n_img_1 ** 2, dim=3, keepdim=True))
147
+ n_img1_norm = n_img_1 / (n_img1_L2 + 1e-8)
148
+
149
+ n_img2_L2 = torch.sqrt(torch.sum(n_img_2 ** 2, dim=3, keepdim=True))
150
+ n_img2_norm = n_img_2 / (n_img2_L2 + 1e-8)
151
+
152
+ # average 2 norms
153
+ n_img_aver = n_img1_norm + n_img2_norm
154
+ n_img_aver_L2 = torch.sqrt(torch.sum(n_img_aver ** 2, dim=3, keepdim=True))
155
+ n_img_aver_norm = n_img_aver / (n_img_aver_L2 + 1e-8)
156
+ # re-orient normals consistently
157
+ orient_mask = torch.sum(n_img_aver_norm * xyz, dim=3) > 0
158
+ n_img_aver_norm[orient_mask] *= -1
159
+ n_img_aver_norm_out = n_img_aver_norm.permute((1, 2, 3, 0)) # [h, w, c, b]
160
+
161
+ # a = torch.sum(n_img1_norm_out*n_img2_norm_out, dim=2).cpu().numpy().squeeze()
162
+ # plt.imshow(np.abs(a), cmap='rainbow')
163
+ # plt.show()
164
+ return n_img_aver_norm_out#n_img1_norm.permute((1, 2, 3, 0))
165
+
166
+ def surface_normal_from_depth(depth, focal_length, valid_mask=None):
167
+ # para depth: depth map, [b, c, h, w]
168
+ b, c, h, w = depth.shape
169
+ focal_length = focal_length[:, None, None, None]
170
+ depth_filter = nn.functional.avg_pool2d(depth, kernel_size=3, stride=1, padding=1)
171
+ #depth_filter = nn.functional.avg_pool2d(depth_filter, kernel_size=3, stride=1, padding=1)
172
+ xyz = depth_to_xyz(depth_filter, focal_length)
173
+ sn_batch = []
174
+ for i in range(b):
175
+ xyz_i = xyz[i, :][None, :, :, :]
176
+ #normal = get_surface_normalv2(xyz_i)
177
+ normal = get_surface_normal(xyz_i)
178
+ sn_batch.append(normal)
179
+ sn_batch = torch.cat(sn_batch, dim=3).permute((3, 2, 0, 1)) # [b, c, h, w]
180
+
181
+ if valid_mask != None:
182
+ mask_invalid = (~valid_mask).repeat(1, 3, 1, 1)
183
+ sn_batch[mask_invalid] = 0.0
184
+
185
+ return sn_batch
186
+
GeoWizard/geowizard/utils/depth_ensemble.py ADDED
@@ -0,0 +1,115 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # A reimplemented version in public environments by Xiao Fu and Mu Hu
2
+
3
+ import numpy as np
4
+ import torch
5
+
6
+ from scipy.optimize import minimize
7
+
8
+ def inter_distances(tensors: torch.Tensor):
9
+ """
10
+ To calculate the distance between each two depth maps.
11
+ """
12
+ distances = []
13
+ for i, j in torch.combinations(torch.arange(tensors.shape[0])):
14
+ arr1 = tensors[i : i + 1]
15
+ arr2 = tensors[j : j + 1]
16
+ distances.append(arr1 - arr2)
17
+ dist = torch.concat(distances, dim=0)
18
+ return dist
19
+
20
+
21
+ def ensemble_depths(input_images:torch.Tensor,
22
+ regularizer_strength: float =0.02,
23
+ max_iter: int =2,
24
+ tol:float =1e-3,
25
+ reduction: str='median',
26
+ max_res: int=None):
27
+ """
28
+ To ensemble multiple affine-invariant depth images (up to scale and shift),
29
+ by aligning estimating the scale and shift
30
+ """
31
+
32
+ device = input_images.device
33
+ dtype = input_images.dtype
34
+ np_dtype = np.float32
35
+
36
+
37
+ original_input = input_images.clone()
38
+ n_img = input_images.shape[0]
39
+ ori_shape = input_images.shape
40
+
41
+ if max_res is not None:
42
+ scale_factor = torch.min(max_res / torch.tensor(ori_shape[-2:]))
43
+ if scale_factor < 1:
44
+ downscaler = torch.nn.Upsample(scale_factor=scale_factor, mode="nearest")
45
+ input_images = downscaler(torch.from_numpy(input_images)).numpy()
46
+
47
+ # init guess
48
+ _min = np.min(input_images.reshape((n_img, -1)).cpu().numpy(), axis=1) # get the min value of each possible depth
49
+ _max = np.max(input_images.reshape((n_img, -1)).cpu().numpy(), axis=1) # get the max value of each possible depth
50
+ s_init = 1.0 / (_max - _min).reshape((-1, 1, 1)) #(10,1,1) : re-scale'f scale
51
+ t_init = (-1 * s_init.flatten() * _min.flatten()).reshape((-1, 1, 1)) #(10,1,1)
52
+
53
+ x = np.concatenate([s_init, t_init]).reshape(-1).astype(np_dtype) #(20,)
54
+
55
+ input_images = input_images.to(device)
56
+
57
+ # objective function
58
+ def closure(x):
59
+ l = len(x)
60
+ s = x[: int(l / 2)]
61
+ t = x[int(l / 2) :]
62
+ s = torch.from_numpy(s).to(dtype=dtype).to(device)
63
+ t = torch.from_numpy(t).to(dtype=dtype).to(device)
64
+
65
+ transformed_arrays = input_images * s.view((-1, 1, 1)) + t.view((-1, 1, 1))
66
+ dists = inter_distances(transformed_arrays)
67
+ sqrt_dist = torch.sqrt(torch.mean(dists**2))
68
+
69
+ if "mean" == reduction:
70
+ pred = torch.mean(transformed_arrays, dim=0)
71
+ elif "median" == reduction:
72
+ pred = torch.median(transformed_arrays, dim=0).values
73
+ else:
74
+ raise ValueError
75
+
76
+ near_err = torch.sqrt((0 - torch.min(pred)) ** 2)
77
+ far_err = torch.sqrt((1 - torch.max(pred)) ** 2)
78
+
79
+ err = sqrt_dist + (near_err + far_err) * regularizer_strength
80
+ err = err.detach().cpu().numpy().astype(np_dtype)
81
+ return err
82
+
83
+ res = minimize(
84
+ closure, x, method="BFGS", tol=tol, options={"maxiter": max_iter, "disp": False}
85
+ )
86
+ x = res.x
87
+ l = len(x)
88
+ s = x[: int(l / 2)]
89
+ t = x[int(l / 2) :]
90
+
91
+ # Prediction
92
+ s = torch.from_numpy(s).to(dtype=dtype).to(device)
93
+ t = torch.from_numpy(t).to(dtype=dtype).to(device)
94
+ transformed_arrays = original_input * s.view(-1, 1, 1) + t.view(-1, 1, 1) #[10,H,W]
95
+
96
+
97
+ if "mean" == reduction:
98
+ aligned_images = torch.mean(transformed_arrays, dim=0)
99
+ std = torch.std(transformed_arrays, dim=0)
100
+ uncertainty = std
101
+
102
+ elif "median" == reduction:
103
+ aligned_images = torch.median(transformed_arrays, dim=0).values
104
+ # MAD (median absolute deviation) as uncertainty indicator
105
+ abs_dev = torch.abs(transformed_arrays - aligned_images)
106
+ mad = torch.median(abs_dev, dim=0).values
107
+ uncertainty = mad
108
+
109
+ # Scale and shift to [0, 1]
110
+ _min = torch.min(aligned_images)
111
+ _max = torch.max(aligned_images)
112
+ aligned_images = (aligned_images - _min) / (_max - _min)
113
+ uncertainty /= _max - _min
114
+
115
+ return aligned_images, uncertainty
GeoWizard/geowizard/utils/depth_transform.py ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Author: Bingxin Ke
2
+ # Last modified: 2024-02-08
3
+
4
+ import torch
5
+
6
+
7
+ def get_depth_normalizer(cfg_normalizer):
8
+ if cfg_normalizer is None:
9
+
10
+ def identical(x):
11
+ return x
12
+
13
+ depth_transform = identical
14
+
15
+ elif "near_far_metric" == cfg_normalizer.type:
16
+ depth_transform = NearFarMetricNormalizer(
17
+ norm_min=cfg_normalizer.norm_min,
18
+ norm_max=cfg_normalizer.norm_max,
19
+ min_max_quantile=cfg_normalizer.min_max_quantile,
20
+ clip=cfg_normalizer.clip,
21
+ )
22
+ else:
23
+ raise NotImplementedError
24
+ return depth_transform
25
+
26
+
27
+ class DepthNormalizerBase:
28
+ is_relative = None
29
+ far_plane_at_max = None
30
+
31
+ def __init__(
32
+ self,
33
+ norm_min=-1.0,
34
+ norm_max=1.0,
35
+ ) -> None:
36
+ self.norm_min = norm_min
37
+ self.norm_max = norm_max
38
+ raise NotImplementedError
39
+
40
+ def __call__(self, depth, valid_mask=None, clip=None):
41
+ raise NotImplementedError
42
+
43
+ def denormalize(self, depth_norm, **kwargs):
44
+ # For metric depth: convert prediction back to metric depth
45
+ # For relative depth: convert prediction to [0, 1]
46
+ raise NotImplementedError
47
+
48
+
49
+ class NearFarMetricNormalizer(DepthNormalizerBase):
50
+ """
51
+ depth in [0, d_max] -> [-1, 1]
52
+ """
53
+
54
+ is_relative = True
55
+ far_plane_at_max = True
56
+
57
+ def __init__(
58
+ self, norm_min=-1.0, norm_max=1.0, min_max_quantile=0.02, clip=True
59
+ ) -> None:
60
+ self.norm_min = norm_min
61
+ self.norm_max = norm_max
62
+ self.norm_range = self.norm_max - self.norm_min
63
+ self.min_quantile = min_max_quantile
64
+ self.max_quantile = 1.0 - self.min_quantile
65
+ self.clip = clip
66
+
67
+ def __call__(self, depth_linear, valid_mask=None, clip=None):
68
+ clip = clip if clip is not None else self.clip
69
+
70
+ if valid_mask is None:
71
+ valid_mask = torch.ones_like(depth_linear).bool()
72
+ valid_mask = valid_mask & (depth_linear > 0)
73
+
74
+ # Take quantiles as min and max
75
+ _min, _max = torch.quantile(
76
+ depth_linear[valid_mask],
77
+ torch.tensor([self.min_quantile, self.max_quantile]),
78
+ )
79
+
80
+ # scale and shift
81
+ depth_norm_linear = (depth_linear - _min) / (
82
+ _max - _min
83
+ ) * self.norm_range + self.norm_min
84
+
85
+ if clip:
86
+ depth_norm_linear = torch.clip(
87
+ depth_norm_linear, self.norm_min, self.norm_max
88
+ )
89
+
90
+ return depth_norm_linear
91
+
92
+ def scale_back(self, depth_norm):
93
+ # scale to [0, 1]
94
+ depth_linear = (depth_norm - self.norm_min) / self.norm_range
95
+ return depth_linear
96
+
97
+ def denormalize(self, depth_norm, **kwargs):
98
+ return self.scale_back(depth_norm=depth_norm)
99
+
GeoWizard/geowizard/utils/image_util.py ADDED
@@ -0,0 +1,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # A reimplemented version in public environments by Xiao Fu and Mu Hu
2
+
3
+ import matplotlib
4
+ import numpy as np
5
+ import torch
6
+ from PIL import Image
7
+
8
+
9
+
10
+
11
+ def resize_max_res(img: Image.Image, max_edge_resolution: int) -> Image.Image:
12
+ """
13
+ Resize image to limit maximum edge length while keeping aspect ratio.
14
+ Args:
15
+ img (`Image.Image`):
16
+ Image to be resized.
17
+ max_edge_resolution (`int`):
18
+ Maximum edge length (pixel).
19
+ Returns:
20
+ `Image.Image`: Resized image.
21
+ """
22
+
23
+ original_width, original_height = img.size
24
+
25
+ downscale_factor = min(
26
+ max_edge_resolution / original_width, max_edge_resolution / original_height
27
+ )
28
+
29
+ new_width = int(original_width * downscale_factor)
30
+ new_height = int(original_height * downscale_factor)
31
+
32
+ resized_img = img.resize((new_width, new_height))
33
+ return resized_img
34
+
35
+
36
+ def colorize_depth_maps(
37
+ depth_map, min_depth, max_depth, cmap="Spectral", valid_mask=None
38
+ ):
39
+ """
40
+ Colorize depth maps.
41
+ """
42
+ assert len(depth_map.shape) >= 2, "Invalid dimension"
43
+
44
+ if isinstance(depth_map, torch.Tensor):
45
+ depth = depth_map.detach().clone().squeeze().numpy()
46
+ elif isinstance(depth_map, np.ndarray):
47
+ depth = depth_map.copy().squeeze()
48
+ # reshape to [ (B,) H, W ]
49
+ if depth.ndim < 3:
50
+ depth = depth[np.newaxis, :, :]
51
+
52
+ # colorize
53
+ cm = matplotlib.colormaps[cmap]
54
+ depth = ((depth - min_depth) / (max_depth - min_depth)).clip(0, 1)
55
+ img_colored_np = cm(depth, bytes=False)[:, :, :, 0:3] # value from 0 to 1
56
+ img_colored_np = np.rollaxis(img_colored_np, 3, 1)
57
+
58
+ if valid_mask is not None:
59
+ if isinstance(depth_map, torch.Tensor):
60
+ valid_mask = valid_mask.detach().numpy()
61
+ valid_mask = valid_mask.squeeze() # [H, W] or [B, H, W]
62
+ if valid_mask.ndim < 3:
63
+ valid_mask = valid_mask[np.newaxis, np.newaxis, :, :]
64
+ else:
65
+ valid_mask = valid_mask[:, np.newaxis, :, :]
66
+ valid_mask = np.repeat(valid_mask, 3, axis=1)
67
+ img_colored_np[~valid_mask] = 0
68
+
69
+ if isinstance(depth_map, torch.Tensor):
70
+ img_colored = torch.from_numpy(img_colored_np).float()
71
+ elif isinstance(depth_map, np.ndarray):
72
+ img_colored = img_colored_np
73
+
74
+ return img_colored
75
+
76
+
77
+ def chw2hwc(chw):
78
+ assert 3 == len(chw.shape)
79
+ if isinstance(chw, torch.Tensor):
80
+ hwc = torch.permute(chw, (1, 2, 0))
81
+ elif isinstance(chw, np.ndarray):
82
+ hwc = np.moveaxis(chw, 0, -1)
83
+ return hwc
GeoWizard/geowizard/utils/metric.py ADDED
@@ -0,0 +1,158 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Author: Bingxin Ke
2
+ # Last modified: 2024-02-15
3
+
4
+
5
+ import pandas as pd
6
+ import torch
7
+
8
+
9
+ # Adapted from: https://github.com/victoresque/pytorch-template/blob/master/utils/util.py
10
+ class MetricTracker:
11
+ def __init__(self, *keys, writer=None):
12
+ self.writer = writer
13
+ self._data = pd.DataFrame(index=keys, columns=["total", "counts", "average"])
14
+ self.reset()
15
+
16
+ def reset(self):
17
+ for col in self._data.columns:
18
+ self._data[col].values[:] = 0
19
+
20
+ def update(self, key, value, n=1):
21
+ if self.writer is not None:
22
+ self.writer.add_scalar(key, value)
23
+ self._data.loc[key, "total"] += value * n
24
+ self._data.loc[key, "counts"] += n
25
+ self._data.loc[key, "average"] = self._data.total[key] / self._data.counts[key]
26
+
27
+ def avg(self, key):
28
+ return self._data.average[key]
29
+
30
+ def result(self):
31
+ return dict(self._data.average)
32
+
33
+
34
+ def abs_relative_difference(output, target, valid_mask=None):
35
+ actual_output = output
36
+ actual_target = target
37
+ abs_relative_diff = torch.abs(actual_output - actual_target) / actual_target
38
+ if valid_mask is not None:
39
+ abs_relative_diff[~valid_mask] = 0
40
+ n = valid_mask.sum((-1, -2))
41
+ else:
42
+ n = output.shape[-1] * output.shape[-2]
43
+ abs_relative_diff = torch.sum(abs_relative_diff, (-1, -2)) / n
44
+ return abs_relative_diff.mean()
45
+
46
+
47
+ def squared_relative_difference(output, target, valid_mask=None):
48
+ actual_output = output
49
+ actual_target = target
50
+ square_relative_diff = (
51
+ torch.pow(torch.abs(actual_output - actual_target), 2) / actual_target
52
+ )
53
+ if valid_mask is not None:
54
+ square_relative_diff[~valid_mask] = 0
55
+ n = valid_mask.sum((-1, -2))
56
+ else:
57
+ n = output.shape[-1] * output.shape[-2]
58
+ square_relative_diff = torch.sum(square_relative_diff, (-1, -2)) / n
59
+ return square_relative_diff.mean()
60
+
61
+
62
+ def rmse_linear(output, target, valid_mask=None):
63
+ actual_output = output
64
+ actual_target = target
65
+ diff = actual_output - actual_target
66
+ if valid_mask is not None:
67
+ diff[~valid_mask] = 0
68
+ n = valid_mask.sum((-1, -2))
69
+ else:
70
+ n = output.shape[-1] * output.shape[-2]
71
+ diff2 = torch.pow(diff, 2)
72
+ mse = torch.sum(diff2, (-1, -2)) / n
73
+ rmse = torch.sqrt(mse)
74
+ return rmse.mean()
75
+
76
+
77
+ def rmse_log(output, target, valid_mask=None):
78
+ diff = torch.log(output) - torch.log(target)
79
+ if valid_mask is not None:
80
+ diff[~valid_mask] = 0
81
+ n = valid_mask.sum((-1, -2))
82
+ else:
83
+ n = output.shape[-1] * output.shape[-2]
84
+ diff2 = torch.pow(diff, 2)
85
+ mse = torch.sum(diff2, (-1, -2)) / n # [B]
86
+ rmse = torch.sqrt(mse)
87
+ return rmse.mean()
88
+
89
+
90
+ def log10(output, target, valid_mask=None):
91
+ if valid_mask is not None:
92
+ diff = torch.abs(
93
+ torch.log10(output[valid_mask]) - torch.log10(target[valid_mask])
94
+ )
95
+ else:
96
+ diff = torch.abs(torch.log10(output) - torch.log10(target))
97
+ return diff.mean()
98
+
99
+
100
+ # adapt from: https://github.com/imran3180/depth-map-prediction/blob/master/main.py
101
+ def threshold_percentage(output, target, threshold_val, valid_mask=None):
102
+ d1 = output / target
103
+ d2 = target / output
104
+ max_d1_d2 = torch.max(d1, d2)
105
+ zero = torch.zeros(*output.shape)
106
+ one = torch.ones(*output.shape)
107
+ bit_mat = torch.where(max_d1_d2.cpu() < threshold_val, one, zero)
108
+ if valid_mask is not None:
109
+ bit_mat[~valid_mask] = 0
110
+ n = valid_mask.sum((-1, -2))
111
+ else:
112
+ n = output.shape[-1] * output.shape[-2]
113
+ count_mat = torch.sum(bit_mat, (-1, -2))
114
+ threshold_mat = count_mat / n.cpu()
115
+ return threshold_mat.mean()
116
+
117
+
118
+ def delta1_acc(pred, gt, valid_mask):
119
+ return threshold_percentage(pred, gt, 1.25, valid_mask)
120
+
121
+
122
+ def delta2_acc(pred, gt, valid_mask):
123
+ return threshold_percentage(pred, gt, 1.25**2, valid_mask)
124
+
125
+
126
+ def delta3_acc(pred, gt, valid_mask):
127
+ return threshold_percentage(pred, gt, 1.25**3, valid_mask)
128
+
129
+
130
+ def i_rmse(output, target, valid_mask=None):
131
+ output_inv = 1.0 / output
132
+ target_inv = 1.0 / target
133
+ diff = output_inv - target_inv
134
+ if valid_mask is not None:
135
+ diff[~valid_mask] = 0
136
+ n = valid_mask.sum((-1, -2))
137
+ else:
138
+ n = output.shape[-1] * output.shape[-2]
139
+ diff2 = torch.pow(diff, 2)
140
+ mse = torch.sum(diff2, (-1, -2)) / n # [B]
141
+ rmse = torch.sqrt(mse)
142
+ return rmse.mean()
143
+
144
+
145
+ def silog_rmse(depth_pred, depth_gt, valid_mask=None):
146
+ diff = torch.log(depth_pred) - torch.log(depth_gt)
147
+ if valid_mask is not None:
148
+ diff[~valid_mask] = 0
149
+ n = valid_mask.sum((-1, -2))
150
+ else:
151
+ n = depth_gt.shape[-2] * depth_gt.shape[-1]
152
+
153
+ diff2 = torch.pow(diff, 2)
154
+
155
+ first_term = torch.sum(diff2, (-1, -2)) / n
156
+ second_term = torch.pow(torch.sum(diff, (-1, -2)), 2) / (n**2)
157
+ loss = torch.sqrt(torch.mean(first_term - second_term)) * 100
158
+ return loss
GeoWizard/geowizard/utils/noise.py ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # add
2
+ import torch
3
+ import numpy as np
4
+ from torch import nn
5
+
6
+
7
+ # Pyramid Noise implementation from GeoWizard training code
8
+ # https://github.com/fuxiao0719/GeoWizard/blob/5b25910f5ceaecb4f5f3db000153052628611c9d/geowizard/training/training/train_depth_normal.py#L299
9
+ def pyramid_noise_like(x, timesteps, discount=0.9):
10
+ b, c, w_ori, h_ori = x.shape
11
+ u = nn.Upsample(size=(w_ori, h_ori), mode='bilinear')
12
+ noise = torch.randn_like(x)
13
+ scale = 1.5
14
+ for i in range(10):
15
+ r = np.random.random()*scale + scale # Rather than always going 2x,
16
+ w, h = max(1, int(w_ori/(r**i))), max(1, int(h_ori/(r**i)))
17
+ noise += u(torch.randn(b, c, w, h).to(x)) * (timesteps[...,None,None,None]/1000) * discount**i
18
+ if w==1 or h==1: break # Lowest resolution is 1x1
19
+ return noise/noise.std() # Scaled back to roughly unit variance
GeoWizard/geowizard/utils/normal_ensemble.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # A reimplemented version in public environments by Xiao Fu and Mu Hu
2
+
3
+ import numpy as np
4
+ import torch
5
+
6
+ def ensemble_normals(input_images:torch.Tensor):
7
+ normal_preds = input_images
8
+
9
+ bsz, d, h, w = normal_preds.shape
10
+ normal_preds = normal_preds / (torch.norm(normal_preds, p=2, dim=1).unsqueeze(1)+1e-5)
11
+
12
+ phi = torch.atan2(normal_preds[:,1,:,:], normal_preds[:,0,:,:]).mean(dim=0)
13
+ theta = torch.atan2(torch.norm(normal_preds[:,:2,:,:], p=2, dim=1), normal_preds[:,2,:,:]).mean(dim=0)
14
+ normal_pred = torch.zeros((d,h,w)).to(normal_preds)
15
+ normal_pred[0,:,:] = torch.sin(theta) * torch.cos(phi)
16
+ normal_pred[1,:,:] = torch.sin(theta) * torch.sin(phi)
17
+ normal_pred[2,:,:] = torch.cos(theta)
18
+
19
+ angle_error = torch.acos(torch.clip(torch.cosine_similarity(normal_pred[None], normal_preds, dim=1),-0.999, 0.999))
20
+ normal_idx = torch.argmin(angle_error.reshape(bsz,-1).sum(-1))
21
+
22
+ return normal_preds[normal_idx]
GeoWizard/geowizard/utils/seed_all.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2023 Bingxin Ke, ETH Zurich. 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
+ # If you find this code useful, we kindly ask you to cite our paper in your work.
16
+ # Please find bibtex at: https://github.com/prs-eth/Marigold#-citation
17
+ # More information about the method can be found at https://marigoldmonodepth.github.io
18
+ # --------------------------------------------------------------------------
19
+
20
+
21
+ import numpy as np
22
+ import random
23
+ import torch
24
+
25
+
26
+ def seed_all(seed: int = 0):
27
+ """
28
+ Set random seeds of all components.
29
+ """
30
+ random.seed(seed)
31
+ np.random.seed(seed)
32
+ torch.manual_seed(seed)
33
+ torch.cuda.manual_seed_all(seed)
GeoWizard/geowizard/utils/surface_normal.py ADDED
@@ -0,0 +1,213 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # A reimplemented version in public environments by Xiao Fu and Mu Hu
2
+
3
+ import torch
4
+ import numpy as np
5
+ import torch.nn as nn
6
+
7
+
8
+ def init_image_coor(height, width):
9
+ x_row = np.arange(0, width)
10
+ x = np.tile(x_row, (height, 1))
11
+ x = x[np.newaxis, :, :]
12
+ x = x.astype(np.float32)
13
+ x = torch.from_numpy(x.copy()).cuda()
14
+ u_u0 = x - width/2.0
15
+
16
+ y_col = np.arange(0, height) # y_col = np.arange(0, height)
17
+ y = np.tile(y_col, (width, 1)).T
18
+ y = y[np.newaxis, :, :]
19
+ y = y.astype(np.float32)
20
+ y = torch.from_numpy(y.copy()).cuda()
21
+ v_v0 = y - height/2.0
22
+ return u_u0, v_v0
23
+
24
+
25
+ def depth_to_xyz(depth, focal_length):
26
+ b, c, h, w = depth.shape
27
+ u_u0, v_v0 = init_image_coor(h, w)
28
+ x = u_u0 * depth / focal_length
29
+ y = v_v0 * depth / focal_length
30
+ z = depth
31
+ pw = torch.cat([x, y, z], 1).permute(0, 2, 3, 1) # [b, h, w, c]
32
+ return pw
33
+
34
+
35
+ def get_surface_normal(xyz, patch_size=3):
36
+ # xyz: [1, h, w, 3]
37
+ x, y, z = torch.unbind(xyz, dim=3)
38
+ x = torch.unsqueeze(x, 0)
39
+ y = torch.unsqueeze(y, 0)
40
+ z = torch.unsqueeze(z, 0)
41
+
42
+ xx = x * x
43
+ yy = y * y
44
+ zz = z * z
45
+ xy = x * y
46
+ xz = x * z
47
+ yz = y * z
48
+ patch_weight = torch.ones((1, 1, patch_size, patch_size), requires_grad=False).cuda()
49
+ xx_patch = nn.functional.conv2d(xx, weight=patch_weight, padding=int(patch_size / 2))
50
+ yy_patch = nn.functional.conv2d(yy, weight=patch_weight, padding=int(patch_size / 2))
51
+ zz_patch = nn.functional.conv2d(zz, weight=patch_weight, padding=int(patch_size / 2))
52
+ xy_patch = nn.functional.conv2d(xy, weight=patch_weight, padding=int(patch_size / 2))
53
+ xz_patch = nn.functional.conv2d(xz, weight=patch_weight, padding=int(patch_size / 2))
54
+ yz_patch = nn.functional.conv2d(yz, weight=patch_weight, padding=int(patch_size / 2))
55
+ ATA = torch.stack([xx_patch, xy_patch, xz_patch, xy_patch, yy_patch, yz_patch, xz_patch, yz_patch, zz_patch],
56
+ dim=4)
57
+ ATA = torch.squeeze(ATA)
58
+ ATA = torch.reshape(ATA, (ATA.size(0), ATA.size(1), 3, 3))
59
+ eps_identity = 1e-6 * torch.eye(3, device=ATA.device, dtype=ATA.dtype)[None, None, :, :].repeat([ATA.size(0), ATA.size(1), 1, 1])
60
+ ATA = ATA + eps_identity
61
+ x_patch = nn.functional.conv2d(x, weight=patch_weight, padding=int(patch_size / 2))
62
+ y_patch = nn.functional.conv2d(y, weight=patch_weight, padding=int(patch_size / 2))
63
+ z_patch = nn.functional.conv2d(z, weight=patch_weight, padding=int(patch_size / 2))
64
+ AT1 = torch.stack([x_patch, y_patch, z_patch], dim=4)
65
+ AT1 = torch.squeeze(AT1)
66
+ AT1 = torch.unsqueeze(AT1, 3)
67
+
68
+ patch_num = 4
69
+ patch_x = int(AT1.size(1) / patch_num)
70
+ patch_y = int(AT1.size(0) / patch_num)
71
+ n_img = torch.randn(AT1.shape).cuda()
72
+ overlap = patch_size // 2 + 1
73
+ for x in range(int(patch_num)):
74
+ for y in range(int(patch_num)):
75
+ left_flg = 0 if x == 0 else 1
76
+ right_flg = 0 if x == patch_num -1 else 1
77
+ top_flg = 0 if y == 0 else 1
78
+ btm_flg = 0 if y == patch_num - 1 else 1
79
+ at1 = AT1[y * patch_y - top_flg * overlap:(y + 1) * patch_y + btm_flg * overlap,
80
+ x * patch_x - left_flg * overlap:(x + 1) * patch_x + right_flg * overlap]
81
+ ata = ATA[y * patch_y - top_flg * overlap:(y + 1) * patch_y + btm_flg * overlap,
82
+ x * patch_x - left_flg * overlap:(x + 1) * patch_x + right_flg * overlap]
83
+ n_img_tmp, _ = torch.solve(at1, ata)
84
+
85
+ n_img_tmp_select = n_img_tmp[top_flg * overlap:patch_y + top_flg * overlap, left_flg * overlap:patch_x + left_flg * overlap, :, :]
86
+ n_img[y * patch_y:y * patch_y + patch_y, x * patch_x:x * patch_x + patch_x, :, :] = n_img_tmp_select
87
+
88
+ n_img_L2 = torch.sqrt(torch.sum(n_img ** 2, dim=2, keepdim=True))
89
+ n_img_norm = n_img / n_img_L2
90
+
91
+ # re-orient normals consistently
92
+ orient_mask = torch.sum(torch.squeeze(n_img_norm) * torch.squeeze(xyz), dim=2) > 0
93
+ n_img_norm[orient_mask] *= -1
94
+ return n_img_norm
95
+
96
+ def get_surface_normalv2(xyz, patch_size=3):
97
+ """
98
+ xyz: xyz coordinates
99
+ patch: [p1, p2, p3,
100
+ p4, p5, p6,
101
+ p7, p8, p9]
102
+ surface_normal = [(p9-p1) x (p3-p7)] + [(p6-p4) - (p8-p2)]
103
+ return: normal [h, w, 3, b]
104
+ """
105
+ b, h, w, c = xyz.shape
106
+ half_patch = patch_size // 2
107
+ xyz_pad = torch.zeros((b, h + patch_size - 1, w + patch_size - 1, c), dtype=xyz.dtype, device=xyz.device)
108
+ xyz_pad[:, half_patch:-half_patch, half_patch:-half_patch, :] = xyz
109
+
110
+ # xyz_left_top = xyz_pad[:, :h, :w, :] # p1
111
+ # xyz_right_bottom = xyz_pad[:, -h:, -w:, :]# p9
112
+ # xyz_left_bottom = xyz_pad[:, -h:, :w, :] # p7
113
+ # xyz_right_top = xyz_pad[:, :h, -w:, :] # p3
114
+ # xyz_cross1 = xyz_left_top - xyz_right_bottom # p1p9
115
+ # xyz_cross2 = xyz_left_bottom - xyz_right_top # p7p3
116
+
117
+ xyz_left = xyz_pad[:, half_patch:half_patch + h, :w, :] # p4
118
+ xyz_right = xyz_pad[:, half_patch:half_patch + h, -w:, :] # p6
119
+ xyz_top = xyz_pad[:, :h, half_patch:half_patch + w, :] # p2
120
+ xyz_bottom = xyz_pad[:, -h:, half_patch:half_patch + w, :] # p8
121
+ xyz_horizon = xyz_left - xyz_right # p4p6
122
+ xyz_vertical = xyz_top - xyz_bottom # p2p8
123
+
124
+ xyz_left_in = xyz_pad[:, half_patch:half_patch + h, 1:w+1, :] # p4
125
+ xyz_right_in = xyz_pad[:, half_patch:half_patch + h, patch_size-1:patch_size-1+w, :] # p6
126
+ xyz_top_in = xyz_pad[:, 1:h+1, half_patch:half_patch + w, :] # p2
127
+ xyz_bottom_in = xyz_pad[:, patch_size-1:patch_size-1+h, half_patch:half_patch + w, :] # p8
128
+ xyz_horizon_in = xyz_left_in - xyz_right_in # p4p6
129
+ xyz_vertical_in = xyz_top_in - xyz_bottom_in # p2p8
130
+
131
+ n_img_1 = torch.cross(xyz_horizon_in, xyz_vertical_in, dim=3)
132
+ n_img_2 = torch.cross(xyz_horizon, xyz_vertical, dim=3)
133
+
134
+ # re-orient normals consistently
135
+ orient_mask = torch.sum(n_img_1 * xyz, dim=3) > 0
136
+ n_img_1[orient_mask] *= -1
137
+ orient_mask = torch.sum(n_img_2 * xyz, dim=3) > 0
138
+ n_img_2[orient_mask] *= -1
139
+
140
+ n_img1_L2 = torch.sqrt(torch.sum(n_img_1 ** 2, dim=3, keepdim=True))
141
+ n_img1_norm = n_img_1 / (n_img1_L2 + 1e-8)
142
+
143
+ n_img2_L2 = torch.sqrt(torch.sum(n_img_2 ** 2, dim=3, keepdim=True))
144
+ n_img2_norm = n_img_2 / (n_img2_L2 + 1e-8)
145
+
146
+ # average 2 norms
147
+ n_img_aver = n_img1_norm + n_img2_norm
148
+ n_img_aver_L2 = torch.sqrt(torch.sum(n_img_aver ** 2, dim=3, keepdim=True))
149
+ n_img_aver_norm = n_img_aver / (n_img_aver_L2 + 1e-8)
150
+ # re-orient normals consistently
151
+ orient_mask = torch.sum(n_img_aver_norm * xyz, dim=3) > 0
152
+ n_img_aver_norm[orient_mask] *= -1
153
+ n_img_aver_norm_out = n_img_aver_norm.permute((1, 2, 3, 0)) # [h, w, c, b]
154
+
155
+ # a = torch.sum(n_img1_norm_out*n_img2_norm_out, dim=2).cpu().numpy().squeeze()
156
+ # plt.imshow(np.abs(a), cmap='rainbow')
157
+ # plt.show()
158
+ return n_img_aver_norm_out#n_img1_norm.permute((1, 2, 3, 0))
159
+
160
+ def surface_normal_from_depth(depth, focal_length, valid_mask=None):
161
+ # para depth: depth map, [b, c, h, w]
162
+ b, c, h, w = depth.shape
163
+ focal_length = focal_length[:, None, None, None]
164
+ depth_filter = nn.functional.avg_pool2d(depth, kernel_size=3, stride=1, padding=1)
165
+ depth_filter = nn.functional.avg_pool2d(depth_filter, kernel_size=3, stride=1, padding=1)
166
+ xyz = depth_to_xyz(depth_filter, focal_length)
167
+ sn_batch = []
168
+ for i in range(b):
169
+ xyz_i = xyz[i, :][None, :, :, :]
170
+ normal = get_surface_normalv2(xyz_i)
171
+ sn_batch.append(normal)
172
+ sn_batch = torch.cat(sn_batch, dim=3).permute((3, 2, 0, 1)) # [b, c, h, w]
173
+ mask_invalid = (~valid_mask).repeat(1, 3, 1, 1)
174
+ sn_batch[mask_invalid] = 0.0
175
+
176
+ return sn_batch
177
+
178
+
179
+ def vis_normal(normal):
180
+ """
181
+ Visualize surface normal. Transfer surface normal value from [-1, 1] to [0, 255]
182
+ @para normal: surface normal, [h, w, 3], numpy.array
183
+ """
184
+ n_img_L2 = np.sqrt(np.sum(normal ** 2, axis=2, keepdims=True))
185
+ n_img_norm = normal / (n_img_L2 + 1e-8)
186
+ normal_vis = n_img_norm * 127
187
+ normal_vis += 128
188
+ normal_vis = normal_vis.astype(np.uint8)
189
+ return normal_vis
190
+
191
+ def vis_normal2(normals):
192
+ '''
193
+ Montage of normal maps. Vectors are unit length and backfaces thresholded.
194
+ '''
195
+ x = normals[:, :, 0] # horizontal; pos right
196
+ y = normals[:, :, 1] # depth; pos far
197
+ z = normals[:, :, 2] # vertical; pos up
198
+ backfacing = (z > 0)
199
+ norm = np.sqrt(np.sum(normals**2, axis=2))
200
+ zero = (norm < 1e-5)
201
+ x += 1.0; x *= 0.5
202
+ y += 1.0; y *= 0.5
203
+ z = np.abs(z)
204
+ x[zero] = 0.0
205
+ y[zero] = 0.0
206
+ z[zero] = 0.0
207
+ normals[:, :, 0] = x # horizontal; pos right
208
+ normals[:, :, 1] = y # depth; pos far
209
+ normals[:, :, 2] = z # vertical; pos up
210
+ return normals
211
+
212
+ if __name__ == '__main__':
213
+ import cv2, os