fred-dev/comfy_ui_ali
0
1"""2 This file is part of ComfyUI.3 Copyright (C) 2024 Comfy4 5 This program is free software: you can redistribute it and/or modify6 it under the terms of the GNU General Public License as published by7 the Free Software Foundation, either version 3 of the License, or8 (at your option) any later version.9 10 This program is distributed in the hope that it will be useful,11 but WITHOUT ANY WARRANTY; without even the implied warranty of12 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the13 GNU General Public License for more details.14 15 You should have received a copy of the GNU General Public License16 along with this program. If not, see <https://www.gnu.org/licenses/>.17"""18 19 20import torch21import math22import struct23import comfy.checkpoint_pickle24import safetensors.torch25import numpy as np26from PIL import Image27import logging28import itertools29from torch.nn.functional import interpolate30from einops import rearrange31 32ALWAYS_SAFE_LOAD = False33if hasattr(torch.serialization, "add_safe_globals"): # TODO: this was added in pytorch 2.4, the unsafe path should be removed once earlier versions are deprecated34 class ModelCheckpoint:35 pass36 ModelCheckpoint.__module__ = "pytorch_lightning.callbacks.model_checkpoint"37 38 from numpy.core.multiarray import scalar39 from numpy import dtype40 from numpy.dtypes import Float64DType41 from _codecs import encode42 43 torch.serialization.add_safe_globals([ModelCheckpoint, scalar, dtype, Float64DType, encode])44 ALWAYS_SAFE_LOAD = True45 logging.info("Checkpoint files will always be loaded safely.")46else:47 logging.info("Warning, you are using an old pytorch version and some ckpt/pt files might be loaded unsafely. Upgrading to 2.4 or above is recommended.")48 49def load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False):50 if device is None:51 device = torch.device("cpu")52 metadata = None53 if ckpt.lower().endswith(".safetensors") or ckpt.lower().endswith(".sft"):54 try:55 with safetensors.safe_open(ckpt, framework="pt", device=device.type) as f:56 sd = {}57 for k in f.keys():58 sd[k] = f.get_tensor(k)59 if return_metadata:60 metadata = f.metadata()61 except Exception as e:62 if len(e.args) > 0:63 message = e.args[0]64 if "HeaderTooLarge" in message:65 raise ValueError("{}\n\nFile path: {}\n\nThe safetensors file is corrupt or invalid. Make sure this is actually a safetensors file and not a ckpt or pt or other filetype.".format(message, ckpt))66 if "MetadataIncompleteBuffer" in message:67 raise ValueError("{}\n\nFile path: {}\n\nThe safetensors file is corrupt/incomplete. Check the file size and make sure you have copied/downloaded it correctly.".format(message, ckpt))68 raise e69 else:70 if safe_load or ALWAYS_SAFE_LOAD:71 pl_sd = torch.load(ckpt, map_location=device, weights_only=True)72 else:73 pl_sd = torch.load(ckpt, map_location=device, pickle_module=comfy.checkpoint_pickle)74 if "global_step" in pl_sd:75 logging.debug(f"Global Step: {pl_sd['global_step']}")76 if "state_dict" in pl_sd:77 sd = pl_sd["state_dict"]78 else:79 if len(pl_sd) == 1:80 key = list(pl_sd.keys())[0]81 sd = pl_sd[key]82 if not isinstance(sd, dict):83 sd = pl_sd84 else:85 sd = pl_sd86 return (sd, metadata) if return_metadata else sd87 88def save_torch_file(sd, ckpt, metadata=None):89 if metadata is not None:90 safetensors.torch.save_file(sd, ckpt, metadata=metadata)91 else:92 safetensors.torch.save_file(sd, ckpt)93 94def calculate_parameters(sd, prefix=""):95 params = 096 for k in sd.keys():97 if k.startswith(prefix):98 w = sd[k]99 params += w.nelement()100 return params101 102def weight_dtype(sd, prefix=""):103 dtypes = {}104 for k in sd.keys():105 if k.startswith(prefix):106 w = sd[k]107 dtypes[w.dtype] = dtypes.get(w.dtype, 0) + w.numel()108 109 if len(dtypes) == 0:110 return None111 112 return max(dtypes, key=dtypes.get)113 114def state_dict_key_replace(state_dict, keys_to_replace):115 for x in keys_to_replace:116 if x in state_dict:117 state_dict[keys_to_replace[x]] = state_dict.pop(x)118 return state_dict119 120def state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=False):121 if filter_keys:122 out = {}123 else:124 out = state_dict125 for rp in replace_prefix:126 replace = list(map(lambda a: (a, "{}{}".format(replace_prefix[rp], a[len(rp):])), filter(lambda a: a.startswith(rp), state_dict.keys())))127 for x in replace:128 w = state_dict.pop(x[0])129 out[x[1]] = w130 return out131 132 133def transformers_convert(sd, prefix_from, prefix_to, number):134 keys_to_replace = {135 "{}positional_embedding": "{}embeddings.position_embedding.weight",136 "{}token_embedding.weight": "{}embeddings.token_embedding.weight",137 "{}ln_final.weight": "{}final_layer_norm.weight",138 "{}ln_final.bias": "{}final_layer_norm.bias",139 }140 141 for k in keys_to_replace:142 x = k.format(prefix_from)143 if x in sd:144 sd[keys_to_replace[k].format(prefix_to)] = sd.pop(x)145 146 resblock_to_replace = {147 "ln_1": "layer_norm1",148 "ln_2": "layer_norm2",149 "mlp.c_fc": "mlp.fc1",150 "mlp.c_proj": "mlp.fc2",151 "attn.out_proj": "self_attn.out_proj",152 }153 154 for resblock in range(number):155 for x in resblock_to_replace:156 for y in ["weight", "bias"]:157 k = "{}transformer.resblocks.{}.{}.{}".format(prefix_from, resblock, x, y)158 k_to = "{}encoder.layers.{}.{}.{}".format(prefix_to, resblock, resblock_to_replace[x], y)159 if k in sd:160 sd[k_to] = sd.pop(k)161 162 for y in ["weight", "bias"]:163 k_from = "{}transformer.resblocks.{}.attn.in_proj_{}".format(prefix_from, resblock, y)164 if k_from in sd:165 weights = sd.pop(k_from)166 shape_from = weights.shape[0] // 3167 for x in range(3):168 p = ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"]169 k_to = "{}encoder.layers.{}.{}.{}".format(prefix_to, resblock, p[x], y)170 sd[k_to] = weights[shape_from*x:shape_from*(x + 1)]171 172 return sd173 174def clip_text_transformers_convert(sd, prefix_from, prefix_to):175 sd = transformers_convert(sd, prefix_from, "{}text_model.".format(prefix_to), 32)176 177 tp = "{}text_projection.weight".format(prefix_from)178 if tp in sd:179 sd["{}text_projection.weight".format(prefix_to)] = sd.pop(tp)180 181 tp = "{}text_projection".format(prefix_from)182 if tp in sd:183 sd["{}text_projection.weight".format(prefix_to)] = sd.pop(tp).transpose(0, 1).contiguous()184 return sd185 186 187UNET_MAP_ATTENTIONS = {188 "proj_in.weight",189 "proj_in.bias",190 "proj_out.weight",191 "proj_out.bias",192 "norm.weight",193 "norm.bias",194}195 196TRANSFORMER_BLOCKS = {197 "norm1.weight",198 "norm1.bias",199 "norm2.weight",200 "norm2.bias",201 "norm3.weight",202 "norm3.bias",203 "attn1.to_q.weight",204 "attn1.to_k.weight",205 "attn1.to_v.weight",206 "attn1.to_out.0.weight",207 "attn1.to_out.0.bias",208 "attn2.to_q.weight",209 "attn2.to_k.weight",210 "attn2.to_v.weight",211 "attn2.to_out.0.weight",212 "attn2.to_out.0.bias",213 "ff.net.0.proj.weight",214 "ff.net.0.proj.bias",215 "ff.net.2.weight",216 "ff.net.2.bias",217}218 219UNET_MAP_RESNET = {220 "in_layers.2.weight": "conv1.weight",221 "in_layers.2.bias": "conv1.bias",222 "emb_layers.1.weight": "time_emb_proj.weight",223 "emb_layers.1.bias": "time_emb_proj.bias",224 "out_layers.3.weight": "conv2.weight",225 "out_layers.3.bias": "conv2.bias",226 "skip_connection.weight": "conv_shortcut.weight",227 "skip_connection.bias": "conv_shortcut.bias",228 "in_layers.0.weight": "norm1.weight",229 "in_layers.0.bias": "norm1.bias",230 "out_layers.0.weight": "norm2.weight",231 "out_layers.0.bias": "norm2.bias",232}233 234UNET_MAP_BASIC = {235 ("label_emb.0.0.weight", "class_embedding.linear_1.weight"),236 ("label_emb.0.0.bias", "class_embedding.linear_1.bias"),237 ("label_emb.0.2.weight", "class_embedding.linear_2.weight"),238 ("label_emb.0.2.bias", "class_embedding.linear_2.bias"),239 ("label_emb.0.0.weight", "add_embedding.linear_1.weight"),240 ("label_emb.0.0.bias", "add_embedding.linear_1.bias"),241 ("label_emb.0.2.weight", "add_embedding.linear_2.weight"),242 ("label_emb.0.2.bias", "add_embedding.linear_2.bias"),243 ("input_blocks.0.0.weight", "conv_in.weight"),244 ("input_blocks.0.0.bias", "conv_in.bias"),245 ("out.0.weight", "conv_norm_out.weight"),246 ("out.0.bias", "conv_norm_out.bias"),247 ("out.2.weight", "conv_out.weight"),248 ("out.2.bias", "conv_out.bias"),249 ("time_embed.0.weight", "time_embedding.linear_1.weight"),250 ("time_embed.0.bias", "time_embedding.linear_1.bias"),251 ("time_embed.2.weight", "time_embedding.linear_2.weight"),252 ("time_embed.2.bias", "time_embedding.linear_2.bias")253}254 255def unet_to_diffusers(unet_config):256 if "num_res_blocks" not in unet_config:257 return {}258 num_res_blocks = unet_config["num_res_blocks"]259 channel_mult = unet_config["channel_mult"]260 transformer_depth = unet_config["transformer_depth"][:]261 transformer_depth_output = unet_config["transformer_depth_output"][:]262 num_blocks = len(channel_mult)263 264 transformers_mid = unet_config.get("transformer_depth_middle", None)265 266 diffusers_unet_map = {}267 for x in range(num_blocks):268 n = 1 + (num_res_blocks[x] + 1) * x269 for i in range(num_res_blocks[x]):270 for b in UNET_MAP_RESNET:271 diffusers_unet_map["down_blocks.{}.resnets.{}.{}".format(x, i, UNET_MAP_RESNET[b])] = "input_blocks.{}.0.{}".format(n, b)272 num_transformers = transformer_depth.pop(0)273 if num_transformers > 0:274 for b in UNET_MAP_ATTENTIONS:275 diffusers_unet_map["down_blocks.{}.attentions.{}.{}".format(x, i, b)] = "input_blocks.{}.1.{}".format(n, b)276 for t in range(num_transformers):277 for b in TRANSFORMER_BLOCKS:278 diffusers_unet_map["down_blocks.{}.attentions.{}.transformer_blocks.{}.{}".format(x, i, t, b)] = "input_blocks.{}.1.transformer_blocks.{}.{}".format(n, t, b)279 n += 1280 for k in ["weight", "bias"]:281 diffusers_unet_map["down_blocks.{}.downsamplers.0.conv.{}".format(x, k)] = "input_blocks.{}.0.op.{}".format(n, k)282 283 i = 0284 for b in UNET_MAP_ATTENTIONS:285 diffusers_unet_map["mid_block.attentions.{}.{}".format(i, b)] = "middle_block.1.{}".format(b)286 for t in range(transformers_mid):287 for b in TRANSFORMER_BLOCKS:288 diffusers_unet_map["mid_block.attentions.{}.transformer_blocks.{}.{}".format(i, t, b)] = "middle_block.1.transformer_blocks.{}.{}".format(t, b)289 290 for i, n in enumerate([0, 2]):291 for b in UNET_MAP_RESNET:292 diffusers_unet_map["mid_block.resnets.{}.{}".format(i, UNET_MAP_RESNET[b])] = "middle_block.{}.{}".format(n, b)293 294 num_res_blocks = list(reversed(num_res_blocks))295 for x in range(num_blocks):296 n = (num_res_blocks[x] + 1) * x297 l = num_res_blocks[x] + 1298 for i in range(l):299 c = 0300 for b in UNET_MAP_RESNET:301 diffusers_unet_map["up_blocks.{}.resnets.{}.{}".format(x, i, UNET_MAP_RESNET[b])] = "output_blocks.{}.0.{}".format(n, b)302 c += 1303 num_transformers = transformer_depth_output.pop()304 if num_transformers > 0:305 c += 1306 for b in UNET_MAP_ATTENTIONS:307 diffusers_unet_map["up_blocks.{}.attentions.{}.{}".format(x, i, b)] = "output_blocks.{}.1.{}".format(n, b)308 for t in range(num_transformers):309 for b in TRANSFORMER_BLOCKS:310 diffusers_unet_map["up_blocks.{}.attentions.{}.transformer_blocks.{}.{}".format(x, i, t, b)] = "output_blocks.{}.1.transformer_blocks.{}.{}".format(n, t, b)311 if i == l - 1:312 for k in ["weight", "bias"]:313 diffusers_unet_map["up_blocks.{}.upsamplers.0.conv.{}".format(x, k)] = "output_blocks.{}.{}.conv.{}".format(n, c, k)314 n += 1315 316 for k in UNET_MAP_BASIC:317 diffusers_unet_map[k[1]] = k[0]318 319 return diffusers_unet_map320 321def swap_scale_shift(weight):322 shift, scale = weight.chunk(2, dim=0)323 new_weight = torch.cat([scale, shift], dim=0)324 return new_weight325 326MMDIT_MAP_BASIC = {327 ("context_embedder.bias", "context_embedder.bias"),328 ("context_embedder.weight", "context_embedder.weight"),329 ("t_embedder.mlp.0.bias", "time_text_embed.timestep_embedder.linear_1.bias"),330 ("t_embedder.mlp.0.weight", "time_text_embed.timestep_embedder.linear_1.weight"),331 ("t_embedder.mlp.2.bias", "time_text_embed.timestep_embedder.linear_2.bias"),332 ("t_embedder.mlp.2.weight", "time_text_embed.timestep_embedder.linear_2.weight"),333 ("x_embedder.proj.bias", "pos_embed.proj.bias"),334 ("x_embedder.proj.weight", "pos_embed.proj.weight"),335 ("y_embedder.mlp.0.bias", "time_text_embed.text_embedder.linear_1.bias"),336 ("y_embedder.mlp.0.weight", "time_text_embed.text_embedder.linear_1.weight"),337 ("y_embedder.mlp.2.bias", "time_text_embed.text_embedder.linear_2.bias"),338 ("y_embedder.mlp.2.weight", "time_text_embed.text_embedder.linear_2.weight"),339 ("pos_embed", "pos_embed.pos_embed"),340 ("final_layer.adaLN_modulation.1.bias", "norm_out.linear.bias", swap_scale_shift),341 ("final_layer.adaLN_modulation.1.weight", "norm_out.linear.weight", swap_scale_shift),342 ("final_layer.linear.bias", "proj_out.bias"),343 ("final_layer.linear.weight", "proj_out.weight"),344}345 346MMDIT_MAP_BLOCK = {347 ("context_block.adaLN_modulation.1.bias", "norm1_context.linear.bias"),348 ("context_block.adaLN_modulation.1.weight", "norm1_context.linear.weight"),349 ("context_block.attn.proj.bias", "attn.to_add_out.bias"),350 ("context_block.attn.proj.weight", "attn.to_add_out.weight"),351 ("context_block.mlp.fc1.bias", "ff_context.net.0.proj.bias"),352 ("context_block.mlp.fc1.weight", "ff_context.net.0.proj.weight"),353 ("context_block.mlp.fc2.bias", "ff_context.net.2.bias"),354 ("context_block.mlp.fc2.weight", "ff_context.net.2.weight"),355 ("context_block.attn.ln_q.weight", "attn.norm_added_q.weight"),356 ("context_block.attn.ln_k.weight", "attn.norm_added_k.weight"),357 ("x_block.adaLN_modulation.1.bias", "norm1.linear.bias"),358 ("x_block.adaLN_modulation.1.weight", "norm1.linear.weight"),359 ("x_block.attn.proj.bias", "attn.to_out.0.bias"),360 ("x_block.attn.proj.weight", "attn.to_out.0.weight"),361 ("x_block.attn.ln_q.weight", "attn.norm_q.weight"),362 ("x_block.attn.ln_k.weight", "attn.norm_k.weight"),363 ("x_block.attn2.proj.bias", "attn2.to_out.0.bias"),364 ("x_block.attn2.proj.weight", "attn2.to_out.0.weight"),365 ("x_block.attn2.ln_q.weight", "attn2.norm_q.weight"),366 ("x_block.attn2.ln_k.weight", "attn2.norm_k.weight"),367 ("x_block.mlp.fc1.bias", "ff.net.0.proj.bias"),368 ("x_block.mlp.fc1.weight", "ff.net.0.proj.weight"),369 ("x_block.mlp.fc2.bias", "ff.net.2.bias"),370 ("x_block.mlp.fc2.weight", "ff.net.2.weight"),371}372 373def mmdit_to_diffusers(mmdit_config, output_prefix=""):374 key_map = {}375 376 depth = mmdit_config.get("depth", 0)377 num_blocks = mmdit_config.get("num_blocks", depth)378 for i in range(num_blocks):379 block_from = "transformer_blocks.{}".format(i)380 block_to = "{}joint_blocks.{}".format(output_prefix, i)381 382 offset = depth * 64383 384 for end in ("weight", "bias"):385 k = "{}.attn.".format(block_from)386 qkv = "{}.x_block.attn.qkv.{}".format(block_to, end)387 key_map["{}to_q.{}".format(k, end)] = (qkv, (0, 0, offset))388 key_map["{}to_k.{}".format(k, end)] = (qkv, (0, offset, offset))389 key_map["{}to_v.{}".format(k, end)] = (qkv, (0, offset * 2, offset))390 391 qkv = "{}.context_block.attn.qkv.{}".format(block_to, end)392 key_map["{}add_q_proj.{}".format(k, end)] = (qkv, (0, 0, offset))393 key_map["{}add_k_proj.{}".format(k, end)] = (qkv, (0, offset, offset))394 key_map["{}add_v_proj.{}".format(k, end)] = (qkv, (0, offset * 2, offset))395 396 k = "{}.attn2.".format(block_from)397 qkv = "{}.x_block.attn2.qkv.{}".format(block_to, end)398 key_map["{}to_q.{}".format(k, end)] = (qkv, (0, 0, offset))399 key_map["{}to_k.{}".format(k, end)] = (qkv, (0, offset, offset))400 key_map["{}to_v.{}".format(k, end)] = (qkv, (0, offset * 2, offset))401 402 for k in MMDIT_MAP_BLOCK:403 key_map["{}.{}".format(block_from, k[1])] = "{}.{}".format(block_to, k[0])404 405 map_basic = MMDIT_MAP_BASIC.copy()406 map_basic.add(("joint_blocks.{}.context_block.adaLN_modulation.1.bias".format(depth - 1), "transformer_blocks.{}.norm1_context.linear.bias".format(depth - 1), swap_scale_shift))407 map_basic.add(("joint_blocks.{}.context_block.adaLN_modulation.1.weight".format(depth - 1), "transformer_blocks.{}.norm1_context.linear.weight".format(depth - 1), swap_scale_shift))408 409 for k in map_basic:410 if len(k) > 2:411 key_map[k[1]] = ("{}{}".format(output_prefix, k[0]), None, k[2])412 else:413 key_map[k[1]] = "{}{}".format(output_prefix, k[0])414 415 return key_map416 417PIXART_MAP_BASIC = {418 ("csize_embedder.mlp.0.weight", "adaln_single.emb.resolution_embedder.linear_1.weight"),419 ("csize_embedder.mlp.0.bias", "adaln_single.emb.resolution_embedder.linear_1.bias"),420 ("csize_embedder.mlp.2.weight", "adaln_single.emb.resolution_embedder.linear_2.weight"),421 ("csize_embedder.mlp.2.bias", "adaln_single.emb.resolution_embedder.linear_2.bias"),422 ("ar_embedder.mlp.0.weight", "adaln_single.emb.aspect_ratio_embedder.linear_1.weight"),423 ("ar_embedder.mlp.0.bias", "adaln_single.emb.aspect_ratio_embedder.linear_1.bias"),424 ("ar_embedder.mlp.2.weight", "adaln_single.emb.aspect_ratio_embedder.linear_2.weight"),425 ("ar_embedder.mlp.2.bias", "adaln_single.emb.aspect_ratio_embedder.linear_2.bias"),426 ("x_embedder.proj.weight", "pos_embed.proj.weight"),427 ("x_embedder.proj.bias", "pos_embed.proj.bias"),428 ("y_embedder.y_embedding", "caption_projection.y_embedding"),429 ("y_embedder.y_proj.fc1.weight", "caption_projection.linear_1.weight"),430 ("y_embedder.y_proj.fc1.bias", "caption_projection.linear_1.bias"),431 ("y_embedder.y_proj.fc2.weight", "caption_projection.linear_2.weight"),432 ("y_embedder.y_proj.fc2.bias", "caption_projection.linear_2.bias"),433 ("t_embedder.mlp.0.weight", "adaln_single.emb.timestep_embedder.linear_1.weight"),434 ("t_embedder.mlp.0.bias", "adaln_single.emb.timestep_embedder.linear_1.bias"),435 ("t_embedder.mlp.2.weight", "adaln_single.emb.timestep_embedder.linear_2.weight"),436 ("t_embedder.mlp.2.bias", "adaln_single.emb.timestep_embedder.linear_2.bias"),437 ("t_block.1.weight", "adaln_single.linear.weight"),438 ("t_block.1.bias", "adaln_single.linear.bias"),439 ("final_layer.linear.weight", "proj_out.weight"),440 ("final_layer.linear.bias", "proj_out.bias"),441 ("final_layer.scale_shift_table", "scale_shift_table"),442}443 444PIXART_MAP_BLOCK = {445 ("scale_shift_table", "scale_shift_table"),446 ("attn.proj.weight", "attn1.to_out.0.weight"),447 ("attn.proj.bias", "attn1.to_out.0.bias"),448 ("mlp.fc1.weight", "ff.net.0.proj.weight"),449 ("mlp.fc1.bias", "ff.net.0.proj.bias"),450 ("mlp.fc2.weight", "ff.net.2.weight"),451 ("mlp.fc2.bias", "ff.net.2.bias"),452 ("cross_attn.proj.weight" ,"attn2.to_out.0.weight"),453 ("cross_attn.proj.bias" ,"attn2.to_out.0.bias"),454}455 456def pixart_to_diffusers(mmdit_config, output_prefix=""):457 key_map = {}458 459 depth = mmdit_config.get("depth", 0)460 offset = mmdit_config.get("hidden_size", 1152)461 462 for i in range(depth):463 block_from = "transformer_blocks.{}".format(i)464 block_to = "{}blocks.{}".format(output_prefix, i)465 466 for end in ("weight", "bias"):467 s = "{}.attn1.".format(block_from)468 qkv = "{}.attn.qkv.{}".format(block_to, end)469 key_map["{}to_q.{}".format(s, end)] = (qkv, (0, 0, offset))470 key_map["{}to_k.{}".format(s, end)] = (qkv, (0, offset, offset))471 key_map["{}to_v.{}".format(s, end)] = (qkv, (0, offset * 2, offset))472 473 s = "{}.attn2.".format(block_from)474 q = "{}.cross_attn.q_linear.{}".format(block_to, end)475 kv = "{}.cross_attn.kv_linear.{}".format(block_to, end)476 477 key_map["{}to_q.{}".format(s, end)] = q478 key_map["{}to_k.{}".format(s, end)] = (kv, (0, 0, offset))479 key_map["{}to_v.{}".format(s, end)] = (kv, (0, offset, offset))480 481 for k in PIXART_MAP_BLOCK:482 key_map["{}.{}".format(block_from, k[1])] = "{}.{}".format(block_to, k[0])483 484 for k in PIXART_MAP_BASIC:485 key_map[k[1]] = "{}{}".format(output_prefix, k[0])486 487 return key_map488 489def auraflow_to_diffusers(mmdit_config, output_prefix=""):490 n_double_layers = mmdit_config.get("n_double_layers", 0)491 n_layers = mmdit_config.get("n_layers", 0)492 493 key_map = {}494 for i in range(n_layers):495 if i < n_double_layers:496 index = i497 prefix_from = "joint_transformer_blocks"498 prefix_to = "{}double_layers".format(output_prefix)499 block_map = {500 "attn.to_q.weight": "attn.w2q.weight",501 "attn.to_k.weight": "attn.w2k.weight",502 "attn.to_v.weight": "attn.w2v.weight",503 "attn.to_out.0.weight": "attn.w2o.weight",504 "attn.add_q_proj.weight": "attn.w1q.weight",505 "attn.add_k_proj.weight": "attn.w1k.weight",506 "attn.add_v_proj.weight": "attn.w1v.weight",507 "attn.to_add_out.weight": "attn.w1o.weight",508 "ff.linear_1.weight": "mlpX.c_fc1.weight",509 "ff.linear_2.weight": "mlpX.c_fc2.weight",510 "ff.out_projection.weight": "mlpX.c_proj.weight",511 "ff_context.linear_1.weight": "mlpC.c_fc1.weight",512 "ff_context.linear_2.weight": "mlpC.c_fc2.weight",513 "ff_context.out_projection.weight": "mlpC.c_proj.weight",514 "norm1.linear.weight": "modX.1.weight",515 "norm1_context.linear.weight": "modC.1.weight",516 }517 else:518 index = i - n_double_layers519 prefix_from = "single_transformer_blocks"520 prefix_to = "{}single_layers".format(output_prefix)521 522 block_map = {523 "attn.to_q.weight": "attn.w1q.weight",524 "attn.to_k.weight": "attn.w1k.weight",525 "attn.to_v.weight": "attn.w1v.weight",526 "attn.to_out.0.weight": "attn.w1o.weight",527 "norm1.linear.weight": "modCX.1.weight",528 "ff.linear_1.weight": "mlp.c_fc1.weight",529 "ff.linear_2.weight": "mlp.c_fc2.weight",530 "ff.out_projection.weight": "mlp.c_proj.weight"531 }532 533 for k in block_map:534 key_map["{}.{}.{}".format(prefix_from, index, k)] = "{}.{}.{}".format(prefix_to, index, block_map[k])535 536 MAP_BASIC = {537 ("positional_encoding", "pos_embed.pos_embed"),538 ("register_tokens", "register_tokens"),539 ("t_embedder.mlp.0.weight", "time_step_proj.linear_1.weight"),540 ("t_embedder.mlp.0.bias", "time_step_proj.linear_1.bias"),541 ("t_embedder.mlp.2.weight", "time_step_proj.linear_2.weight"),542 ("t_embedder.mlp.2.bias", "time_step_proj.linear_2.bias"),543 ("cond_seq_linear.weight", "context_embedder.weight"),544 ("init_x_linear.weight", "pos_embed.proj.weight"),545 ("init_x_linear.bias", "pos_embed.proj.bias"),546 ("final_linear.weight", "proj_out.weight"),547 ("modF.1.weight", "norm_out.linear.weight", swap_scale_shift),548 }549 550 for k in MAP_BASIC:551 if len(k) > 2:552 key_map[k[1]] = ("{}{}".format(output_prefix, k[0]), None, k[2])553 else:554 key_map[k[1]] = "{}{}".format(output_prefix, k[0])555 556 return key_map557 558def flux_to_diffusers(mmdit_config, output_prefix=""):559 n_double_layers = mmdit_config.get("depth", 0)560 n_single_layers = mmdit_config.get("depth_single_blocks", 0)561 hidden_size = mmdit_config.get("hidden_size", 0)562 563 key_map = {}564 for index in range(n_double_layers):565 prefix_from = "transformer_blocks.{}".format(index)566 prefix_to = "{}double_blocks.{}".format(output_prefix, index)567 568 for end in ("weight", "bias"):569 k = "{}.attn.".format(prefix_from)570 qkv = "{}.img_attn.qkv.{}".format(prefix_to, end)571 key_map["{}to_q.{}".format(k, end)] = (qkv, (0, 0, hidden_size))572 key_map["{}to_k.{}".format(k, end)] = (qkv, (0, hidden_size, hidden_size))573 key_map["{}to_v.{}".format(k, end)] = (qkv, (0, hidden_size * 2, hidden_size))574 575 k = "{}.attn.".format(prefix_from)576 qkv = "{}.txt_attn.qkv.{}".format(prefix_to, end)577 key_map["{}add_q_proj.{}".format(k, end)] = (qkv, (0, 0, hidden_size))578 key_map["{}add_k_proj.{}".format(k, end)] = (qkv, (0, hidden_size, hidden_size))579 key_map["{}add_v_proj.{}".format(k, end)] = (qkv, (0, hidden_size * 2, hidden_size))580 581 block_map = {582 "attn.to_out.0.weight": "img_attn.proj.weight",583 "attn.to_out.0.bias": "img_attn.proj.bias",584 "norm1.linear.weight": "img_mod.lin.weight",585 "norm1.linear.bias": "img_mod.lin.bias",586 "norm1_context.linear.weight": "txt_mod.lin.weight",587 "norm1_context.linear.bias": "txt_mod.lin.bias",588 "attn.to_add_out.weight": "txt_attn.proj.weight",589 "attn.to_add_out.bias": "txt_attn.proj.bias",590 "ff.net.0.proj.weight": "img_mlp.0.weight",591 "ff.net.0.proj.bias": "img_mlp.0.bias",592 "ff.net.2.weight": "img_mlp.2.weight",593 "ff.net.2.bias": "img_mlp.2.bias",594 "ff_context.net.0.proj.weight": "txt_mlp.0.weight",595 "ff_context.net.0.proj.bias": "txt_mlp.0.bias",596 "ff_context.net.2.weight": "txt_mlp.2.weight",597 "ff_context.net.2.bias": "txt_mlp.2.bias",598 "attn.norm_q.weight": "img_attn.norm.query_norm.scale",599 "attn.norm_k.weight": "img_attn.norm.key_norm.scale",600 "attn.norm_added_q.weight": "txt_attn.norm.query_norm.scale",601 "attn.norm_added_k.weight": "txt_attn.norm.key_norm.scale",602 }603 604 for k in block_map:605 key_map["{}.{}".format(prefix_from, k)] = "{}.{}".format(prefix_to, block_map[k])606 607 for index in range(n_single_layers):608 prefix_from = "single_transformer_blocks.{}".format(index)609 prefix_to = "{}single_blocks.{}".format(output_prefix, index)610 611 for end in ("weight", "bias"):612 k = "{}.attn.".format(prefix_from)613 qkv = "{}.linear1.{}".format(prefix_to, end)614 key_map["{}to_q.{}".format(k, end)] = (qkv, (0, 0, hidden_size))615 key_map["{}to_k.{}".format(k, end)] = (qkv, (0, hidden_size, hidden_size))616 key_map["{}to_v.{}".format(k, end)] = (qkv, (0, hidden_size * 2, hidden_size))617 key_map["{}.proj_mlp.{}".format(prefix_from, end)] = (qkv, (0, hidden_size * 3, hidden_size * 4))618 619 block_map = {620 "norm.linear.weight": "modulation.lin.weight",621 "norm.linear.bias": "modulation.lin.bias",622 "proj_out.weight": "linear2.weight",623 "proj_out.bias": "linear2.bias",624 "attn.norm_q.weight": "norm.query_norm.scale",625 "attn.norm_k.weight": "norm.key_norm.scale",626 }627 628 for k in block_map:629 key_map["{}.{}".format(prefix_from, k)] = "{}.{}".format(prefix_to, block_map[k])630 631 MAP_BASIC = {632 ("final_layer.linear.bias", "proj_out.bias"),633 ("final_layer.linear.weight", "proj_out.weight"),634 ("img_in.bias", "x_embedder.bias"),635 ("img_in.weight", "x_embedder.weight"),636 ("time_in.in_layer.bias", "time_text_embed.timestep_embedder.linear_1.bias"),637 ("time_in.in_layer.weight", "time_text_embed.timestep_embedder.linear_1.weight"),638 ("time_in.out_layer.bias", "time_text_embed.timestep_embedder.linear_2.bias"),639 ("time_in.out_layer.weight", "time_text_embed.timestep_embedder.linear_2.weight"),640 ("txt_in.bias", "context_embedder.bias"),641 ("txt_in.weight", "context_embedder.weight"),642 ("vector_in.in_layer.bias", "time_text_embed.text_embedder.linear_1.bias"),643 ("vector_in.in_layer.weight", "time_text_embed.text_embedder.linear_1.weight"),644 ("vector_in.out_layer.bias", "time_text_embed.text_embedder.linear_2.bias"),645 ("vector_in.out_layer.weight", "time_text_embed.text_embedder.linear_2.weight"),646 ("guidance_in.in_layer.bias", "time_text_embed.guidance_embedder.linear_1.bias"),647 ("guidance_in.in_layer.weight", "time_text_embed.guidance_embedder.linear_1.weight"),648 ("guidance_in.out_layer.bias", "time_text_embed.guidance_embedder.linear_2.bias"),649 ("guidance_in.out_layer.weight", "time_text_embed.guidance_embedder.linear_2.weight"),650 ("final_layer.adaLN_modulation.1.bias", "norm_out.linear.bias", swap_scale_shift),651 ("final_layer.adaLN_modulation.1.weight", "norm_out.linear.weight", swap_scale_shift),652 ("pos_embed_input.bias", "controlnet_x_embedder.bias"),653 ("pos_embed_input.weight", "controlnet_x_embedder.weight"),654 }655 656 for k in MAP_BASIC:657 if len(k) > 2:658 key_map[k[1]] = ("{}{}".format(output_prefix, k[0]), None, k[2])659 else:660 key_map[k[1]] = "{}{}".format(output_prefix, k[0])661 662 return key_map663 664def repeat_to_batch_size(tensor, batch_size, dim=0):665 if tensor.shape[dim] > batch_size:666 return tensor.narrow(dim, 0, batch_size)667 elif tensor.shape[dim] < batch_size:668 return tensor.repeat(dim * [1] + [math.ceil(batch_size / tensor.shape[dim])] + [1] * (len(tensor.shape) - 1 - dim)).narrow(dim, 0, batch_size)669 return tensor670 671def resize_to_batch_size(tensor, batch_size):672 in_batch_size = tensor.shape[0]673 if in_batch_size == batch_size:674 return tensor675 676 if batch_size <= 1:677 return tensor[:batch_size]678 679 output = torch.empty([batch_size] + list(tensor.shape)[1:], dtype=tensor.dtype, device=tensor.device)680 if batch_size < in_batch_size:681 scale = (in_batch_size - 1) / (batch_size - 1)682 for i in range(batch_size):683 output[i] = tensor[min(round(i * scale), in_batch_size - 1)]684 else:685 scale = in_batch_size / batch_size686 for i in range(batch_size):687 output[i] = tensor[min(math.floor((i + 0.5) * scale), in_batch_size - 1)]688 689 return output690 691def convert_sd_to(state_dict, dtype):692 keys = list(state_dict.keys())693 for k in keys:694 state_dict[k] = state_dict[k].to(dtype)695 return state_dict696 697def safetensors_header(safetensors_path, max_size=100*1024*1024):698 with open(safetensors_path, "rb") as f:699 header = f.read(8)700 length_of_header = struct.unpack('<Q', header)[0]701 if length_of_header > max_size:702 return None703 return f.read(length_of_header)704 705def set_attr(obj, attr, value):706 attrs = attr.split(".")707 for name in attrs[:-1]:708 obj = getattr(obj, name)709 prev = getattr(obj, attrs[-1])710 setattr(obj, attrs[-1], value)711 return prev712 713def set_attr_param(obj, attr, value):714 return set_attr(obj, attr, torch.nn.Parameter(value, requires_grad=False))715 716def copy_to_param(obj, attr, value):717 # inplace update tensor instead of replacing it718 attrs = attr.split(".")719 for name in attrs[:-1]:720 obj = getattr(obj, name)721 prev = getattr(obj, attrs[-1])722 prev.data.copy_(value)723 724def get_attr(obj, attr: str):725 """Retrieves a nested attribute from an object using dot notation.726 727 Args:728 obj: The object to get the attribute from729 attr (str): The attribute path using dot notation (e.g. "model.layer.weight")730 731 Returns:732 The value of the requested attribute733 734 Example:735 model = MyModel()736 weight = get_attr(model, "layer1.conv.weight")737 # Equivalent to: model.layer1.conv.weight738 739 Important:740 Always prefer `comfy.model_patcher.ModelPatcher.get_model_object` when741 accessing nested model objects under `ModelPatcher.model`.742 """743 attrs = attr.split(".")744 for name in attrs:745 obj = getattr(obj, name)746 return obj747 748def bislerp(samples, width, height):749 def slerp(b1, b2, r):750 '''slerps batches b1, b2 according to ratio r, batches should be flat e.g. NxC'''751 752 c = b1.shape[-1]753 754 #norms755 b1_norms = torch.norm(b1, dim=-1, keepdim=True)756 b2_norms = torch.norm(b2, dim=-1, keepdim=True)757 758 #normalize759 b1_normalized = b1 / b1_norms760 b2_normalized = b2 / b2_norms761 762 #zero when norms are zero763 b1_normalized[b1_norms.expand(-1,c) == 0.0] = 0.0764 b2_normalized[b2_norms.expand(-1,c) == 0.0] = 0.0765 766 #slerp767 dot = (b1_normalized*b2_normalized).sum(1)768 omega = torch.acos(dot)769 so = torch.sin(omega)770 771 #technically not mathematically correct, but more pleasing?772 res = (torch.sin((1.0-r.squeeze(1))*omega)/so).unsqueeze(1)*b1_normalized + (torch.sin(r.squeeze(1)*omega)/so).unsqueeze(1) * b2_normalized773 res *= (b1_norms * (1.0-r) + b2_norms * r).expand(-1,c)774 775 #edge cases for same or polar opposites776 res[dot > 1 - 1e-5] = b1[dot > 1 - 1e-5]777 res[dot < 1e-5 - 1] = (b1 * (1.0-r) + b2 * r)[dot < 1e-5 - 1]778 return res779 780 def generate_bilinear_data(length_old, length_new, device):781 coords_1 = torch.arange(length_old, dtype=torch.float32, device=device).reshape((1,1,1,-1))782 coords_1 = torch.nn.functional.interpolate(coords_1, size=(1, length_new), mode="bilinear")783 ratios = coords_1 - coords_1.floor()784 coords_1 = coords_1.to(torch.int64)785 786 coords_2 = torch.arange(length_old, dtype=torch.float32, device=device).reshape((1,1,1,-1)) + 1787 coords_2[:,:,:,-1] -= 1788 coords_2 = torch.nn.functional.interpolate(coords_2, size=(1, length_new), mode="bilinear")789 coords_2 = coords_2.to(torch.int64)790 return ratios, coords_1, coords_2791 792 orig_dtype = samples.dtype793 samples = samples.float()794 n,c,h,w = samples.shape795 h_new, w_new = (height, width)796 797 #linear w798 ratios, coords_1, coords_2 = generate_bilinear_data(w, w_new, samples.device)799 coords_1 = coords_1.expand((n, c, h, -1))800 coords_2 = coords_2.expand((n, c, h, -1))801 ratios = ratios.expand((n, 1, h, -1))802 803 pass_1 = samples.gather(-1,coords_1).movedim(1, -1).reshape((-1,c))804 pass_2 = samples.gather(-1,coords_2).movedim(1, -1).reshape((-1,c))805 ratios = ratios.movedim(1, -1).reshape((-1,1))806 807 result = slerp(pass_1, pass_2, ratios)808 result = result.reshape(n, h, w_new, c).movedim(-1, 1)809 810 #linear h811 ratios, coords_1, coords_2 = generate_bilinear_data(h, h_new, samples.device)812 coords_1 = coords_1.reshape((1,1,-1,1)).expand((n, c, -1, w_new))813 coords_2 = coords_2.reshape((1,1,-1,1)).expand((n, c, -1, w_new))814 ratios = ratios.reshape((1,1,-1,1)).expand((n, 1, -1, w_new))815 816 pass_1 = result.gather(-2,coords_1).movedim(1, -1).reshape((-1,c))817 pass_2 = result.gather(-2,coords_2).movedim(1, -1).reshape((-1,c))818 ratios = ratios.movedim(1, -1).reshape((-1,1))819 820 result = slerp(pass_1, pass_2, ratios)821 result = result.reshape(n, h_new, w_new, c).movedim(-1, 1)822 return result.to(orig_dtype)823 824def lanczos(samples, width, height):825 images = [Image.fromarray(np.clip(255. * image.movedim(0, -1).cpu().numpy(), 0, 255).astype(np.uint8)) for image in samples]826 images = [image.resize((width, height), resample=Image.Resampling.LANCZOS) for image in images]827 images = [torch.from_numpy(np.array(image).astype(np.float32) / 255.0).movedim(-1, 0) for image in images]828 result = torch.stack(images)829 return result.to(samples.device, samples.dtype)830 831def common_upscale(samples, width, height, upscale_method, crop):832 orig_shape = tuple(samples.shape)833 if len(orig_shape) > 4:834 samples = samples.reshape(samples.shape[0], samples.shape[1], -1, samples.shape[-2], samples.shape[-1])835 samples = samples.movedim(2, 1)836 samples = samples.reshape(-1, orig_shape[1], orig_shape[-2], orig_shape[-1])837 if crop == "center":838 old_width = samples.shape[-1]839 old_height = samples.shape[-2]840 old_aspect = old_width / old_height841 new_aspect = width / height842 x = 0843 y = 0844 if old_aspect > new_aspect:845 x = round((old_width - old_width * (new_aspect / old_aspect)) / 2)846 elif old_aspect < new_aspect:847 y = round((old_height - old_height * (old_aspect / new_aspect)) / 2)848 s = samples.narrow(-2, y, old_height - y * 2).narrow(-1, x, old_width - x * 2)849 else:850 s = samples851 852 if upscale_method == "bislerp":853 out = bislerp(s, width, height)854 elif upscale_method == "lanczos":855 out = lanczos(s, width, height)856 else:857 out = torch.nn.functional.interpolate(s, size=(height, width), mode=upscale_method)858 859 if len(orig_shape) == 4:860 return out861 862 out = out.reshape((orig_shape[0], -1, orig_shape[1]) + (height, width))863 return out.movedim(2, 1).reshape(orig_shape[:-2] + (height, width))864 865def get_tiled_scale_steps(width, height, tile_x, tile_y, overlap):866 rows = 1 if height <= tile_y else math.ceil((height - overlap) / (tile_y - overlap))867 cols = 1 if width <= tile_x else math.ceil((width - overlap) / (tile_x - overlap))868 return rows * cols869 870@torch.inference_mode()871def tiled_scale_multidim(samples, function, tile=(64, 64), overlap=8, upscale_amount=4, out_channels=3, output_device="cpu", downscale=False, index_formulas=None, pbar=None):872 dims = len(tile)873 874 if not (isinstance(upscale_amount, (tuple, list))):875 upscale_amount = [upscale_amount] * dims876 877 if not (isinstance(overlap, (tuple, list))):878 overlap = [overlap] * dims879 880 if index_formulas is None:881 index_formulas = upscale_amount882 883 if not (isinstance(index_formulas, (tuple, list))):884 index_formulas = [index_formulas] * dims885 886 def get_upscale(dim, val):887 up = upscale_amount[dim]888 if callable(up):889 return up(val)890 else:891 return up * val892 893 def get_downscale(dim, val):894 up = upscale_amount[dim]895 if callable(up):896 return up(val)897 else:898 return val / up899 900 def get_upscale_pos(dim, val):901 up = index_formulas[dim]902 if callable(up):903 return up(val)904 else:905 return up * val906 907 def get_downscale_pos(dim, val):908 up = index_formulas[dim]909 if callable(up):910 return up(val)911 else:912 return val / up913 914 if downscale:915 get_scale = get_downscale916 get_pos = get_downscale_pos917 else:918 get_scale = get_upscale919 get_pos = get_upscale_pos920 921 def mult_list_upscale(a):922 out = []923 for i in range(len(a)):924 out.append(round(get_scale(i, a[i])))925 return out926 927 output = torch.empty([samples.shape[0], out_channels] + mult_list_upscale(samples.shape[2:]), device=output_device)928 929 for b in range(samples.shape[0]):930 s = samples[b:b+1]931 932 # handle entire input fitting in a single tile933 if all(s.shape[d+2] <= tile[d] for d in range(dims)):934 output[b:b+1] = function(s).to(output_device)935 if pbar is not None:936 pbar.update(1)937 continue938 939 out = torch.zeros([s.shape[0], out_channels] + mult_list_upscale(s.shape[2:]), device=output_device)940 out_div = torch.zeros([s.shape[0], out_channels] + mult_list_upscale(s.shape[2:]), device=output_device)941 942 positions = [range(0, s.shape[d+2] - overlap[d], tile[d] - overlap[d]) if s.shape[d+2] > tile[d] else [0] for d in range(dims)]943 944 for it in itertools.product(*positions):945 s_in = s946 upscaled = []947 948 for d in range(dims):949 pos = max(0, min(s.shape[d + 2] - overlap[d], it[d]))950 l = min(tile[d], s.shape[d + 2] - pos)951 s_in = s_in.narrow(d + 2, pos, l)952 upscaled.append(round(get_pos(d, pos)))953 954 ps = function(s_in).to(output_device)955 mask = torch.ones_like(ps)956 957 for d in range(2, dims + 2):958 feather = round(get_scale(d - 2, overlap[d - 2]))959 if feather >= mask.shape[d]:960 continue961 for t in range(feather):962 a = (t + 1) / feather963 mask.narrow(d, t, 1).mul_(a)964 mask.narrow(d, mask.shape[d] - 1 - t, 1).mul_(a)965 966 o = out967 o_d = out_div968 for d in range(dims):969 o = o.narrow(d + 2, upscaled[d], mask.shape[d + 2])970 o_d = o_d.narrow(d + 2, upscaled[d], mask.shape[d + 2])971 972 o.add_(ps * mask)973 o_d.add_(mask)974 975 if pbar is not None:976 pbar.update(1)977 978 output[b:b+1] = out/out_div979 return output980 981def tiled_scale(samples, function, tile_x=64, tile_y=64, overlap = 8, upscale_amount = 4, out_channels = 3, output_device="cpu", pbar = None):982 return tiled_scale_multidim(samples, function, (tile_y, tile_x), overlap=overlap, upscale_amount=upscale_amount, out_channels=out_channels, output_device=output_device, pbar=pbar)983 984PROGRESS_BAR_ENABLED = True985def set_progress_bar_enabled(enabled):986 global PROGRESS_BAR_ENABLED987 PROGRESS_BAR_ENABLED = enabled988 989PROGRESS_BAR_HOOK = None990def set_progress_bar_global_hook(function):991 global PROGRESS_BAR_HOOK992 PROGRESS_BAR_HOOK = function993 994class ProgressBar:995 def __init__(self, total):996 global PROGRESS_BAR_HOOK997 self.total = total998 self.current = 0999 self.hook = PROGRESS_BAR_HOOK1000 1001 def update_absolute(self, value, total=None, preview=None):1002 if total is not None:1003 self.total = total1004 if value > self.total:1005 value = self.total1006 self.current = value1007 if self.hook is not None:1008 self.hook(self.current, self.total, preview)1009 1010 def update(self, value):1011 self.update_absolute(self.current + value)1012 1013def reshape_mask(input_mask, output_shape):1014 dims = len(output_shape) - 21015 1016 if dims == 1:1017 scale_mode = "linear"1018 1019 if dims == 2:1020 input_mask = input_mask.reshape((-1, 1, input_mask.shape[-2], input_mask.shape[-1]))1021 scale_mode = "bilinear"1022 1023 if dims == 3:1024 if len(input_mask.shape) < 5:1025 input_mask = input_mask.reshape((1, 1, -1, input_mask.shape[-2], input_mask.shape[-1]))1026 scale_mode = "trilinear"1027 1028 mask = torch.nn.functional.interpolate(input_mask, size=output_shape[2:], mode=scale_mode)1029 if mask.shape[1] < output_shape[1]:1030 mask = mask.repeat((1, output_shape[1]) + (1,) * dims)[:,:output_shape[1]]1031 mask = repeat_to_batch_size(mask, output_shape[0])1032 return mask1033 1034def upscale_dit_mask(mask: torch.Tensor, img_size_in, img_size_out):1035 hi, wi = img_size_in1036 ho, wo = img_size_out1037 # if it's already the correct size, no need to do anything1038 if (hi, wi) == (ho, wo):1039 return mask1040 if mask.ndim == 2:1041 mask = mask.unsqueeze(0)1042 if mask.ndim != 3:1043 raise ValueError(f"Got a mask of shape {list(mask.shape)}, expected [b, q, k] or [q, k]")1044 txt_tokens = mask.shape[1] - (hi * wi)1045 # quadrants of the mask1046 txt_to_txt = mask[:, :txt_tokens, :txt_tokens]1047 txt_to_img = mask[:, :txt_tokens, txt_tokens:]1048 img_to_img = mask[:, txt_tokens:, txt_tokens:]1049 img_to_txt = mask[:, txt_tokens:, :txt_tokens]1050 1051 # convert to 1d x 2d, interpolate, then back to 1d x 1d1052 txt_to_img = rearrange (txt_to_img, "b t (h w) -> b t h w", h=hi, w=wi)1053 txt_to_img = interpolate(txt_to_img, size=img_size_out, mode="bilinear")1054 txt_to_img = rearrange (txt_to_img, "b t h w -> b t (h w)")1055 # this one is hard because we have to do it twice1056 # convert to 1d x 2d, interpolate, then to 2d x 1d, interpolate, then 1d x 1d1057 img_to_img = rearrange (img_to_img, "b hw (h w) -> b hw h w", h=hi, w=wi)1058 img_to_img = interpolate(img_to_img, size=img_size_out, mode="bilinear")1059 img_to_img = rearrange (img_to_img, "b (hk wk) hq wq -> b (hq wq) hk wk", hk=hi, wk=wi)1060 img_to_img = interpolate(img_to_img, size=img_size_out, mode="bilinear")1061 img_to_img = rearrange (img_to_img, "b (hq wq) hk wk -> b (hk wk) (hq wq)", hq=ho, wq=wo)1062 # convert to 2d x 1d, interpolate, then back to 1d x 1d1063 img_to_txt = rearrange (img_to_txt, "b (h w) t -> b t h w", h=hi, w=wi)1064 img_to_txt = interpolate(img_to_txt, size=img_size_out, mode="bilinear")1065 img_to_txt = rearrange (img_to_txt, "b t h w -> b (h w) t")1066 1067 # reassemble the mask from blocks1068 out = torch.cat([1069 torch.cat([txt_to_txt, txt_to_img], dim=2),1070 torch.cat([img_to_txt, img_to_img], dim=2)],1071 dim=11072 )1073 return out1074