From 61196d88576c95c1cd8535e881af48172d5af525 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 25 Nov 2024 05:00:23 -0500 Subject: [PATCH 1/6] Add option to inference the diffusion model in fp32 and fp64. --- comfy/cli_args.py | 6 ++++-- comfy/model_management.py | 6 +++++- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 20b9f4749..847f35abd 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -60,8 +60,10 @@ fp_group.add_argument("--force-fp32", action="store_true", help="Force fp32 (If fp_group.add_argument("--force-fp16", action="store_true", help="Force fp16.") fpunet_group = parser.add_mutually_exclusive_group() -fpunet_group.add_argument("--bf16-unet", action="store_true", help="Run the UNET in bf16. This should only be used for testing stuff.") -fpunet_group.add_argument("--fp16-unet", action="store_true", help="Store unet weights in fp16.") +fpunet_group.add_argument("--fp32-unet", action="store_true", help="Run the diffusion model in fp32.") +fpunet_group.add_argument("--fp64-unet", action="store_true", help="Run the diffusion model in fp64.") +fpunet_group.add_argument("--bf16-unet", action="store_true", help="Run the diffusion model in bf16.") +fpunet_group.add_argument("--fp16-unet", action="store_true", help="Run the diffusion model in fp16") fpunet_group.add_argument("--fp8_e4m3fn-unet", action="store_true", help="Store unet weights in fp8_e4m3fn.") fpunet_group.add_argument("--fp8_e5m2-unet", action="store_true", help="Store unet weights in fp8_e5m2.") diff --git a/comfy/model_management.py b/comfy/model_management.py index fd493aff0..a793cab3e 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -628,6 +628,10 @@ def maximum_vram_for_weights(device=None): def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]): if model_params < 0: model_params = 1000000000000000000000 + if args.fp32_unet: + return torch.float32 + if args.fp64_unet: + return torch.float64 if args.bf16_unet: return torch.bfloat16 if args.fp16_unet: @@ -674,7 +678,7 @@ def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, tor # None means no manual cast def unet_manual_cast(weight_dtype, inference_device, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]): - if weight_dtype == torch.float32: + if weight_dtype == torch.float32 or weight_dtype == torch.float64: return None fp16_supported = should_use_fp16(inference_device, prioritize_performance=False) From b7143b74ce6df191004fab4a683cb2792c82248b Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 26 Nov 2024 01:33:01 -0500 Subject: [PATCH 2/6] Flux inpaint model does not work in fp16. --- comfy/supported_models.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/comfy/supported_models.py b/comfy/supported_models.py index caa1a1846..b426d3085 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -659,6 +659,15 @@ class Flux(supported_models_base.BASE): t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref)) return supported_models_base.ClipTarget(comfy.text_encoders.flux.FluxTokenizer, comfy.text_encoders.flux.flux_clip(**t5_detect)) +class FluxInpaint(Flux): + unet_config = { + "image_model": "flux", + "guidance_embed": True, + "in_channels": 96, + } + + supported_inference_dtypes = [torch.bfloat16, torch.float32] + class FluxSchnell(Flux): unet_config = { "image_model": "flux", @@ -731,6 +740,6 @@ class LTXV(supported_models_base.BASE): t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref)) return supported_models_base.ClipTarget(comfy.text_encoders.lt.LTXVT5Tokenizer, comfy.text_encoders.lt.ltxv_te(**t5_detect)) -models = [Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, HunyuanDiT, HunyuanDiT1, Flux, FluxSchnell, GenmoMochi, LTXV] +models = [Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, FluxSchnell, GenmoMochi, LTXV] models += [SVD_img2vid] From 15c39ea75774b5d36065eb8caa5ecbbfcd168b71 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 26 Nov 2024 03:34:36 -0500 Subject: [PATCH 3/6] Support for the official mochi lora format. --- comfy/lora.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/comfy/lora.py b/comfy/lora.py index 18602f24c..1080169b1 100644 --- a/comfy/lora.py +++ b/comfy/lora.py @@ -62,6 +62,7 @@ def load_lora(lora, to_load): diffusers_lora = "{}_lora.up.weight".format(x) diffusers2_lora = "{}.lora_B.weight".format(x) diffusers3_lora = "{}.lora.up.weight".format(x) + mochi_lora = "{}.lora_B".format(x) transformers_lora = "{}.lora_linear_layer.up.weight".format(x) A_name = None @@ -81,6 +82,10 @@ def load_lora(lora, to_load): A_name = diffusers3_lora B_name = "{}.lora.down.weight".format(x) mid_name = None + elif mochi_lora in lora.keys(): + A_name = mochi_lora + B_name = "{}.lora_A".format(x) + mid_name = None elif transformers_lora in lora.keys(): A_name = transformers_lora B_name ="{}.lora_linear_layer.down.weight".format(x) @@ -362,6 +367,12 @@ def model_lora_keys_unet(model, key_map={}): key_map["lycoris_{}".format(k[:-len(".weight")].replace(".", "_"))] = to #simpletrainer lycoris key_map["lora_transformer_{}".format(k[:-len(".weight")].replace(".", "_"))] = to #onetrainer + if isinstance(model, comfy.model_base.GenmoMochi): + for k in sdk: + if k.startswith("diffusion_model.") and k.endswith(".weight"): #Official Mochi lora format + key_lora = k[len("diffusion_model."):-len(".weight")] + key_map["{}".format(key_lora)] = k + return key_map From 4c82741b545c6cedcfa397034f56ce1377b3675a Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 26 Nov 2024 11:31:25 -0500 Subject: [PATCH 4/6] Support official SD3.5 Controlnets. --- comfy/cldm/dit_embedder.py | 122 +++++++++++++++++++++++++++++++++++++ comfy/controlnet.py | 91 ++++++++++++++++++++++++++- 2 files changed, 210 insertions(+), 3 deletions(-) create mode 100644 comfy/cldm/dit_embedder.py diff --git a/comfy/cldm/dit_embedder.py b/comfy/cldm/dit_embedder.py new file mode 100644 index 000000000..e9cdd4991 --- /dev/null +++ b/comfy/cldm/dit_embedder.py @@ -0,0 +1,122 @@ +import math +from typing import List, Optional, Tuple + +import numpy as np +import torch +import torch.nn as nn +from einops import rearrange +from torch import Tensor + +from comfy.ldm.modules.diffusionmodules.mmdit import DismantledBlock, PatchEmbed, VectorEmbedder, TimestepEmbedder, get_2d_sincos_pos_embed_torch + + +class ControlNetEmbedder(nn.Module): + + def __init__( + self, + img_size: int, + patch_size: int, + in_chans: int, + attention_head_dim: int, + num_attention_heads: int, + adm_in_channels: int, + num_layers: int, + main_model_double: int, + double_y_emb: bool, + device: torch.device, + dtype: torch.dtype, + pos_embed_max_size: Optional[int] = None, + operations = None, + ): + super().__init__() + self.main_model_double = main_model_double + self.dtype = dtype + self.hidden_size = num_attention_heads * attention_head_dim + self.patch_size = patch_size + self.x_embedder = PatchEmbed( + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=self.hidden_size, + strict_img_size=pos_embed_max_size is None, + device=device, + dtype=dtype, + operations=operations, + ) + + self.t_embedder = TimestepEmbedder(self.hidden_size, dtype=dtype, device=device, operations=operations) + + self.double_y_emb = double_y_emb + if self.double_y_emb: + self.orig_y_embedder = VectorEmbedder( + adm_in_channels, self.hidden_size, dtype, device, operations=operations + ) + self.y_embedder = VectorEmbedder( + self.hidden_size, self.hidden_size, dtype, device, operations=operations + ) + else: + self.y_embedder = VectorEmbedder( + adm_in_channels, self.hidden_size, dtype, device, operations=operations + ) + + self.transformer_blocks = nn.ModuleList( + DismantledBlock( + hidden_size=self.hidden_size, num_heads=num_attention_heads, qkv_bias=True, + dtype=dtype, device=device, operations=operations + ) + for _ in range(num_layers) + ) + + # self.use_y_embedder = pooled_projection_dim != self.time_text_embed.text_embedder.linear_1.in_features + # TODO double check this logic when 8b + self.use_y_embedder = True + + self.controlnet_blocks = nn.ModuleList([]) + for _ in range(len(self.transformer_blocks)): + controlnet_block = operations.Linear(self.hidden_size, self.hidden_size, dtype=dtype, device=device) + self.controlnet_blocks.append(controlnet_block) + + self.pos_embed_input = PatchEmbed( + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=self.hidden_size, + strict_img_size=False, + device=device, + dtype=dtype, + operations=operations, + ) + + def forward( + self, + x: torch.Tensor, + timesteps: torch.Tensor, + y: Optional[torch.Tensor] = None, + context: Optional[torch.Tensor] = None, + hint = None, + ) -> Tuple[Tensor, List[Tensor]]: + x_shape = list(x.shape) + x = self.x_embedder(x) + if not self.double_y_emb: + h = (x_shape[-2] + 1) // self.patch_size + w = (x_shape[-1] + 1) // self.patch_size + x += get_2d_sincos_pos_embed_torch(self.hidden_size, w, h, device=x.device) + c = self.t_embedder(timesteps, dtype=x.dtype) + if y is not None and self.y_embedder is not None: + if self.double_y_emb: + y = self.orig_y_embedder(y) + y = self.y_embedder(y) + c = c + y + + x = x + self.pos_embed_input(hint) + + block_out = () + + repeat = math.ceil(self.main_model_double / len(self.transformer_blocks)) + for i in range(len(self.transformer_blocks)): + out = self.transformer_blocks[i](x, c) + if not self.double_y_emb: + x = out + block_out += (self.controlnet_blocks[i](out),) * repeat + + return {"output": block_out} diff --git a/comfy/controlnet.py b/comfy/controlnet.py index d2744e427..cf9f894ca 100644 --- a/comfy/controlnet.py +++ b/comfy/controlnet.py @@ -35,7 +35,7 @@ import comfy.ldm.cascade.controlnet import comfy.cldm.mmdit import comfy.ldm.hydit.controlnet import comfy.ldm.flux.controlnet - +import comfy.cldm.dit_embedder def broadcast_image_to(tensor, target_batch_size, batched_number): current_batch_size = tensor.shape[0] @@ -78,6 +78,7 @@ class ControlBase: self.concat_mask = False self.extra_concat_orig = [] self.extra_concat = None + self.preprocess_image = lambda a: a def set_cond_hint(self, cond_hint, strength=1.0, timestep_percent_range=(0.0, 1.0), vae=None, extra_concat=[]): self.cond_hint_original = cond_hint @@ -129,6 +130,7 @@ class ControlBase: c.strength_type = self.strength_type c.concat_mask = self.concat_mask c.extra_concat_orig = self.extra_concat_orig.copy() + c.preprocess_image = self.preprocess_image def inference_memory_requirements(self, dtype): if self.previous_controlnet is not None: @@ -181,7 +183,7 @@ class ControlBase: class ControlNet(ControlBase): - 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): + 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): super().__init__() self.control_model = control_model self.load_device = load_device @@ -196,6 +198,7 @@ class ControlNet(ControlBase): self.extra_conds += extra_conds self.strength_type = strength_type self.concat_mask = concat_mask + self.preprocess_image = preprocess_image def get_control(self, x_noisy, t, cond, batched_number): control_prev = None @@ -224,6 +227,7 @@ class ControlNet(ControlBase): if self.latent_format is not None: raise ValueError("This Controlnet needs a VAE but none was provided, please use a ControlNetApply node with a VAE input and connect it.") 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") + self.cond_hint = self.preprocess_image(self.cond_hint) if self.vae is not None: loaded_models = comfy.model_management.loaded_models(only_currently_used=True) self.cond_hint = self.vae.encode(self.cond_hint.movedim(1, -1)) @@ -427,6 +431,7 @@ def controlnet_load_state_dict(control_model, sd): logging.debug("unexpected controlnet keys: {}".format(unexpected)) return control_model + def load_controlnet_mmdit(sd, model_options={}): new_sd = comfy.model_detection.convert_diffusers_mmdit(sd, "") model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(new_sd, model_options=model_options) @@ -448,6 +453,83 @@ def load_controlnet_mmdit(sd, model_options={}): return control +class ControlNetSD35(ControlNet): + def pre_run(self, model, percent_to_timestep_function): + if self.control_model.double_y_emb: + missing, unexpected = self.control_model.orig_y_embedder.load_state_dict(model.diffusion_model.y_embedder.state_dict(), strict=False) + else: + missing, unexpected = self.control_model.x_embedder.load_state_dict(model.diffusion_model.x_embedder.state_dict(), strict=False) + super().pre_run(model, percent_to_timestep_function) + + def copy(self): + c = ControlNetSD35(None, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype) + c.control_model = self.control_model + c.control_model_wrapped = self.control_model_wrapped + self.copy_to(c) + return c + +def load_controlnet_sd35(sd, model_options={}): + control_type = -1 + if "control_type" in sd: + control_type = round(sd.pop("control_type").item()) + + # blur_cnet = control_type == 0 + canny_cnet = control_type == 1 + depth_cnet = control_type == 2 + + print(control_type, canny_cnet, depth_cnet) + new_sd = {} + for k in comfy.utils.MMDIT_MAP_BASIC: + if k[1] in sd: + new_sd[k[0]] = sd.pop(k[1]) + for k in sd: + new_sd[k] = sd[k] + sd = new_sd + + y_emb_shape = sd["y_embedder.mlp.0.weight"].shape + depth = y_emb_shape[0] // 64 + hidden_size = 64 * depth + num_heads = depth + head_dim = hidden_size // num_heads + num_blocks = comfy.model_detection.count_blocks(new_sd, 'transformer_blocks.{}.') + + load_device = comfy.model_management.get_torch_device() + offload_device = comfy.model_management.unet_offload_device() + unet_dtype = comfy.model_management.unet_dtype(model_params=-1) + + manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device) + + operations = model_options.get("custom_operations", None) + if operations is None: + operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype, disable_fast_fp8=True) + + control_model = comfy.cldm.dit_embedder.ControlNetEmbedder(img_size=None, + patch_size=2, + in_chans=16, + num_layers=num_blocks, + main_model_double=depth, + double_y_emb=y_emb_shape[0] == y_emb_shape[1], + attention_head_dim=head_dim, + num_attention_heads=num_heads, + adm_in_channels=2048, + device=offload_device, + dtype=unet_dtype, + operations=operations) + + control_model = controlnet_load_state_dict(control_model, sd) + + latent_format = comfy.latent_formats.SD3() + preprocess_image = lambda a: a + if canny_cnet: + preprocess_image = lambda a: (a * 255 * 0.5 + 0.5) + elif depth_cnet: + preprocess_image = lambda a: 1.0 - a + + 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) + return control + + + def load_controlnet_hunyuandit(controlnet_data, model_options={}): model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(controlnet_data, model_options=model_options) @@ -560,7 +642,10 @@ def load_controlnet_state_dict(state_dict, model=None, model_options={}): if "double_blocks.0.img_attn.norm.key_norm.scale" in controlnet_data: return load_controlnet_flux_xlabs_mistoline(controlnet_data, model_options=model_options) elif "pos_embed_input.proj.weight" in controlnet_data: - return load_controlnet_mmdit(controlnet_data, model_options=model_options) #SD3 diffusers controlnet + if "transformer_blocks.0.adaLN_modulation.1.bias" in controlnet_data: + return load_controlnet_sd35(controlnet_data, model_options=model_options) #Stability sd3.5 format + else: + return load_controlnet_mmdit(controlnet_data, model_options=model_options) #SD3 diffusers controlnet elif "controlnet_x_embedder.weight" in controlnet_data: return load_controlnet_flux_instantx(controlnet_data, model_options=model_options) elif "controlnet_blocks.0.linear.weight" in controlnet_data: #mistoline flux From 24dc581dc3c3bfb12d4a2910609425b4b295d731 Mon Sep 17 00:00:00 2001 From: lky <32677614+lyksdu@users.noreply.github.com> Date: Wed, 27 Nov 2024 04:34:19 +0800 Subject: [PATCH 5/6] fix multi add makedirs error (#5786) try to start multiple comfyui server at the same time, and this got error --- app/user_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app/user_manager.py b/app/user_manager.py index bb6217995..e863b93dd 100644 --- a/app/user_manager.py +++ b/app/user_manager.py @@ -36,7 +36,7 @@ class UserManager(): self.settings = AppSettings(self) if not os.path.exists(user_directory): - os.mkdir(user_directory) + os.makedirs(user_directory, exist_ok=True) if not args.multi_user: print("****** User settings have been changed to be stored on the server instead of browser storage. ******") print("****** For multi-user setups add the --multi-user CLI argument to enable multiple user profiles. ******") From 497db6212faea643d3aac1253ccfa23834109881 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 26 Nov 2024 17:51:40 -0500 Subject: [PATCH 6/6] Alternative fix for #5767 --- comfy/model_base.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/comfy/model_base.py b/comfy/model_base.py index 3c385cda9..c305014a4 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -712,7 +712,13 @@ class Flux(BaseModel): super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.flux.model.Flux) def concat_cond(self, **kwargs): - num_channels = self.diffusion_model.img_in.weight.shape[1] // (self.diffusion_model.patch_size * self.diffusion_model.patch_size) + try: + #Handle Flux control loras dynamically changing the img_in weight. + num_channels = self.diffusion_model.img_in.weight.shape[1] // (self.diffusion_model.patch_size * self.diffusion_model.patch_size) + except: + #Some cases like tensorrt might not have the weights accessible + num_channels = self.model_config.unet_config["in_channels"] + out_channels = self.model_config.unet_config["out_channels"] if num_channels <= out_channels: