Spaces:
Running
on
Zero
Running
on
Zero
Upload 22 files
Browse files- GeoWizard/README.md +1 -0
- GeoWizard/geowizard/models/attention.py +777 -0
- GeoWizard/geowizard/models/geowizard_pipeline.py +383 -0
- GeoWizard/geowizard/models/transformer_2d.py +463 -0
- GeoWizard/geowizard/models/unet_2d_blocks.py +0 -0
- GeoWizard/geowizard/models/unet_2d_condition.py +1222 -0
- GeoWizard/geowizard/utils/README.txt +2 -0
- GeoWizard/geowizard/utils/alignment.py +73 -0
- GeoWizard/geowizard/utils/batch_size.py +63 -0
- GeoWizard/geowizard/utils/colormap.py +45 -0
- GeoWizard/geowizard/utils/common.py +42 -0
- GeoWizard/geowizard/utils/dataset_configuration.py +81 -0
- GeoWizard/geowizard/utils/de_normalized.py +31 -0
- GeoWizard/geowizard/utils/depth2normal.py +186 -0
- GeoWizard/geowizard/utils/depth_ensemble.py +115 -0
- GeoWizard/geowizard/utils/depth_transform.py +99 -0
- GeoWizard/geowizard/utils/image_util.py +83 -0
- GeoWizard/geowizard/utils/metric.py +158 -0
- GeoWizard/geowizard/utils/noise.py +19 -0
- GeoWizard/geowizard/utils/normal_ensemble.py +22 -0
- GeoWizard/geowizard/utils/seed_all.py +33 -0
- GeoWizard/geowizard/utils/surface_normal.py +213 -0
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
|