ashawkey commited on
Commit
a8a63dd
0 Parent(s):
README.md ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # MVDream-hf
2
+
3
+ modified from https://github.com/KokeCacao/mvdream-hf.
4
+
5
+ ### convert weights
6
+ ```bash
7
+ # download original ckpt
8
+ wget https://huggingface.co/MVDream/MVDream/resolve/main/sd-v2.1-base-4view.pt
9
+ wget https://raw.githubusercontent.com/bytedance/MVDream/main/mvdream/configs/sd-v2-base.yaml
10
+
11
+ # convert
12
+ python convert_mvdream_to_diffusers.py --checkpoint_path ./sd-v2.1-base-4view.pt --dump_path ./weights --original_config_file ./sd-v2-base.yaml --half --to_safetensors --test
13
+ ```
14
+
15
+ ### run pipeline
16
+ ```python
17
+ import torch
18
+ import kiui
19
+ from mvdream.pipeline_mvdream import MVDreamStableDiffusionPipeline
20
+
21
+ pipe = MVDreamStableDiffusionPipeline.from_pretrained('./weights', torch_dtype=torch.float16)
22
+ pipe = pipe.to("cuda")
23
+
24
+ prompt = "a photo of an astronaut riding a horse on mars"
25
+ image = pipe(prompt) # np.ndarray [4, 256, 256, 3]
26
+
27
+ kiui.vis.plot_image(image)
28
+ ```
convert_mvdream_to_diffusers.py ADDED
@@ -0,0 +1,408 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Modified from https://github.com/huggingface/diffusers/blob/bc691231360a4cbc7d19a58742ebb8ed0f05e027/scripts/convert_original_stable_diffusion_to_diffusers.py
2
+
3
+ import argparse
4
+ import torch
5
+ import sys
6
+
7
+ sys.path.insert(0, '.')
8
+
9
+ from diffusers.models import (
10
+ AutoencoderKL,
11
+ )
12
+ from omegaconf import OmegaConf
13
+ from diffusers.schedulers import DDIMScheduler
14
+ from diffusers.utils import logging
15
+ from typing import Any
16
+ from accelerate import init_empty_weights
17
+ from accelerate.utils import set_module_tensor_to_device
18
+ from mvdream.models import MultiViewUNetWrapperModel
19
+ from mvdream.pipeline_mvdream import MVDreamStableDiffusionPipeline
20
+ from transformers import CLIPTokenizer, CLIPTextModel
21
+
22
+ logger = logging.get_logger(__name__)
23
+
24
+
25
+ def assign_to_checkpoint(paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None):
26
+ """
27
+ This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits
28
+ attention layers, and takes into account additional replacements that may arise.
29
+ Assigns the weights to the new checkpoint.
30
+ """
31
+ assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys."
32
+
33
+ # Splits the attention layers into three variables.
34
+ if attention_paths_to_split is not None:
35
+ for path, path_map in attention_paths_to_split.items():
36
+ old_tensor = old_checkpoint[path]
37
+ channels = old_tensor.shape[0] // 3
38
+
39
+ target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1)
40
+
41
+ assert config is not None
42
+ num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3
43
+
44
+ old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:])
45
+ query, key, value = old_tensor.split(channels // num_heads, dim=1)
46
+
47
+ checkpoint[path_map["query"]] = query.reshape(target_shape)
48
+ checkpoint[path_map["key"]] = key.reshape(target_shape)
49
+ checkpoint[path_map["value"]] = value.reshape(target_shape)
50
+
51
+ for path in paths:
52
+ new_path = path["new"]
53
+
54
+ # These have already been assigned
55
+ if attention_paths_to_split is not None and new_path in attention_paths_to_split:
56
+ continue
57
+
58
+ # Global renaming happens here
59
+ new_path = new_path.replace("middle_block.0", "mid_block.resnets.0")
60
+ new_path = new_path.replace("middle_block.1", "mid_block.attentions.0")
61
+ new_path = new_path.replace("middle_block.2", "mid_block.resnets.1")
62
+
63
+ if additional_replacements is not None:
64
+ for replacement in additional_replacements:
65
+ new_path = new_path.replace(replacement["old"], replacement["new"])
66
+
67
+ # proj_attn.weight has to be converted from conv 1D to linear
68
+ is_attn_weight = "proj_attn.weight" in new_path or ("attentions" in new_path and "to_" in new_path)
69
+ shape = old_checkpoint[path["old"]].shape
70
+ if is_attn_weight and len(shape) == 3:
71
+ checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0]
72
+ elif is_attn_weight and len(shape) == 4:
73
+ checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0, 0]
74
+ else:
75
+ checkpoint[new_path] = old_checkpoint[path["old"]]
76
+
77
+
78
+ def shave_segments(path, n_shave_prefix_segments=1):
79
+ """
80
+ Removes segments. Positive values shave the first segments, negative shave the last segments.
81
+ """
82
+ if n_shave_prefix_segments >= 0:
83
+ return ".".join(path.split(".")[n_shave_prefix_segments:])
84
+ else:
85
+ return ".".join(path.split(".")[:n_shave_prefix_segments])
86
+
87
+
88
+ def create_vae_diffusers_config(original_config, image_size: int):
89
+ """
90
+ Creates a config for the diffusers based on the config of the LDM model.
91
+ """
92
+ vae_params = original_config.model.params.first_stage_config.params.ddconfig
93
+ _ = original_config.model.params.first_stage_config.params.embed_dim
94
+
95
+ block_out_channels = [vae_params.ch * mult for mult in vae_params.ch_mult]
96
+ down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels)
97
+ up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels)
98
+
99
+ config = {
100
+ "sample_size": image_size,
101
+ "in_channels": vae_params.in_channels,
102
+ "out_channels": vae_params.out_ch,
103
+ "down_block_types": tuple(down_block_types),
104
+ "up_block_types": tuple(up_block_types),
105
+ "block_out_channels": tuple(block_out_channels),
106
+ "latent_channels": vae_params.z_channels,
107
+ "layers_per_block": vae_params.num_res_blocks,
108
+ }
109
+ return config
110
+
111
+
112
+ def convert_ldm_vae_checkpoint(checkpoint, config):
113
+ # extract state dict for VAE
114
+ vae_state_dict = {}
115
+ vae_key = "first_stage_model."
116
+ keys = list(checkpoint.keys())
117
+ for key in keys:
118
+ if key.startswith(vae_key):
119
+ vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key)
120
+
121
+ new_checkpoint = {}
122
+
123
+ new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"]
124
+ new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"]
125
+ new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"]
126
+ new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"]
127
+ new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"]
128
+ new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"]
129
+
130
+ new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"]
131
+ new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"]
132
+ new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"]
133
+ new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"]
134
+ new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"]
135
+ new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"]
136
+
137
+ new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"]
138
+ new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"]
139
+ new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"]
140
+ new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"]
141
+
142
+ # Retrieves the keys for the encoder down blocks only
143
+ num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer})
144
+ down_blocks = {layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks)}
145
+
146
+ # Retrieves the keys for the decoder up blocks only
147
+ num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer})
148
+ up_blocks = {layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks)}
149
+
150
+ for i in range(num_down_blocks):
151
+ resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key]
152
+
153
+ if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict:
154
+ new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop(f"encoder.down.{i}.downsample.conv.weight")
155
+ new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop(f"encoder.down.{i}.downsample.conv.bias")
156
+
157
+ paths = renew_vae_resnet_paths(resnets)
158
+ meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"}
159
+ assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
160
+
161
+ mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key]
162
+ num_mid_res_blocks = 2
163
+ for i in range(1, num_mid_res_blocks + 1):
164
+ resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key]
165
+
166
+ paths = renew_vae_resnet_paths(resnets)
167
+ meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
168
+ assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
169
+
170
+ mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key]
171
+ paths = renew_vae_attention_paths(mid_attentions)
172
+ meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
173
+ assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
174
+ conv_attn_to_linear(new_checkpoint)
175
+
176
+ for i in range(num_up_blocks):
177
+ block_id = num_up_blocks - 1 - i
178
+ resnets = [key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key]
179
+
180
+ if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict:
181
+ new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[f"decoder.up.{block_id}.upsample.conv.weight"]
182
+ new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[f"decoder.up.{block_id}.upsample.conv.bias"]
183
+
184
+ paths = renew_vae_resnet_paths(resnets)
185
+ meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"}
186
+ assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
187
+
188
+ mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key]
189
+ num_mid_res_blocks = 2
190
+ for i in range(1, num_mid_res_blocks + 1):
191
+ resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key]
192
+
193
+ paths = renew_vae_resnet_paths(resnets)
194
+ meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
195
+ assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
196
+
197
+ mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key]
198
+ paths = renew_vae_attention_paths(mid_attentions)
199
+ meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
200
+ assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
201
+ conv_attn_to_linear(new_checkpoint)
202
+ return new_checkpoint
203
+
204
+
205
+ def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0):
206
+ """
207
+ Updates paths inside resnets to the new naming scheme (local renaming)
208
+ """
209
+ mapping = []
210
+ for old_item in old_list:
211
+ new_item = old_item
212
+
213
+ new_item = new_item.replace("nin_shortcut", "conv_shortcut")
214
+ new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
215
+
216
+ mapping.append({"old": old_item, "new": new_item})
217
+
218
+ return mapping
219
+
220
+
221
+ def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0):
222
+ """
223
+ Updates paths inside attentions to the new naming scheme (local renaming)
224
+ """
225
+ mapping = []
226
+ for old_item in old_list:
227
+ new_item = old_item
228
+
229
+ new_item = new_item.replace("norm.weight", "group_norm.weight")
230
+ new_item = new_item.replace("norm.bias", "group_norm.bias")
231
+
232
+ new_item = new_item.replace("q.weight", "to_q.weight")
233
+ new_item = new_item.replace("q.bias", "to_q.bias")
234
+
235
+ new_item = new_item.replace("k.weight", "to_k.weight")
236
+ new_item = new_item.replace("k.bias", "to_k.bias")
237
+
238
+ new_item = new_item.replace("v.weight", "to_v.weight")
239
+ new_item = new_item.replace("v.bias", "to_v.bias")
240
+
241
+ new_item = new_item.replace("proj_out.weight", "to_out.0.weight")
242
+ new_item = new_item.replace("proj_out.bias", "to_out.0.bias")
243
+
244
+ new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
245
+
246
+ mapping.append({"old": old_item, "new": new_item})
247
+
248
+ return mapping
249
+
250
+
251
+ def conv_attn_to_linear(checkpoint):
252
+ keys = list(checkpoint.keys())
253
+ attn_keys = ["query.weight", "key.weight", "value.weight"]
254
+ for key in keys:
255
+ if ".".join(key.split(".")[-2:]) in attn_keys:
256
+ if checkpoint[key].ndim > 2:
257
+ checkpoint[key] = checkpoint[key][:, :, 0, 0]
258
+ elif "proj_attn.weight" in key:
259
+ if checkpoint[key].ndim > 2:
260
+ checkpoint[key] = checkpoint[key][:, :, 0]
261
+
262
+ def create_unet_config(original_config) -> Any:
263
+ return OmegaConf.to_container(original_config.model.params.unet_config.params, resolve=True)
264
+
265
+ def convert_from_original_mvdream_ckpt(checkpoint_path, original_config_file, device):
266
+ checkpoint = torch.load(checkpoint_path, map_location=device)
267
+ # print(f"Checkpoint: {checkpoint.keys()}")
268
+ torch.cuda.empty_cache()
269
+
270
+ original_config = OmegaConf.load(original_config_file)
271
+ # print(f"Original Config: {original_config}")
272
+ prediction_type = "epsilon"
273
+ image_size = 256
274
+ num_train_timesteps = getattr(original_config.model.params, "timesteps", None) or 1000
275
+ beta_start = getattr(original_config.model.params, "linear_start", None) or 0.02
276
+ beta_end = getattr(original_config.model.params, "linear_end", None) or 0.085
277
+ scheduler = DDIMScheduler(
278
+ beta_end=beta_end,
279
+ beta_schedule="scaled_linear",
280
+ beta_start=beta_start,
281
+ num_train_timesteps=num_train_timesteps,
282
+ steps_offset=1,
283
+ clip_sample=False,
284
+ set_alpha_to_one=False,
285
+ prediction_type=prediction_type,
286
+ )
287
+ scheduler.register_to_config(clip_sample=False)
288
+
289
+ # Convert the UNet2DConditionModel model.
290
+ # upcast_attention = None
291
+ # unet_config = create_unet_diffusers_config(original_config, image_size=image_size)
292
+ # unet_config["upcast_attention"] = upcast_attention
293
+ # with init_empty_weights():
294
+ # unet = UNet2DConditionModel(**unet_config)
295
+ # converted_unet_checkpoint = convert_ldm_unet_checkpoint(
296
+ # checkpoint, unet_config, path=None, extract_ema=extract_ema
297
+ # )
298
+ # print(f"Unet Config: {original_config.model.params.unet_config.params}")
299
+ unet_config = create_unet_config(original_config)
300
+ unet: MultiViewUNetWrapperModel = MultiViewUNetWrapperModel(**unet_config)
301
+ unet.register_to_config(**unet_config)
302
+ # print(f"Unet State Dict: {unet.state_dict().keys()}")
303
+ unet.load_state_dict({key.replace("model.diffusion_model.", "unet."): value for key, value in checkpoint.items() if key.replace("model.diffusion_model.", "unet.") in unet.state_dict()})
304
+ for param_name, param in unet.state_dict().items():
305
+ set_module_tensor_to_device(unet, param_name, device=device, value=param)
306
+
307
+ # Convert the VAE model.
308
+ vae_config = create_vae_diffusers_config(original_config, image_size=image_size)
309
+ converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config)
310
+
311
+ if ("model" in original_config and "params" in original_config.model and "scale_factor" in original_config.model.params):
312
+ vae_scaling_factor = original_config.model.params.scale_factor
313
+ else:
314
+ vae_scaling_factor = 0.18215 # default SD scaling factor
315
+
316
+ vae_config["scaling_factor"] = vae_scaling_factor
317
+
318
+ with init_empty_weights():
319
+ vae = AutoencoderKL(**vae_config)
320
+
321
+ for param_name, param in converted_vae_checkpoint.items():
322
+ set_module_tensor_to_device(vae, param_name, device=device, value=param)
323
+
324
+ if original_config.model.params.unet_config.params.context_dim == 768:
325
+ tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
326
+ text_encoder: CLIPTextModel = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to(device=device) # type: ignore
327
+ elif original_config.model.params.unet_config.params.context_dim == 1024:
328
+ tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="tokenizer")
329
+ text_encoder: CLIPTextModel = CLIPTextModel.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="text_encoder").to(device=device) # type: ignore
330
+ else:
331
+ raise ValueError(f"Unknown context_dim: {original_config.model.paams.unet_config.params.context_dim}")
332
+
333
+ pipe = MVDreamStableDiffusionPipeline(
334
+ vae=vae,
335
+ unet=unet,
336
+ tokenizer=tokenizer,
337
+ text_encoder=text_encoder,
338
+ scheduler=scheduler,
339
+ )
340
+
341
+ return pipe
342
+
343
+
344
+ if __name__ == "__main__":
345
+ parser = argparse.ArgumentParser()
346
+
347
+ parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert.")
348
+ parser.add_argument(
349
+ "--original_config_file",
350
+ default=None,
351
+ type=str,
352
+ help="The YAML config file corresponding to the original architecture.",
353
+ )
354
+ parser.add_argument(
355
+ "--to_safetensors",
356
+ action="store_true",
357
+ help="Whether to store pipeline in safetensors format or not.",
358
+ )
359
+ parser.add_argument("--half", action="store_true", help="Save weights in half precision.")
360
+ parser.add_argument("--test", action="store_true", help="Whether to test inference after convertion.")
361
+ parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.")
362
+ parser.add_argument("--device", type=str, help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)")
363
+ args = parser.parse_args()
364
+
365
+ args.device = torch.device(args.device if args.device is not None else "cuda" if torch.cuda.is_available() else "cpu")
366
+
367
+ pipe = convert_from_original_mvdream_ckpt(
368
+ checkpoint_path=args.checkpoint_path,
369
+ original_config_file=args.original_config_file,
370
+ device=args.device,
371
+ )
372
+
373
+ if args.half:
374
+ pipe.to(torch_dtype=torch.float16)
375
+
376
+ print(f"Saving pipeline to {args.dump_path}...")
377
+ pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)
378
+
379
+ if args.test:
380
+ try:
381
+ print(f"Testing each subcomponent of the pipeline...")
382
+ images = pipe(
383
+ prompt="Head of Hatsune Miku",
384
+ negative_prompt="painting, bad quality, flat",
385
+ output_type="pil",
386
+ guidance_scale=7.5,
387
+ num_inference_steps=50,
388
+ device=args.device,
389
+ )
390
+ for i, image in enumerate(images):
391
+ image.save(f"image_{i}.png") # type: ignore
392
+
393
+ print(f"Testing entire pipeline...")
394
+ loaded_pipe: MVDreamStableDiffusionPipeline = MVDreamStableDiffusionPipeline.from_pretrained(args.dump_path, safe_serialization=args.to_safetensors) # type: ignore
395
+ images = loaded_pipe(
396
+ prompt="Head of Hatsune Miku",
397
+ negative_prompt="painting, bad quality, flat",
398
+ output_type="pil",
399
+ guidance_scale=7.5,
400
+ num_inference_steps=50,
401
+ device=args.device,
402
+ )
403
+ for i, image in enumerate(images):
404
+ image.save(f"image_{i}.png") # type: ignore
405
+ except Exception as e:
406
+ print(f"Failed to test inference: {e}")
407
+ raise e from e
408
+ print("Inference test passed!")
main.py ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import kiui
3
+ from mvdream.pipeline_mvdream import MVDreamStableDiffusionPipeline
4
+
5
+ pipe = MVDreamStableDiffusionPipeline.from_pretrained('./weights', torch_dtype=torch.float16)
6
+ pipe = pipe.to("cuda")
7
+
8
+ prompt = "a photo of an astronaut riding a horse on mars"
9
+ image = pipe(prompt)
10
+
11
+ kiui.vis.plot_image(image)
mvdream/attention.py ADDED
@@ -0,0 +1,352 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # obtained and modified from https://github.com/bytedance/MVDream
2
+
3
+ import math
4
+ import torch
5
+ import torch.nn.functional as F
6
+
7
+ from inspect import isfunction
8
+ from torch import nn, einsum
9
+ from torch.amp.autocast_mode import autocast
10
+ from einops import rearrange, repeat
11
+ from typing import Optional, Any
12
+ from .util import checkpoint
13
+
14
+ try:
15
+ import xformers # type: ignore
16
+ import xformers.ops # type: ignore
17
+ XFORMERS_IS_AVAILBLE = True
18
+ except:
19
+ XFORMERS_IS_AVAILBLE = False
20
+
21
+ # CrossAttn precision handling
22
+ import os
23
+
24
+ _ATTN_PRECISION = os.environ.get("ATTN_PRECISION", "fp32")
25
+
26
+
27
+ def uniq(arr):
28
+ return {el: True for el in arr}.keys()
29
+
30
+
31
+ def default(val, d):
32
+ if val is not None:
33
+ return val
34
+ return d() if isfunction(d) else d
35
+
36
+
37
+ def max_neg_value(t):
38
+ return -torch.finfo(t.dtype).max
39
+
40
+
41
+ def init_(tensor):
42
+ dim = tensor.shape[-1]
43
+ std = 1 / math.sqrt(dim)
44
+ tensor.uniform_(-std, std)
45
+ return tensor
46
+
47
+
48
+ # feedforward
49
+ class GEGLU(nn.Module):
50
+
51
+ def __init__(self, dim_in, dim_out):
52
+ super().__init__()
53
+ self.proj = nn.Linear(dim_in, dim_out * 2)
54
+
55
+ def forward(self, x):
56
+ x, gate = self.proj(x).chunk(2, dim=-1)
57
+ return x * F.gelu(gate)
58
+
59
+
60
+ class FeedForward(nn.Module):
61
+
62
+ def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
63
+ super().__init__()
64
+ inner_dim = int(dim * mult)
65
+ dim_out = default(dim_out, dim)
66
+ project_in = nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU()) if not glu else GEGLU(dim, inner_dim)
67
+
68
+ self.net = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))
69
+
70
+ def forward(self, x):
71
+ return self.net(x)
72
+
73
+
74
+ def zero_module(module):
75
+ """
76
+ Zero out the parameters of a module and return it.
77
+ """
78
+ for p in module.parameters():
79
+ p.detach().zero_()
80
+ return module
81
+
82
+
83
+ def Normalize(in_channels):
84
+ return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
85
+
86
+
87
+ class SpatialSelfAttention(nn.Module):
88
+
89
+ def __init__(self, in_channels):
90
+ super().__init__()
91
+ self.in_channels = in_channels
92
+
93
+ self.norm = Normalize(in_channels)
94
+ self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
95
+ self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
96
+ self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
97
+ self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
98
+
99
+ def forward(self, x):
100
+ h_ = x
101
+ h_ = self.norm(h_)
102
+ q = self.q(h_)
103
+ k = self.k(h_)
104
+ v = self.v(h_)
105
+
106
+ # compute attention
107
+ b, c, h, w = q.shape
108
+ q = rearrange(q, 'b c h w -> b (h w) c')
109
+ k = rearrange(k, 'b c h w -> b c (h w)')
110
+ w_ = torch.einsum('bij,bjk->bik', q, k)
111
+
112
+ w_ = w_ * (int(c)**(-0.5))
113
+ w_ = torch.nn.functional.softmax(w_, dim=2)
114
+
115
+ # attend to values
116
+ v = rearrange(v, 'b c h w -> b c (h w)')
117
+ w_ = rearrange(w_, 'b i j -> b j i')
118
+ h_ = torch.einsum('bij,bjk->bik', v, w_)
119
+ h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h)
120
+ h_ = self.proj_out(h_)
121
+
122
+ return x + h_
123
+
124
+
125
+ class CrossAttention(nn.Module):
126
+
127
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):
128
+ super().__init__()
129
+ inner_dim = dim_head * heads
130
+ context_dim = default(context_dim, query_dim)
131
+
132
+ self.scale = dim_head**-0.5
133
+ self.heads = heads
134
+
135
+ self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
136
+ self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
137
+ self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
138
+
139
+ self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))
140
+
141
+ def forward(self, x, context=None, mask=None):
142
+ h = self.heads
143
+
144
+ q = self.to_q(x)
145
+ context = default(context, x)
146
+ k = self.to_k(context)
147
+ v = self.to_v(context)
148
+
149
+ q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
150
+
151
+ # force cast to fp32 to avoid overflowing
152
+ if _ATTN_PRECISION == "fp32":
153
+ with autocast(enabled=False, device_type='cuda'):
154
+ q, k = q.float(), k.float()
155
+ sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
156
+ else:
157
+ sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
158
+
159
+ del q, k
160
+
161
+ if mask is not None:
162
+ mask = rearrange(mask, 'b ... -> b (...)')
163
+ max_neg_value = -torch.finfo(sim.dtype).max
164
+ mask = repeat(mask, 'b j -> (b h) () j', h=h)
165
+ sim.masked_fill_(~mask, max_neg_value)
166
+
167
+ # attention, what we cannot get enough of
168
+ sim = sim.softmax(dim=-1)
169
+
170
+ out = einsum('b i j, b j d -> b i d', sim, v)
171
+ out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
172
+ return self.to_out(out)
173
+
174
+
175
+ class MemoryEfficientCrossAttention(nn.Module):
176
+ # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
177
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):
178
+ super().__init__()
179
+ # print(f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using {heads} heads.")
180
+ inner_dim = dim_head * heads
181
+ context_dim = default(context_dim, query_dim)
182
+
183
+ self.heads = heads
184
+ self.dim_head = dim_head
185
+
186
+ self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
187
+ self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
188
+ self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
189
+
190
+ self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))
191
+ self.attention_op: Optional[Any] = None
192
+
193
+ def forward(self, x, context=None, mask=None):
194
+ q = self.to_q(x)
195
+ context = default(context, x)
196
+ k = self.to_k(context)
197
+ v = self.to_v(context)
198
+
199
+ b, _, _ = q.shape
200
+ q, k, v = map(
201
+ lambda t: t.unsqueeze(3).reshape(b, t.shape[1], self.heads, self.dim_head).permute(0, 2, 1, 3).reshape(b * self.heads, t.shape[1], self.dim_head).contiguous(),
202
+ (q, k, v),
203
+ )
204
+
205
+ # actually compute the attention, what we cannot get enough of
206
+ out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
207
+
208
+ if mask is not None:
209
+ raise NotImplementedError
210
+ out = (out.unsqueeze(0).reshape(b, self.heads, out.shape[1], self.dim_head).permute(0, 2, 1, 3).reshape(b, out.shape[1], self.heads * self.dim_head))
211
+ return self.to_out(out)
212
+
213
+
214
+ class BasicTransformerBlock(nn.Module):
215
+ ATTENTION_MODES = {
216
+ "softmax": CrossAttention, # vanilla attention
217
+ "softmax-xformers": MemoryEfficientCrossAttention
218
+ }
219
+
220
+ def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True, disable_self_attn=False):
221
+ super().__init__()
222
+ attn_mode = "softmax-xformers" if XFORMERS_IS_AVAILBLE else "softmax"
223
+ assert attn_mode in self.ATTENTION_MODES
224
+ attn_cls = self.ATTENTION_MODES[attn_mode]
225
+ self.disable_self_attn = disable_self_attn
226
+ self.attn1 = attn_cls(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout, context_dim=context_dim if self.disable_self_attn else None) # is a self-attention if not self.disable_self_attn
227
+ self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
228
+ self.attn2 = attn_cls(query_dim=dim, context_dim=context_dim, heads=n_heads, dim_head=d_head, dropout=dropout) # is self-attn if context is none
229
+ self.norm1 = nn.LayerNorm(dim)
230
+ self.norm2 = nn.LayerNorm(dim)
231
+ self.norm3 = nn.LayerNorm(dim)
232
+ self.checkpoint = checkpoint
233
+
234
+ def forward(self, x, context=None):
235
+ return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
236
+
237
+ def _forward(self, x, context=None):
238
+ x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None) + x
239
+ x = self.attn2(self.norm2(x), context=context) + x
240
+ x = self.ff(self.norm3(x)) + x
241
+ return x
242
+
243
+
244
+ class SpatialTransformer(nn.Module):
245
+ """
246
+ Transformer block for image-like data.
247
+ First, project the input (aka embedding)
248
+ and reshape to b, t, d.
249
+ Then apply standard transformer action.
250
+ Finally, reshape to image
251
+ NEW: use_linear for more efficiency instead of the 1x1 convs
252
+ """
253
+
254
+ def __init__(self, in_channels, n_heads, d_head, depth=1, dropout=0., context_dim=None, disable_self_attn=False, use_linear=False, use_checkpoint=True):
255
+ super().__init__()
256
+ assert context_dim is not None
257
+ if not isinstance(context_dim, list):
258
+ context_dim = [context_dim]
259
+ self.in_channels = in_channels
260
+ inner_dim = n_heads * d_head
261
+ self.norm = Normalize(in_channels)
262
+ if not use_linear:
263
+ self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
264
+ else:
265
+ self.proj_in = nn.Linear(in_channels, inner_dim)
266
+
267
+ self.transformer_blocks = nn.ModuleList([BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim[d], disable_self_attn=disable_self_attn, checkpoint=use_checkpoint) for d in range(depth)])
268
+ if not use_linear:
269
+ self.proj_out = zero_module(nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0))
270
+ else:
271
+ self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
272
+ self.use_linear = use_linear
273
+
274
+ def forward(self, x, context=None):
275
+ # note: if no context is given, cross-attention defaults to self-attention
276
+ if not isinstance(context, list):
277
+ context = [context]
278
+ b, c, h, w = x.shape
279
+ x_in = x
280
+ x = self.norm(x)
281
+ if not self.use_linear:
282
+ x = self.proj_in(x)
283
+ x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
284
+ if self.use_linear:
285
+ x = self.proj_in(x)
286
+ for i, block in enumerate(self.transformer_blocks):
287
+ x = block(x, context=context[i])
288
+ if self.use_linear:
289
+ x = self.proj_out(x)
290
+ x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
291
+ if not self.use_linear:
292
+ x = self.proj_out(x)
293
+ return x + x_in
294
+
295
+
296
+ class BasicTransformerBlock3D(BasicTransformerBlock):
297
+
298
+ def forward(self, x, context=None, num_frames=1):
299
+ return checkpoint(self._forward, (x, context, num_frames), self.parameters(), self.checkpoint)
300
+
301
+ def _forward(self, x, context=None, num_frames=1):
302
+ x = rearrange(x, "(b f) l c -> b (f l) c", f=num_frames).contiguous()
303
+ x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None) + x
304
+ x = rearrange(x, "b (f l) c -> (b f) l c", f=num_frames).contiguous()
305
+ x = self.attn2(self.norm2(x), context=context) + x
306
+ x = self.ff(self.norm3(x)) + x
307
+ return x
308
+
309
+
310
+ class SpatialTransformer3D(nn.Module):
311
+ ''' 3D self-attention '''
312
+
313
+ def __init__(self, in_channels, n_heads, d_head, depth=1, dropout=0., context_dim=None, disable_self_attn=False, use_linear=False, use_checkpoint=True):
314
+ super().__init__()
315
+ assert context_dim is not None
316
+ if not isinstance(context_dim, list):
317
+ context_dim = [context_dim]
318
+ self.in_channels = in_channels
319
+ inner_dim = n_heads * d_head
320
+ self.norm = Normalize(in_channels)
321
+ if not use_linear:
322
+ self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
323
+ else:
324
+ self.proj_in = nn.Linear(in_channels, inner_dim)
325
+
326
+ self.transformer_blocks = nn.ModuleList([BasicTransformerBlock3D(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim[d], disable_self_attn=disable_self_attn, checkpoint=use_checkpoint) for d in range(depth)])
327
+ if not use_linear:
328
+ self.proj_out = zero_module(nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0))
329
+ else:
330
+ self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
331
+ self.use_linear = use_linear
332
+
333
+ def forward(self, x, context=None, num_frames=1):
334
+ # note: if no context is given, cross-attention defaults to self-attention
335
+ if not isinstance(context, list):
336
+ context = [context]
337
+ b, c, h, w = x.shape
338
+ x_in = x
339
+ x = self.norm(x)
340
+ if not self.use_linear:
341
+ x = self.proj_in(x)
342
+ x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
343
+ if self.use_linear:
344
+ x = self.proj_in(x)
345
+ for i, block in enumerate(self.transformer_blocks):
346
+ x = block(x, context=context[i], num_frames=num_frames)
347
+ if self.use_linear:
348
+ x = self.proj_out(x)
349
+ x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
350
+ if not self.use_linear:
351
+ x = self.proj_out(x)
352
+ return x + x_in
mvdream/models.py ADDED
@@ -0,0 +1,775 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # obtained and modified from https://github.com/bytedance/MVDream
2
+
3
+ import math
4
+ import numpy as np
5
+ import torch as th
6
+ import torch.nn as nn
7
+ import torch.nn.functional as F
8
+
9
+ from abc import abstractmethod
10
+ from .util import (
11
+ checkpoint,
12
+ conv_nd,
13
+ linear,
14
+ avg_pool_nd,
15
+ zero_module,
16
+ normalization,
17
+ timestep_embedding,
18
+ )
19
+ from .attention import SpatialTransformer, SpatialTransformer3D
20
+ from diffusers.configuration_utils import ConfigMixin
21
+ from diffusers.models.modeling_utils import ModelMixin
22
+ from typing import Any, List, Optional
23
+ from torch import Tensor
24
+
25
+
26
+ class MultiViewUNetWrapperModel(ModelMixin, ConfigMixin):
27
+
28
+ def __init__(self,
29
+ image_size,
30
+ in_channels,
31
+ model_channels,
32
+ out_channels,
33
+ num_res_blocks,
34
+ attention_resolutions,
35
+ dropout=0,
36
+ channel_mult=(1, 2, 4, 8),
37
+ conv_resample=True,
38
+ dims=2,
39
+ num_classes=None,
40
+ use_checkpoint=False,
41
+ num_heads=-1,
42
+ num_head_channels=-1,
43
+ num_heads_upsample=-1,
44
+ use_scale_shift_norm=False,
45
+ resblock_updown=False,
46
+ use_new_attention_order=False,
47
+ use_spatial_transformer=False, # custom transformer support
48
+ transformer_depth=1, # custom transformer support
49
+ context_dim=None, # custom transformer support
50
+ n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
51
+ legacy=True,
52
+ disable_self_attentions=None,
53
+ num_attention_blocks=None,
54
+ disable_middle_self_attn=False,
55
+ use_linear_in_transformer=False,
56
+ adm_in_channels=None,
57
+ camera_dim=None,):
58
+ super().__init__()
59
+ self.unet: MultiViewUNetModel = MultiViewUNetModel(
60
+ image_size=image_size,
61
+ in_channels=in_channels,
62
+ model_channels=model_channels,
63
+ out_channels=out_channels,
64
+ num_res_blocks=num_res_blocks,
65
+ attention_resolutions=attention_resolutions,
66
+ dropout=dropout,
67
+ channel_mult=channel_mult,
68
+ conv_resample=conv_resample,
69
+ dims=dims,
70
+ num_classes=num_classes,
71
+ use_checkpoint=use_checkpoint,
72
+ num_heads=num_heads,
73
+ num_head_channels=num_head_channels,
74
+ num_heads_upsample=num_heads_upsample,
75
+ use_scale_shift_norm=use_scale_shift_norm,
76
+ resblock_updown=resblock_updown,
77
+ use_new_attention_order=use_new_attention_order,
78
+ use_spatial_transformer=use_spatial_transformer,
79
+ transformer_depth=transformer_depth,
80
+ context_dim=context_dim,
81
+ n_embed=n_embed,
82
+ legacy=legacy,
83
+ disable_self_attentions=disable_self_attentions,
84
+ num_attention_blocks=num_attention_blocks,
85
+ disable_middle_self_attn=disable_middle_self_attn,
86
+ use_linear_in_transformer=use_linear_in_transformer,
87
+ adm_in_channels=adm_in_channels,
88
+ camera_dim=camera_dim,
89
+ )
90
+
91
+ def forward(self, *args, **kwargs):
92
+ return self.unet(*args, **kwargs)
93
+
94
+
95
+ class TimestepBlock(nn.Module):
96
+ """
97
+ Any module where forward() takes timestep embeddings as a second argument.
98
+ """
99
+
100
+ @abstractmethod
101
+ def forward(self, x, emb):
102
+ """
103
+ Apply the module to `x` given `emb` timestep embeddings.
104
+ """
105
+
106
+
107
+ class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
108
+ """
109
+ A sequential module that passes timestep embeddings to the children that
110
+ support it as an extra input.
111
+ """
112
+
113
+ def forward(self, x, emb, context=None, num_frames=1):
114
+ for layer in self:
115
+ if isinstance(layer, TimestepBlock):
116
+ x = layer(x, emb)
117
+ elif isinstance(layer, SpatialTransformer3D):
118
+ x = layer(x, context, num_frames=num_frames)
119
+ elif isinstance(layer, SpatialTransformer):
120
+ x = layer(x, context)
121
+ else:
122
+ x = layer(x)
123
+ return x
124
+
125
+
126
+ class Upsample(nn.Module):
127
+ """
128
+ An upsampling layer with an optional convolution.
129
+ :param channels: channels in the inputs and outputs.
130
+ :param use_conv: a bool determining if a convolution is applied.
131
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
132
+ upsampling occurs in the inner-two dimensions.
133
+ """
134
+
135
+ def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
136
+ super().__init__()
137
+ self.channels = channels
138
+ self.out_channels = out_channels or channels
139
+ self.use_conv = use_conv
140
+ self.dims = dims
141
+ if use_conv:
142
+ self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=padding)
143
+
144
+ def forward(self, x):
145
+ assert x.shape[1] == self.channels
146
+ if self.dims == 3:
147
+ x = F.interpolate(x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest")
148
+ else:
149
+ x = F.interpolate(x, scale_factor=2, mode="nearest")
150
+ if self.use_conv:
151
+ x = self.conv(x)
152
+ return x
153
+
154
+
155
+ class Downsample(nn.Module):
156
+ """
157
+ A downsampling layer with an optional convolution.
158
+ :param channels: channels in the inputs and outputs.
159
+ :param use_conv: a bool determining if a convolution is applied.
160
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
161
+ downsampling occurs in the inner-two dimensions.
162
+ """
163
+
164
+ def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
165
+ super().__init__()
166
+ self.channels = channels
167
+ self.out_channels = out_channels or channels
168
+ self.use_conv = use_conv
169
+ self.dims = dims
170
+ stride = 2 if dims != 3 else (1, 2, 2)
171
+ if use_conv:
172
+ self.op = conv_nd(dims, self.channels, self.out_channels, 3, stride=stride, padding=padding)
173
+ else:
174
+ assert self.channels == self.out_channels
175
+ self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
176
+
177
+ def forward(self, x):
178
+ assert x.shape[1] == self.channels
179
+ return self.op(x)
180
+
181
+
182
+ class ResBlock(TimestepBlock):
183
+ """
184
+ A residual block that can optionally change the number of channels.
185
+ :param channels: the number of input channels.
186
+ :param emb_channels: the number of timestep embedding channels.
187
+ :param dropout: the rate of dropout.
188
+ :param out_channels: if specified, the number of out channels.
189
+ :param use_conv: if True and out_channels is specified, use a spatial
190
+ convolution instead of a smaller 1x1 convolution to change the
191
+ channels in the skip connection.
192
+ :param dims: determines if the signal is 1D, 2D, or 3D.
193
+ :param use_checkpoint: if True, use gradient checkpointing on this module.
194
+ :param up: if True, use this block for upsampling.
195
+ :param down: if True, use this block for downsampling.
196
+ """
197
+
198
+ def __init__(
199
+ self,
200
+ channels,
201
+ emb_channels,
202
+ dropout,
203
+ out_channels=None,
204
+ use_conv=False,
205
+ use_scale_shift_norm=False,
206
+ dims=2,
207
+ use_checkpoint=False,
208
+ up=False,
209
+ down=False,
210
+ ):
211
+ super().__init__()
212
+ self.channels = channels
213
+ self.emb_channels = emb_channels
214
+ self.dropout = dropout
215
+ self.out_channels = out_channels or channels
216
+ self.use_conv = use_conv
217
+ self.use_checkpoint = use_checkpoint
218
+ self.use_scale_shift_norm = use_scale_shift_norm
219
+
220
+ self.in_layers = nn.Sequential(
221
+ normalization(channels),
222
+ nn.SiLU(),
223
+ conv_nd(dims, channels, self.out_channels, 3, padding=1),
224
+ )
225
+
226
+ self.updown = up or down
227
+
228
+ if up:
229
+ self.h_upd = Upsample(channels, False, dims)
230
+ self.x_upd = Upsample(channels, False, dims)
231
+ elif down:
232
+ self.h_upd = Downsample(channels, False, dims)
233
+ self.x_upd = Downsample(channels, False, dims)
234
+ else:
235
+ self.h_upd = self.x_upd = nn.Identity()
236
+
237
+ self.emb_layers = nn.Sequential(
238
+ nn.SiLU(),
239
+ linear(
240
+ emb_channels,
241
+ 2 * self.out_channels if use_scale_shift_norm else self.out_channels,
242
+ ),
243
+ )
244
+ self.out_layers = nn.Sequential(
245
+ normalization(self.out_channels),
246
+ nn.SiLU(),
247
+ nn.Dropout(p=dropout),
248
+ zero_module(conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)),
249
+ )
250
+
251
+ if self.out_channels == channels:
252
+ self.skip_connection = nn.Identity()
253
+ elif use_conv:
254
+ self.skip_connection = conv_nd(dims, channels, self.out_channels, 3, padding=1)
255
+ else:
256
+ self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
257
+
258
+ def forward(self, x, emb):
259
+ """
260
+ Apply the block to a Tensor, conditioned on a timestep embedding.
261
+ :param x: an [N x C x ...] Tensor of features.
262
+ :param emb: an [N x emb_channels] Tensor of timestep embeddings.
263
+ :return: an [N x C x ...] Tensor of outputs.
264
+ """
265
+ return checkpoint(self._forward, (x, emb), self.parameters(), self.use_checkpoint)
266
+
267
+ def _forward(self, x, emb):
268
+ if self.updown:
269
+ in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
270
+ h = in_rest(x)
271
+ h = self.h_upd(h)
272
+ x = self.x_upd(x)
273
+ h = in_conv(h)
274
+ else:
275
+ h = self.in_layers(x)
276
+ emb_out = self.emb_layers(emb).type(h.dtype)
277
+ while len(emb_out.shape) < len(h.shape):
278
+ emb_out = emb_out[..., None]
279
+ if self.use_scale_shift_norm:
280
+ out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
281
+ scale, shift = th.chunk(emb_out, 2, dim=1)
282
+ h = out_norm(h) * (1 + scale) + shift
283
+ h = out_rest(h)
284
+ else:
285
+ h = h + emb_out
286
+ h = self.out_layers(h)
287
+ return self.skip_connection(x) + h
288
+
289
+
290
+ class AttentionBlock(nn.Module):
291
+ """
292
+ An attention block that allows spatial positions to attend to each other.
293
+ Originally ported from here, but adapted to the N-d case.
294
+ https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
295
+ """
296
+
297
+ def __init__(
298
+ self,
299
+ channels,
300
+ num_heads=1,
301
+ num_head_channels=-1,
302
+ use_checkpoint=False,
303
+ use_new_attention_order=False,
304
+ ):
305
+ super().__init__()
306
+ self.channels = channels
307
+ if num_head_channels == -1:
308
+ self.num_heads = num_heads
309
+ else:
310
+ assert (channels % num_head_channels == 0), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
311
+ self.num_heads = channels // num_head_channels
312
+ self.use_checkpoint = use_checkpoint
313
+ self.norm = normalization(channels)
314
+ self.qkv = conv_nd(1, channels, channels * 3, 1)
315
+ if use_new_attention_order:
316
+ # split qkv before split heads
317
+ self.attention = QKVAttention(self.num_heads)
318
+ else:
319
+ # split heads before split qkv
320
+ self.attention = QKVAttentionLegacy(self.num_heads)
321
+
322
+ self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
323
+
324
+ def forward(self, x):
325
+ return checkpoint(self._forward, (x,), self.parameters(), True) # TODO: check checkpoint usage, is True # TODO: fix the .half call!!!
326
+ #return pt_checkpoint(self._forward, x) # pytorch
327
+
328
+ def _forward(self, x):
329
+ b, c, *spatial = x.shape
330
+ x = x.reshape(b, c, -1)
331
+ qkv = self.qkv(self.norm(x))
332
+ h = self.attention(qkv)
333
+ h = self.proj_out(h)
334
+ return (x + h).reshape(b, c, *spatial)
335
+
336
+
337
+ def count_flops_attn(model, _x, y):
338
+ """
339
+ A counter for the `thop` package to count the operations in an
340
+ attention operation.
341
+ Meant to be used like:
342
+ macs, params = thop.profile(
343
+ model,
344
+ inputs=(inputs, timestamps),
345
+ custom_ops={QKVAttention: QKVAttention.count_flops},
346
+ )
347
+ """
348
+ b, c, *spatial = y[0].shape
349
+ num_spatial = int(np.prod(spatial))
350
+ # We perform two matmuls with the same number of ops.
351
+ # The first computes the weight matrix, the second computes
352
+ # the combination of the value vectors.
353
+ matmul_ops = 2 * b * (num_spatial**2) * c
354
+ model.total_ops += th.DoubleTensor([matmul_ops])
355
+
356
+
357
+ class QKVAttentionLegacy(nn.Module):
358
+ """
359
+ A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping
360
+ """
361
+
362
+ def __init__(self, n_heads):
363
+ super().__init__()
364
+ self.n_heads = n_heads
365
+
366
+ def forward(self, qkv):
367
+ """
368
+ Apply QKV attention.
369
+ :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
370
+ :return: an [N x (H * C) x T] tensor after attention.
371
+ """
372
+ bs, width, length = qkv.shape
373
+ assert width % (3 * self.n_heads) == 0
374
+ ch = width // (3 * self.n_heads)
375
+ q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
376
+ scale = 1 / math.sqrt(math.sqrt(ch))
377
+ weight = th.einsum("bct,bcs->bts", q * scale, k * scale) # More stable with f16 than dividing afterwards
378
+ weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
379
+ a = th.einsum("bts,bcs->bct", weight, v)
380
+ return a.reshape(bs, -1, length)
381
+
382
+ @staticmethod
383
+ def count_flops(model, _x, y):
384
+ return count_flops_attn(model, _x, y)
385
+
386
+
387
+ class QKVAttention(nn.Module):
388
+ """
389
+ A module which performs QKV attention and splits in a different order.
390
+ """
391
+
392
+ def __init__(self, n_heads):
393
+ super().__init__()
394
+ self.n_heads = n_heads
395
+
396
+ def forward(self, qkv):
397
+ """
398
+ Apply QKV attention.
399
+ :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.
400
+ :return: an [N x (H * C) x T] tensor after attention.
401
+ """
402
+ bs, width, length = qkv.shape
403
+ assert width % (3 * self.n_heads) == 0
404
+ ch = width // (3 * self.n_heads)
405
+ q, k, v = qkv.chunk(3, dim=1)
406
+ scale = 1 / math.sqrt(math.sqrt(ch))
407
+ weight = th.einsum(
408
+ "bct,bcs->bts",
409
+ (q * scale).view(bs * self.n_heads, ch, length),
410
+ (k * scale).view(bs * self.n_heads, ch, length),
411
+ ) # More stable with f16 than dividing afterwards
412
+ weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
413
+ a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length))
414
+ return a.reshape(bs, -1, length)
415
+
416
+ @staticmethod
417
+ def count_flops(model, _x, y):
418
+ return count_flops_attn(model, _x, y)
419
+
420
+
421
+ class Timestep(nn.Module):
422
+
423
+ def __init__(self, dim):
424
+ super().__init__()
425
+ self.dim = dim
426
+
427
+ def forward(self, t):
428
+ return timestep_embedding(t, self.dim)
429
+
430
+
431
+ class MultiViewUNetModel(nn.Module):
432
+ """
433
+ The full multi-view UNet model with attention, timestep embedding and camera embedding.
434
+ :param in_channels: channels in the input Tensor.
435
+ :param model_channels: base channel count for the model.
436
+ :param out_channels: channels in the output Tensor.
437
+ :param num_res_blocks: number of residual blocks per downsample.
438
+ :param attention_resolutions: a collection of downsample rates at which
439
+ attention will take place. May be a set, list, or tuple.
440
+ For example, if this contains 4, then at 4x downsampling, attention
441
+ will be used.
442
+ :param dropout: the dropout probability.
443
+ :param channel_mult: channel multiplier for each level of the UNet.
444
+ :param conv_resample: if True, use learned convolutions for upsampling and
445
+ downsampling.
446
+ :param dims: determines if the signal is 1D, 2D, or 3D.
447
+ :param num_classes: if specified (as an int), then this model will be
448
+ class-conditional with `num_classes` classes.
449
+ :param use_checkpoint: use gradient checkpointing to reduce memory usage.
450
+ :param num_heads: the number of attention heads in each attention layer.
451
+ :param num_heads_channels: if specified, ignore num_heads and instead use
452
+ a fixed channel width per attention head.
453
+ :param num_heads_upsample: works with num_heads to set a different number
454
+ of heads for upsampling. Deprecated.
455
+ :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
456
+ :param resblock_updown: use residual blocks for up/downsampling.
457
+ :param use_new_attention_order: use a different attention pattern for potentially
458
+ increased efficiency.
459
+ :param camera_dim: dimensionality of camera input.
460
+ """
461
+
462
+ def __init__(
463
+ self,
464
+ image_size,
465
+ in_channels,
466
+ model_channels,
467
+ out_channels,
468
+ num_res_blocks,
469
+ attention_resolutions,
470
+ dropout=0,
471
+ channel_mult=(1, 2, 4, 8),
472
+ conv_resample=True,
473
+ dims=2,
474
+ num_classes=None,
475
+ use_checkpoint=False,
476
+ num_heads=-1,
477
+ num_head_channels=-1,
478
+ num_heads_upsample=-1,
479
+ use_scale_shift_norm=False,
480
+ resblock_updown=False,
481
+ use_new_attention_order=False,
482
+ use_spatial_transformer=False, # custom transformer support
483
+ transformer_depth=1, # custom transformer support
484
+ context_dim=None, # custom transformer support
485
+ n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
486
+ legacy=True,
487
+ disable_self_attentions=None,
488
+ num_attention_blocks=None,
489
+ disable_middle_self_attn=False,
490
+ use_linear_in_transformer=False,
491
+ adm_in_channels=None,
492
+ camera_dim=None,
493
+ ):
494
+ super().__init__()
495
+ if use_spatial_transformer:
496
+ assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
497
+
498
+ if context_dim is not None:
499
+ assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
500
+ from omegaconf.listconfig import ListConfig
501
+ if type(context_dim) == ListConfig:
502
+ context_dim = list(context_dim)
503
+
504
+ if num_heads_upsample == -1:
505
+ num_heads_upsample = num_heads
506
+
507
+ if num_heads == -1:
508
+ assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
509
+
510
+ if num_head_channels == -1:
511
+ assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
512
+
513
+ self.image_size = image_size
514
+ self.in_channels = in_channels
515
+ self.model_channels = model_channels
516
+ self.out_channels = out_channels
517
+ if isinstance(num_res_blocks, int):
518
+ self.num_res_blocks = len(channel_mult) * [num_res_blocks]
519
+ else:
520
+ if len(num_res_blocks) != len(channel_mult):
521
+ raise ValueError("provide num_res_blocks either as an int (globally constant) or "
522
+ "as a list/tuple (per-level) with the same length as channel_mult")
523
+ self.num_res_blocks = num_res_blocks
524
+ if disable_self_attentions is not None:
525
+ # should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
526
+ assert len(disable_self_attentions) == len(channel_mult)
527
+ if num_attention_blocks is not None:
528
+ assert len(num_attention_blocks) == len(self.num_res_blocks)
529
+ assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks))))
530
+ print(f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
531
+ f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
532
+ f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
533
+ f"attention will still not be set.")
534
+
535
+ self.attention_resolutions = attention_resolutions
536
+ self.dropout = dropout
537
+ self.channel_mult = channel_mult
538
+ self.conv_resample = conv_resample
539
+ self.num_classes = num_classes
540
+ self.use_checkpoint = use_checkpoint
541
+ self.num_heads = num_heads
542
+ self.num_head_channels = num_head_channels
543
+ self.num_heads_upsample = num_heads_upsample
544
+ self.predict_codebook_ids = n_embed is not None
545
+
546
+ time_embed_dim = model_channels * 4
547
+ self.time_embed = nn.Sequential(
548
+ linear(model_channels, time_embed_dim),
549
+ nn.SiLU(),
550
+ linear(time_embed_dim, time_embed_dim),
551
+ )
552
+
553
+ if camera_dim is not None:
554
+ time_embed_dim = model_channels * 4
555
+ self.camera_embed = nn.Sequential(
556
+ linear(camera_dim, time_embed_dim),
557
+ nn.SiLU(),
558
+ linear(time_embed_dim, time_embed_dim),
559
+ )
560
+
561
+ if self.num_classes is not None:
562
+ if isinstance(self.num_classes, int):
563
+ self.label_emb = nn.Embedding(self.num_classes, time_embed_dim)
564
+ elif self.num_classes == "continuous":
565
+ # print("setting up linear c_adm embedding layer")
566
+ self.label_emb = nn.Linear(1, time_embed_dim)
567
+ elif self.num_classes == "sequential":
568
+ assert adm_in_channels is not None
569
+ self.label_emb = nn.Sequential(nn.Sequential(
570
+ linear(adm_in_channels, time_embed_dim),
571
+ nn.SiLU(),
572
+ linear(time_embed_dim, time_embed_dim),
573
+ ))
574
+ else:
575
+ raise ValueError()
576
+
577
+ self.input_blocks = nn.ModuleList([TimestepEmbedSequential(conv_nd(dims, in_channels, model_channels, 3, padding=1))])
578
+ self._feature_size = model_channels
579
+ input_block_chans = [model_channels]
580
+ ch = model_channels
581
+ ds = 1
582
+ for level, mult in enumerate(channel_mult):
583
+ for nr in range(self.num_res_blocks[level]):
584
+ layers: List[Any] = [ResBlock(
585
+ ch,
586
+ time_embed_dim,
587
+ dropout,
588
+ out_channels=mult * model_channels,
589
+ dims=dims,
590
+ use_checkpoint=use_checkpoint,
591
+ use_scale_shift_norm=use_scale_shift_norm,
592
+ )]
593
+ ch = mult * model_channels
594
+ if ds in attention_resolutions:
595
+ if num_head_channels == -1:
596
+ dim_head = ch // num_heads
597
+ else:
598
+ num_heads = ch // num_head_channels
599
+ dim_head = num_head_channels
600
+ if legacy:
601
+ #num_heads = 1
602
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
603
+ if disable_self_attentions is not None:
604
+ disabled_sa = disable_self_attentions[level]
605
+ else:
606
+ disabled_sa = False
607
+
608
+ if num_attention_blocks is None or nr < num_attention_blocks[level]:
609
+ layers.append(AttentionBlock(
610
+ ch,
611
+ use_checkpoint=use_checkpoint,
612
+ num_heads=num_heads,
613
+ num_head_channels=dim_head,
614
+ use_new_attention_order=use_new_attention_order,
615
+ ) if not use_spatial_transformer else SpatialTransformer3D(ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim, disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer, use_checkpoint=use_checkpoint))
616
+ self.input_blocks.append(TimestepEmbedSequential(*layers))
617
+ self._feature_size += ch
618
+ input_block_chans.append(ch)
619
+ if level != len(channel_mult) - 1:
620
+ out_ch = ch
621
+ self.input_blocks.append(TimestepEmbedSequential(ResBlock(
622
+ ch,
623
+ time_embed_dim,
624
+ dropout,
625
+ out_channels=out_ch,
626
+ dims=dims,
627
+ use_checkpoint=use_checkpoint,
628
+ use_scale_shift_norm=use_scale_shift_norm,
629
+ down=True,
630
+ ) if resblock_updown else Downsample(ch, conv_resample, dims=dims, out_channels=out_ch)))
631
+ ch = out_ch
632
+ input_block_chans.append(ch)
633
+ ds *= 2
634
+ self._feature_size += ch
635
+
636
+ if num_head_channels == -1:
637
+ dim_head = ch // num_heads
638
+ else:
639
+ num_heads = ch // num_head_channels
640
+ dim_head = num_head_channels
641
+ if legacy:
642
+ #num_heads = 1
643
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
644
+ self.middle_block = TimestepEmbedSequential(
645
+ ResBlock(
646
+ ch,
647
+ time_embed_dim,
648
+ dropout,
649
+ dims=dims,
650
+ use_checkpoint=use_checkpoint,
651
+ use_scale_shift_norm=use_scale_shift_norm,
652
+ ),
653
+ AttentionBlock(
654
+ ch,
655
+ use_checkpoint=use_checkpoint,
656
+ num_heads=num_heads,
657
+ num_head_channels=dim_head,
658
+ use_new_attention_order=use_new_attention_order,
659
+ ) if not use_spatial_transformer else SpatialTransformer3D( # always uses a self-attn
660
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim, disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer, use_checkpoint=use_checkpoint),
661
+ ResBlock(
662
+ ch,
663
+ time_embed_dim,
664
+ dropout,
665
+ dims=dims,
666
+ use_checkpoint=use_checkpoint,
667
+ use_scale_shift_norm=use_scale_shift_norm,
668
+ ),
669
+ )
670
+ self._feature_size += ch
671
+
672
+ self.output_blocks = nn.ModuleList([])
673
+ for level, mult in list(enumerate(channel_mult))[::-1]:
674
+ for i in range(self.num_res_blocks[level] + 1):
675
+ ich = input_block_chans.pop()
676
+ layers = [ResBlock(
677
+ ch + ich,
678
+ time_embed_dim,
679
+ dropout,
680
+ out_channels=model_channels * mult,
681
+ dims=dims,
682
+ use_checkpoint=use_checkpoint,
683
+ use_scale_shift_norm=use_scale_shift_norm,
684
+ )]
685
+ ch = model_channels * mult
686
+ if ds in attention_resolutions:
687
+ if num_head_channels == -1:
688
+ dim_head = ch // num_heads
689
+ else:
690
+ num_heads = ch // num_head_channels
691
+ dim_head = num_head_channels
692
+ if legacy:
693
+ #num_heads = 1
694
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
695
+ if disable_self_attentions is not None:
696
+ disabled_sa = disable_self_attentions[level]
697
+ else:
698
+ disabled_sa = False
699
+
700
+ if num_attention_blocks is None or i < num_attention_blocks[level]:
701
+ layers.append(AttentionBlock(
702
+ ch,
703
+ use_checkpoint=use_checkpoint,
704
+ num_heads=num_heads_upsample,
705
+ num_head_channels=dim_head,
706
+ use_new_attention_order=use_new_attention_order,
707
+ ) if not use_spatial_transformer else SpatialTransformer3D(ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim, disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer, use_checkpoint=use_checkpoint))
708
+ if level and i == self.num_res_blocks[level]:
709
+ out_ch = ch
710
+ layers.append(ResBlock(
711
+ ch,
712
+ time_embed_dim,
713
+ dropout,
714
+ out_channels=out_ch,
715
+ dims=dims,
716
+ use_checkpoint=use_checkpoint,
717
+ use_scale_shift_norm=use_scale_shift_norm,
718
+ up=True,
719
+ ) if resblock_updown else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch))
720
+ ds //= 2
721
+ self.output_blocks.append(TimestepEmbedSequential(*layers))
722
+ self._feature_size += ch
723
+
724
+ self.out = nn.Sequential(
725
+ normalization(ch),
726
+ nn.SiLU(),
727
+ zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
728
+ )
729
+ if self.predict_codebook_ids:
730
+ self.id_predictor = nn.Sequential(
731
+ normalization(ch),
732
+ conv_nd(dims, model_channels, n_embed, 1),
733
+ #nn.LogSoftmax(dim=1) # change to cross_entropy and produce non-normalized logits
734
+ )
735
+
736
+ def forward(self, x, timesteps=None, context=None, y: Optional[Tensor] = None, camera=None, num_frames=1, **kwargs):
737
+ """
738
+ Apply the model to an input batch.
739
+ :param x: an [(N x F) x C x ...] Tensor of inputs. F is the number of frames (views).
740
+ :param timesteps: a 1-D batch of timesteps.
741
+ :param context: conditioning plugged in via crossattn
742
+ :param y: an [N] Tensor of labels, if class-conditional.
743
+ :param num_frames: a integer indicating number of frames for tensor reshaping.
744
+ :return: an [(N x F) x C x ...] Tensor of outputs. F is the number of frames (views).
745
+ """
746
+ assert x.shape[0] % num_frames == 0, "[UNet] input batch size must be dividable by num_frames!"
747
+ assert (y is not None) == (self.num_classes is not None), "must specify y if and only if the model is class-conditional"
748
+ hs = []
749
+ t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
750
+
751
+ emb = self.time_embed(t_emb)
752
+
753
+ if self.num_classes is not None:
754
+ assert y is not None
755
+ assert y.shape[0] == x.shape[0]
756
+ emb = emb + self.label_emb(y)
757
+
758
+ # Add camera embeddings
759
+ if camera is not None:
760
+ assert camera.shape[0] == emb.shape[0]
761
+ emb = emb + self.camera_embed(camera)
762
+
763
+ h = x
764
+ for module in self.input_blocks:
765
+ h = module(h, emb, context, num_frames=num_frames)
766
+ hs.append(h)
767
+ h = self.middle_block(h, emb, context, num_frames=num_frames)
768
+ for module in self.output_blocks:
769
+ h = th.cat([h, hs.pop()], dim=1)
770
+ h = module(h, emb, context, num_frames=num_frames)
771
+ h = h.type(x.dtype)
772
+ if self.predict_codebook_ids:
773
+ return self.id_predictor(h)
774
+ else:
775
+ return self.out(h)
mvdream/pipeline_mvdream.py ADDED
@@ -0,0 +1,484 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+ import inspect
4
+ from typing import Callable, List, Optional, Union
5
+ from transformers import CLIPTextModel, CLIPTokenizer
6
+ from diffusers import AutoencoderKL, DiffusionPipeline
7
+ from diffusers.utils import (
8
+ deprecate,
9
+ is_accelerate_available,
10
+ is_accelerate_version,
11
+ logging,
12
+ )
13
+ from diffusers.configuration_utils import FrozenDict
14
+ from diffusers.schedulers import DDIMScheduler
15
+ try:
16
+ from diffusers import randn_tensor # old import # type: ignore
17
+ except ImportError:
18
+ from diffusers.utils.torch_utils import randn_tensor # new import # type: ignore
19
+
20
+ from .models import MultiViewUNetWrapperModel
21
+ from accelerate.utils import set_module_tensor_to_device
22
+
23
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
24
+
25
+ def create_camera_to_world_matrix(elevation, azimuth):
26
+ elevation = np.radians(elevation)
27
+ azimuth = np.radians(azimuth)
28
+ # Convert elevation and azimuth angles to Cartesian coordinates on a unit sphere
29
+ x = np.cos(elevation) * np.sin(azimuth)
30
+ y = np.sin(elevation)
31
+ z = np.cos(elevation) * np.cos(azimuth)
32
+
33
+ # Calculate camera position, target, and up vectors
34
+ camera_pos = np.array([x, y, z])
35
+ target = np.array([0, 0, 0])
36
+ up = np.array([0, 1, 0])
37
+
38
+ # Construct view matrix
39
+ forward = target - camera_pos
40
+ forward /= np.linalg.norm(forward)
41
+ right = np.cross(forward, up)
42
+ right /= np.linalg.norm(right)
43
+ new_up = np.cross(right, forward)
44
+ new_up /= np.linalg.norm(new_up)
45
+ cam2world = np.eye(4)
46
+ cam2world[:3, :3] = np.array([right, new_up, -forward]).T
47
+ cam2world[:3, 3] = camera_pos
48
+ return cam2world
49
+
50
+
51
+ def convert_opengl_to_blender(camera_matrix):
52
+ if isinstance(camera_matrix, np.ndarray):
53
+ # Construct transformation matrix to convert from OpenGL space to Blender space
54
+ flip_yz = np.array([[1, 0, 0, 0], [0, 0, -1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
55
+ camera_matrix_blender = np.dot(flip_yz, camera_matrix)
56
+ else:
57
+ # Construct transformation matrix to convert from OpenGL space to Blender space
58
+ flip_yz = torch.tensor([[1, 0, 0, 0], [0, 0, -1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
59
+ if camera_matrix.ndim == 3:
60
+ flip_yz = flip_yz.unsqueeze(0)
61
+ camera_matrix_blender = torch.matmul(flip_yz.to(camera_matrix), camera_matrix)
62
+ return camera_matrix_blender
63
+
64
+
65
+ def get_camera(num_frames, elevation=15, azimuth_start=0, azimuth_span=360, blender_coord=True):
66
+ angle_gap = azimuth_span / num_frames
67
+ cameras = []
68
+ for azimuth in np.arange(azimuth_start, azimuth_span + azimuth_start, angle_gap):
69
+ camera_matrix = create_camera_to_world_matrix(elevation, azimuth)
70
+ if blender_coord:
71
+ camera_matrix = convert_opengl_to_blender(camera_matrix)
72
+ cameras.append(camera_matrix.flatten())
73
+ return torch.tensor(np.stack(cameras, 0)).float()
74
+
75
+
76
+ class MVDreamStableDiffusionPipeline(DiffusionPipeline):
77
+
78
+ def __init__(
79
+ self,
80
+ vae: AutoencoderKL,
81
+ unet: MultiViewUNetWrapperModel,
82
+ tokenizer: CLIPTokenizer,
83
+ text_encoder: CLIPTextModel,
84
+ scheduler: DDIMScheduler,
85
+ requires_safety_checker: bool = False,
86
+ ):
87
+ super().__init__()
88
+
89
+ if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1: # type: ignore
90
+ deprecation_message = (f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
91
+ f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " # type: ignore
92
+ "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
93
+ " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
94
+ " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
95
+ " file")
96
+ deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)
97
+ new_config = dict(scheduler.config)
98
+ new_config["steps_offset"] = 1
99
+ scheduler._internal_dict = FrozenDict(new_config)
100
+
101
+ if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True: # type: ignore
102
+ deprecation_message = (f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
103
+ " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
104
+ " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
105
+ " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
106
+ " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file")
107
+ deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)
108
+ new_config = dict(scheduler.config)
109
+ new_config["clip_sample"] = False
110
+ scheduler._internal_dict = FrozenDict(new_config)
111
+
112
+ self.register_modules(
113
+ vae=vae,
114
+ unet=unet,
115
+ scheduler=scheduler,
116
+ tokenizer=tokenizer,
117
+ text_encoder=text_encoder,
118
+ )
119
+ self.vae_scale_factor = 2**(len(self.vae.config.block_out_channels) - 1)
120
+ self.register_to_config(requires_safety_checker=requires_safety_checker)
121
+
122
+ def enable_vae_slicing(self):
123
+ r"""
124
+ Enable sliced VAE decoding.
125
+
126
+ When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several
127
+ steps. This is useful to save some memory and allow larger batch sizes.
128
+ """
129
+ self.vae.enable_slicing()
130
+
131
+ def disable_vae_slicing(self):
132
+ r"""
133
+ Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to
134
+ computing decoding in one step.
135
+ """
136
+ self.vae.disable_slicing()
137
+
138
+ def enable_vae_tiling(self):
139
+ r"""
140
+ Enable tiled VAE decoding.
141
+
142
+ When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in
143
+ several steps. This is useful to save a large amount of memory and to allow the processing of larger images.
144
+ """
145
+ self.vae.enable_tiling()
146
+
147
+ def disable_vae_tiling(self):
148
+ r"""
149
+ Disable tiled VAE decoding. If `enable_vae_tiling` was previously invoked, this method will go back to
150
+ computing decoding in one step.
151
+ """
152
+ self.vae.disable_tiling()
153
+
154
+ def enable_sequential_cpu_offload(self, gpu_id=0):
155
+ r"""
156
+ Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
157
+ text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
158
+ `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
159
+ Note that offloading happens on a submodule basis. Memory savings are higher than with
160
+ `enable_model_cpu_offload`, but performance is lower.
161
+ """
162
+ if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):
163
+ from accelerate import cpu_offload
164
+ else:
165
+ raise ImportError("`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher")
166
+
167
+ device = torch.device(f"cuda:{gpu_id}")
168
+
169
+ if self.device.type != "cpu":
170
+ self.to("cpu", silence_dtype_warnings=True)
171
+ torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
172
+
173
+ for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
174
+ cpu_offload(cpu_offloaded_model, device)
175
+
176
+ def enable_model_cpu_offload(self, gpu_id=0):
177
+ r"""
178
+ Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
179
+ to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
180
+ method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
181
+ `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
182
+ """
183
+ if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
184
+ from accelerate import cpu_offload_with_hook
185
+ else:
186
+ raise ImportError("`enable_model_offload` requires `accelerate v0.17.0` or higher.")
187
+
188
+ device = torch.device(f"cuda:{gpu_id}")
189
+
190
+ if self.device.type != "cpu":
191
+ self.to("cpu", silence_dtype_warnings=True)
192
+ torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
193
+
194
+ hook = None
195
+ for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:
196
+ _, hook = cpu_offload_with_hook(cpu_offloaded_model, device, prev_module_hook=hook)
197
+
198
+ # We'll offload the last model manually.
199
+ self.final_offload_hook = hook
200
+
201
+ @property
202
+ def _execution_device(self):
203
+ r"""
204
+ Returns the device on which the pipeline's models will be executed. After calling
205
+ `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
206
+ hooks.
207
+ """
208
+ if not hasattr(self.unet, "_hf_hook"):
209
+ return self.device
210
+ for module in self.unet.modules():
211
+ if (hasattr(module, "_hf_hook") and hasattr(module._hf_hook, "execution_device") and module._hf_hook.execution_device is not None):
212
+ return torch.device(module._hf_hook.execution_device)
213
+ return self.device
214
+
215
+ def _encode_prompt(
216
+ self,
217
+ prompt,
218
+ device,
219
+ num_images_per_prompt,
220
+ do_classifier_free_guidance: bool,
221
+ negative_prompt=None,
222
+ ):
223
+ r"""
224
+ Encodes the prompt into text encoder hidden states.
225
+
226
+ Args:
227
+ prompt (`str` or `List[str]`, *optional*):
228
+ prompt to be encoded
229
+ device: (`torch.device`):
230
+ torch device
231
+ num_images_per_prompt (`int`):
232
+ number of images that should be generated per prompt
233
+ do_classifier_free_guidance (`bool`):
234
+ whether to use classifier free guidance or not
235
+ negative_prompt (`str` or `List[str]`, *optional*):
236
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
237
+ `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
238
+ Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
239
+ prompt_embeds (`torch.FloatTensor`, *optional*):
240
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
241
+ provided, text embeddings will be generated from `prompt` input argument.
242
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
243
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
244
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
245
+ argument.
246
+ """
247
+ if prompt is not None and isinstance(prompt, str):
248
+ batch_size = 1
249
+ elif prompt is not None and isinstance(prompt, list):
250
+ batch_size = len(prompt)
251
+ else:
252
+ raise ValueError(f"`prompt` should be either a string or a list of strings, but got {type(prompt)}.")
253
+
254
+ text_inputs = self.tokenizer(
255
+ prompt,
256
+ padding="max_length",
257
+ max_length=self.tokenizer.model_max_length,
258
+ truncation=True,
259
+ return_tensors="pt",
260
+ )
261
+ text_input_ids = text_inputs.input_ids
262
+ untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
263
+
264
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
265
+ removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1:-1])
266
+ logger.warning("The following part of your input was truncated because CLIP can only handle sequences up to"
267
+ f" {self.tokenizer.model_max_length} tokens: {removed_text}")
268
+
269
+ if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
270
+ attention_mask = text_inputs.attention_mask.to(device)
271
+ else:
272
+ attention_mask = None
273
+
274
+ prompt_embeds = self.text_encoder(
275
+ text_input_ids.to(device),
276
+ attention_mask=attention_mask,
277
+ )
278
+ prompt_embeds = prompt_embeds[0]
279
+
280
+ prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
281
+
282
+ bs_embed, seq_len, _ = prompt_embeds.shape
283
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
284
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
285
+ prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
286
+
287
+ # get unconditional embeddings for classifier free guidance
288
+ if do_classifier_free_guidance:
289
+ uncond_tokens: List[str]
290
+ if negative_prompt is None:
291
+ uncond_tokens = [""] * batch_size
292
+ elif type(prompt) is not type(negative_prompt):
293
+ raise TypeError(f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
294
+ f" {type(prompt)}.")
295
+ elif isinstance(negative_prompt, str):
296
+ uncond_tokens = [negative_prompt]
297
+ elif batch_size != len(negative_prompt):
298
+ raise ValueError(f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
299
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
300
+ " the batch size of `prompt`.")
301
+ else:
302
+ uncond_tokens = negative_prompt
303
+
304
+ max_length = prompt_embeds.shape[1]
305
+ uncond_input = self.tokenizer(
306
+ uncond_tokens,
307
+ padding="max_length",
308
+ max_length=max_length,
309
+ truncation=True,
310
+ return_tensors="pt",
311
+ )
312
+
313
+ if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
314
+ attention_mask = uncond_input.attention_mask.to(device)
315
+ else:
316
+ attention_mask = None
317
+
318
+ negative_prompt_embeds = self.text_encoder(
319
+ uncond_input.input_ids.to(device),
320
+ attention_mask=attention_mask,
321
+ )
322
+ negative_prompt_embeds = negative_prompt_embeds[0]
323
+
324
+ # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
325
+ seq_len = negative_prompt_embeds.shape[1]
326
+
327
+ negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
328
+
329
+ negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
330
+ negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
331
+
332
+ # For classifier free guidance, we need to do two forward passes.
333
+ # Here we concatenate the unconditional and text embeddings into a single batch
334
+ # to avoid doing two forward passes
335
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
336
+
337
+ return prompt_embeds
338
+
339
+ def decode_latents(self, latents):
340
+ latents = 1 / self.vae.config.scaling_factor * latents
341
+ image = self.vae.decode(latents).sample
342
+ image = (image / 2 + 0.5).clamp(0, 1)
343
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
344
+ image = image.cpu().permute(0, 2, 3, 1).float().numpy()
345
+ return image
346
+
347
+ def prepare_extra_step_kwargs(self, generator, eta):
348
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
349
+ # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
350
+ # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
351
+ # and should be between [0, 1]
352
+
353
+ accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
354
+ extra_step_kwargs = {}
355
+ if accepts_eta:
356
+ extra_step_kwargs["eta"] = eta
357
+
358
+ # check if the scheduler accepts generator
359
+ accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
360
+ if accepts_generator:
361
+ extra_step_kwargs["generator"] = generator
362
+ return extra_step_kwargs
363
+
364
+ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
365
+ shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
366
+ if isinstance(generator, list) and len(generator) != batch_size:
367
+ raise ValueError(f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
368
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators.")
369
+
370
+ if latents is None:
371
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
372
+ else:
373
+ latents = latents.to(device)
374
+
375
+ # scale the initial noise by the standard deviation required by the scheduler
376
+ latents = latents * self.scheduler.init_noise_sigma
377
+ return latents
378
+
379
+ @torch.no_grad()
380
+ def __call__(
381
+ self,
382
+ prompt: str = "a car",
383
+ height: int = 256,
384
+ width: int = 256,
385
+ num_inference_steps: int = 50,
386
+ guidance_scale: float = 7.0,
387
+ negative_prompt: str = "bad quality",
388
+ num_images_per_prompt: int = 1,
389
+ eta: float = 0.0,
390
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
391
+ output_type: Optional[str] = "image",
392
+ callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
393
+ callback_steps: int = 1,
394
+ batch_size: int = 4,
395
+ device = torch.device("cuda:0"),
396
+ ):
397
+ self.unet = self.unet.to(device=device)
398
+ self.vae = self.vae.to(device=device)
399
+
400
+ self.text_encoder = self.text_encoder.to(device=device)
401
+
402
+
403
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
404
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
405
+ # corresponds to doing no classifier free guidance.
406
+ do_classifier_free_guidance = guidance_scale > 1.0
407
+
408
+ # Prepare timesteps
409
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
410
+ timesteps = self.scheduler.timesteps
411
+
412
+ _prompt_embeds: torch.Tensor = self._encode_prompt(
413
+ prompt=prompt,
414
+ device=device,
415
+ num_images_per_prompt=num_images_per_prompt,
416
+ do_classifier_free_guidance=do_classifier_free_guidance,
417
+ negative_prompt=negative_prompt,
418
+ ) # type: ignore
419
+ prompt_embeds_neg, prompt_embeds_pos = _prompt_embeds.chunk(2)
420
+
421
+ # Prepare latent variables
422
+ latents: torch.Tensor = self.prepare_latents(
423
+ batch_size * num_images_per_prompt,
424
+ 4,
425
+ height,
426
+ width,
427
+ prompt_embeds_pos.dtype,
428
+ device,
429
+ generator,
430
+ None,
431
+ )
432
+
433
+ camera = get_camera(batch_size).to(dtype=latents.dtype, device=device)
434
+
435
+ # Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
436
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
437
+
438
+ # Denoising loop
439
+ num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
440
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
441
+ for i, t in enumerate(timesteps):
442
+ # expand the latents if we are doing classifier free guidance
443
+ multiplier = 2 if do_classifier_free_guidance else 1
444
+ latent_model_input = torch.cat([latents] * multiplier)
445
+ latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
446
+
447
+ # predict the noise residual
448
+ noise_pred = self.unet.forward(
449
+ x=latent_model_input,
450
+ timesteps=torch.tensor([t] * 4 * multiplier, dtype=latent_model_input.dtype, device=device),
451
+ context=torch.cat([prompt_embeds_neg] * 4 + [prompt_embeds_pos] * 4),
452
+ num_frames=4,
453
+ camera=torch.cat([camera] * multiplier),
454
+ )
455
+
456
+ # perform guidance
457
+ if do_classifier_free_guidance:
458
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
459
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
460
+
461
+ # compute the previous noisy sample x_t -> x_t-1
462
+ # latents = self.scheduler.step(noise_pred.to(dtype=torch.float32), t, latents.to(dtype=torch.float32)).prev_sample.to(prompt_embeds.dtype)
463
+ latents: torch.Tensor = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
464
+
465
+ # call the callback, if provided
466
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
467
+ progress_bar.update()
468
+ if callback is not None and i % callback_steps == 0:
469
+ callback(i, t, latents) # type: ignore
470
+
471
+ # Post-processing
472
+ if output_type == "latent":
473
+ image = latents
474
+ elif output_type == "pil":
475
+ image = self.decode_latents(latents)
476
+ image = self.numpy_to_pil(image)
477
+ else:
478
+ image = self.decode_latents(latents)
479
+
480
+ # Offload last model to CPU
481
+ if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
482
+ self.final_offload_hook.offload()
483
+
484
+ return image
mvdream/util.py ADDED
@@ -0,0 +1,320 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # adopted from
2
+ # https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
3
+ # and
4
+ # https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
5
+ # and
6
+ # https://github.com/openai/guided-diffusion/blob/0ba878e517b276c45d1195eb29f6f5f72659a05b/guided_diffusion/nn.py
7
+ #
8
+ # thanks!
9
+
10
+ import math
11
+ import torch
12
+ import torch.nn as nn
13
+ import numpy as np
14
+ import importlib
15
+ from einops import repeat
16
+ from typing import Any
17
+
18
+
19
+ def instantiate_from_config(config):
20
+ if not "target" in config:
21
+ if config == '__is_first_stage__':
22
+ return None
23
+ elif config == "__is_unconditional__":
24
+ return None
25
+ raise KeyError("Expected key `target` to instantiate.")
26
+ return get_obj_from_str(config["target"])(**config.get("params", dict()))
27
+
28
+
29
+ def get_obj_from_str(string, reload=False):
30
+ module, cls = string.rsplit(".", 1)
31
+ if reload:
32
+ module_imp = importlib.import_module(module)
33
+ importlib.reload(module_imp)
34
+ return getattr(importlib.import_module(module, package=None), cls)
35
+
36
+
37
+ def make_beta_schedule(schedule,
38
+ n_timestep,
39
+ linear_start=1e-4,
40
+ linear_end=2e-2,
41
+ cosine_s=8e-3):
42
+ if schedule == "linear":
43
+ betas = (torch.linspace(linear_start**0.5,
44
+ linear_end**0.5,
45
+ n_timestep,
46
+ dtype=torch.float64)**2)
47
+
48
+ elif schedule == "cosine":
49
+ timesteps = (
50
+ torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep +
51
+ cosine_s)
52
+ alphas = timesteps / (1 + cosine_s) * np.pi / 2
53
+ alphas = torch.cos(alphas).pow(2)
54
+ alphas = alphas / alphas[0]
55
+ betas = 1 - alphas[1:] / alphas[:-1]
56
+ betas = np.clip(betas, a_min=0, a_max=0.999)
57
+
58
+ elif schedule == "sqrt_linear":
59
+ betas = torch.linspace(linear_start,
60
+ linear_end,
61
+ n_timestep,
62
+ dtype=torch.float64)
63
+ elif schedule == "sqrt":
64
+ betas = torch.linspace(linear_start,
65
+ linear_end,
66
+ n_timestep,
67
+ dtype=torch.float64)**0.5
68
+ else:
69
+ raise ValueError(f"schedule '{schedule}' unknown.")
70
+ return betas.numpy() # type: ignore
71
+
72
+
73
+ def make_ddim_timesteps(ddim_discr_method,
74
+ num_ddim_timesteps,
75
+ num_ddpm_timesteps,
76
+ verbose=True):
77
+ if ddim_discr_method == 'uniform':
78
+ c = num_ddpm_timesteps // num_ddim_timesteps
79
+ ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))
80
+ elif ddim_discr_method == 'quad':
81
+ ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8),
82
+ num_ddim_timesteps))**2).astype(int)
83
+ else:
84
+ raise NotImplementedError(
85
+ f'There is no ddim discretization method called "{ddim_discr_method}"'
86
+ )
87
+
88
+ # assert ddim_timesteps.shape[0] == num_ddim_timesteps
89
+ # add one to get the final alpha values right (the ones from first scale to data during sampling)
90
+ steps_out = ddim_timesteps + 1
91
+ if verbose:
92
+ print(f'Selected timesteps for ddim sampler: {steps_out}')
93
+ return steps_out
94
+
95
+
96
+ def make_ddim_sampling_parameters(alphacums,
97
+ ddim_timesteps,
98
+ eta,
99
+ verbose=True):
100
+ # select alphas for computing the variance schedule
101
+ alphas = alphacums[ddim_timesteps]
102
+ alphas_prev = np.asarray([alphacums[0]] +
103
+ alphacums[ddim_timesteps[:-1]].tolist())
104
+
105
+ # according the the formula provided in https://arxiv.org/abs/2010.02502
106
+ sigmas = eta * np.sqrt(
107
+ (1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev))
108
+ if verbose:
109
+ print(
110
+ f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}'
111
+ )
112
+ print(
113
+ f'For the chosen value of eta, which is {eta}, '
114
+ f'this results in the following sigma_t schedule for ddim sampler {sigmas}'
115
+ )
116
+ return sigmas, alphas, alphas_prev
117
+
118
+
119
+ def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
120
+ """
121
+ Create a beta schedule that discretizes the given alpha_t_bar function,
122
+ which defines the cumulative product of (1-beta) over time from t = [0,1].
123
+ :param num_diffusion_timesteps: the number of betas to produce.
124
+ :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
125
+ produces the cumulative product of (1-beta) up to that
126
+ part of the diffusion process.
127
+ :param max_beta: the maximum beta to use; use values lower than 1 to
128
+ prevent singularities.
129
+ """
130
+ betas = []
131
+ for i in range(num_diffusion_timesteps):
132
+ t1 = i / num_diffusion_timesteps
133
+ t2 = (i + 1) / num_diffusion_timesteps
134
+ betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
135
+ return np.array(betas)
136
+
137
+
138
+ def extract_into_tensor(a, t, x_shape):
139
+ b, *_ = t.shape
140
+ out = a.gather(-1, t)
141
+ return out.reshape(b, *((1, ) * (len(x_shape) - 1)))
142
+
143
+
144
+ def checkpoint(func, inputs, params, flag):
145
+ """
146
+ Evaluate a function without caching intermediate activations, allowing for
147
+ reduced memory at the expense of extra compute in the backward pass.
148
+ :param func: the function to evaluate.
149
+ :param inputs: the argument sequence to pass to `func`.
150
+ :param params: a sequence of parameters `func` depends on but does not
151
+ explicitly take as arguments.
152
+ :param flag: if False, disable gradient checkpointing.
153
+ """
154
+ if flag:
155
+ args = tuple(inputs) + tuple(params)
156
+ return CheckpointFunction.apply(func, len(inputs), *args)
157
+ else:
158
+ return func(*inputs)
159
+
160
+
161
+ class CheckpointFunction(torch.autograd.Function):
162
+
163
+ @staticmethod
164
+ def forward(ctx, run_function, length, *args):
165
+ ctx.run_function = run_function
166
+ ctx.input_tensors = list(args[:length])
167
+ ctx.input_params = list(args[length:])
168
+
169
+ with torch.no_grad():
170
+ output_tensors = ctx.run_function(*ctx.input_tensors)
171
+ return output_tensors
172
+
173
+ @staticmethod
174
+ def backward(ctx, *output_grads):
175
+ ctx.input_tensors = [
176
+ x.detach().requires_grad_(True) for x in ctx.input_tensors
177
+ ]
178
+ with torch.enable_grad():
179
+ # Fixes a bug where the first op in run_function modifies the
180
+ # Tensor storage in place, which is not allowed for detach()'d
181
+ # Tensors.
182
+ shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
183
+ output_tensors = ctx.run_function(*shallow_copies)
184
+ input_grads = torch.autograd.grad(
185
+ output_tensors,
186
+ ctx.input_tensors + ctx.input_params,
187
+ output_grads,
188
+ allow_unused=True,
189
+ )
190
+ del ctx.input_tensors
191
+ del ctx.input_params
192
+ del output_tensors
193
+ return (None, None) + input_grads
194
+
195
+
196
+ def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
197
+ """
198
+ Create sinusoidal timestep embeddings.
199
+ :param timesteps: a 1-D Tensor of N indices, one per batch element.
200
+ These may be fractional.
201
+ :param dim: the dimension of the output.
202
+ :param max_period: controls the minimum frequency of the embeddings.
203
+ :return: an [N x dim] Tensor of positional embeddings.
204
+ """
205
+ if not repeat_only:
206
+ half = dim // 2
207
+ freqs = torch.exp(
208
+ -math.log(max_period) *
209
+ torch.arange(start=0, end=half, dtype=torch.float32) /
210
+ half).to(device=timesteps.device)
211
+ args = timesteps[:, None] * freqs[None]
212
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
213
+ if dim % 2:
214
+ embedding = torch.cat(
215
+ [embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
216
+ else:
217
+ embedding = repeat(timesteps, 'b -> b d', d=dim)
218
+ # import pdb; pdb.set_trace()
219
+ return embedding
220
+
221
+
222
+ def zero_module(module):
223
+ """
224
+ Zero out the parameters of a module and return it.
225
+ """
226
+ for p in module.parameters():
227
+ p.detach().zero_()
228
+ return module
229
+
230
+
231
+ def scale_module(module, scale):
232
+ """
233
+ Scale the parameters of a module and return it.
234
+ """
235
+ for p in module.parameters():
236
+ p.detach().mul_(scale)
237
+ return module
238
+
239
+
240
+ def mean_flat(tensor):
241
+ """
242
+ Take the mean over all non-batch dimensions.
243
+ """
244
+ return tensor.mean(dim=list(range(1, len(tensor.shape))))
245
+
246
+
247
+ def normalization(channels):
248
+ """
249
+ Make a standard normalization layer.
250
+ :param channels: number of input channels.
251
+ :return: an nn.Module for normalization.
252
+ """
253
+ return GroupNorm32(32, channels)
254
+
255
+
256
+ # PyTorch 1.7 has SiLU, but we support PyTorch 1.5.
257
+ class SiLU(nn.Module):
258
+
259
+ def forward(self, x):
260
+ return x * torch.sigmoid(x)
261
+
262
+
263
+ class GroupNorm32(nn.GroupNorm):
264
+
265
+ def forward(self, x):
266
+ return super().forward(x)
267
+
268
+
269
+ def conv_nd(dims, *args, **kwargs):
270
+ """
271
+ Create a 1D, 2D, or 3D convolution module.
272
+ """
273
+ if dims == 1:
274
+ return nn.Conv1d(*args, **kwargs)
275
+ elif dims == 2:
276
+ return nn.Conv2d(*args, **kwargs)
277
+ elif dims == 3:
278
+ return nn.Conv3d(*args, **kwargs)
279
+ raise ValueError(f"unsupported dimensions: {dims}")
280
+
281
+
282
+ def linear(*args, **kwargs):
283
+ """
284
+ Create a linear module.
285
+ """
286
+ return nn.Linear(*args, **kwargs)
287
+
288
+
289
+ def avg_pool_nd(dims, *args, **kwargs):
290
+ """
291
+ Create a 1D, 2D, or 3D average pooling module.
292
+ """
293
+ if dims == 1:
294
+ return nn.AvgPool1d(*args, **kwargs)
295
+ elif dims == 2:
296
+ return nn.AvgPool2d(*args, **kwargs)
297
+ elif dims == 3:
298
+ return nn.AvgPool3d(*args, **kwargs)
299
+ raise ValueError(f"unsupported dimensions: {dims}")
300
+
301
+
302
+ class HybridConditioner(nn.Module):
303
+
304
+ def __init__(self, c_concat_config, c_crossattn_config):
305
+ super().__init__()
306
+ self.concat_conditioner: Any = instantiate_from_config(c_concat_config)
307
+ self.crossattn_conditioner: Any = instantiate_from_config(
308
+ c_crossattn_config)
309
+
310
+ def forward(self, c_concat, c_crossattn):
311
+ c_concat = self.concat_conditioner(c_concat)
312
+ c_crossattn = self.crossattn_conditioner(c_crossattn)
313
+ return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]}
314
+
315
+
316
+ def noise_like(shape, device, repeat=False):
317
+ repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(
318
+ shape[0], *((1, ) * (len(shape) - 1)))
319
+ noise = lambda: torch.randn(shape, device=device)
320
+ return repeat_noise() if repeat else noise()