Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| from torch import Tensor | |
| import comfy.model_detection | |
| from comfy.utils import UNET_MAP_BASIC, UNET_MAP_RESNET, UNET_MAP_ATTENTIONS, TRANSFORMER_BLOCKS | |
| import torch | |
| from comfy.ldm.modules.diffusionmodules.util import ( | |
| zero_module, | |
| timestep_embedding, | |
| ) | |
| from comfy.ldm.modules.attention import SpatialVideoTransformer | |
| from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, VideoResBlock, Downsample | |
| from comfy.ldm.util import exists | |
| import comfy.ops | |
| class SVDControlNet(nn.Module): | |
| def __init__( | |
| self, | |
| image_size, | |
| in_channels, | |
| model_channels, | |
| hint_channels, | |
| num_res_blocks, | |
| dropout=0, | |
| channel_mult=(1, 2, 4, 8), | |
| conv_resample=True, | |
| dims=2, | |
| num_classes=None, | |
| use_checkpoint=False, | |
| dtype=torch.float32, | |
| num_heads=-1, | |
| num_head_channels=-1, | |
| num_heads_upsample=-1, | |
| use_scale_shift_norm=False, | |
| resblock_updown=False, | |
| use_new_attention_order=False, | |
| use_spatial_transformer=False, # custom transformer support | |
| transformer_depth=1, # custom transformer support | |
| context_dim=None, # custom transformer support | |
| n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model | |
| legacy=True, | |
| disable_self_attentions=None, | |
| num_attention_blocks=None, | |
| disable_middle_self_attn=False, | |
| use_linear_in_transformer=False, | |
| adm_in_channels=None, | |
| transformer_depth_middle=None, | |
| transformer_depth_output=None, | |
| use_spatial_context=False, | |
| extra_ff_mix_layer=False, | |
| merge_strategy="fixed", | |
| merge_factor=0.5, | |
| video_kernel_size=3, | |
| device=None, | |
| operations=comfy.ops.disable_weight_init, | |
| **kwargs, | |
| ): | |
| super().__init__() | |
| assert use_spatial_transformer == True, "use_spatial_transformer has to be true" | |
| if use_spatial_transformer: | |
| assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...' | |
| if context_dim is not None: | |
| assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...' | |
| # from omegaconf.listconfig import ListConfig | |
| # if type(context_dim) == ListConfig: | |
| # context_dim = list(context_dim) | |
| if num_heads_upsample == -1: | |
| num_heads_upsample = num_heads | |
| if num_heads == -1: | |
| assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set' | |
| if num_head_channels == -1: | |
| assert num_heads != -1, 'Either num_heads or num_head_channels has to be set' | |
| self.dims = dims | |
| self.image_size = image_size | |
| self.in_channels = in_channels | |
| self.model_channels = model_channels | |
| if isinstance(num_res_blocks, int): | |
| self.num_res_blocks = len(channel_mult) * [num_res_blocks] | |
| else: | |
| if len(num_res_blocks) != len(channel_mult): | |
| raise ValueError("provide num_res_blocks either as an int (globally constant) or " | |
| "as a list/tuple (per-level) with the same length as channel_mult") | |
| self.num_res_blocks = num_res_blocks | |
| if disable_self_attentions is not None: | |
| # should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not | |
| assert len(disable_self_attentions) == len(channel_mult) | |
| if num_attention_blocks is not None: | |
| assert len(num_attention_blocks) == len(self.num_res_blocks) | |
| assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks)))) | |
| transformer_depth = transformer_depth[:] | |
| self.dropout = dropout | |
| self.channel_mult = channel_mult | |
| self.conv_resample = conv_resample | |
| self.num_classes = num_classes | |
| self.use_checkpoint = use_checkpoint | |
| self.dtype = dtype | |
| self.num_heads = num_heads | |
| self.num_head_channels = num_head_channels | |
| self.num_heads_upsample = num_heads_upsample | |
| self.predict_codebook_ids = n_embed is not None | |
| time_embed_dim = model_channels * 4 | |
| self.time_embed = nn.Sequential( | |
| operations.Linear(model_channels, time_embed_dim, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device), | |
| ) | |
| if self.num_classes is not None: | |
| if isinstance(self.num_classes, int): | |
| self.label_emb = nn.Embedding(num_classes, time_embed_dim) | |
| elif self.num_classes == "continuous": | |
| print("setting up linear c_adm embedding layer") | |
| self.label_emb = nn.Linear(1, time_embed_dim) | |
| elif self.num_classes == "sequential": | |
| assert adm_in_channels is not None | |
| self.label_emb = nn.Sequential( | |
| nn.Sequential( | |
| operations.Linear(adm_in_channels, time_embed_dim, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device), | |
| ) | |
| ) | |
| else: | |
| raise ValueError() | |
| self.input_blocks = nn.ModuleList( | |
| [ | |
| TimestepEmbedSequential( | |
| operations.conv_nd(dims, in_channels, model_channels, 3, padding=1, dtype=self.dtype, device=device) | |
| ) | |
| ] | |
| ) | |
| self.zero_convs = nn.ModuleList([self.make_zero_conv(model_channels, operations=operations, dtype=self.dtype, device=device)]) | |
| self.input_hint_block = TimestepEmbedSequential( | |
| operations.conv_nd(dims, hint_channels, 16, 3, padding=1, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.conv_nd(dims, 16, 16, 3, padding=1, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.conv_nd(dims, 16, 32, 3, padding=1, stride=2, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.conv_nd(dims, 32, 32, 3, padding=1, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.conv_nd(dims, 32, 96, 3, padding=1, stride=2, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.conv_nd(dims, 96, 96, 3, padding=1, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.conv_nd(dims, 96, 256, 3, padding=1, stride=2, dtype=self.dtype, device=device), | |
| nn.SiLU(), | |
| operations.conv_nd(dims, 256, model_channels, 3, padding=1, dtype=self.dtype, device=device) | |
| ) | |
| self._feature_size = model_channels | |
| input_block_chans = [model_channels] | |
| ch = model_channels | |
| ds = 1 | |
| for level, mult in enumerate(channel_mult): | |
| for nr in range(self.num_res_blocks[level]): | |
| layers = [ | |
| VideoResBlock( | |
| ch, | |
| time_embed_dim, | |
| dropout, | |
| out_channels=mult * model_channels, | |
| dims=dims, | |
| use_checkpoint=use_checkpoint, | |
| use_scale_shift_norm=use_scale_shift_norm, | |
| dtype=self.dtype, | |
| device=device, | |
| operations=operations, | |
| video_kernel_size=video_kernel_size, | |
| merge_strategy=merge_strategy, merge_factor=merge_factor, | |
| ) | |
| ] | |
| ch = mult * model_channels | |
| num_transformers = transformer_depth.pop(0) | |
| if num_transformers > 0: | |
| if num_head_channels == -1: | |
| dim_head = ch // num_heads | |
| else: | |
| num_heads = ch // num_head_channels | |
| dim_head = num_head_channels | |
| if legacy: | |
| #num_heads = 1 | |
| dim_head = ch // num_heads if use_spatial_transformer else num_head_channels | |
| if exists(disable_self_attentions): | |
| disabled_sa = disable_self_attentions[level] | |
| else: | |
| disabled_sa = False | |
| if not exists(num_attention_blocks) or nr < num_attention_blocks[level]: | |
| layers.append( | |
| SpatialVideoTransformer( | |
| ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim, | |
| disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer, | |
| checkpoint=use_checkpoint, dtype=self.dtype, device=device, operations=operations, | |
| use_spatial_context=use_spatial_context, ff_in=extra_ff_mix_layer, | |
| merge_strategy=merge_strategy, merge_factor=merge_factor, | |
| ) | |
| ) | |
| self.input_blocks.append(TimestepEmbedSequential(*layers)) | |
| self.zero_convs.append(self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device)) | |
| self._feature_size += ch | |
| input_block_chans.append(ch) | |
| if level != len(channel_mult) - 1: | |
| out_ch = ch | |
| self.input_blocks.append( | |
| TimestepEmbedSequential( | |
| VideoResBlock( | |
| ch, | |
| time_embed_dim, | |
| dropout, | |
| out_channels=out_ch, | |
| dims=dims, | |
| use_checkpoint=use_checkpoint, | |
| use_scale_shift_norm=use_scale_shift_norm, | |
| down=True, | |
| dtype=self.dtype, | |
| device=device, | |
| operations=operations, | |
| video_kernel_size=video_kernel_size, | |
| merge_strategy=merge_strategy, merge_factor=merge_factor, | |
| ) | |
| if resblock_updown | |
| else Downsample( | |
| ch, conv_resample, dims=dims, out_channels=out_ch, dtype=self.dtype, device=device, operations=operations | |
| ) | |
| ) | |
| ) | |
| ch = out_ch | |
| input_block_chans.append(ch) | |
| self.zero_convs.append(self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device)) | |
| ds *= 2 | |
| self._feature_size += ch | |
| if num_head_channels == -1: | |
| dim_head = ch // num_heads | |
| else: | |
| num_heads = ch // num_head_channels | |
| dim_head = num_head_channels | |
| if legacy: | |
| #num_heads = 1 | |
| dim_head = ch // num_heads if use_spatial_transformer else num_head_channels | |
| mid_block = [ | |
| VideoResBlock( | |
| ch, | |
| time_embed_dim, | |
| dropout, | |
| dims=dims, | |
| use_checkpoint=use_checkpoint, | |
| use_scale_shift_norm=use_scale_shift_norm, | |
| dtype=self.dtype, | |
| device=device, | |
| operations=operations, | |
| video_kernel_size=video_kernel_size, | |
| merge_strategy=merge_strategy, merge_factor=merge_factor, | |
| )] | |
| if transformer_depth_middle >= 0: | |
| mid_block += [SpatialVideoTransformer( # always uses a self-attn | |
| ch, num_heads, dim_head, depth=transformer_depth_middle, context_dim=context_dim, | |
| disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer, | |
| checkpoint=use_checkpoint, dtype=self.dtype, device=device, operations=operations, | |
| use_spatial_context=use_spatial_context, ff_in=extra_ff_mix_layer, | |
| merge_strategy=merge_strategy, merge_factor=merge_factor, | |
| ), | |
| VideoResBlock( | |
| ch, | |
| time_embed_dim, | |
| dropout, | |
| dims=dims, | |
| use_checkpoint=use_checkpoint, | |
| use_scale_shift_norm=use_scale_shift_norm, | |
| dtype=self.dtype, | |
| device=device, | |
| operations=operations, | |
| video_kernel_size=video_kernel_size, | |
| merge_strategy=merge_strategy, merge_factor=merge_factor, | |
| )] | |
| self.middle_block = TimestepEmbedSequential(*mid_block) | |
| self.middle_block_out = self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device) | |
| self._feature_size += ch | |
| def make_zero_conv(self, channels, operations=None, dtype=None, device=None): | |
| return TimestepEmbedSequential(operations.conv_nd(self.dims, channels, channels, 1, padding=0, dtype=dtype, device=device)) | |
| def forward(self, x, hint, timesteps, context, y=None, **kwargs): | |
| t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype) | |
| emb = self.time_embed(t_emb) | |
| cond = kwargs["cond"] | |
| num_video_frames = cond["num_video_frames"] | |
| image_only_indicator = cond.get("image_only_indicator", None) | |
| time_context = cond.get("time_context", None) | |
| del cond | |
| guided_hint = self.input_hint_block(hint, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator) | |
| out_output = [] | |
| out_middle = [] | |
| hs = [] | |
| if self.num_classes is not None: | |
| assert y.shape[0] == x.shape[0] | |
| emb = emb + self.label_emb(y) | |
| h = x | |
| for module, zero_conv in zip(self.input_blocks, self.zero_convs): | |
| if guided_hint is not None: | |
| h = module(h, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator) | |
| h += guided_hint | |
| guided_hint = None | |
| else: | |
| h = module(h, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator) | |
| out_output.append(zero_conv(h, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)) | |
| h = self.middle_block(h, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator) | |
| out_middle.append(self.middle_block_out(h, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)) | |
| return {"middle": out_middle, "output": out_output} | |
| TEMPORAL_TRANSFORMER_BLOCKS = { | |
| "norm_in.weight", | |
| "norm_in.bias", | |
| "ff_in.net.0.proj.weight", | |
| "ff_in.net.0.proj.bias", | |
| "ff_in.net.2.weight", | |
| "ff_in.net.2.bias", | |
| } | |
| TEMPORAL_TRANSFORMER_BLOCKS.update(TRANSFORMER_BLOCKS) | |
| TEMPORAL_UNET_MAP_ATTENTIONS = { | |
| "time_mixer.mix_factor", | |
| } | |
| TEMPORAL_UNET_MAP_ATTENTIONS.update(UNET_MAP_ATTENTIONS) | |
| TEMPORAL_TRANSFORMER_MAP = { | |
| "time_pos_embed.0.weight": "time_pos_embed.linear_1.weight", | |
| "time_pos_embed.0.bias": "time_pos_embed.linear_1.bias", | |
| "time_pos_embed.2.weight": "time_pos_embed.linear_2.weight", | |
| "time_pos_embed.2.bias": "time_pos_embed.linear_2.bias", | |
| } | |
| TEMPORAL_RESNET = { | |
| "time_mixer.mix_factor", | |
| } | |
| def svd_unet_config_from_diffusers_unet(state_dict: dict[str, Tensor], dtype): | |
| match = {} | |
| transformer_depth = [] | |
| attn_res = 1 | |
| down_blocks = comfy.model_detection.count_blocks(state_dict, "down_blocks.{}") | |
| for i in range(down_blocks): | |
| attn_blocks = comfy.model_detection.count_blocks(state_dict, "down_blocks.{}.attentions.".format(i) + '{}') | |
| for ab in range(attn_blocks): | |
| transformer_count = comfy.model_detection.count_blocks(state_dict, "down_blocks.{}.attentions.{}.transformer_blocks.".format(i, ab) + '{}') | |
| transformer_depth.append(transformer_count) | |
| if transformer_count > 0: | |
| match["context_dim"] = state_dict["down_blocks.{}.attentions.{}.transformer_blocks.0.attn2.to_k.weight".format(i, ab)].shape[1] | |
| attn_res *= 2 | |
| if attn_blocks == 0: | |
| transformer_depth.append(0) | |
| transformer_depth.append(0) | |
| match["transformer_depth"] = transformer_depth | |
| match["model_channels"] = state_dict["conv_in.weight"].shape[0] | |
| match["in_channels"] = state_dict["conv_in.weight"].shape[1] | |
| match["adm_in_channels"] = None | |
| if "class_embedding.linear_1.weight" in state_dict: | |
| match["adm_in_channels"] = state_dict["class_embedding.linear_1.weight"].shape[1] | |
| elif "add_embedding.linear_1.weight" in state_dict: | |
| match["adm_in_channels"] = state_dict["add_embedding.linear_1.weight"].shape[1] | |
| # based on unet_config of SVD | |
| SVD = { | |
| 'use_checkpoint': False, | |
| 'image_size': 32, | |
| 'use_spatial_transformer': True, | |
| 'legacy': False, | |
| 'num_classes': 'sequential', | |
| 'adm_in_channels': 768, | |
| 'dtype': dtype, | |
| 'in_channels': 8, | |
| 'out_channels': 4, | |
| 'model_channels': 320, | |
| 'num_res_blocks': [2, 2, 2, 2], | |
| 'transformer_depth': [1, 1, 1, 1, 1, 1, 0, 0], | |
| 'transformer_depth_output': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0], | |
| 'channel_mult': [1, 2, 4, 4], | |
| 'transformer_depth_middle': 1, | |
| 'use_linear_in_transformer': True, | |
| 'context_dim': 1024, | |
| 'extra_ff_mix_layer': True, | |
| 'use_spatial_context': True, | |
| 'merge_strategy': 'learned_with_images', | |
| 'merge_factor': 0.0, | |
| 'video_kernel_size': [3, 1, 1], | |
| 'use_temporal_attention': True, | |
| 'use_temporal_resblock': True, | |
| 'num_heads': -1, | |
| 'num_head_channels': 64, | |
| } | |
| supported_models = [SVD] | |
| for unet_config in supported_models: | |
| matches = True | |
| for k in match: | |
| if match[k] != unet_config[k]: | |
| matches = False | |
| break | |
| if matches: | |
| return comfy.model_detection.convert_config(unet_config) | |
| return None | |
| def svd_unet_to_diffusers(unet_config): | |
| num_res_blocks = unet_config["num_res_blocks"] | |
| channel_mult = unet_config["channel_mult"] | |
| transformer_depth = unet_config["transformer_depth"][:] | |
| transformer_depth_output = unet_config["transformer_depth_output"][:] | |
| num_blocks = len(channel_mult) | |
| transformers_mid = unet_config.get("transformer_depth_middle", None) | |
| diffusers_unet_map = {} | |
| for x in range(num_blocks): | |
| n = 1 + (num_res_blocks[x] + 1) * x | |
| for i in range(num_res_blocks[x]): | |
| for b in TEMPORAL_RESNET: | |
| diffusers_unet_map["down_blocks.{}.resnets.{}.{}".format(x, i, b)] = "input_blocks.{}.0.{}".format(n, b) | |
| for b in UNET_MAP_RESNET: | |
| diffusers_unet_map["down_blocks.{}.resnets.{}.spatial_res_block.{}".format(x, i, UNET_MAP_RESNET[b])] = "input_blocks.{}.0.{}".format(n, b) | |
| diffusers_unet_map["down_blocks.{}.resnets.{}.temporal_res_block.{}".format(x, i, UNET_MAP_RESNET[b])] = "input_blocks.{}.0.time_stack.{}".format(n, b) | |
| #diffusers_unet_map["down_blocks.{}.resnets.{}.{}".format(x, i, UNET_MAP_RESNET[b])] = "input_blocks.{}.0.{}".format(n, b) | |
| num_transformers = transformer_depth.pop(0) | |
| if num_transformers > 0: | |
| for b in TEMPORAL_UNET_MAP_ATTENTIONS: | |
| diffusers_unet_map["down_blocks.{}.attentions.{}.{}".format(x, i, b)] = "input_blocks.{}.1.{}".format(n, b) | |
| for b in TEMPORAL_TRANSFORMER_MAP: | |
| diffusers_unet_map["down_blocks.{}.attentions.{}.{}".format(x, i, TEMPORAL_TRANSFORMER_MAP[b])] = "input_blocks.{}.1.{}".format(n, b) | |
| for t in range(num_transformers): | |
| for b in TRANSFORMER_BLOCKS: | |
| diffusers_unet_map["down_blocks.{}.attentions.{}.transformer_blocks.{}.{}".format(x, i, t, b)] = "input_blocks.{}.1.transformer_blocks.{}.{}".format(n, t, b) | |
| for b in TEMPORAL_TRANSFORMER_BLOCKS: | |
| diffusers_unet_map["down_blocks.{}.attentions.{}.temporal_transformer_blocks.{}.{}".format(x, i, t, b)] = "input_blocks.{}.1.time_stack.{}.{}".format(n, t, b) | |
| n += 1 | |
| for k in ["weight", "bias"]: | |
| diffusers_unet_map["down_blocks.{}.downsamplers.0.conv.{}".format(x, k)] = "input_blocks.{}.0.op.{}".format(n, k) | |
| i = 0 | |
| for b in TEMPORAL_UNET_MAP_ATTENTIONS: | |
| diffusers_unet_map["mid_block.attentions.{}.{}".format(i, b)] = "middle_block.1.{}".format(b) | |
| for b in TEMPORAL_TRANSFORMER_MAP: | |
| diffusers_unet_map["mid_block.attentions.{}.{}".format(i, TEMPORAL_TRANSFORMER_MAP[b])] = "middle_block.1.{}".format(b) | |
| for t in range(transformers_mid): | |
| for b in TRANSFORMER_BLOCKS: | |
| diffusers_unet_map["mid_block.attentions.{}.transformer_blocks.{}.{}".format(i, t, b)] = "middle_block.1.transformer_blocks.{}.{}".format(t, b) | |
| for b in TEMPORAL_TRANSFORMER_BLOCKS: | |
| diffusers_unet_map["mid_block.attentions.{}.temporal_transformer_blocks.{}.{}".format(i, t, b)] = "middle_block.1.time_stack.{}.{}".format(t, b) | |
| for i, n in enumerate([0, 2]): | |
| for b in TEMPORAL_RESNET: | |
| diffusers_unet_map["mid_block.resnets.{}.{}".format(i, b)] = "middle_block.{}.{}".format(n, b) | |
| for b in UNET_MAP_RESNET: | |
| diffusers_unet_map["mid_block.resnets.{}.spatial_res_block.{}".format(i, UNET_MAP_RESNET[b])] = "middle_block.{}.{}".format(n, b) | |
| diffusers_unet_map["mid_block.resnets.{}.temporal_res_block.{}".format(i, UNET_MAP_RESNET[b])] = "middle_block.{}.time_stack.{}".format(n, b) | |
| #diffusers_unet_map["mid_block.resnets.{}.{}".format(i, UNET_MAP_RESNET[b])] = "middle_block.{}.{}".format(n, b) | |
| num_res_blocks = list(reversed(num_res_blocks)) | |
| for x in range(num_blocks): | |
| n = (num_res_blocks[x] + 1) * x | |
| l = num_res_blocks[x] + 1 | |
| for i in range(l): | |
| c = 0 | |
| for b in UNET_MAP_RESNET: | |
| diffusers_unet_map["up_blocks.{}.resnets.{}.{}".format(x, i, UNET_MAP_RESNET[b])] = "output_blocks.{}.0.{}".format(n, b) | |
| c += 1 | |
| num_transformers = transformer_depth_output.pop() | |
| if num_transformers > 0: | |
| c += 1 | |
| for b in UNET_MAP_ATTENTIONS: | |
| diffusers_unet_map["up_blocks.{}.attentions.{}.{}".format(x, i, b)] = "output_blocks.{}.1.{}".format(n, b) | |
| for t in range(num_transformers): | |
| for b in TRANSFORMER_BLOCKS: | |
| diffusers_unet_map["up_blocks.{}.attentions.{}.transformer_blocks.{}.{}".format(x, i, t, b)] = "output_blocks.{}.1.transformer_blocks.{}.{}".format(n, t, b) | |
| if i == l - 1: | |
| for k in ["weight", "bias"]: | |
| diffusers_unet_map["up_blocks.{}.upsamplers.0.conv.{}".format(x, k)] = "output_blocks.{}.{}.conv.{}".format(n, c, k) | |
| n += 1 | |
| for k in UNET_MAP_BASIC: | |
| diffusers_unet_map[k[1]] = k[0] | |
| return diffusers_unet_map | |