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 torch21from enum import Enum22import math23import os24import logging25import comfy.utils26import comfy.model_management27import comfy.model_detection28import comfy.model_patcher29import comfy.ops30import comfy.latent_formats31 32import comfy.cldm.cldm33import comfy.t2i_adapter.adapter34import comfy.ldm.cascade.controlnet35import comfy.cldm.mmdit36import comfy.ldm.hydit.controlnet37import comfy.ldm.flux.controlnet38import comfy.cldm.dit_embedder39from typing import TYPE_CHECKING40if TYPE_CHECKING:41 from comfy.hooks import HookGroup42 43 44def broadcast_image_to(tensor, target_batch_size, batched_number):45 current_batch_size = tensor.shape[0]46 #print(current_batch_size, target_batch_size)47 if current_batch_size == 1:48 return tensor49 50 per_batch = target_batch_size // batched_number51 tensor = tensor[:per_batch]52 53 if per_batch > tensor.shape[0]:54 tensor = torch.cat([tensor] * (per_batch // tensor.shape[0]) + [tensor[:(per_batch % tensor.shape[0])]], dim=0)55 56 current_batch_size = tensor.shape[0]57 if current_batch_size == target_batch_size:58 return tensor59 else:60 return torch.cat([tensor] * batched_number, dim=0)61 62class StrengthType(Enum):63 CONSTANT = 164 LINEAR_UP = 265 66class ControlBase:67 def __init__(self):68 self.cond_hint_original = None69 self.cond_hint = None70 self.strength = 1.071 self.timestep_percent_range = (0.0, 1.0)72 self.latent_format = None73 self.vae = None74 self.global_average_pooling = False75 self.timestep_range = None76 self.compression_ratio = 877 self.upscale_algorithm = 'nearest-exact'78 self.extra_args = {}79 self.previous_controlnet = None80 self.extra_conds = []81 self.strength_type = StrengthType.CONSTANT82 self.concat_mask = False83 self.extra_concat_orig = []84 self.extra_concat = None85 self.extra_hooks: HookGroup = None86 self.preprocess_image = lambda a: a87 88 def set_cond_hint(self, cond_hint, strength=1.0, timestep_percent_range=(0.0, 1.0), vae=None, extra_concat=[]):89 self.cond_hint_original = cond_hint90 self.strength = strength91 self.timestep_percent_range = timestep_percent_range92 if self.latent_format is not None:93 if vae is None:94 logging.warning("WARNING: no VAE provided to the controlnet apply node when this controlnet requires one.")95 self.vae = vae96 self.extra_concat_orig = extra_concat.copy()97 if self.concat_mask and len(self.extra_concat_orig) == 0:98 self.extra_concat_orig.append(torch.tensor([[[[1.0]]]]))99 return self100 101 def pre_run(self, model, percent_to_timestep_function):102 self.timestep_range = (percent_to_timestep_function(self.timestep_percent_range[0]), percent_to_timestep_function(self.timestep_percent_range[1]))103 if self.previous_controlnet is not None:104 self.previous_controlnet.pre_run(model, percent_to_timestep_function)105 106 def set_previous_controlnet(self, controlnet):107 self.previous_controlnet = controlnet108 return self109 110 def cleanup(self):111 if self.previous_controlnet is not None:112 self.previous_controlnet.cleanup()113 114 self.cond_hint = None115 self.extra_concat = None116 self.timestep_range = None117 118 def get_models(self):119 out = []120 if self.previous_controlnet is not None:121 out += self.previous_controlnet.get_models()122 return out123 124 def get_extra_hooks(self):125 out = []126 if self.extra_hooks is not None:127 out.append(self.extra_hooks)128 if self.previous_controlnet is not None:129 out += self.previous_controlnet.get_extra_hooks()130 return out131 132 def copy_to(self, c):133 c.cond_hint_original = self.cond_hint_original134 c.strength = self.strength135 c.timestep_percent_range = self.timestep_percent_range136 c.global_average_pooling = self.global_average_pooling137 c.compression_ratio = self.compression_ratio138 c.upscale_algorithm = self.upscale_algorithm139 c.latent_format = self.latent_format140 c.extra_args = self.extra_args.copy()141 c.vae = self.vae142 c.extra_conds = self.extra_conds.copy()143 c.strength_type = self.strength_type144 c.concat_mask = self.concat_mask145 c.extra_concat_orig = self.extra_concat_orig.copy()146 c.extra_hooks = self.extra_hooks.clone() if self.extra_hooks else None147 c.preprocess_image = self.preprocess_image148 149 def inference_memory_requirements(self, dtype):150 if self.previous_controlnet is not None:151 return self.previous_controlnet.inference_memory_requirements(dtype)152 return 0153 154 def control_merge(self, control, control_prev, output_dtype):155 out = {'input':[], 'middle':[], 'output': []}156 157 for key in control:158 control_output = control[key]159 applied_to = set()160 for i in range(len(control_output)):161 x = control_output[i]162 if x is not None:163 if self.global_average_pooling:164 x = torch.mean(x, dim=(2, 3), keepdim=True).repeat(1, 1, x.shape[2], x.shape[3])165 166 if x not in applied_to: #memory saving strategy, allow shared tensors and only apply strength to shared tensors once167 applied_to.add(x)168 if self.strength_type == StrengthType.CONSTANT:169 x *= self.strength170 elif self.strength_type == StrengthType.LINEAR_UP:171 x *= (self.strength ** float(len(control_output) - i))172 173 if output_dtype is not None and x.dtype != output_dtype:174 x = x.to(output_dtype)175 176 out[key].append(x)177 178 if control_prev is not None:179 for x in ['input', 'middle', 'output']:180 o = out[x]181 for i in range(len(control_prev[x])):182 prev_val = control_prev[x][i]183 if i >= len(o):184 o.append(prev_val)185 elif prev_val is not None:186 if o[i] is None:187 o[i] = prev_val188 else:189 if o[i].shape[0] < prev_val.shape[0]:190 o[i] = prev_val + o[i]191 else:192 o[i] = prev_val + o[i] #TODO: change back to inplace add if shared tensors stop being an issue193 return out194 195 def set_extra_arg(self, argument, value=None):196 self.extra_args[argument] = value197 198 199class ControlNet(ControlBase):200 def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT, concat_mask=False, preprocess_image=lambda a: a):201 super().__init__()202 self.control_model = control_model203 self.load_device = load_device204 if control_model is not None:205 self.control_model_wrapped = comfy.model_patcher.ModelPatcher(self.control_model, load_device=load_device, offload_device=comfy.model_management.unet_offload_device())206 207 self.compression_ratio = compression_ratio208 self.global_average_pooling = global_average_pooling209 self.model_sampling_current = None210 self.manual_cast_dtype = manual_cast_dtype211 self.latent_format = latent_format212 self.extra_conds += extra_conds213 self.strength_type = strength_type214 self.concat_mask = concat_mask215 self.preprocess_image = preprocess_image216 217 def get_control(self, x_noisy, t, cond, batched_number, transformer_options):218 control_prev = None219 if self.previous_controlnet is not None:220 control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number, transformer_options)221 222 if self.timestep_range is not None:223 if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:224 if control_prev is not None:225 return control_prev226 else:227 return None228 229 dtype = self.control_model.dtype230 if self.manual_cast_dtype is not None:231 dtype = self.manual_cast_dtype232 233 if self.cond_hint is None or x_noisy.shape[2] * self.compression_ratio != self.cond_hint.shape[2] or x_noisy.shape[3] * self.compression_ratio != self.cond_hint.shape[3]:234 if self.cond_hint is not None:235 del self.cond_hint236 self.cond_hint = None237 compression_ratio = self.compression_ratio238 if self.vae is not None:239 compression_ratio *= self.vae.downscale_ratio240 else:241 if self.latent_format is not None:242 raise ValueError("This Controlnet needs a VAE but none was provided, please use a ControlNetApply node with a VAE input and connect it.")243 self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, self.upscale_algorithm, "center")244 self.cond_hint = self.preprocess_image(self.cond_hint)245 if self.vae is not None:246 loaded_models = comfy.model_management.loaded_models(only_currently_used=True)247 self.cond_hint = self.vae.encode(self.cond_hint.movedim(1, -1))248 comfy.model_management.load_models_gpu(loaded_models)249 if self.latent_format is not None:250 self.cond_hint = self.latent_format.process_in(self.cond_hint)251 if len(self.extra_concat_orig) > 0:252 to_concat = []253 for c in self.extra_concat_orig:254 c = c.to(self.cond_hint.device)255 c = comfy.utils.common_upscale(c, self.cond_hint.shape[3], self.cond_hint.shape[2], self.upscale_algorithm, "center")256 to_concat.append(comfy.utils.repeat_to_batch_size(c, self.cond_hint.shape[0]))257 self.cond_hint = torch.cat([self.cond_hint] + to_concat, dim=1)258 259 self.cond_hint = self.cond_hint.to(device=x_noisy.device, dtype=dtype)260 if x_noisy.shape[0] != self.cond_hint.shape[0]:261 self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)262 263 context = cond.get('crossattn_controlnet', cond['c_crossattn'])264 extra = self.extra_args.copy()265 for c in self.extra_conds:266 temp = cond.get(c, None)267 if temp is not None:268 extra[c] = temp.to(dtype)269 270 timestep = self.model_sampling_current.timestep(t)271 x_noisy = self.model_sampling_current.calculate_input(t, x_noisy)272 273 control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.to(dtype), context=context.to(dtype), **extra)274 return self.control_merge(control, control_prev, output_dtype=None)275 276 def copy(self):277 c = ControlNet(None, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype)278 c.control_model = self.control_model279 c.control_model_wrapped = self.control_model_wrapped280 self.copy_to(c)281 return c282 283 def get_models(self):284 out = super().get_models()285 out.append(self.control_model_wrapped)286 return out287 288 def pre_run(self, model, percent_to_timestep_function):289 super().pre_run(model, percent_to_timestep_function)290 self.model_sampling_current = model.model_sampling291 292 def cleanup(self):293 self.model_sampling_current = None294 super().cleanup()295 296class ControlLoraOps:297 class Linear(torch.nn.Module, comfy.ops.CastWeightBiasOp):298 def __init__(self, in_features: int, out_features: int, bias: bool = True,299 device=None, dtype=None) -> None:300 super().__init__()301 self.in_features = in_features302 self.out_features = out_features303 self.weight = None304 self.up = None305 self.down = None306 self.bias = None307 308 def forward(self, input):309 weight, bias = comfy.ops.cast_bias_weight(self, input)310 if self.up is not None:311 return torch.nn.functional.linear(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias)312 else:313 return torch.nn.functional.linear(input, weight, bias)314 315 class Conv2d(torch.nn.Module, comfy.ops.CastWeightBiasOp):316 def __init__(317 self,318 in_channels,319 out_channels,320 kernel_size,321 stride=1,322 padding=0,323 dilation=1,324 groups=1,325 bias=True,326 padding_mode='zeros',327 device=None,328 dtype=None329 ):330 super().__init__()331 self.in_channels = in_channels332 self.out_channels = out_channels333 self.kernel_size = kernel_size334 self.stride = stride335 self.padding = padding336 self.dilation = dilation337 self.transposed = False338 self.output_padding = 0339 self.groups = groups340 self.padding_mode = padding_mode341 342 self.weight = None343 self.bias = None344 self.up = None345 self.down = None346 347 348 def forward(self, input):349 weight, bias = comfy.ops.cast_bias_weight(self, input)350 if self.up is not None:351 return torch.nn.functional.conv2d(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias, self.stride, self.padding, self.dilation, self.groups)352 else:353 return torch.nn.functional.conv2d(input, weight, bias, self.stride, self.padding, self.dilation, self.groups)354 355 356class ControlLora(ControlNet):357 def __init__(self, control_weights, global_average_pooling=False, model_options={}): #TODO? model_options358 ControlBase.__init__(self)359 self.control_weights = control_weights360 self.global_average_pooling = global_average_pooling361 self.extra_conds += ["y"]362 363 def pre_run(self, model, percent_to_timestep_function):364 super().pre_run(model, percent_to_timestep_function)365 controlnet_config = model.model_config.unet_config.copy()366 controlnet_config.pop("out_channels")367 controlnet_config["hint_channels"] = self.control_weights["input_hint_block.0.weight"].shape[1]368 self.manual_cast_dtype = model.manual_cast_dtype369 dtype = model.get_dtype()370 if self.manual_cast_dtype is None:371 class control_lora_ops(ControlLoraOps, comfy.ops.disable_weight_init):372 pass373 else:374 class control_lora_ops(ControlLoraOps, comfy.ops.manual_cast):375 pass376 dtype = self.manual_cast_dtype377 378 controlnet_config["operations"] = control_lora_ops379 controlnet_config["dtype"] = dtype380 self.control_model = comfy.cldm.cldm.ControlNet(**controlnet_config)381 self.control_model.to(comfy.model_management.get_torch_device())382 diffusion_model = model.diffusion_model383 sd = diffusion_model.state_dict()384 385 for k in sd:386 weight = sd[k]387 try:388 comfy.utils.set_attr_param(self.control_model, k, weight)389 except:390 pass391 392 for k in self.control_weights:393 if k not in {"lora_controlnet"}:394 comfy.utils.set_attr_param(self.control_model, k, self.control_weights[k].to(dtype).to(comfy.model_management.get_torch_device()))395 396 def copy(self):397 c = ControlLora(self.control_weights, global_average_pooling=self.global_average_pooling)398 self.copy_to(c)399 return c400 401 def cleanup(self):402 del self.control_model403 self.control_model = None404 super().cleanup()405 406 def get_models(self):407 out = ControlBase.get_models(self)408 return out409 410 def inference_memory_requirements(self, dtype):411 return comfy.utils.calculate_parameters(self.control_weights) * comfy.model_management.dtype_size(dtype) + ControlBase.inference_memory_requirements(self, dtype)412 413def controlnet_config(sd, model_options={}):414 model_config = comfy.model_detection.model_config_from_unet(sd, "", True)415 416 unet_dtype = model_options.get("dtype", None)417 if unet_dtype is None:418 weight_dtype = comfy.utils.weight_dtype(sd)419 420 supported_inference_dtypes = list(model_config.supported_inference_dtypes)421 unet_dtype = comfy.model_management.unet_dtype(model_params=-1, supported_dtypes=supported_inference_dtypes, weight_dtype=weight_dtype)422 423 load_device = comfy.model_management.get_torch_device()424 manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device)425 426 operations = model_options.get("custom_operations", None)427 if operations is None:428 operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype, disable_fast_fp8=True)429 430 offload_device = comfy.model_management.unet_offload_device()431 return model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device432 433def controlnet_load_state_dict(control_model, sd):434 missing, unexpected = control_model.load_state_dict(sd, strict=False)435 436 if len(missing) > 0:437 logging.warning("missing controlnet keys: {}".format(missing))438 439 if len(unexpected) > 0:440 logging.debug("unexpected controlnet keys: {}".format(unexpected))441 return control_model442 443 444def load_controlnet_mmdit(sd, model_options={}):445 new_sd = comfy.model_detection.convert_diffusers_mmdit(sd, "")446 model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(new_sd, model_options=model_options)447 num_blocks = comfy.model_detection.count_blocks(new_sd, 'joint_blocks.{}.')448 for k in sd:449 new_sd[k] = sd[k]450 451 concat_mask = False452 control_latent_channels = new_sd.get("pos_embed_input.proj.weight").shape[1]453 if control_latent_channels == 17: #inpaint controlnet454 concat_mask = True455 456 control_model = comfy.cldm.mmdit.ControlNet(num_blocks=num_blocks, control_latent_channels=control_latent_channels, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config)457 control_model = controlnet_load_state_dict(control_model, new_sd)458 459 latent_format = comfy.latent_formats.SD3()460 latent_format.shift_factor = 0 #SD3 controlnet weirdness461 control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype)462 return control463 464 465class ControlNetSD35(ControlNet):466 def pre_run(self, model, percent_to_timestep_function):467 if self.control_model.double_y_emb:468 missing, unexpected = self.control_model.orig_y_embedder.load_state_dict(model.diffusion_model.y_embedder.state_dict(), strict=False)469 else:470 missing, unexpected = self.control_model.x_embedder.load_state_dict(model.diffusion_model.x_embedder.state_dict(), strict=False)471 super().pre_run(model, percent_to_timestep_function)472 473 def copy(self):474 c = ControlNetSD35(None, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype)475 c.control_model = self.control_model476 c.control_model_wrapped = self.control_model_wrapped477 self.copy_to(c)478 return c479 480def load_controlnet_sd35(sd, model_options={}):481 control_type = -1482 if "control_type" in sd:483 control_type = round(sd.pop("control_type").item())484 485 # blur_cnet = control_type == 0486 canny_cnet = control_type == 1487 depth_cnet = control_type == 2488 489 new_sd = {}490 for k in comfy.utils.MMDIT_MAP_BASIC:491 if k[1] in sd:492 new_sd[k[0]] = sd.pop(k[1])493 for k in sd:494 new_sd[k] = sd[k]495 sd = new_sd496 497 y_emb_shape = sd["y_embedder.mlp.0.weight"].shape498 depth = y_emb_shape[0] // 64499 hidden_size = 64 * depth500 num_heads = depth501 head_dim = hidden_size // num_heads502 num_blocks = comfy.model_detection.count_blocks(new_sd, 'transformer_blocks.{}.')503 504 load_device = comfy.model_management.get_torch_device()505 offload_device = comfy.model_management.unet_offload_device()506 unet_dtype = comfy.model_management.unet_dtype(model_params=-1)507 508 manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device)509 510 operations = model_options.get("custom_operations", None)511 if operations is None:512 operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype, disable_fast_fp8=True)513 514 control_model = comfy.cldm.dit_embedder.ControlNetEmbedder(img_size=None,515 patch_size=2,516 in_chans=16,517 num_layers=num_blocks,518 main_model_double=depth,519 double_y_emb=y_emb_shape[0] == y_emb_shape[1],520 attention_head_dim=head_dim,521 num_attention_heads=num_heads,522 adm_in_channels=2048,523 device=offload_device,524 dtype=unet_dtype,525 operations=operations)526 527 control_model = controlnet_load_state_dict(control_model, sd)528 529 latent_format = comfy.latent_formats.SD3()530 preprocess_image = lambda a: a531 if canny_cnet:532 preprocess_image = lambda a: (a * 255 * 0.5 + 0.5)533 elif depth_cnet:534 preprocess_image = lambda a: 1.0 - a535 536 control = ControlNetSD35(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype, preprocess_image=preprocess_image)537 return control538 539 540 541def load_controlnet_hunyuandit(controlnet_data, model_options={}):542 model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(controlnet_data, model_options=model_options)543 544 control_model = comfy.ldm.hydit.controlnet.HunYuanControlNet(operations=operations, device=offload_device, dtype=unet_dtype)545 control_model = controlnet_load_state_dict(control_model, controlnet_data)546 547 latent_format = comfy.latent_formats.SDXL()548 extra_conds = ['text_embedding_mask', 'encoder_hidden_states_t5', 'text_embedding_mask_t5', 'image_meta_size', 'style', 'cos_cis_img', 'sin_cis_img']549 control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds, strength_type=StrengthType.CONSTANT)550 return control551 552def load_controlnet_flux_xlabs_mistoline(sd, mistoline=False, model_options={}):553 model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(sd, model_options=model_options)554 control_model = comfy.ldm.flux.controlnet.ControlNetFlux(mistoline=mistoline, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config)555 control_model = controlnet_load_state_dict(control_model, sd)556 extra_conds = ['y', 'guidance']557 control = ControlNet(control_model, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds)558 return control559 560def load_controlnet_flux_instantx(sd, model_options={}):561 new_sd = comfy.model_detection.convert_diffusers_mmdit(sd, "")562 model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(new_sd, model_options=model_options)563 for k in sd:564 new_sd[k] = sd[k]565 566 num_union_modes = 0567 union_cnet = "controlnet_mode_embedder.weight"568 if union_cnet in new_sd:569 num_union_modes = new_sd[union_cnet].shape[0]570 571 control_latent_channels = new_sd.get("pos_embed_input.weight").shape[1] // 4572 concat_mask = False573 if control_latent_channels == 17:574 concat_mask = True575 576 control_model = comfy.ldm.flux.controlnet.ControlNetFlux(latent_input=True, num_union_modes=num_union_modes, control_latent_channels=control_latent_channels, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config)577 control_model = controlnet_load_state_dict(control_model, new_sd)578 579 latent_format = comfy.latent_formats.Flux()580 extra_conds = ['y', 'guidance']581 control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds)582 return control583 584def convert_mistoline(sd):585 return comfy.utils.state_dict_prefix_replace(sd, {"single_controlnet_blocks.": "controlnet_single_blocks."})586 587 588def load_controlnet_state_dict(state_dict, model=None, model_options={}):589 controlnet_data = state_dict590 if 'after_proj_list.18.bias' in controlnet_data.keys(): #Hunyuan DiT591 return load_controlnet_hunyuandit(controlnet_data, model_options=model_options)592 593 if "lora_controlnet" in controlnet_data:594 return ControlLora(controlnet_data, model_options=model_options)595 596 controlnet_config = None597 supported_inference_dtypes = None598 599 if "controlnet_cond_embedding.conv_in.weight" in controlnet_data: #diffusers format600 controlnet_config = comfy.model_detection.unet_config_from_diffusers_unet(controlnet_data)601 diffusers_keys = comfy.utils.unet_to_diffusers(controlnet_config)602 diffusers_keys["controlnet_mid_block.weight"] = "middle_block_out.0.weight"603 diffusers_keys["controlnet_mid_block.bias"] = "middle_block_out.0.bias"604 605 count = 0606 loop = True607 while loop:608 suffix = [".weight", ".bias"]609 for s in suffix:610 k_in = "controlnet_down_blocks.{}{}".format(count, s)611 k_out = "zero_convs.{}.0{}".format(count, s)612 if k_in not in controlnet_data:613 loop = False614 break615 diffusers_keys[k_in] = k_out616 count += 1617 618 count = 0619 loop = True620 while loop:621 suffix = [".weight", ".bias"]622 for s in suffix:623 if count == 0:624 k_in = "controlnet_cond_embedding.conv_in{}".format(s)625 else:626 k_in = "controlnet_cond_embedding.blocks.{}{}".format(count - 1, s)627 k_out = "input_hint_block.{}{}".format(count * 2, s)628 if k_in not in controlnet_data:629 k_in = "controlnet_cond_embedding.conv_out{}".format(s)630 loop = False631 diffusers_keys[k_in] = k_out632 count += 1633 634 new_sd = {}635 for k in diffusers_keys:636 if k in controlnet_data:637 new_sd[diffusers_keys[k]] = controlnet_data.pop(k)638 639 if "control_add_embedding.linear_1.bias" in controlnet_data: #Union Controlnet640 controlnet_config["union_controlnet_num_control_type"] = controlnet_data["task_embedding"].shape[0]641 for k in list(controlnet_data.keys()):642 new_k = k.replace('.attn.in_proj_', '.attn.in_proj.')643 new_sd[new_k] = controlnet_data.pop(k)644 645 leftover_keys = controlnet_data.keys()646 if len(leftover_keys) > 0:647 logging.warning("leftover keys: {}".format(leftover_keys))648 controlnet_data = new_sd649 elif "controlnet_blocks.0.weight" in controlnet_data:650 if "double_blocks.0.img_attn.norm.key_norm.scale" in controlnet_data:651 return load_controlnet_flux_xlabs_mistoline(controlnet_data, model_options=model_options)652 elif "pos_embed_input.proj.weight" in controlnet_data:653 if "transformer_blocks.0.adaLN_modulation.1.bias" in controlnet_data:654 return load_controlnet_sd35(controlnet_data, model_options=model_options) #Stability sd3.5 format655 else:656 return load_controlnet_mmdit(controlnet_data, model_options=model_options) #SD3 diffusers controlnet657 elif "controlnet_x_embedder.weight" in controlnet_data:658 return load_controlnet_flux_instantx(controlnet_data, model_options=model_options)659 elif "controlnet_blocks.0.linear.weight" in controlnet_data: #mistoline flux660 return load_controlnet_flux_xlabs_mistoline(convert_mistoline(controlnet_data), mistoline=True, model_options=model_options)661 662 pth_key = 'control_model.zero_convs.0.0.weight'663 pth = False664 key = 'zero_convs.0.0.weight'665 if pth_key in controlnet_data:666 pth = True667 key = pth_key668 prefix = "control_model."669 elif key in controlnet_data:670 prefix = ""671 else:672 net = load_t2i_adapter(controlnet_data, model_options=model_options)673 if net is None:674 logging.error("error could not detect control model type.")675 return net676 677 if controlnet_config is None:678 model_config = comfy.model_detection.model_config_from_unet(controlnet_data, prefix, True)679 supported_inference_dtypes = list(model_config.supported_inference_dtypes)680 controlnet_config = model_config.unet_config681 682 unet_dtype = model_options.get("dtype", None)683 if unet_dtype is None:684 weight_dtype = comfy.utils.weight_dtype(controlnet_data)685 686 if supported_inference_dtypes is None:687 supported_inference_dtypes = [comfy.model_management.unet_dtype()]688 689 unet_dtype = comfy.model_management.unet_dtype(model_params=-1, supported_dtypes=supported_inference_dtypes, weight_dtype=weight_dtype)690 691 load_device = comfy.model_management.get_torch_device()692 693 manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device)694 operations = model_options.get("custom_operations", None)695 if operations is None:696 operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype)697 698 controlnet_config["operations"] = operations699 controlnet_config["dtype"] = unet_dtype700 controlnet_config["device"] = comfy.model_management.unet_offload_device()701 controlnet_config.pop("out_channels")702 controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1]703 control_model = comfy.cldm.cldm.ControlNet(**controlnet_config)704 705 if pth:706 if 'difference' in controlnet_data:707 if model is not None:708 comfy.model_management.load_models_gpu([model])709 model_sd = model.model_state_dict()710 for x in controlnet_data:711 c_m = "control_model."712 if x.startswith(c_m):713 sd_key = "diffusion_model.{}".format(x[len(c_m):])714 if sd_key in model_sd:715 cd = controlnet_data[x]716 cd += model_sd[sd_key].type(cd.dtype).to(cd.device)717 else:718 logging.warning("WARNING: Loaded a diff controlnet without a model. It will very likely not work.")719 720 class WeightsLoader(torch.nn.Module):721 pass722 w = WeightsLoader()723 w.control_model = control_model724 missing, unexpected = w.load_state_dict(controlnet_data, strict=False)725 else:726 missing, unexpected = control_model.load_state_dict(controlnet_data, strict=False)727 728 if len(missing) > 0:729 logging.warning("missing controlnet keys: {}".format(missing))730 731 if len(unexpected) > 0:732 logging.debug("unexpected controlnet keys: {}".format(unexpected))733 734 global_average_pooling = model_options.get("global_average_pooling", False)735 control = ControlNet(control_model, global_average_pooling=global_average_pooling, load_device=load_device, manual_cast_dtype=manual_cast_dtype)736 return control737 738def load_controlnet(ckpt_path, model=None, model_options={}):739 if "global_average_pooling" not in model_options:740 filename = os.path.splitext(ckpt_path)[0]741 if filename.endswith("_shuffle") or filename.endswith("_shuffle_fp16"): #TODO: smarter way of enabling global_average_pooling742 model_options["global_average_pooling"] = True743 744 cnet = load_controlnet_state_dict(comfy.utils.load_torch_file(ckpt_path, safe_load=True), model=model, model_options=model_options)745 if cnet is None:746 logging.error("error checkpoint does not contain controlnet or t2i adapter data {}".format(ckpt_path))747 return cnet748 749class T2IAdapter(ControlBase):750 def __init__(self, t2i_model, channels_in, compression_ratio, upscale_algorithm, device=None):751 super().__init__()752 self.t2i_model = t2i_model753 self.channels_in = channels_in754 self.control_input = None755 self.compression_ratio = compression_ratio756 self.upscale_algorithm = upscale_algorithm757 if device is None:758 device = comfy.model_management.get_torch_device()759 self.device = device760 761 def scale_image_to(self, width, height):762 unshuffle_amount = self.t2i_model.unshuffle_amount763 width = math.ceil(width / unshuffle_amount) * unshuffle_amount764 height = math.ceil(height / unshuffle_amount) * unshuffle_amount765 return width, height766 767 def get_control(self, x_noisy, t, cond, batched_number, transformer_options):768 control_prev = None769 if self.previous_controlnet is not None:770 control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number, transformer_options)771 772 if self.timestep_range is not None:773 if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:774 if control_prev is not None:775 return control_prev776 else:777 return None778 779 if self.cond_hint is None or x_noisy.shape[2] * self.compression_ratio != self.cond_hint.shape[2] or x_noisy.shape[3] * self.compression_ratio != self.cond_hint.shape[3]:780 if self.cond_hint is not None:781 del self.cond_hint782 self.control_input = None783 self.cond_hint = None784 width, height = self.scale_image_to(x_noisy.shape[3] * self.compression_ratio, x_noisy.shape[2] * self.compression_ratio)785 self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, width, height, self.upscale_algorithm, "center").float().to(self.device)786 if self.channels_in == 1 and self.cond_hint.shape[1] > 1:787 self.cond_hint = torch.mean(self.cond_hint, 1, keepdim=True)788 if x_noisy.shape[0] != self.cond_hint.shape[0]:789 self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)790 if self.control_input is None:791 self.t2i_model.to(x_noisy.dtype)792 self.t2i_model.to(self.device)793 self.control_input = self.t2i_model(self.cond_hint.to(x_noisy.dtype))794 self.t2i_model.cpu()795 796 control_input = {}797 for k in self.control_input:798 control_input[k] = list(map(lambda a: None if a is None else a.clone(), self.control_input[k]))799 800 return self.control_merge(control_input, control_prev, x_noisy.dtype)801 802 def copy(self):803 c = T2IAdapter(self.t2i_model, self.channels_in, self.compression_ratio, self.upscale_algorithm)804 self.copy_to(c)805 return c806 807def load_t2i_adapter(t2i_data, model_options={}): #TODO: model_options808 compression_ratio = 8809 upscale_algorithm = 'nearest-exact'810 811 if 'adapter' in t2i_data:812 t2i_data = t2i_data['adapter']813 if 'adapter.body.0.resnets.0.block1.weight' in t2i_data: #diffusers format814 prefix_replace = {}815 for i in range(4):816 for j in range(2):817 prefix_replace["adapter.body.{}.resnets.{}.".format(i, j)] = "body.{}.".format(i * 2 + j)818 prefix_replace["adapter.body.{}.".format(i, )] = "body.{}.".format(i * 2)819 prefix_replace["adapter."] = ""820 t2i_data = comfy.utils.state_dict_prefix_replace(t2i_data, prefix_replace)821 keys = t2i_data.keys()822 823 if "body.0.in_conv.weight" in keys:824 cin = t2i_data['body.0.in_conv.weight'].shape[1]825 model_ad = comfy.t2i_adapter.adapter.Adapter_light(cin=cin, channels=[320, 640, 1280, 1280], nums_rb=4)826 elif 'conv_in.weight' in keys:827 cin = t2i_data['conv_in.weight'].shape[1]828 channel = t2i_data['conv_in.weight'].shape[0]829 ksize = t2i_data['body.0.block2.weight'].shape[2]830 use_conv = False831 down_opts = list(filter(lambda a: a.endswith("down_opt.op.weight"), keys))832 if len(down_opts) > 0:833 use_conv = True834 xl = False835 if cin == 256 or cin == 768:836 xl = True837 model_ad = comfy.t2i_adapter.adapter.Adapter(cin=cin, channels=[channel, channel*2, channel*4, channel*4][:4], nums_rb=2, ksize=ksize, sk=True, use_conv=use_conv, xl=xl)838 elif "backbone.0.0.weight" in keys:839 model_ad = comfy.ldm.cascade.controlnet.ControlNet(c_in=t2i_data['backbone.0.0.weight'].shape[1], proj_blocks=[0, 4, 8, 12, 51, 55, 59, 63])840 compression_ratio = 32841 upscale_algorithm = 'bilinear'842 elif "backbone.10.blocks.0.weight" in keys:843 model_ad = comfy.ldm.cascade.controlnet.ControlNet(c_in=t2i_data['backbone.0.weight'].shape[1], bottleneck_mode="large", proj_blocks=[0, 4, 8, 12, 51, 55, 59, 63])844 compression_ratio = 1845 upscale_algorithm = 'nearest-exact'846 else:847 return None848 849 missing, unexpected = model_ad.load_state_dict(t2i_data)850 if len(missing) > 0:851 logging.warning("t2i missing {}".format(missing))852 853 if len(unexpected) > 0:854 logging.debug("t2i unexpected {}".format(unexpected))855 856 return T2IAdapter(model_ad, model_ad.input_channels, compression_ratio, upscale_algorithm)857