Bethie commited on
Commit
67f650b
·
verified ·
1 Parent(s): 7dfb762

Code convert ONNX

Browse files
code_convert/convert_cnext.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from models.controlnet import ControlNetModel
2
+ from safetensors.torch import load_file
3
+ import torch
4
+
5
+ def load_safetensors(model, safetensors_path, strict=True, load_weight_increasement=False):
6
+ if not load_weight_increasement:
7
+ state_dict = load_file(safetensors_path)
8
+ model.load_state_dict(state_dict, strict=strict)
9
+ else:
10
+ state_dict = load_file(safetensors_path)
11
+ pretrained_state_dict = model.state_dict()
12
+ for k in state_dict.keys():
13
+ state_dict[k] = state_dict[k] + pretrained_state_dict[k]
14
+ model.load_state_dict(state_dict, strict=False)
15
+
16
+ controlnet = ControlNetModel()
17
+
18
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
19
+ controlnet.to(device)
20
+
21
+ load_safetensors(controlnet, '/home/ControlNeXt/ControlNeXt-SDXL/controlnet.safetensors')
22
+
23
+ image = torch.randn((1, 3, 1024, 1024), dtype=torch.float32).to(device)
24
+ timestep = torch.rand(1, dtype= torch.float32).to(device)
25
+
26
+ dummy_inputs = (image, timestep)
27
+
28
+ onnx_output_path = '/home/new_onnx/cnext/model.onnx'
29
+ torch.onnx.export(
30
+ controlnet,
31
+ dummy_inputs,
32
+ onnx_output_path,
33
+ export_params=True,
34
+ opset_version=18,
35
+ do_constant_folding=True,
36
+ input_names=['controlnext_image', 'timestep'],
37
+ output_names=['sample'],
38
+ dynamic_axes={
39
+ 'controlnext_image': {0: 'batch_size', 2: 'height', 3: 'width'},
40
+ 'sample': {0: 'batch_size'},
41
+ }
42
+ )
code_convert/convert_image_encoder.py ADDED
@@ -0,0 +1,187 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # import os
2
+ # from pathlib import Path
3
+ # from PIL import Image
4
+ # import onnx
5
+ # import onnx_graphsurgeon as gs
6
+ # import torch
7
+ # from onnx import shape_inference
8
+ # from packaging import version
9
+ # from polygraphy.backend.onnx.loader import fold_constants
10
+ # from torch.onnx import export
11
+
12
+ # from transformers import CLIPVisionModelWithProjection
13
+ # from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor
14
+
15
+ # is_torch_less_than_1_11 = version.parse(version.parse(torch.__version__).base_version) < version.parse("1.11")
16
+ # is_torch_2_0_1 = version.parse(version.parse(torch.__version__).base_version) == version.parse("2.0.1")
17
+
18
+ # class Optimizer:
19
+ # def __init__(self, onnx_graph, verbose=False):
20
+ # self.graph = gs.import_onnx(onnx_graph)
21
+ # self.verbose = verbose
22
+
23
+ # def info(self, prefix):
24
+ # if self.verbose:
25
+ # print(
26
+ # f"{prefix} .. {len(self.graph.nodes)} nodes, {len(self.graph.tensors().keys())} tensors, {len(self.graph.inputs)} inputs, {len(self.graph.outputs)} outputs"
27
+ # )
28
+
29
+ # def cleanup(self, return_onnx=False):
30
+ # self.graph.cleanup().toposort()
31
+ # if return_onnx:
32
+ # return gs.export_onnx(self.graph)
33
+
34
+ # def select_outputs(self, keep, names=None):
35
+ # self.graph.outputs = [self.graph.outputs[o] for o in keep]
36
+ # if names:
37
+ # for i, name in enumerate(names):
38
+ # self.graph.outputs[i].name = name
39
+
40
+ # def fold_constants(self, return_onnx=False):
41
+ # onnx_graph = fold_constants(gs.export_onnx(self.graph), allow_onnxruntime_shape_inference=True)
42
+ # self.graph = gs.import_onnx(onnx_graph)
43
+ # if return_onnx:
44
+ # return onnx_graph
45
+
46
+ # def infer_shapes(self, return_onnx=False):
47
+ # onnx_graph = gs.export_onnx(self.graph)
48
+ # if onnx_graph.ByteSize() > 4147483648:
49
+ # raise TypeError("ERROR: model size exceeds supported 2GB limit")
50
+ # else:
51
+ # onnx_graph = shape_inference.infer_shapes(onnx_graph)
52
+
53
+ # self.graph = gs.import_onnx(onnx_graph)
54
+ # if return_onnx:
55
+ # return onnx_graph
56
+
57
+
58
+ # def optimize(onnx_graph, name, verbose):
59
+ # opt = Optimizer(onnx_graph, verbose=verbose)
60
+ # opt.info(name + ": original")
61
+ # opt.cleanup()
62
+ # opt.info(name + ": cleanup")
63
+ # opt.fold_constants()
64
+ # opt.info(name + ": fold constants")
65
+ # # opt.infer_shapes()
66
+ # # opt.info(name + ': shape inference')
67
+ # onnx_opt_graph = opt.cleanup(return_onnx=True)
68
+ # opt.info(name + ": finished")
69
+ # return onnx_opt_graph
70
+
71
+
72
+ # class CLIPVisionProj(torch.nn.Module):
73
+ # def __init__(self, clip_model) -> None:
74
+ # super().__init__()
75
+ # self.clip_model = clip_model
76
+
77
+ # def forward(self, image_embedding):
78
+ # result = self.clip_model(image_embedding,return_dict = False)
79
+ # return result[0]
80
+
81
+ # def onnx_export(
82
+ # model,
83
+ # model_args: tuple,
84
+ # output_path: Path,
85
+ # ordered_input_names,
86
+ # output_names,
87
+ # dynamic_axes,
88
+ # opset: int,
89
+ # use_external_data_format=False,
90
+ # verbose=False, # Thêm tham số verbose
91
+ # ):
92
+ # output_path.parent.mkdir(parents=True, exist_ok=True)
93
+ # with torch.inference_mode(), torch.autocast("cuda"):
94
+ # if is_torch_less_than_1_11:
95
+ # export(
96
+ # model,
97
+ # model_args,
98
+ # f=output_path.as_posix(),
99
+ # input_names=ordered_input_names,
100
+ # output_names=output_names,
101
+ # dynamic_axes=dynamic_axes,
102
+ # do_constant_folding=True,
103
+ # use_external_data_format=use_external_data_format,
104
+ # enable_onnx_checker=True,
105
+ # opset_version=opset,
106
+ # verbose=verbose, # Thêm verbose ở đây
107
+ # )
108
+ # else:
109
+ # export(
110
+ # model,
111
+ # model_args,
112
+ # f=output_path.as_posix(),
113
+ # input_names=ordered_input_names,
114
+ # output_names=output_names,
115
+ # dynamic_axes=dynamic_axes,
116
+ # do_constant_folding=True,
117
+ # opset_version=opset,
118
+ # verbose=verbose, # Thêm verbose ở đây
119
+ # )
120
+
121
+ # def convert_models(
122
+ # image_path:str,
123
+ # output_path:str,
124
+ # opset:int=16,
125
+ # ):
126
+ # dtype = torch.float32
127
+ # device = 'cpu'
128
+ # image = Image.open(image_path)
129
+ # image_encoder_processor = CLIPImageProcessor()
130
+ # image_embedding = image_encoder_processor(image, return_tensors="pt").pixel_values
131
+ # clip_model= CLIPVisionModelWithProjection.from_pretrained("h94/IP-Adapter", subfolder = 'sdxl_models/image_encoder')
132
+ # image_encoder = CLIPVisionProj(clip_model).to(device=device)
133
+ # output_path = Path(output_path)
134
+
135
+ # clip_path = output_path / "clip_vision_proj" / "model.onnx"
136
+ # clip_optimize = output_path / 'clip_vision_proj' / 'optimize' / 'model.onnx'
137
+ # #create folder for optimize clip
138
+ # os.makedirs(output_path / 'optimize', exist_ok= True)
139
+ # onnx_export(image_encoder,
140
+ # model_args= (image_embedding).to(dtype = torch.float32, device = device),
141
+ # output_path =clip_path,
142
+ # ordered_input_names= ['image_embedding'],
143
+ # output_names=["image_encoder"],
144
+ # dynamic_axes={'image_embedding': {0: 'Batch_size', 1: 'channel', 2: 'height', 3:'width'},
145
+ # 'image_encoder': {0:'Batch_size', 1: 'sequence_length'} },
146
+ # opset=opset,
147
+ # verbose=True,
148
+ # use_external_data_format=True,
149
+ # )
150
+ # clip_opt_graph = onnx.load(clip_path)
151
+ # onnx.save_model(
152
+ # clip_opt_graph,
153
+ # clip_optimize,
154
+ # save_as_external_data=True,
155
+ # all_tensors_to_one_file=True,
156
+ # convert_attribute=False,
157
+ # location="weights.pb",
158
+ # )
159
+
160
+ # if __name__ == "__main__":
161
+ # convert_models(image_path= '/home/SDXL_CNextAnimeCanny_IPAdapter_ONNX/code_inference/image_condition/control_canny_edge/condition_0.jpg',
162
+ # output_path='/home/new_onnx/image_encoder',
163
+ # opset=18,
164
+ # )
165
+
166
+
167
+ import onnx
168
+
169
+ def optimize_onnx_model(clip_path, output_model_path, weight_file="weights.pb"):
170
+ # Load the ONNX model từ đường dẫn clip_path
171
+ clip_opt_graph = onnx.load(clip_path)
172
+
173
+ # Save optimized model, với weights được lưu riêng vào weight_file
174
+ onnx.save_model(
175
+ clip_opt_graph,
176
+ output_model_path,
177
+ save_as_external_data=True, # Lưu dữ liệu tensor lớn ra ngoài
178
+ all_tensors_to_one_file=True, # Gom tất cả tensor vào một file duy nhất
179
+ location=weight_file, # Đường dẫn đến file lưu weights
180
+ )
181
+
182
+ # Sử dụng hàm
183
+ clip_path = "/home/new_onnx/image_encoder/clip_vision_proj/model.onnx" # Đường dẫn model đã convert
184
+ output_model_path = "/home/new_onnx/image_encoder/optimize/model.onnx" # Đường dẫn file model bạn muốn lưu
185
+ weight_file = "weights.pb" # Tên file chứa weights
186
+ optimize_onnx_model(clip_path, output_model_path, weight_file)
187
+
code_convert/convert_proj.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from huggingface_hub import hf_hub_download
3
+
4
+ class ImageProjModel(torch.nn.Module):
5
+ def __init__(self, cross_attention_dim=2048, clip_embeddings_dim=1280, clip_extra_context_tokens=4):
6
+ super().__init__()
7
+
8
+ self.generator = None
9
+ self.cross_attention_dim = cross_attention_dim
10
+ self.clip_extra_context_tokens = clip_extra_context_tokens
11
+ self.proj = torch.nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim)
12
+ self.norm = torch.nn.LayerNorm(cross_attention_dim)
13
+
14
+ def forward(self, image_embeds):
15
+ embeds = image_embeds
16
+ clip_extra_context_tokens = self.proj(embeds).reshape(
17
+ -1, self.clip_extra_context_tokens, self.cross_attention_dim
18
+ )
19
+ clip_extra_context_tokens = self.norm(clip_extra_context_tokens)
20
+ return clip_extra_context_tokens
21
+
22
+ image_proj_model = ImageProjModel()
23
+
24
+ model_filename = hf_hub_download(repo_id="h94/IP-Adapter", filename="sdxl_models/ip-adapter_sdxl.bin")
25
+ state_dict = torch.load(model_filename, map_location="cpu", weights_only=True)
26
+
27
+ image_proj_model.load_state_dict(state_dict["image_proj"])
28
+
29
+ clip_image_embeds = torch.rand((1, 1280))
30
+
31
+ onnx_output_path = '/home/new_onnx/proj/model.onnx'
32
+ torch.onnx.export(
33
+ image_proj_model,
34
+ clip_image_embeds,
35
+ onnx_output_path,
36
+ export_params=True,
37
+ opset_version=16,
38
+ do_constant_folding=True,
39
+ input_names=['clip_image_embeds'],
40
+ output_names=['image_prompt_embeds'],
41
+ dynamic_axes={
42
+ 'clip_image_embeds': {0: 'batch_size', 1:'embed_size'},
43
+ 'image_prompt_embeds': {0: 'batch_size'},
44
+ })
code_convert/convert_unet.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from models.unet import UNet2DConditionModel
2
+ import torch
3
+ from ip_adapter import IPAdapterXL
4
+ from safetensors.torch import load_file
5
+ import onnx
6
+ from pathlib import Path
7
+
8
+ output_path = '/home/new_onnx/unet'
9
+ output_path = Path(output_path)
10
+
11
+ unet = UNet2DConditionModel.from_pretrained(
12
+ "neta-art/neta-xl-2.0",
13
+ subfolder="unet",
14
+ )
15
+
16
+ state_dict = load_file('/home/ControlNeXt/ControlNeXt-SDXL/unet.safetensors')
17
+ unet.load_state_dict(state_dict, strict=False)
18
+
19
+ image_encoder_path = "h94/IP-Adapter"
20
+ ip_ckpt = "h94/IP-Adapter"
21
+ device = 'cpu'
22
+ ip_model = IPAdapterXL(unet, image_encoder_path, ip_ckpt, device, num_tokens=4)
23
+
24
+ unet = ip_model.unet
25
+
26
+ sample = torch.randn((1, 4, 128, 128))
27
+ timestep = torch.rand(1, dtype=torch.float32)
28
+ encoder_hidden_state = torch.randn((1, 81, 2048))
29
+ mid_block_additional_residual_scale = torch.tensor([1], dtype=torch.float32)
30
+ mid_block_additional_residual = torch.randn((1, 320, 128, 128), dtype=torch.float32)
31
+
32
+ dummy_inputs = (sample, timestep, encoder_hidden_state, mid_block_additional_residual, mid_block_additional_residual_scale)
33
+
34
+ onnx_output_path = output_path / "unet" / "model.onnx"
35
+ torch.onnx.export(
36
+ unet,
37
+ dummy_inputs,
38
+ str(onnx_output_path), # Đường dẫn dưới dạng chuỗi để đảm bảo tương thích
39
+ export_params=True,
40
+ opset_version=18,
41
+ do_constant_folding=True,
42
+ input_names=['sample', 'timestep', 'encoder_hidden_state', 'control_out', 'control_scale'],
43
+ output_names=['predict_noise'],
44
+ dynamic_axes={
45
+ "sample": {0: "B"},
46
+ "encoder_hidden_state": {0: "B", 1: "1B", 2: '2B'},
47
+ "control_out": {0: "B"},
48
+ "predict_noise": {0: 'B'}
49
+ }
50
+ )
51
+
52
+ unet_opt_graph = onnx.load(str(onnx_output_path))
53
+ unet_optimize = output_path / "unet_optimize" / "model.onnx"
54
+ onnx.save_model(
55
+ unet_opt_graph,
56
+ str(unet_optimize),
57
+ save_as_external_data=True,
58
+ all_tensors_to_one_file=True,
59
+ location="weights.pb",
60
+ )
code_convert/models/controlnet.py ADDED
@@ -0,0 +1,492 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2023 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ from dataclasses import dataclass
15
+ from typing import Any, Dict, List, Optional, Tuple, Union
16
+
17
+ import torch
18
+ from torch import nn
19
+
20
+ from diffusers.configuration_utils import ConfigMixin, register_to_config
21
+ from diffusers.utils import BaseOutput, logging
22
+ from diffusers.models.embeddings import TimestepEmbedding, Timesteps
23
+ from diffusers.models.modeling_utils import ModelMixin
24
+ from diffusers.models.resnet import Downsample2D, ResnetBlock2D
25
+ from einops import rearrange
26
+
27
+
28
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
29
+
30
+
31
+ @dataclass
32
+ class ControlNetOutput(BaseOutput):
33
+ """
34
+ The output of [`ControlNetModel`].
35
+
36
+ Args:
37
+ down_block_res_samples (`tuple[torch.Tensor]`):
38
+ A tuple of downsample activations at different resolutions for each downsampling block. Each tensor should
39
+ be of shape `(batch_size, channel * resolution, height //resolution, width // resolution)`. Output can be
40
+ used to condition the original UNet's downsampling activations.
41
+ mid_down_block_re_sample (`torch.Tensor`):
42
+ The activation of the midde block (the lowest sample resolution). Each tensor should be of shape
43
+ `(batch_size, channel * lowest_resolution, height // lowest_resolution, width // lowest_resolution)`.
44
+ Output can be used to condition the original UNet's middle block activation.
45
+ """
46
+
47
+ down_block_res_samples: Tuple[torch.Tensor]
48
+ mid_block_res_sample: torch.Tensor
49
+
50
+
51
+ class Block2D(nn.Module):
52
+ def __init__(
53
+ self,
54
+ in_channels: int,
55
+ out_channels: int,
56
+ temb_channels: int,
57
+ dropout: float = 0.0,
58
+ num_layers: int = 1,
59
+ resnet_eps: float = 1e-6,
60
+ resnet_time_scale_shift: str = "default",
61
+ resnet_act_fn: str = "swish",
62
+ resnet_groups: int = 32,
63
+ resnet_pre_norm: bool = True,
64
+ output_scale_factor: float = 1.0,
65
+ add_downsample: bool = True,
66
+ downsample_padding: int = 1,
67
+ ):
68
+ super().__init__()
69
+ resnets = []
70
+
71
+ for i in range(num_layers):
72
+ in_channels = in_channels if i == 0 else out_channels
73
+ resnets.append(
74
+ ResnetBlock2D(
75
+ in_channels=in_channels,
76
+ out_channels=out_channels,
77
+ temb_channels=temb_channels,
78
+ eps=resnet_eps,
79
+ groups=resnet_groups,
80
+ dropout=dropout,
81
+ time_embedding_norm=resnet_time_scale_shift,
82
+ non_linearity=resnet_act_fn,
83
+ output_scale_factor=output_scale_factor,
84
+ pre_norm=resnet_pre_norm,
85
+ )
86
+ )
87
+
88
+ self.resnets = nn.ModuleList(resnets)
89
+
90
+ if add_downsample:
91
+ self.downsamplers = nn.ModuleList(
92
+ [
93
+ Downsample2D(
94
+ out_channels,
95
+ use_conv=True,
96
+ out_channels=out_channels,
97
+ padding=downsample_padding,
98
+ name="op",
99
+ )
100
+ ]
101
+ )
102
+ else:
103
+ self.downsamplers = None
104
+
105
+ self.gradient_checkpointing = False
106
+
107
+ def forward(
108
+ self,
109
+ hidden_states: torch.FloatTensor,
110
+ temb: Optional[torch.FloatTensor] = None,
111
+ ) -> Union[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
112
+ output_states = ()
113
+
114
+ for resnet in zip(self.resnets):
115
+ hidden_states = resnet(hidden_states, temb)
116
+ output_states += (hidden_states,)
117
+
118
+ if self.downsamplers is not None:
119
+ for downsampler in self.downsamplers:
120
+ hidden_states = downsampler(hidden_states)
121
+
122
+ output_states += (hidden_states,)
123
+
124
+ return hidden_states, output_states
125
+
126
+
127
+ class IdentityModule(nn.Module):
128
+ def __init__(self):
129
+ super(IdentityModule, self).__init__()
130
+
131
+ def forward(self, *args):
132
+ if len(args) > 0:
133
+ return args[0]
134
+ else:
135
+ return None
136
+
137
+
138
+ class BasicBlock(nn.Module):
139
+ def __init__(self,
140
+ in_channels: int,
141
+ out_channels: Optional[int] = None,
142
+ stride=1,
143
+ conv_shortcut: bool = False,
144
+ dropout: float = 0.0,
145
+ temb_channels: int = 512,
146
+ groups: int = 32,
147
+ groups_out: Optional[int] = None,
148
+ pre_norm: bool = True,
149
+ eps: float = 1e-6,
150
+ non_linearity: str = "swish",
151
+ skip_time_act: bool = False,
152
+ time_embedding_norm: str = "default", # default, scale_shift, ada_group, spatial
153
+ kernel: Optional[torch.FloatTensor] = None,
154
+ output_scale_factor: float = 1.0,
155
+ use_in_shortcut: Optional[bool] = None,
156
+ up: bool = False,
157
+ down: bool = False,
158
+ conv_shortcut_bias: bool = True,
159
+ conv_2d_out_channels: Optional[int] = None,):
160
+ super(BasicBlock, self).__init__()
161
+ self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False)
162
+ self.bn1 = nn.BatchNorm2d(out_channels)
163
+ self.relu = nn.ReLU(inplace=True)
164
+ self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
165
+ self.bn2 = nn.BatchNorm2d(out_channels)
166
+
167
+ self.downsample = None
168
+ if stride != 1 or in_channels != out_channels:
169
+ self.downsample = nn.Sequential(
170
+ nn.Conv2d(in_channels,
171
+ out_channels,
172
+ kernel_size=3 if stride != 1 else 1,
173
+ stride=stride,
174
+ padding=1 if stride != 1 else 0,
175
+ bias=False),
176
+ nn.BatchNorm2d(out_channels)
177
+ )
178
+
179
+ def forward(self, x, *args):
180
+ residual = x
181
+ out = self.conv1(x)
182
+ out = self.bn1(out)
183
+ out = self.relu(out)
184
+
185
+ out = self.conv2(out)
186
+ out = self.bn2(out)
187
+
188
+ if self.downsample is not None:
189
+ residual = self.downsample(x)
190
+
191
+ out += residual
192
+ out = self.relu(out)
193
+
194
+ return out
195
+
196
+
197
+ class Block2D(nn.Module):
198
+ def __init__(
199
+ self,
200
+ in_channels: int,
201
+ out_channels: int,
202
+ temb_channels: int,
203
+ dropout: float = 0.0,
204
+ num_layers: int = 1,
205
+ resnet_eps: float = 1e-6,
206
+ resnet_time_scale_shift: str = "default",
207
+ resnet_act_fn: str = "swish",
208
+ resnet_groups: int = 32,
209
+ resnet_pre_norm: bool = True,
210
+ output_scale_factor: float = 1.0,
211
+ add_downsample: bool = True,
212
+ downsample_padding: int = 1,
213
+ ):
214
+ super().__init__()
215
+ resnets = []
216
+
217
+ for i in range(num_layers):
218
+ # in_channels = in_channels if i == 0 else out_channels
219
+ resnets.append(
220
+ # ResnetBlock2D(
221
+ # in_channels=in_channels,
222
+ # out_channels=out_channels,
223
+ # temb_channels=temb_channels,
224
+ # eps=resnet_eps,
225
+ # groups=resnet_groups,
226
+ # dropout=dropout,
227
+ # time_embedding_norm=resnet_time_scale_shift,
228
+ # non_linearity=resnet_act_fn,
229
+ # output_scale_factor=output_scale_factor,
230
+ # pre_norm=resnet_pre_norm,
231
+ BasicBlock(
232
+ in_channels=in_channels,
233
+ out_channels=out_channels,
234
+ temb_channels=temb_channels,
235
+ eps=resnet_eps,
236
+ groups=resnet_groups,
237
+ dropout=dropout,
238
+ time_embedding_norm=resnet_time_scale_shift,
239
+ non_linearity=resnet_act_fn,
240
+ output_scale_factor=output_scale_factor,
241
+ pre_norm=resnet_pre_norm,
242
+ ) if i == num_layers - 1 else \
243
+ IdentityModule()
244
+ )
245
+
246
+ self.resnets = nn.ModuleList(resnets)
247
+
248
+ if add_downsample:
249
+ self.downsamplers = nn.ModuleList(
250
+ [
251
+ # Downsample2D(
252
+ # out_channels,
253
+ # use_conv=True,
254
+ # out_channels=out_channels,
255
+ # padding=downsample_padding,
256
+ # name="op",
257
+ # )
258
+ BasicBlock(
259
+ in_channels=out_channels,
260
+ out_channels=out_channels,
261
+ temb_channels=temb_channels,
262
+ stride=2,
263
+ eps=resnet_eps,
264
+ groups=resnet_groups,
265
+ dropout=dropout,
266
+ time_embedding_norm=resnet_time_scale_shift,
267
+ non_linearity=resnet_act_fn,
268
+ output_scale_factor=output_scale_factor,
269
+ pre_norm=resnet_pre_norm,
270
+ )
271
+ ]
272
+ )
273
+ else:
274
+ self.downsamplers = None
275
+
276
+ self.gradient_checkpointing = False
277
+
278
+ def forward(
279
+ self,
280
+ hidden_states: torch.FloatTensor,
281
+ temb: Optional[torch.FloatTensor] = None,
282
+ ) -> Union[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
283
+ output_states = ()
284
+
285
+ for resnet in self.resnets:
286
+ hidden_states = resnet(hidden_states, temb)
287
+ output_states += (hidden_states,)
288
+
289
+ if self.downsamplers is not None:
290
+ for downsampler in self.downsamplers:
291
+ hidden_states = downsampler(hidden_states)
292
+
293
+ output_states += (hidden_states,)
294
+
295
+ return hidden_states, output_states
296
+
297
+
298
+ class ControlProject(nn.Module):
299
+ def __init__(self, num_channels, scale=8, is_empty=False) -> None:
300
+ super().__init__()
301
+ assert scale and scale & (scale - 1) == 0
302
+ self.is_empty = is_empty
303
+ self.scale = scale
304
+ if not is_empty:
305
+ if scale > 1:
306
+ self.down_scale = nn.AvgPool2d(scale, scale)
307
+ else:
308
+ self.down_scale = nn.Identity()
309
+ self.out = nn.Conv2d(num_channels, num_channels, kernel_size=1, stride=1, bias=False)
310
+ for p in self.out.parameters():
311
+ nn.init.zeros_(p)
312
+
313
+ def forward(
314
+ self,
315
+ hidden_states: torch.FloatTensor):
316
+ if self.is_empty:
317
+ shape = list(hidden_states.shape)
318
+ shape[-2] = shape[-2] // self.scale
319
+ shape[-1] = shape[-1] // self.scale
320
+ return torch.zeros(shape).to(hidden_states)
321
+
322
+ if len(hidden_states.shape) == 5:
323
+ B, F, C, H, W = hidden_states.shape
324
+ hidden_states = rearrange(hidden_states, "B F C H W -> (B F) C H W")
325
+ hidden_states = self.down_scale(hidden_states)
326
+ hidden_states = self.out(hidden_states)
327
+ hidden_states = rearrange(hidden_states, "(B F) C H W -> B F C H W", F=F)
328
+ else:
329
+ hidden_states = self.down_scale(hidden_states)
330
+ hidden_states = self.out(hidden_states)
331
+ return hidden_states
332
+
333
+
334
+ class ControlNetModel(ModelMixin, ConfigMixin):
335
+
336
+ _supports_gradient_checkpointing = True
337
+
338
+ @register_to_config
339
+ def __init__(
340
+ self,
341
+ in_channels: List[int] = [128, 128],
342
+ out_channels: List[int] = [128, 256],
343
+ groups: List[int] = [4, 8],
344
+ time_embed_dim: int = 256,
345
+ final_out_channels: int = 320,
346
+ ):
347
+ super().__init__()
348
+
349
+ self.time_proj = Timesteps(128, True, downscale_freq_shift=0)
350
+ self.time_embedding = TimestepEmbedding(128, time_embed_dim)
351
+
352
+ self.embedding = nn.Sequential(
353
+ nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1),
354
+ nn.GroupNorm(2, 64),
355
+ nn.ReLU(),
356
+ nn.Conv2d(64, 64, kernel_size=3, padding=1),
357
+ nn.GroupNorm(2, 64),
358
+ nn.ReLU(),
359
+ nn.Conv2d(64, 128, kernel_size=3, padding=1),
360
+ nn.GroupNorm(2, 128),
361
+ nn.ReLU(),
362
+ )
363
+
364
+ self.down_res = nn.ModuleList()
365
+ self.down_sample = nn.ModuleList()
366
+ for i in range(len(in_channels)):
367
+ self.down_res.append(
368
+ ResnetBlock2D(
369
+ in_channels=in_channels[i],
370
+ out_channels=out_channels[i],
371
+ temb_channels=time_embed_dim,
372
+ groups=groups[i]
373
+ ),
374
+ )
375
+ self.down_sample.append(
376
+ Downsample2D(
377
+ out_channels[i],
378
+ use_conv=True,
379
+ out_channels=out_channels[i],
380
+ padding=1,
381
+ name="op",
382
+ )
383
+ )
384
+
385
+ self.mid_convs = nn.ModuleList()
386
+ self.mid_convs.append(nn.Sequential(
387
+ nn.Conv2d(
388
+ in_channels=out_channels[-1],
389
+ out_channels=out_channels[-1],
390
+ kernel_size=3,
391
+ stride=1,
392
+ padding=1
393
+ ),
394
+ nn.ReLU(),
395
+ nn.GroupNorm(8, out_channels[-1]),
396
+ nn.Conv2d(
397
+ in_channels=out_channels[-1],
398
+ out_channels=out_channels[-1],
399
+ kernel_size=3,
400
+ stride=1,
401
+ padding=1
402
+ ),
403
+ nn.GroupNorm(8, out_channels[-1]),
404
+ ))
405
+ self.mid_convs.append(
406
+ nn.Conv2d(
407
+ in_channels=out_channels[-1],
408
+ out_channels=final_out_channels,
409
+ kernel_size=1,
410
+ stride=1,
411
+ ))
412
+ self.scale = 1.0 # nn.Parameter(torch.tensor(1.))
413
+
414
+ def _set_gradient_checkpointing(self, module, value=False):
415
+ if hasattr(module, "gradient_checkpointing"):
416
+ module.gradient_checkpointing = value
417
+
418
+ # Copied from diffusers.models.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
419
+ def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
420
+ """
421
+ Sets the attention processor to use [feed forward
422
+ chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
423
+
424
+ Parameters:
425
+ chunk_size (`int`, *optional*):
426
+ The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually
427
+ over each tensor of dim=`dim`.
428
+ dim (`int`, *optional*, defaults to `0`):
429
+ The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch)
430
+ or dim=1 (sequence length).
431
+ """
432
+ if dim not in [0, 1]:
433
+ raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}")
434
+
435
+ # By default chunk size is 1
436
+ chunk_size = chunk_size or 1
437
+
438
+ def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
439
+ if hasattr(module, "set_chunk_feed_forward"):
440
+ module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
441
+
442
+ for child in module.children():
443
+ fn_recursive_feed_forward(child, chunk_size, dim)
444
+
445
+ for module in self.children():
446
+ fn_recursive_feed_forward(module, chunk_size, dim)
447
+
448
+ def forward(
449
+ self,
450
+ sample: torch.FloatTensor,
451
+ timestep: Union[torch.Tensor, float, int],
452
+ ) -> Union[ControlNetOutput, Tuple]:
453
+
454
+ timesteps = timestep
455
+ if not torch.is_tensor(timesteps):
456
+ # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
457
+ # This would be a good case for the `match` statement (Python 3.10+)
458
+ is_mps = sample.device.type == "mps"
459
+ if isinstance(timestep, float):
460
+ dtype = torch.float32 if is_mps else torch.float64
461
+ else:
462
+ dtype = torch.int32 if is_mps else torch.int64
463
+ timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
464
+ elif len(timesteps.shape) == 0:
465
+ timesteps = timesteps[None].to(sample.device)
466
+
467
+ # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
468
+ batch_size = sample.shape[0]
469
+ timesteps = timesteps.expand(batch_size)
470
+ t_emb = self.time_proj(timesteps)
471
+ # `Timesteps` does not contain any weights and will always return f32 tensors
472
+ # but time_embedding might actually be running in fp16. so we need to cast here.
473
+ # there might be better ways to encapsulate this.
474
+ t_emb = t_emb.to(dtype=sample.dtype)
475
+ emb_batch = self.time_embedding(t_emb)
476
+
477
+ # Repeat the embeddings num_video_frames times
478
+ # emb: [batch, channels] -> [batch * frames, channels]
479
+ emb = emb_batch
480
+ sample = self.embedding(sample)
481
+ for res, downsample in zip(self.down_res, self.down_sample):
482
+ sample = res(sample, emb)
483
+ sample = downsample(sample, emb)
484
+ sample = self.mid_convs[0](sample) + sample
485
+ sample = self.mid_convs[1](sample)
486
+ return sample
487
+
488
+
489
+ def zero_module(module):
490
+ for p in module.parameters():
491
+ nn.init.zeros_(p)
492
+ return module
code_convert/models/unet.py ADDED
@@ -0,0 +1,1308 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from dataclasses import dataclass
2
+ from typing import Any, Dict, List, Optional, Tuple, Union
3
+
4
+ import torch
5
+ import torch.nn as nn
6
+ import torch.utils.checkpoint
7
+
8
+ from diffusers.configuration_utils import ConfigMixin, register_to_config
9
+ from diffusers.loaders import PeftAdapterMixin, UNet2DConditionLoadersMixin
10
+ from diffusers.loaders.single_file_model import FromOriginalModelMixin
11
+ from diffusers.utils import USE_PEFT_BACKEND, BaseOutput, deprecate, logging, scale_lora_layers, unscale_lora_layers
12
+ from diffusers.models.activations import get_activation
13
+ from diffusers.models.attention_processor import (
14
+ ADDED_KV_ATTENTION_PROCESSORS,
15
+ CROSS_ATTENTION_PROCESSORS,
16
+ Attention,
17
+ AttentionProcessor,
18
+ AttnAddedKVProcessor,
19
+ AttnProcessor,
20
+ )
21
+ from diffusers.models.embeddings import (
22
+ GaussianFourierProjection,
23
+ GLIGENTextBoundingboxProjection,
24
+ ImageHintTimeEmbedding,
25
+ ImageProjection,
26
+ ImageTimeEmbedding,
27
+ TextImageProjection,
28
+ TextImageTimeEmbedding,
29
+ TextTimeEmbedding,
30
+ TimestepEmbedding,
31
+ Timesteps,
32
+ )
33
+ from diffusers.models.modeling_utils import ModelMixin
34
+ from diffusers.models.unets.unet_2d_blocks import (
35
+ get_down_block,
36
+ get_mid_block,
37
+ get_up_block,
38
+ )
39
+
40
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
41
+
42
+
43
+ @dataclass
44
+ class UNet2DConditionOutput(BaseOutput):
45
+ """
46
+ The output of [`UNet2DConditionModel`].
47
+
48
+ Args:
49
+ sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)`):
50
+ The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model.
51
+ """
52
+
53
+ sample: torch.Tensor = None
54
+
55
+
56
+ class UNet2DConditionModel(
57
+ ModelMixin, ConfigMixin, FromOriginalModelMixin, UNet2DConditionLoadersMixin, PeftAdapterMixin
58
+ ):
59
+ r"""
60
+ A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample
61
+ shaped output.
62
+
63
+ This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
64
+ for all models (such as downloading or saving).
65
+
66
+ Parameters:
67
+ sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
68
+ Height and width of input/output sample.
69
+ in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample.
70
+ out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
71
+ center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample.
72
+ flip_sin_to_cos (`bool`, *optional*, defaults to `True`):
73
+ Whether to flip the sin to cos in the time embedding.
74
+ freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding.
75
+ down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
76
+ The tuple of downsample blocks to use.
77
+ mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`):
78
+ Block type for middle of UNet, it can be one of `UNetMidBlock2DCrossAttn`, `UNetMidBlock2D`, or
79
+ `UNetMidBlock2DSimpleCrossAttn`. If `None`, the mid block layer is skipped.
80
+ up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`):
81
+ The tuple of upsample blocks to use.
82
+ only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`):
83
+ Whether to include self-attention in the basic transformer blocks, see
84
+ [`~models.attention.BasicTransformerBlock`].
85
+ block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
86
+ The tuple of output channels for each block.
87
+ layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
88
+ downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.
89
+ mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block.
90
+ dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
91
+ act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
92
+ norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.
93
+ If `None`, normalization and activation layers is skipped in post-processing.
94
+ norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.
95
+ cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
96
+ The dimension of the cross attention features.
97
+ transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1):
98
+ The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
99
+ [`~models.unets.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unets.unet_2d_blocks.CrossAttnUpBlock2D`],
100
+ [`~models.unets.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
101
+ reverse_transformer_layers_per_block : (`Tuple[Tuple]`, *optional*, defaults to None):
102
+ The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`], in the upsampling
103
+ blocks of the U-Net. Only relevant if `transformer_layers_per_block` is of type `Tuple[Tuple]` and for
104
+ [`~models.unets.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unets.unet_2d_blocks.CrossAttnUpBlock2D`],
105
+ [`~models.unets.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
106
+ encoder_hid_dim (`int`, *optional*, defaults to None):
107
+ If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim`
108
+ dimension to `cross_attention_dim`.
109
+ encoder_hid_dim_type (`str`, *optional*, defaults to `None`):
110
+ If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text
111
+ embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`.
112
+ attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.
113
+ num_attention_heads (`int`, *optional*):
114
+ The number of attention heads. If not defined, defaults to `attention_head_dim`
115
+ resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config
116
+ for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`.
117
+ class_embed_type (`str`, *optional*, defaults to `None`):
118
+ The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`,
119
+ `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`.
120
+ addition_embed_type (`str`, *optional*, defaults to `None`):
121
+ Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or
122
+ "text". "text" will use the `TextTimeEmbedding` layer.
123
+ addition_time_embed_dim: (`int`, *optional*, defaults to `None`):
124
+ Dimension for the timestep embeddings.
125
+ num_class_embeds (`int`, *optional*, defaults to `None`):
126
+ Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing
127
+ class conditioning with `class_embed_type` equal to `None`.
128
+ time_embedding_type (`str`, *optional*, defaults to `positional`):
129
+ The type of position embedding to use for timesteps. Choose from `positional` or `fourier`.
130
+ time_embedding_dim (`int`, *optional*, defaults to `None`):
131
+ An optional override for the dimension of the projected time embedding.
132
+ time_embedding_act_fn (`str`, *optional*, defaults to `None`):
133
+ Optional activation function to use only once on the time embeddings before they are passed to the rest of
134
+ the UNet. Choose from `silu`, `mish`, `gelu`, and `swish`.
135
+ timestep_post_act (`str`, *optional*, defaults to `None`):
136
+ The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`.
137
+ time_cond_proj_dim (`int`, *optional*, defaults to `None`):
138
+ The dimension of `cond_proj` layer in the timestep embedding.
139
+ conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer.
140
+ conv_out_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_out` layer.
141
+ projection_class_embeddings_input_dim (`int`, *optional*): The dimension of the `class_labels` input when
142
+ `class_embed_type="projection"`. Required when `class_embed_type="projection"`.
143
+ class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time
144
+ embeddings with the class embeddings.
145
+ mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`):
146
+ Whether to use cross attention with the mid block when using the `UNetMidBlock2DSimpleCrossAttn`. If
147
+ `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is `None`, the
148
+ `only_cross_attention` value is used as the value for `mid_block_only_cross_attention`. Default to `False`
149
+ otherwise.
150
+ """
151
+
152
+ _supports_gradient_checkpointing = True
153
+ _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D", "CrossAttnUpBlock2D"]
154
+
155
+ @register_to_config
156
+ def __init__(
157
+ self,
158
+ sample_size: Optional[int] = None,
159
+ in_channels: int = 4,
160
+ out_channels: int = 4,
161
+ center_input_sample: bool = False,
162
+ flip_sin_to_cos: bool = True,
163
+ freq_shift: int = 0,
164
+ down_block_types: Tuple[str] = (
165
+ "CrossAttnDownBlock2D",
166
+ "CrossAttnDownBlock2D",
167
+ "CrossAttnDownBlock2D",
168
+ "DownBlock2D",
169
+ ),
170
+ mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
171
+ up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
172
+ only_cross_attention: Union[bool, Tuple[bool]] = False,
173
+ block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
174
+ layers_per_block: Union[int, Tuple[int]] = 2,
175
+ downsample_padding: int = 1,
176
+ mid_block_scale_factor: float = 1,
177
+ dropout: float = 0.0,
178
+ act_fn: str = "silu",
179
+ norm_num_groups: Optional[int] = 32,
180
+ norm_eps: float = 1e-5,
181
+ cross_attention_dim: Union[int, Tuple[int]] = 1280,
182
+ transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
183
+ reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None,
184
+ encoder_hid_dim: Optional[int] = None,
185
+ encoder_hid_dim_type: Optional[str] = None,
186
+ attention_head_dim: Union[int, Tuple[int]] = 8,
187
+ num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
188
+ dual_cross_attention: bool = False,
189
+ use_linear_projection: bool = False,
190
+ class_embed_type: Optional[str] = None,
191
+ addition_embed_type: Optional[str] = None,
192
+ addition_time_embed_dim: Optional[int] = None,
193
+ num_class_embeds: Optional[int] = None,
194
+ upcast_attention: bool = False,
195
+ resnet_time_scale_shift: str = "default",
196
+ resnet_skip_time_act: bool = False,
197
+ resnet_out_scale_factor: float = 1.0,
198
+ time_embedding_type: str = "positional",
199
+ time_embedding_dim: Optional[int] = None,
200
+ time_embedding_act_fn: Optional[str] = None,
201
+ timestep_post_act: Optional[str] = None,
202
+ time_cond_proj_dim: Optional[int] = None,
203
+ conv_in_kernel: int = 3,
204
+ conv_out_kernel: int = 3,
205
+ projection_class_embeddings_input_dim: Optional[int] = None,
206
+ attention_type: str = "default",
207
+ class_embeddings_concat: bool = False,
208
+ mid_block_only_cross_attention: Optional[bool] = None,
209
+ cross_attention_norm: Optional[str] = None,
210
+ addition_embed_type_num_heads: int = 64,
211
+ ):
212
+ super().__init__()
213
+
214
+ self.sample_size = sample_size
215
+
216
+ if num_attention_heads is not None:
217
+ raise ValueError(
218
+ "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."
219
+ )
220
+
221
+ # If `num_attention_heads` is not defined (which is the case for most models)
222
+ # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
223
+ # The reason for this behavior is to correct for incorrectly named variables that were introduced
224
+ # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
225
+ # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
226
+ # which is why we correct for the naming here.
227
+ num_attention_heads = num_attention_heads or attention_head_dim
228
+
229
+ # Check inputs
230
+ self._check_config(
231
+ down_block_types=down_block_types,
232
+ up_block_types=up_block_types,
233
+ only_cross_attention=only_cross_attention,
234
+ block_out_channels=block_out_channels,
235
+ layers_per_block=layers_per_block,
236
+ cross_attention_dim=cross_attention_dim,
237
+ transformer_layers_per_block=transformer_layers_per_block,
238
+ reverse_transformer_layers_per_block=reverse_transformer_layers_per_block,
239
+ attention_head_dim=attention_head_dim,
240
+ num_attention_heads=num_attention_heads,
241
+ )
242
+
243
+ # input
244
+ conv_in_padding = (conv_in_kernel - 1) // 2
245
+ self.conv_in = nn.Conv2d(
246
+ in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding
247
+ )
248
+
249
+ # time
250
+ time_embed_dim, timestep_input_dim = self._set_time_proj(
251
+ time_embedding_type,
252
+ block_out_channels=block_out_channels,
253
+ flip_sin_to_cos=flip_sin_to_cos,
254
+ freq_shift=freq_shift,
255
+ time_embedding_dim=time_embedding_dim,
256
+ )
257
+
258
+ self.time_embedding = TimestepEmbedding(
259
+ timestep_input_dim,
260
+ time_embed_dim,
261
+ act_fn=act_fn,
262
+ post_act_fn=timestep_post_act,
263
+ cond_proj_dim=time_cond_proj_dim,
264
+ )
265
+
266
+ self._set_encoder_hid_proj(
267
+ encoder_hid_dim_type,
268
+ cross_attention_dim=cross_attention_dim,
269
+ encoder_hid_dim=encoder_hid_dim,
270
+ )
271
+
272
+ # class embedding
273
+ self._set_class_embedding(
274
+ class_embed_type,
275
+ act_fn=act_fn,
276
+ num_class_embeds=num_class_embeds,
277
+ projection_class_embeddings_input_dim=projection_class_embeddings_input_dim,
278
+ time_embed_dim=time_embed_dim,
279
+ timestep_input_dim=timestep_input_dim,
280
+ )
281
+
282
+ self._set_add_embedding(
283
+ addition_embed_type,
284
+ addition_embed_type_num_heads=addition_embed_type_num_heads,
285
+ addition_time_embed_dim=addition_time_embed_dim,
286
+ cross_attention_dim=cross_attention_dim,
287
+ encoder_hid_dim=encoder_hid_dim,
288
+ flip_sin_to_cos=flip_sin_to_cos,
289
+ freq_shift=freq_shift,
290
+ projection_class_embeddings_input_dim=projection_class_embeddings_input_dim,
291
+ time_embed_dim=time_embed_dim,
292
+ )
293
+
294
+ if time_embedding_act_fn is None:
295
+ self.time_embed_act = None
296
+ else:
297
+ self.time_embed_act = get_activation(time_embedding_act_fn)
298
+
299
+ self.down_blocks = nn.ModuleList([])
300
+ self.up_blocks = nn.ModuleList([])
301
+
302
+ if isinstance(only_cross_attention, bool):
303
+ if mid_block_only_cross_attention is None:
304
+ mid_block_only_cross_attention = only_cross_attention
305
+
306
+ only_cross_attention = [only_cross_attention] * len(down_block_types)
307
+
308
+ if mid_block_only_cross_attention is None:
309
+ mid_block_only_cross_attention = False
310
+
311
+ if isinstance(num_attention_heads, int):
312
+ num_attention_heads = (num_attention_heads,) * len(down_block_types)
313
+
314
+ if isinstance(attention_head_dim, int):
315
+ attention_head_dim = (attention_head_dim,) * len(down_block_types)
316
+
317
+ if isinstance(cross_attention_dim, int):
318
+ cross_attention_dim = (cross_attention_dim,) * len(down_block_types)
319
+
320
+ if isinstance(layers_per_block, int):
321
+ layers_per_block = [layers_per_block] * len(down_block_types)
322
+
323
+ if isinstance(transformer_layers_per_block, int):
324
+ transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
325
+
326
+ if class_embeddings_concat:
327
+ # The time embeddings are concatenated with the class embeddings. The dimension of the
328
+ # time embeddings passed to the down, middle, and up blocks is twice the dimension of the
329
+ # regular time embeddings
330
+ blocks_time_embed_dim = time_embed_dim * 2
331
+ else:
332
+ blocks_time_embed_dim = time_embed_dim
333
+
334
+ # down
335
+ output_channel = block_out_channels[0]
336
+ for i, down_block_type in enumerate(down_block_types):
337
+ input_channel = output_channel
338
+ output_channel = block_out_channels[i]
339
+ is_final_block = i == len(block_out_channels) - 1
340
+
341
+ down_block = get_down_block(
342
+ down_block_type,
343
+ num_layers=layers_per_block[i],
344
+ transformer_layers_per_block=transformer_layers_per_block[i],
345
+ in_channels=input_channel,
346
+ out_channels=output_channel,
347
+ temb_channels=blocks_time_embed_dim,
348
+ add_downsample=not is_final_block,
349
+ resnet_eps=norm_eps,
350
+ resnet_act_fn=act_fn,
351
+ resnet_groups=norm_num_groups,
352
+ cross_attention_dim=cross_attention_dim[i],
353
+ num_attention_heads=num_attention_heads[i],
354
+ downsample_padding=downsample_padding,
355
+ dual_cross_attention=dual_cross_attention,
356
+ use_linear_projection=use_linear_projection,
357
+ only_cross_attention=only_cross_attention[i],
358
+ upcast_attention=upcast_attention,
359
+ resnet_time_scale_shift=resnet_time_scale_shift,
360
+ attention_type=attention_type,
361
+ resnet_skip_time_act=resnet_skip_time_act,
362
+ resnet_out_scale_factor=resnet_out_scale_factor,
363
+ cross_attention_norm=cross_attention_norm,
364
+ attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
365
+ dropout=dropout,
366
+ )
367
+ self.down_blocks.append(down_block)
368
+
369
+ # mid
370
+ self.mid_block = get_mid_block(
371
+ mid_block_type,
372
+ temb_channels=blocks_time_embed_dim,
373
+ in_channels=block_out_channels[-1],
374
+ resnet_eps=norm_eps,
375
+ resnet_act_fn=act_fn,
376
+ resnet_groups=norm_num_groups,
377
+ output_scale_factor=mid_block_scale_factor,
378
+ transformer_layers_per_block=transformer_layers_per_block[-1],
379
+ num_attention_heads=num_attention_heads[-1],
380
+ cross_attention_dim=cross_attention_dim[-1],
381
+ dual_cross_attention=dual_cross_attention,
382
+ use_linear_projection=use_linear_projection,
383
+ mid_block_only_cross_attention=mid_block_only_cross_attention,
384
+ upcast_attention=upcast_attention,
385
+ resnet_time_scale_shift=resnet_time_scale_shift,
386
+ attention_type=attention_type,
387
+ resnet_skip_time_act=resnet_skip_time_act,
388
+ cross_attention_norm=cross_attention_norm,
389
+ attention_head_dim=attention_head_dim[-1],
390
+ dropout=dropout,
391
+ )
392
+
393
+ # count how many layers upsample the images
394
+ self.num_upsamplers = 0
395
+
396
+ # up
397
+ reversed_block_out_channels = list(reversed(block_out_channels))
398
+ reversed_num_attention_heads = list(reversed(num_attention_heads))
399
+ reversed_layers_per_block = list(reversed(layers_per_block))
400
+ reversed_cross_attention_dim = list(reversed(cross_attention_dim))
401
+ reversed_transformer_layers_per_block = (
402
+ list(reversed(transformer_layers_per_block))
403
+ if reverse_transformer_layers_per_block is None
404
+ else reverse_transformer_layers_per_block
405
+ )
406
+ only_cross_attention = list(reversed(only_cross_attention))
407
+
408
+ output_channel = reversed_block_out_channels[0]
409
+ for i, up_block_type in enumerate(up_block_types):
410
+ is_final_block = i == len(block_out_channels) - 1
411
+
412
+ prev_output_channel = output_channel
413
+ output_channel = reversed_block_out_channels[i]
414
+ input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
415
+
416
+ # add upsample block for all BUT final layer
417
+ if not is_final_block:
418
+ add_upsample = True
419
+ self.num_upsamplers += 1
420
+ else:
421
+ add_upsample = False
422
+
423
+ up_block = get_up_block(
424
+ up_block_type,
425
+ num_layers=reversed_layers_per_block[i] + 1,
426
+ transformer_layers_per_block=reversed_transformer_layers_per_block[i],
427
+ in_channels=input_channel,
428
+ out_channels=output_channel,
429
+ prev_output_channel=prev_output_channel,
430
+ temb_channels=blocks_time_embed_dim,
431
+ add_upsample=add_upsample,
432
+ resnet_eps=norm_eps,
433
+ resnet_act_fn=act_fn,
434
+ resolution_idx=i,
435
+ resnet_groups=norm_num_groups,
436
+ cross_attention_dim=reversed_cross_attention_dim[i],
437
+ num_attention_heads=reversed_num_attention_heads[i],
438
+ dual_cross_attention=dual_cross_attention,
439
+ use_linear_projection=use_linear_projection,
440
+ only_cross_attention=only_cross_attention[i],
441
+ upcast_attention=upcast_attention,
442
+ resnet_time_scale_shift=resnet_time_scale_shift,
443
+ attention_type=attention_type,
444
+ resnet_skip_time_act=resnet_skip_time_act,
445
+ resnet_out_scale_factor=resnet_out_scale_factor,
446
+ cross_attention_norm=cross_attention_norm,
447
+ attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
448
+ dropout=dropout,
449
+ )
450
+ self.up_blocks.append(up_block)
451
+
452
+ # out
453
+ if norm_num_groups is not None:
454
+ self.conv_norm_out = nn.GroupNorm(
455
+ num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps
456
+ )
457
+
458
+ self.conv_act = get_activation(act_fn)
459
+
460
+ else:
461
+ self.conv_norm_out = None
462
+ self.conv_act = None
463
+
464
+ conv_out_padding = (conv_out_kernel - 1) // 2
465
+ self.conv_out = nn.Conv2d(
466
+ block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding
467
+ )
468
+
469
+ self._set_pos_net_if_use_gligen(attention_type=attention_type, cross_attention_dim=cross_attention_dim)
470
+
471
+ def _check_config(
472
+ self,
473
+ down_block_types: Tuple[str],
474
+ up_block_types: Tuple[str],
475
+ only_cross_attention: Union[bool, Tuple[bool]],
476
+ block_out_channels: Tuple[int],
477
+ layers_per_block: Union[int, Tuple[int]],
478
+ cross_attention_dim: Union[int, Tuple[int]],
479
+ transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple[int]]],
480
+ reverse_transformer_layers_per_block: bool,
481
+ attention_head_dim: int,
482
+ num_attention_heads: Optional[Union[int, Tuple[int]]],
483
+ ):
484
+ if len(down_block_types) != len(up_block_types):
485
+ raise ValueError(
486
+ 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}."
487
+ )
488
+
489
+ if len(block_out_channels) != len(down_block_types):
490
+ raise ValueError(
491
+ 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}."
492
+ )
493
+
494
+ if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):
495
+ raise ValueError(
496
+ 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}."
497
+ )
498
+
499
+ if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
500
+ raise ValueError(
501
+ 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}."
502
+ )
503
+
504
+ if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):
505
+ raise ValueError(
506
+ 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}."
507
+ )
508
+
509
+ if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types):
510
+ raise ValueError(
511
+ 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}."
512
+ )
513
+
514
+ if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types):
515
+ raise ValueError(
516
+ 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}."
517
+ )
518
+ if isinstance(transformer_layers_per_block, list) and reverse_transformer_layers_per_block is None:
519
+ for layer_number_per_block in transformer_layers_per_block:
520
+ if isinstance(layer_number_per_block, list):
521
+ raise ValueError("Must provide 'reverse_transformer_layers_per_block` if using asymmetrical UNet.")
522
+
523
+ def _set_time_proj(
524
+ self,
525
+ time_embedding_type: str,
526
+ block_out_channels: int,
527
+ flip_sin_to_cos: bool,
528
+ freq_shift: float,
529
+ time_embedding_dim: int,
530
+ ) -> Tuple[int, int]:
531
+ if time_embedding_type == "fourier":
532
+ time_embed_dim = time_embedding_dim or block_out_channels[0] * 2
533
+ if time_embed_dim % 2 != 0:
534
+ raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.")
535
+ self.time_proj = GaussianFourierProjection(
536
+ time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos
537
+ )
538
+ timestep_input_dim = time_embed_dim
539
+ elif time_embedding_type == "positional":
540
+ time_embed_dim = time_embedding_dim or block_out_channels[0] * 4
541
+
542
+ self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
543
+ timestep_input_dim = block_out_channels[0]
544
+ else:
545
+ raise ValueError(
546
+ f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`."
547
+ )
548
+
549
+ return time_embed_dim, timestep_input_dim
550
+
551
+ def _set_encoder_hid_proj(
552
+ self,
553
+ encoder_hid_dim_type: Optional[str],
554
+ cross_attention_dim: Union[int, Tuple[int]],
555
+ encoder_hid_dim: Optional[int],
556
+ ):
557
+ if encoder_hid_dim_type is None and encoder_hid_dim is not None:
558
+ encoder_hid_dim_type = "text_proj"
559
+ self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type)
560
+ logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.")
561
+
562
+ if encoder_hid_dim is None and encoder_hid_dim_type is not None:
563
+ raise ValueError(
564
+ f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}."
565
+ )
566
+
567
+ if encoder_hid_dim_type == "text_proj":
568
+ self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim)
569
+ elif encoder_hid_dim_type == "text_image_proj":
570
+ # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much
571
+ # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
572
+ # case when `addition_embed_type == "text_image_proj"` (Kandinsky 2.1)`
573
+ self.encoder_hid_proj = TextImageProjection(
574
+ text_embed_dim=encoder_hid_dim,
575
+ image_embed_dim=cross_attention_dim,
576
+ cross_attention_dim=cross_attention_dim,
577
+ )
578
+ elif encoder_hid_dim_type == "image_proj":
579
+ # Kandinsky 2.2
580
+ self.encoder_hid_proj = ImageProjection(
581
+ image_embed_dim=encoder_hid_dim,
582
+ cross_attention_dim=cross_attention_dim,
583
+ )
584
+ elif encoder_hid_dim_type is not None:
585
+ raise ValueError(
586
+ f"`encoder_hid_dim_type`: {encoder_hid_dim_type} must be None, 'text_proj', 'text_image_proj', or 'image_proj'."
587
+ )
588
+ else:
589
+ self.encoder_hid_proj = None
590
+
591
+ def _set_class_embedding(
592
+ self,
593
+ class_embed_type: Optional[str],
594
+ act_fn: str,
595
+ num_class_embeds: Optional[int],
596
+ projection_class_embeddings_input_dim: Optional[int],
597
+ time_embed_dim: int,
598
+ timestep_input_dim: int,
599
+ ):
600
+ if class_embed_type is None and num_class_embeds is not None:
601
+ self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
602
+ elif class_embed_type == "timestep":
603
+ self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim, act_fn=act_fn)
604
+ elif class_embed_type == "identity":
605
+ self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
606
+ elif class_embed_type == "projection":
607
+ if projection_class_embeddings_input_dim is None:
608
+ raise ValueError(
609
+ "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"
610
+ )
611
+ # The projection `class_embed_type` is the same as the timestep `class_embed_type` except
612
+ # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings
613
+ # 2. it projects from an arbitrary input dimension.
614
+ #
615
+ # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.
616
+ # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.
617
+ # As a result, `TimestepEmbedding` can be passed arbitrary vectors.
618
+ self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
619
+ elif class_embed_type == "simple_projection":
620
+ if projection_class_embeddings_input_dim is None:
621
+ raise ValueError(
622
+ "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set"
623
+ )
624
+ self.class_embedding = nn.Linear(projection_class_embeddings_input_dim, time_embed_dim)
625
+ else:
626
+ self.class_embedding = None
627
+
628
+ def _set_add_embedding(
629
+ self,
630
+ addition_embed_type: str,
631
+ addition_embed_type_num_heads: int,
632
+ addition_time_embed_dim: Optional[int],
633
+ flip_sin_to_cos: bool,
634
+ freq_shift: float,
635
+ cross_attention_dim: Optional[int],
636
+ encoder_hid_dim: Optional[int],
637
+ projection_class_embeddings_input_dim: Optional[int],
638
+ time_embed_dim: int,
639
+ ):
640
+ if addition_embed_type == "text":
641
+ if encoder_hid_dim is not None:
642
+ text_time_embedding_from_dim = encoder_hid_dim
643
+ else:
644
+ text_time_embedding_from_dim = cross_attention_dim
645
+
646
+ self.add_embedding = TextTimeEmbedding(
647
+ text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads
648
+ )
649
+ elif addition_embed_type == "text_image":
650
+ # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much
651
+ # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
652
+ # case when `addition_embed_type == "text_image"` (Kandinsky 2.1)`
653
+ self.add_embedding = TextImageTimeEmbedding(
654
+ text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim
655
+ )
656
+ elif addition_embed_type == "text_time":
657
+ self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift)
658
+ self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
659
+ elif addition_embed_type == "image":
660
+ # Kandinsky 2.2
661
+ self.add_embedding = ImageTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim)
662
+ elif addition_embed_type == "image_hint":
663
+ # Kandinsky 2.2 ControlNet
664
+ self.add_embedding = ImageHintTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim)
665
+ elif addition_embed_type is not None:
666
+ raise ValueError(
667
+ f"`addition_embed_type`: {addition_embed_type} must be None, 'text', 'text_image', 'text_time', 'image', or 'image_hint'."
668
+ )
669
+
670
+ def _set_pos_net_if_use_gligen(self, attention_type: str, cross_attention_dim: int):
671
+ if attention_type in ["gated", "gated-text-image"]:
672
+ positive_len = 768
673
+ if isinstance(cross_attention_dim, int):
674
+ positive_len = cross_attention_dim
675
+ elif isinstance(cross_attention_dim, (list, tuple)):
676
+ positive_len = cross_attention_dim[0]
677
+
678
+ feature_type = "text-only" if attention_type == "gated" else "text-image"
679
+ self.position_net = GLIGENTextBoundingboxProjection(
680
+ positive_len=positive_len, out_dim=cross_attention_dim, feature_type=feature_type
681
+ )
682
+
683
+ @property
684
+ def attn_processors(self) -> Dict[str, AttentionProcessor]:
685
+ r"""
686
+ Returns:
687
+ `dict` of attention processors: A dictionary containing all attention processors used in the model with
688
+ indexed by its weight name.
689
+ """
690
+ # set recursively
691
+ processors = {}
692
+
693
+ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
694
+ if hasattr(module, "get_processor"):
695
+ processors[f"{name}.processor"] = module.get_processor()
696
+
697
+ for sub_name, child in module.named_children():
698
+ fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
699
+
700
+ return processors
701
+
702
+ for name, module in self.named_children():
703
+ fn_recursive_add_processors(name, module, processors)
704
+
705
+ return processors
706
+
707
+ def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
708
+ r"""
709
+ Sets the attention processor to use to compute attention.
710
+
711
+ Parameters:
712
+ processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
713
+ The instantiated processor class or a dictionary of processor classes that will be set as the processor
714
+ for **all** `Attention` layers.
715
+
716
+ If `processor` is a dict, the key needs to define the path to the corresponding cross attention
717
+ processor. This is strongly recommended when setting trainable attention processors.
718
+
719
+ """
720
+ count = len(self.attn_processors.keys())
721
+
722
+ if isinstance(processor, dict) and len(processor) != count:
723
+ raise ValueError(
724
+ f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
725
+ f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
726
+ )
727
+
728
+ def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
729
+ if hasattr(module, "set_processor"):
730
+ if not isinstance(processor, dict):
731
+ module.set_processor(processor)
732
+ else:
733
+ module.set_processor(processor.pop(f"{name}.processor"))
734
+
735
+ for sub_name, child in module.named_children():
736
+ fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
737
+
738
+ for name, module in self.named_children():
739
+ fn_recursive_attn_processor(name, module, processor)
740
+
741
+ def set_default_attn_processor(self):
742
+ """
743
+ Disables custom attention processors and sets the default attention implementation.
744
+ """
745
+ if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
746
+ processor = AttnAddedKVProcessor()
747
+ elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
748
+ processor = AttnProcessor()
749
+ else:
750
+ raise ValueError(
751
+ f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
752
+ )
753
+
754
+ self.set_attn_processor(processor)
755
+
756
+ def set_attention_slice(self, slice_size: Union[str, int, List[int]] = "auto"):
757
+ r"""
758
+ Enable sliced attention computation.
759
+
760
+ When this option is enabled, the attention module splits the input tensor in slices to compute attention in
761
+ several steps. This is useful for saving some memory in exchange for a small decrease in speed.
762
+
763
+ Args:
764
+ slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
765
+ When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If
766
+ `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is
767
+ provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
768
+ must be a multiple of `slice_size`.
769
+ """
770
+ sliceable_head_dims = []
771
+
772
+ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
773
+ if hasattr(module, "set_attention_slice"):
774
+ sliceable_head_dims.append(module.sliceable_head_dim)
775
+
776
+ for child in module.children():
777
+ fn_recursive_retrieve_sliceable_dims(child)
778
+
779
+ # retrieve number of attention layers
780
+ for module in self.children():
781
+ fn_recursive_retrieve_sliceable_dims(module)
782
+
783
+ num_sliceable_layers = len(sliceable_head_dims)
784
+
785
+ if slice_size == "auto":
786
+ # half the attention head size is usually a good trade-off between
787
+ # speed and memory
788
+ slice_size = [dim // 2 for dim in sliceable_head_dims]
789
+ elif slice_size == "max":
790
+ # make smallest slice possible
791
+ slice_size = num_sliceable_layers * [1]
792
+
793
+ slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size
794
+
795
+ if len(slice_size) != len(sliceable_head_dims):
796
+ raise ValueError(
797
+ f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
798
+ f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
799
+ )
800
+
801
+ for i in range(len(slice_size)):
802
+ size = slice_size[i]
803
+ dim = sliceable_head_dims[i]
804
+ if size is not None and size > dim:
805
+ raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
806
+
807
+ # Recursively walk through all the children.
808
+ # Any children which exposes the set_attention_slice method
809
+ # gets the message
810
+ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
811
+ if hasattr(module, "set_attention_slice"):
812
+ module.set_attention_slice(slice_size.pop())
813
+
814
+ for child in module.children():
815
+ fn_recursive_set_attention_slice(child, slice_size)
816
+
817
+ reversed_slice_size = list(reversed(slice_size))
818
+ for module in self.children():
819
+ fn_recursive_set_attention_slice(module, reversed_slice_size)
820
+
821
+ def _set_gradient_checkpointing(self, module, value=False):
822
+ if hasattr(module, "gradient_checkpointing"):
823
+ module.gradient_checkpointing = value
824
+
825
+ def enable_freeu(self, s1: float, s2: float, b1: float, b2: float):
826
+ r"""Enables the FreeU mechanism from https://arxiv.org/abs/2309.11497.
827
+
828
+ The suffixes after the scaling factors represent the stage blocks where they are being applied.
829
+
830
+ Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of values that
831
+ are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL.
832
+
833
+ Args:
834
+ s1 (`float`):
835
+ Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to
836
+ mitigate the "oversmoothing effect" in the enhanced denoising process.
837
+ s2 (`float`):
838
+ Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to
839
+ mitigate the "oversmoothing effect" in the enhanced denoising process.
840
+ b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features.
841
+ b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features.
842
+ """
843
+ for i, upsample_block in enumerate(self.up_blocks):
844
+ setattr(upsample_block, "s1", s1)
845
+ setattr(upsample_block, "s2", s2)
846
+ setattr(upsample_block, "b1", b1)
847
+ setattr(upsample_block, "b2", b2)
848
+
849
+ def disable_freeu(self):
850
+ """Disables the FreeU mechanism."""
851
+ freeu_keys = {"s1", "s2", "b1", "b2"}
852
+ for i, upsample_block in enumerate(self.up_blocks):
853
+ for k in freeu_keys:
854
+ if hasattr(upsample_block, k) or getattr(upsample_block, k, None) is not None:
855
+ setattr(upsample_block, k, None)
856
+
857
+ def fuse_qkv_projections(self):
858
+ """
859
+ Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
860
+ are fused. For cross-attention modules, key and value projection matrices are fused.
861
+
862
+ <Tip warning={true}>
863
+
864
+ This API is 🧪 experimental.
865
+
866
+ </Tip>
867
+ """
868
+ self.original_attn_processors = None
869
+
870
+ for _, attn_processor in self.attn_processors.items():
871
+ if "Added" in str(attn_processor.__class__.__name__):
872
+ raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
873
+
874
+ self.original_attn_processors = self.attn_processors
875
+
876
+ for module in self.modules():
877
+ if isinstance(module, Attention):
878
+ module.fuse_projections(fuse=True)
879
+
880
+ self.set_attn_processor(FusedAttnProcessor2_0())
881
+
882
+ def unfuse_qkv_projections(self):
883
+ """Disables the fused QKV projection if enabled.
884
+
885
+ <Tip warning={true}>
886
+
887
+ This API is 🧪 experimental.
888
+
889
+ </Tip>
890
+
891
+ """
892
+ if self.original_attn_processors is not None:
893
+ self.set_attn_processor(self.original_attn_processors)
894
+
895
+ def get_time_embed(
896
+ self, sample: torch.Tensor, timestep: Union[torch.Tensor, float, int]
897
+ ) -> Optional[torch.Tensor]:
898
+ timesteps = timestep
899
+ if not torch.is_tensor(timesteps):
900
+ # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
901
+ # This would be a good case for the `match` statement (Python 3.10+)
902
+ is_mps = sample.device.type == "mps"
903
+ if isinstance(timestep, float):
904
+ dtype = torch.float32 if is_mps else torch.float64
905
+ else:
906
+ dtype = torch.int32 if is_mps else torch.int64
907
+ timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
908
+ elif len(timesteps.shape) == 0:
909
+ timesteps = timesteps[None].to(sample.device)
910
+
911
+ # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
912
+ timesteps = timesteps.expand(sample.shape[0])
913
+
914
+ t_emb = self.time_proj(timesteps)
915
+ # `Timesteps` does not contain any weights and will always return f32 tensors
916
+ # but time_embedding might actually be running in fp16. so we need to cast here.
917
+ # there might be better ways to encapsulate this.
918
+ t_emb = t_emb.to(dtype=sample.dtype)
919
+ return t_emb
920
+
921
+ def get_class_embed(self, sample: torch.Tensor, class_labels: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
922
+ class_emb = None
923
+ if self.class_embedding is not None:
924
+ if class_labels is None:
925
+ raise ValueError("class_labels should be provided when num_class_embeds > 0")
926
+
927
+ if self.config.class_embed_type == "timestep":
928
+ class_labels = self.time_proj(class_labels)
929
+
930
+ # `Timesteps` does not contain any weights and will always return f32 tensors
931
+ # there might be better ways to encapsulate this.
932
+ class_labels = class_labels.to(dtype=sample.dtype)
933
+
934
+ class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype)
935
+ return class_emb
936
+
937
+ def get_aug_embed(
938
+ self, emb: torch.Tensor, encoder_hidden_states: torch.Tensor, added_cond_kwargs: Dict[str, Any]
939
+ ) -> Optional[torch.Tensor]:
940
+ aug_emb = None
941
+ if self.config.addition_embed_type == "text":
942
+ aug_emb = self.add_embedding(encoder_hidden_states)
943
+ elif self.config.addition_embed_type == "text_image":
944
+ # Kandinsky 2.1 - style
945
+ if "image_embeds" not in added_cond_kwargs:
946
+ raise ValueError(
947
+ 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`"
948
+ )
949
+
950
+ image_embs = added_cond_kwargs.get("image_embeds")
951
+ text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states)
952
+ aug_emb = self.add_embedding(text_embs, image_embs)
953
+ elif self.config.addition_embed_type == "text_time":
954
+ # SDXL - style
955
+ if "text_embeds" not in added_cond_kwargs:
956
+ raise ValueError(
957
+ 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`"
958
+ )
959
+ text_embeds = added_cond_kwargs.get("text_embeds")
960
+ if "time_ids" not in added_cond_kwargs:
961
+ raise ValueError(
962
+ 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`"
963
+ )
964
+ time_ids = added_cond_kwargs.get("time_ids")
965
+ time_embeds = self.add_time_proj(time_ids.flatten())
966
+ time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
967
+ add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
968
+ add_embeds = add_embeds.to(emb.dtype)
969
+ aug_emb = self.add_embedding(add_embeds)
970
+ elif self.config.addition_embed_type == "image":
971
+ # Kandinsky 2.2 - style
972
+ if "image_embeds" not in added_cond_kwargs:
973
+ raise ValueError(
974
+ 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`"
975
+ )
976
+ image_embs = added_cond_kwargs.get("image_embeds")
977
+ aug_emb = self.add_embedding(image_embs)
978
+ elif self.config.addition_embed_type == "image_hint":
979
+ # Kandinsky 2.2 ControlNet - style
980
+ if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs:
981
+ raise ValueError(
982
+ 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`"
983
+ )
984
+ image_embs = added_cond_kwargs.get("image_embeds")
985
+ hint = added_cond_kwargs.get("hint")
986
+ aug_emb = self.add_embedding(image_embs, hint)
987
+ return aug_emb
988
+
989
+ def process_encoder_hidden_states(
990
+ self, encoder_hidden_states: torch.Tensor, added_cond_kwargs: Dict[str, Any]
991
+ ) -> torch.Tensor:
992
+ if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj":
993
+ encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states)
994
+ elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj":
995
+ # Kandinsky 2.1 - style
996
+ if "image_embeds" not in added_cond_kwargs:
997
+ raise ValueError(
998
+ 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_cond_kwargs`"
999
+ )
1000
+
1001
+ image_embeds = added_cond_kwargs.get("image_embeds")
1002
+ encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds)
1003
+ elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj":
1004
+ # Kandinsky 2.2 - style
1005
+ if "image_embeds" not in added_cond_kwargs:
1006
+ raise ValueError(
1007
+ 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_cond_kwargs`"
1008
+ )
1009
+ image_embeds = added_cond_kwargs.get("image_embeds")
1010
+ encoder_hidden_states = self.encoder_hid_proj(image_embeds)
1011
+ elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "ip_image_proj":
1012
+ if "image_embeds" not in added_cond_kwargs:
1013
+ raise ValueError(
1014
+ 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_cond_kwargs`"
1015
+ )
1016
+
1017
+ if hasattr(self, "text_encoder_hid_proj") and self.text_encoder_hid_proj is not None:
1018
+ encoder_hidden_states = self.text_encoder_hid_proj(encoder_hidden_states)
1019
+
1020
+ image_embeds = added_cond_kwargs.get("image_embeds")
1021
+ image_embeds = self.encoder_hid_proj(image_embeds)
1022
+ encoder_hidden_states = (encoder_hidden_states, image_embeds)
1023
+ return encoder_hidden_states
1024
+
1025
+ def forward(
1026
+ self,
1027
+ sample: torch.Tensor,
1028
+ timestep: Union[torch.Tensor, float, int],
1029
+ encoder_hidden_states: torch.Tensor,
1030
+ controls: torch.Tensor,
1031
+ controls_scale: torch.Tensor,
1032
+ class_labels: Optional[torch.Tensor] = None,
1033
+ timestep_cond: Optional[torch.Tensor] = None,
1034
+ attention_mask: Optional[torch.Tensor] = None,
1035
+ cross_attention_kwargs: Optional[Dict[str, Any]] = None,
1036
+ added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
1037
+ down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
1038
+ mid_block_additional_residual: Optional[torch.Tensor] = None,
1039
+ down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
1040
+ encoder_attention_mask: Optional[torch.Tensor] = None,
1041
+ #return_dict: bool = True,
1042
+ ) -> Union[UNet2DConditionOutput, Tuple]:
1043
+ r"""
1044
+ The [`UNet2DConditionModel`] forward method.
1045
+
1046
+ Args:
1047
+ sample (`torch.Tensor`):
1048
+ The noisy input tensor with the following shape `(batch, channel, height, width)`.
1049
+ timestep (`torch.Tensor` or `float` or `int`): The number of timesteps to denoise an input.
1050
+ encoder_hidden_states (`torch.Tensor`):
1051
+ The encoder hidden states with shape `(batch, sequence_length, feature_dim)`.
1052
+ class_labels (`torch.Tensor`, *optional*, defaults to `None`):
1053
+ Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
1054
+ timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`):
1055
+ Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed
1056
+ through the `self.time_embedding` layer to obtain the timestep embeddings.
1057
+ attention_mask (`torch.Tensor`, *optional*, defaults to `None`):
1058
+ An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
1059
+ is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
1060
+ negative values to the attention scores corresponding to "discard" tokens.
1061
+ cross_attention_kwargs (`dict`, *optional*):
1062
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
1063
+ `self.processor` in
1064
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
1065
+ added_cond_kwargs: (`dict`, *optional*):
1066
+ A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that
1067
+ are passed along to the UNet blocks.
1068
+ down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*):
1069
+ A tuple of tensors that if specified are added to the residuals of down unet blocks.
1070
+ mid_block_additional_residual: (`torch.Tensor`, *optional*):
1071
+ A tensor that if specified is added to the residual of the middle unet block.
1072
+ down_intrablock_additional_residuals (`tuple` of `torch.Tensor`, *optional*):
1073
+ additional residuals to be added within UNet down blocks, for example from T2I-Adapter side model(s)
1074
+ encoder_attention_mask (`torch.Tensor`):
1075
+ A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If
1076
+ `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias,
1077
+ which adds large negative values to the attention scores corresponding to "discard" tokens.
1078
+ return_dict (`bool`, *optional*, defaults to `True`):
1079
+ Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
1080
+ tuple.
1081
+
1082
+ Returns:
1083
+ [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
1084
+ If `return_dict` is True, an [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] is returned,
1085
+ otherwise a `tuple` is returned where the first element is the sample tensor.
1086
+ """
1087
+ # By default samples have to be AT least a multiple of the overall upsampling factor.
1088
+ # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).
1089
+ # However, the upsampling interpolation output size can be forced to fit any upsampling size
1090
+ # on the fly if necessary.
1091
+ default_overall_up_factor = 2**self.num_upsamplers
1092
+
1093
+ # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
1094
+ forward_upsample_size = False
1095
+ upsample_size = None
1096
+
1097
+ for dim in sample.shape[-2:]:
1098
+ if dim % default_overall_up_factor != 0:
1099
+ # Forward upsample size to force interpolation output size.
1100
+ forward_upsample_size = True
1101
+ break
1102
+
1103
+ # ensure attention_mask is a bias, and give it a singleton query_tokens dimension
1104
+ # expects mask of shape:
1105
+ # [batch, key_tokens]
1106
+ # adds singleton query_tokens dimension:
1107
+ # [batch, 1, key_tokens]
1108
+ # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
1109
+ # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)
1110
+ # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
1111
+ if attention_mask is not None:
1112
+ # assume that mask is expressed as:
1113
+ # (1 = keep, 0 = discard)
1114
+ # convert mask into a bias that can be added to attention scores:
1115
+ # (keep = +0, discard = -10000.0)
1116
+ attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
1117
+ attention_mask = attention_mask.unsqueeze(1)
1118
+
1119
+ # convert encoder_attention_mask to a bias the same way we do for attention_mask
1120
+ if encoder_attention_mask is not None:
1121
+ encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
1122
+ encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
1123
+
1124
+ # 0. center input if necessary
1125
+ if self.config.center_input_sample:
1126
+ sample = 2 * sample - 1.0
1127
+
1128
+ # 1. time
1129
+ t_emb = self.get_time_embed(sample=sample, timestep=timestep)
1130
+ emb = self.time_embedding(t_emb, timestep_cond)
1131
+
1132
+ class_emb = self.get_class_embed(sample=sample, class_labels=class_labels)
1133
+ if class_emb is not None:
1134
+ if self.config.class_embeddings_concat:
1135
+ emb = torch.cat([emb, class_emb], dim=-1)
1136
+ else:
1137
+ emb = emb + class_emb
1138
+
1139
+ aug_emb = None
1140
+
1141
+ if self.config.addition_embed_type == "image_hint":
1142
+ aug_emb, hint = aug_emb
1143
+ sample = torch.cat([sample, hint], dim=1)
1144
+
1145
+ emb = emb + aug_emb if aug_emb is not None else emb
1146
+
1147
+ if self.time_embed_act is not None:
1148
+ emb = self.time_embed_act(emb)
1149
+
1150
+ encoder_hidden_states = self.process_encoder_hidden_states(
1151
+ encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs
1152
+ )
1153
+
1154
+ # 2. pre-process
1155
+ sample = self.conv_in(sample)
1156
+
1157
+ # 2.5 GLIGEN position net
1158
+ if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None:
1159
+ cross_attention_kwargs = cross_attention_kwargs.copy()
1160
+ gligen_args = cross_attention_kwargs.pop("gligen")
1161
+ cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)}
1162
+
1163
+ # 3. down
1164
+ # we're popping the `scale` instead of getting it because otherwise `scale` will be propagated
1165
+ # to the internal blocks and will raise deprecation warnings. this will be confusing for our users.
1166
+ if cross_attention_kwargs is not None:
1167
+ cross_attention_kwargs = cross_attention_kwargs.copy()
1168
+ lora_scale = cross_attention_kwargs.pop("scale", 1.0)
1169
+ else:
1170
+ lora_scale = 1.0
1171
+
1172
+ if USE_PEFT_BACKEND:
1173
+ # weight the lora layers by setting `lora_scale` for each PEFT layer
1174
+ scale_lora_layers(self, lora_scale)
1175
+
1176
+ is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
1177
+ # using new arg down_intrablock_additional_residuals for T2I-Adapters, to distinguish from controlnets
1178
+ is_adapter = down_intrablock_additional_residuals is not None
1179
+ # maintain backward compatibility for legacy usage, where
1180
+ # T2I-Adapter and ControlNet both use down_block_additional_residuals arg
1181
+ # but can only use one or the other
1182
+ if not is_adapter and mid_block_additional_residual is None and down_block_additional_residuals is not None:
1183
+ deprecate(
1184
+ "T2I should not use down_block_additional_residuals",
1185
+ "1.3.0",
1186
+ "Passing intrablock residual connections with `down_block_additional_residuals` is deprecated \
1187
+ and will be removed in diffusers 1.3.0. `down_block_additional_residuals` should only be used \
1188
+ for ControlNet. Please make sure use `down_intrablock_additional_residuals` instead. ",
1189
+ standard_warn=False,
1190
+ )
1191
+ down_intrablock_additional_residuals = down_block_additional_residuals
1192
+ is_adapter = True
1193
+
1194
+ down_block_res_samples = (sample,)
1195
+
1196
+ scale = controls_scale
1197
+ controls = controls.to(sample)
1198
+ mean_latents, std_latents = torch.mean(sample, dim=(1, 2, 3), keepdim=True), torch.std(sample, dim=(1, 2, 3), keepdim=True)
1199
+ mean_control, std_control = torch.mean(controls, dim=(1, 2, 3), keepdim=True), torch.std(controls, dim=(1, 2, 3), keepdim=True)
1200
+ controls = (controls - mean_control) * (std_latents / (std_control + 1e-12)) + mean_latents
1201
+ controls = nn.functional.adaptive_avg_pool2d(controls, (128, 128))
1202
+ sample = sample + controls * scale
1203
+
1204
+ for downsample_block in self.down_blocks:
1205
+ if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
1206
+ # For t2i-adapter CrossAttnDownBlock2D
1207
+ additional_residuals = {}
1208
+ if is_adapter and len(down_intrablock_additional_residuals) > 0:
1209
+ additional_residuals["additional_residuals"] = down_intrablock_additional_residuals.pop(0)
1210
+
1211
+ sample, res_samples = downsample_block(
1212
+ hidden_states=sample,
1213
+ temb=emb,
1214
+ encoder_hidden_states=encoder_hidden_states,
1215
+ attention_mask=attention_mask,
1216
+ cross_attention_kwargs=cross_attention_kwargs,
1217
+ encoder_attention_mask=encoder_attention_mask,
1218
+ **additional_residuals,
1219
+ )
1220
+ else:
1221
+ sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
1222
+ if is_adapter and len(down_intrablock_additional_residuals) > 0:
1223
+ sample += down_intrablock_additional_residuals.pop(0)
1224
+
1225
+ down_block_res_samples += res_samples
1226
+
1227
+ if is_controlnet:
1228
+ new_down_block_res_samples = ()
1229
+
1230
+ for down_block_res_sample, down_block_additional_residual in zip(
1231
+ down_block_res_samples, down_block_additional_residuals
1232
+ ):
1233
+ down_block_res_sample = down_block_res_sample + down_block_additional_residual
1234
+ new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
1235
+
1236
+ down_block_res_samples = new_down_block_res_samples
1237
+
1238
+ # 4. mid
1239
+ if self.mid_block is not None:
1240
+ if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention:
1241
+ sample = self.mid_block(
1242
+ sample,
1243
+ emb,
1244
+ encoder_hidden_states=encoder_hidden_states,
1245
+ attention_mask=attention_mask,
1246
+ cross_attention_kwargs=cross_attention_kwargs,
1247
+ encoder_attention_mask=encoder_attention_mask,
1248
+ )
1249
+ else:
1250
+ sample = self.mid_block(sample, emb)
1251
+
1252
+ # To support T2I-Adapter-XL
1253
+ if (
1254
+ is_adapter
1255
+ and len(down_intrablock_additional_residuals) > 0
1256
+ and sample.shape == down_intrablock_additional_residuals[0].shape
1257
+ ):
1258
+ sample += down_intrablock_additional_residuals.pop(0)
1259
+
1260
+ if is_controlnet:
1261
+ sample = sample + mid_block_additional_residual
1262
+
1263
+ # 5. up
1264
+ for i, upsample_block in enumerate(self.up_blocks):
1265
+ is_final_block = i == len(self.up_blocks) - 1
1266
+
1267
+ res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
1268
+ down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
1269
+
1270
+ # if we have not reached the final block and need to forward the
1271
+ # upsample size, we do it here
1272
+ if not is_final_block and forward_upsample_size:
1273
+ upsample_size = down_block_res_samples[-1].shape[2:]
1274
+
1275
+ if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
1276
+ sample = upsample_block(
1277
+ hidden_states=sample,
1278
+ temb=emb,
1279
+ res_hidden_states_tuple=res_samples,
1280
+ encoder_hidden_states=encoder_hidden_states,
1281
+ cross_attention_kwargs=cross_attention_kwargs,
1282
+ upsample_size=upsample_size,
1283
+ attention_mask=attention_mask,
1284
+ encoder_attention_mask=encoder_attention_mask,
1285
+ )
1286
+ else:
1287
+ sample = upsample_block(
1288
+ hidden_states=sample,
1289
+ temb=emb,
1290
+ res_hidden_states_tuple=res_samples,
1291
+ upsample_size=upsample_size,
1292
+ )
1293
+
1294
+ # 6. post-process
1295
+ if self.conv_norm_out:
1296
+ sample = self.conv_norm_out(sample)
1297
+ sample = self.conv_act(sample)
1298
+ sample = self.conv_out(sample)
1299
+
1300
+ if USE_PEFT_BACKEND:
1301
+ # remove `lora_scale` from each PEFT layer
1302
+ unscale_lora_layers(self, lora_scale)
1303
+
1304
+ #if not return_dict:
1305
+ return (sample,)
1306
+
1307
+ #return UNet2DConditionOutput(sample=sample)
1308
+