From 5c69cde0374207c1d4b6ec1ec033cfd5592d6de0 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 11 Aug 2024 23:50:01 -0400 Subject: [PATCH 01/14] Load TE model straight to vram if certain conditions are met. --- comfy/model_management.py | 11 +++++++++++ comfy/sd.py | 12 ++++++------ 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index a0105131d..6e2738391 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -684,6 +684,17 @@ def text_encoder_device(): else: return torch.device("cpu") +def text_encoder_initial_device(load_device, offload_device, model_size=0): + if load_device == offload_device or model_size <= 1024 * 1024 * 1024: + return offload_device + + mem_l = get_free_memory(load_device) + mem_o = get_free_memory(offload_device) + if mem_l > (mem_o * 0.5) and model_size * 1.2 < mem_l: + return load_device + else: + return offload_device + def text_encoder_dtype(device=None): if args.fp8_e4m3fn_text_enc: return torch.float8_e4m3fn diff --git a/comfy/sd.py b/comfy/sd.py index ee91ad53b..6d729929a 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -62,7 +62,7 @@ def load_lora_for_models(model, clip, lora, strength_model, strength_clip): class CLIP: - def __init__(self, target=None, embedding_directory=None, no_init=False, tokenizer_data={}): + def __init__(self, target=None, embedding_directory=None, no_init=False, tokenizer_data={}, parameters=0): if no_init: return params = target.params.copy() @@ -71,10 +71,9 @@ class CLIP: load_device = model_management.text_encoder_device() offload_device = model_management.text_encoder_offload_device() - params['device'] = offload_device dtype = model_management.text_encoder_dtype(load_device) params['dtype'] = dtype - + params['device'] = model_management.text_encoder_initial_device(load_device, offload_device, parameters * model_management.dtype_size(dtype)) self.cond_stage_model = clip(**(params)) for dt in self.cond_stage_model.dtypes: @@ -84,7 +83,7 @@ class CLIP: self.tokenizer = tokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data) self.patcher = comfy.model_patcher.ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device) self.layer_idx = None - logging.debug("CLIP model load device: {}, offload device: {}".format(load_device, offload_device)) + logging.debug("CLIP model load device: {}, offload device: {}, current: {}".format(load_device, offload_device, params['device'])) def clone(self): n = CLIP(no_init=True) @@ -456,7 +455,7 @@ def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DI clip_target.clip = comfy.text_encoders.sd3_clip.SD3ClipModel clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer - clip = CLIP(clip_target, embedding_directory=embedding_directory) + clip = CLIP(clip_target, embedding_directory=embedding_directory, state_dicts=clip_data) for c in clip_data: m, u = clip.load_sd(c) if len(m) > 0: @@ -554,7 +553,8 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c if clip_target is not None: clip_sd = model_config.process_clip_state_dict(sd) if len(clip_sd) > 0: - clip = CLIP(clip_target, embedding_directory=embedding_directory, tokenizer_data=clip_sd) + parameters = comfy.utils.calculate_parameters(clip_sd) + clip = CLIP(clip_target, embedding_directory=embedding_directory, tokenizer_data=clip_sd, parameters=parameters) m, u = clip.load_sd(clip_sd, full_model=True) if len(m) > 0: m_filter = list(filter(lambda a: ".logit_scale" not in a and ".transformer.text_projection.weight" not in a, m)) From 9829b013eaef91a29e47128d1addf98fb0f1ea48 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 12 Aug 2024 00:00:17 -0400 Subject: [PATCH 02/14] Fix mistake in last commit. --- comfy/sd.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/sd.py b/comfy/sd.py index 6d729929a..bbd9412d0 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -455,7 +455,7 @@ def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DI clip_target.clip = comfy.text_encoders.sd3_clip.SD3ClipModel clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer - clip = CLIP(clip_target, embedding_directory=embedding_directory, state_dicts=clip_data) + clip = CLIP(clip_target, embedding_directory=embedding_directory) for c in clip_data: m, u = clip.load_sd(c) if len(m) > 0: From 9acfe4df41b3b0ad8c600fc2d70a3af5c82cf4a4 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 12 Aug 2024 00:06:01 -0400 Subject: [PATCH 03/14] Support loading directly to vram with CLIPLoader node. --- comfy/sd.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/comfy/sd.py b/comfy/sd.py index bbd9412d0..c88f5a30d 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -82,6 +82,8 @@ class CLIP: self.tokenizer = tokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data) self.patcher = comfy.model_patcher.ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device) + if params['device'] == load_device: + model_management.load_model_gpu(self.patcher) self.layer_idx = None logging.debug("CLIP model load device: {}, offload device: {}, current: {}".format(load_device, offload_device, params['device'])) @@ -455,7 +457,11 @@ def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DI clip_target.clip = comfy.text_encoders.sd3_clip.SD3ClipModel clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer - clip = CLIP(clip_target, embedding_directory=embedding_directory) + parameters = 0 + for c in clip_data: + parameters += comfy.utils.calculate_parameters(c) + + clip = CLIP(clip_target, embedding_directory=embedding_directory, parameters=parameters) for c in clip_data: m, u = clip.load_sd(c) if len(m) > 0: From ad76574cb8b28ee498f3dceafc9d00b56f12f992 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 12 Aug 2024 00:23:29 -0400 Subject: [PATCH 04/14] Fix some potential issues with the previous commits. --- comfy/model_management.py | 3 +++ comfy/sd.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/comfy/model_management.py b/comfy/model_management.py index 6e2738391..686f124c9 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -688,6 +688,9 @@ def text_encoder_initial_device(load_device, offload_device, model_size=0): if load_device == offload_device or model_size <= 1024 * 1024 * 1024: return offload_device + if is_device_mps(load_device): + return offload_device + mem_l = get_free_memory(load_device) mem_o = get_free_memory(offload_device) if mem_l > (mem_o * 0.5) and model_size * 1.2 < mem_l: diff --git a/comfy/sd.py b/comfy/sd.py index c88f5a30d..c8a2f086c 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -79,6 +79,9 @@ class CLIP: for dt in self.cond_stage_model.dtypes: if not model_management.supports_cast(load_device, dt): load_device = offload_device + if params['device'] != offload_device: + self.cond_stage_model.to(offload_device) + logging.warning("Had to shift TE back.") self.tokenizer = tokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data) self.patcher = comfy.model_patcher.ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device) From 52a471c5c7e4af28423b3c690cbb6e1238ea9d60 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 12 Aug 2024 10:35:06 -0400 Subject: [PATCH 05/14] Change name of log. --- comfy/model_patcher.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index ae3d20514..aa32244e8 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -407,7 +407,7 @@ class ModelPatcher: logging.debug("lowvram: loaded module regularly {} {}".format(n, m)) if lowvram_counter > 0: - logging.info("loaded in lowvram mode {}".format(lowvram_model_memory / (1024 * 1024))) + logging.info("loaded partially {} {}".format(lowvram_model_memory / (1024 * 1024), patch_counter)) self.model.model_lowvram = True else: logging.info("loaded completely {} {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024))) From 517f4a94e4a5c45edc64594d70585ec8aeb787e0 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 12 Aug 2024 11:38:06 -0400 Subject: [PATCH 06/14] Fix some lora loading slowdowns. --- comfy/model_patcher.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index aa32244e8..c8009db1a 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -355,13 +355,14 @@ class ModelPatcher: return self.model - def lowvram_load(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False): + def lowvram_load(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False, full_load=False): mem_counter = 0 patch_counter = 0 lowvram_counter = 0 for n, m in self.model.named_modules(): lowvram_weight = False - if hasattr(m, "comfy_cast_weights"): + + if not full_load and hasattr(m, "comfy_cast_weights"): module_mem = comfy.model_management.module_size(m) if mem_counter + module_mem >= lowvram_model_memory: lowvram_weight = True @@ -401,8 +402,11 @@ class ModelPatcher: if weight.device == device_to: continue - self.patch_weight_to_device(weight_key) #TODO: speed this up without OOM - self.patch_weight_to_device(bias_key) + weight_to = None + if full_load:#TODO + weight_to = device_to + self.patch_weight_to_device(weight_key, device_to=weight_to) #TODO: speed this up without OOM + self.patch_weight_to_device(bias_key, device_to=weight_to) m.to(device_to) logging.debug("lowvram: loaded module regularly {} {}".format(n, m)) @@ -665,12 +669,13 @@ class ModelPatcher: return memory_freed def partially_load(self, device_to, extra_memory=0): + full_load = False if self.model.model_lowvram == False: return 0 if self.model.model_loaded_weight_memory + extra_memory > self.model_size(): - pass #TODO: Full load + full_load = True current_used = self.model.model_loaded_weight_memory - self.lowvram_load(device_to, lowvram_model_memory=current_used + extra_memory) + self.lowvram_load(device_to, lowvram_model_memory=current_used + extra_memory, full_load=full_load) return self.model.model_loaded_weight_memory - current_used def current_loaded_device(self): From 5d43e75e5b94c203075e315e4516fee47c4d6950 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 12 Aug 2024 12:27:54 -0400 Subject: [PATCH 07/14] Fix some issues with the model sometimes not getting patched. --- comfy/model_patcher.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index c8009db1a..39d00d82e 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -669,6 +669,8 @@ class ModelPatcher: return memory_freed def partially_load(self, device_to, extra_memory=0): + self.unpatch_model(unpatch_weights=False) + self.patch_model(patch_weights=False) full_load = False if self.model.model_lowvram == False: return 0 From b5c3906b38fdb493b22d113c4e191ef17801652f Mon Sep 17 00:00:00 2001 From: "Alex \"mcmonkey\" Goodwin" <4000772+mcmonkey4eva@users.noreply.github.com> Date: Mon, 12 Aug 2024 09:32:16 -0700 Subject: [PATCH 08/14] Automatically link the Comfy CI page on PRs (#4326) also use_prior_commit so it doesn't get a janked merge commit instead of the real one --- .github/workflows/pullrequest-ci-run.yml | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/.github/workflows/pullrequest-ci-run.yml b/.github/workflows/pullrequest-ci-run.yml index fd1ba16e2..691480bcf 100644 --- a/.github/workflows/pullrequest-ci-run.yml +++ b/.github/workflows/pullrequest-ci-run.yml @@ -35,3 +35,19 @@ jobs: torch_version: ${{ matrix.torch_version }} google_credentials: ${{ secrets.GCS_SERVICE_ACCOUNT_JSON }} comfyui_flags: ${{ matrix.flags }} + use_prior_commit: 'true' + comment: + if: ${{ github.event.label.name == 'Run-CI-Test' }} + runs-on: ubuntu-latest + permissions: + pull-requests: write + steps: + - uses: actions/github-script@v6 + with: + script: | + github.rest.issues.createComment({ + issue_number: context.issue.number, + owner: context.repo.owner, + repo: context.repo.repo, + body: '(Automated Bot Message) CI Tests are running, you can view the results at https://ci.comfy.org/?branch=${{ github.event.pull_request.number }}%2Fmerge' + }) From ce37c11164ebc452592f3b0e67fb63c8c16374c0 Mon Sep 17 00:00:00 2001 From: Vladimir Semyonov <20096510+vovsemenv@users.noreply.github.com> Date: Mon, 12 Aug 2024 19:32:34 +0300 Subject: [PATCH 09/14] add DS_Store to gitignore (#4324) --- .gitignore | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 5092c98f4..b212b0297 100644 --- a/.gitignore +++ b/.gitignore @@ -18,4 +18,5 @@ venv/ /tests-ui/data/object_info.json /user/ *.log -web_custom_versions/ \ No newline at end of file +web_custom_versions/ +.DS_Store From b8ffb2937f9daeaead6e9225f8f5d1dde6afc577 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 12 Aug 2024 15:03:33 -0400 Subject: [PATCH 10/14] Memory tweaks. --- comfy/model_management.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 686f124c9..fdf3308ba 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -438,11 +438,11 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu global vram_state inference_memory = minimum_inference_memory() - extra_mem = max(inference_memory, memory_required) + 100 * 1024 * 1024 + extra_mem = max(inference_memory, memory_required + 300 * 1024 * 1024) if minimum_memory_required is None: minimum_memory_required = extra_mem else: - minimum_memory_required = max(inference_memory, minimum_memory_required) + 100 * 1024 * 1024 + minimum_memory_required = max(inference_memory, minimum_memory_required + 300 * 1024 * 1024) models = set(models) From c032b11e074530a6892c4de8d9b457a3d268698e Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Mon, 12 Aug 2024 21:22:22 -0400 Subject: [PATCH 11/14] xlabs Flux controlnet implementation. (#4260) * xlabs Flux controlnet. * Fix not working on old python. * Remove comment. --- comfy/controlnet.py | 17 ++++- comfy/ldm/flux/controlnet_xlabs.py | 104 +++++++++++++++++++++++++++++ comfy/ldm/flux/model.py | 21 ++++-- 3 files changed, 135 insertions(+), 7 deletions(-) create mode 100644 comfy/ldm/flux/controlnet_xlabs.py diff --git a/comfy/controlnet.py b/comfy/controlnet.py index 354b00367..dcfe492ce 100644 --- a/comfy/controlnet.py +++ b/comfy/controlnet.py @@ -34,6 +34,8 @@ import comfy.t2i_adapter.adapter import comfy.ldm.cascade.controlnet import comfy.cldm.mmdit import comfy.ldm.hydit.controlnet +import comfy.ldm.flux.controlnet_xlabs + def broadcast_image_to(tensor, target_batch_size, batched_number): current_batch_size = tensor.shape[0] @@ -416,6 +418,7 @@ def load_controlnet_mmdit(sd): control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype) return control + def load_controlnet_hunyuandit(controlnet_data): model_config, operations, load_device, unet_dtype, manual_cast_dtype = controlnet_config(controlnet_data) @@ -427,6 +430,15 @@ def load_controlnet_hunyuandit(controlnet_data): 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) return control +def load_controlnet_flux_xlabs(sd): + model_config, operations, load_device, unet_dtype, manual_cast_dtype = controlnet_config(sd) + control_model = comfy.ldm.flux.controlnet_xlabs.ControlNetFlux(operations=operations, device=load_device, dtype=unet_dtype, **model_config.unet_config) + control_model = controlnet_load_state_dict(control_model, sd) + extra_conds = ['y', 'guidance'] + control = ControlNet(control_model, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds) + return control + + def load_controlnet(ckpt_path, model=None): controlnet_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True) if 'after_proj_list.18.bias' in controlnet_data.keys(): #Hunyuan DiT @@ -489,7 +501,10 @@ def load_controlnet(ckpt_path, model=None): logging.warning("leftover keys: {}".format(leftover_keys)) controlnet_data = new_sd elif "controlnet_blocks.0.weight" in controlnet_data: #SD3 diffusers format - return load_controlnet_mmdit(controlnet_data) + if "double_blocks.0.img_attn.norm.key_norm.scale" in controlnet_data: + return load_controlnet_flux_xlabs(controlnet_data) + else: + return load_controlnet_mmdit(controlnet_data) pth_key = 'control_model.zero_convs.0.0.weight' pth = False diff --git a/comfy/ldm/flux/controlnet_xlabs.py b/comfy/ldm/flux/controlnet_xlabs.py new file mode 100644 index 000000000..3f40021b2 --- /dev/null +++ b/comfy/ldm/flux/controlnet_xlabs.py @@ -0,0 +1,104 @@ +#Original code can be found on: https://github.com/XLabs-AI/x-flux/blob/main/src/flux/controlnet.py + +import torch +from torch import Tensor, nn +from einops import rearrange, repeat + +from .layers import (DoubleStreamBlock, EmbedND, LastLayer, + MLPEmbedder, SingleStreamBlock, + timestep_embedding) + +from .model import Flux +import comfy.ldm.common_dit + + +class ControlNetFlux(Flux): + def __init__(self, image_model=None, dtype=None, device=None, operations=None, **kwargs): + super().__init__(final_layer=False, dtype=dtype, device=device, operations=operations, **kwargs) + + # add ControlNet blocks + self.controlnet_blocks = nn.ModuleList([]) + for _ in range(self.params.depth): + controlnet_block = operations.Linear(self.hidden_size, self.hidden_size, dtype=dtype, device=device) + # controlnet_block = zero_module(controlnet_block) + self.controlnet_blocks.append(controlnet_block) + self.pos_embed_input = operations.Linear(self.in_channels, self.hidden_size, bias=True, dtype=dtype, device=device) + self.gradient_checkpointing = False + self.input_hint_block = nn.Sequential( + operations.Conv2d(3, 16, 3, padding=1, dtype=dtype, device=device), + nn.SiLU(), + operations.Conv2d(16, 16, 3, padding=1, dtype=dtype, device=device), + nn.SiLU(), + operations.Conv2d(16, 16, 3, padding=1, stride=2, dtype=dtype, device=device), + nn.SiLU(), + operations.Conv2d(16, 16, 3, padding=1, dtype=dtype, device=device), + nn.SiLU(), + operations.Conv2d(16, 16, 3, padding=1, stride=2, dtype=dtype, device=device), + nn.SiLU(), + operations.Conv2d(16, 16, 3, padding=1, dtype=dtype, device=device), + nn.SiLU(), + operations.Conv2d(16, 16, 3, padding=1, stride=2, dtype=dtype, device=device), + nn.SiLU(), + operations.Conv2d(16, 16, 3, padding=1, dtype=dtype, device=device) + ) + + def forward_orig( + self, + img: Tensor, + img_ids: Tensor, + controlnet_cond: Tensor, + txt: Tensor, + txt_ids: Tensor, + timesteps: Tensor, + y: Tensor, + guidance: Tensor = None, + ) -> Tensor: + if img.ndim != 3 or txt.ndim != 3: + raise ValueError("Input img and txt tensors must have 3 dimensions.") + + # running on sequences img + img = self.img_in(img) + controlnet_cond = self.input_hint_block(controlnet_cond) + controlnet_cond = rearrange(controlnet_cond, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2) + controlnet_cond = self.pos_embed_input(controlnet_cond) + img = img + controlnet_cond + vec = self.time_in(timestep_embedding(timesteps, 256)) + if self.params.guidance_embed: + vec = vec + self.guidance_in(timestep_embedding(guidance, 256)) + vec = vec + self.vector_in(y) + txt = self.txt_in(txt) + + ids = torch.cat((txt_ids, img_ids), dim=1) + pe = self.pe_embedder(ids) + + block_res_samples = () + + for block in self.double_blocks: + img, txt = block(img=img, txt=txt, vec=vec, pe=pe) + block_res_samples = block_res_samples + (img,) + + controlnet_block_res_samples = () + for block_res_sample, controlnet_block in zip(block_res_samples, self.controlnet_blocks): + block_res_sample = controlnet_block(block_res_sample) + controlnet_block_res_samples = controlnet_block_res_samples + (block_res_sample,) + + return {"output": (controlnet_block_res_samples * 10)[:19]} + + def forward(self, x, timesteps, context, y, guidance=None, hint=None, **kwargs): + hint = hint * 2.0 - 1.0 + + bs, c, h, w = x.shape + patch_size = 2 + x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch_size, patch_size)) + + img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size) + + h_len = ((h + (patch_size // 2)) // patch_size) + w_len = ((w + (patch_size // 2)) // patch_size) + img_ids = torch.zeros((h_len, w_len, 3), device=x.device, dtype=x.dtype) + img_ids[..., 1] = img_ids[..., 1] + torch.linspace(0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype)[:, None] + img_ids[..., 2] = img_ids[..., 2] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype)[None, :] + img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs) + + txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype) + return self.forward_orig(img, img_ids, hint, context, txt_ids, timesteps, y, guidance) diff --git a/comfy/ldm/flux/model.py b/comfy/ldm/flux/model.py index db6cf3d22..b5373540a 100644 --- a/comfy/ldm/flux/model.py +++ b/comfy/ldm/flux/model.py @@ -38,7 +38,7 @@ class Flux(nn.Module): Transformer model for flow matching on sequences. """ - def __init__(self, image_model=None, dtype=None, device=None, operations=None, **kwargs): + def __init__(self, image_model=None, final_layer=True, dtype=None, device=None, operations=None, **kwargs): super().__init__() self.dtype = dtype params = FluxParams(**kwargs) @@ -83,7 +83,8 @@ class Flux(nn.Module): ] ) - self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels, dtype=dtype, device=device, operations=operations) + if final_layer: + self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels, dtype=dtype, device=device, operations=operations) def forward_orig( self, @@ -94,6 +95,7 @@ class Flux(nn.Module): timesteps: Tensor, y: Tensor, guidance: Tensor = None, + control=None, ) -> Tensor: if img.ndim != 3 or txt.ndim != 3: raise ValueError("Input img and txt tensors must have 3 dimensions.") @@ -112,8 +114,15 @@ class Flux(nn.Module): ids = torch.cat((txt_ids, img_ids), dim=1) pe = self.pe_embedder(ids) - for block in self.double_blocks: - img, txt = block(img=img, txt=txt, vec=vec, pe=pe) + for i in range(len(self.double_blocks)): + img, txt = self.double_blocks[i](img=img, txt=txt, vec=vec, pe=pe) + + if control is not None: #Controlnet + control_o = control.get("output") + if i < len(control_o): + add = control_o[i] + if add is not None: + img += add img = torch.cat((txt, img), 1) for block in self.single_blocks: @@ -123,7 +132,7 @@ class Flux(nn.Module): img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels) return img - def forward(self, x, timestep, context, y, guidance, **kwargs): + def forward(self, x, timestep, context, y, guidance, control=None, **kwargs): bs, c, h, w = x.shape patch_size = 2 x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch_size, patch_size)) @@ -138,5 +147,5 @@ class Flux(nn.Module): img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs) txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype) - out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance) + out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance, control) return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=2, pw=2)[:,:,:h,:w] From 5942c17d5558e3a6a9065e24e86971db3bce0f7f Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 12 Aug 2024 21:56:18 -0400 Subject: [PATCH 12/14] Order of operations matters. --- comfy/ldm/flux/layers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy/ldm/flux/layers.py b/comfy/ldm/flux/layers.py index a7b0c55fa..da0cf61b1 100644 --- a/comfy/ldm/flux/layers.py +++ b/comfy/ldm/flux/layers.py @@ -170,8 +170,8 @@ class DoubleStreamBlock(nn.Module): txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :] # calculate the img bloks - img += img_mod1.gate * self.img_attn.proj(img_attn) - img += img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift) + img = img + img_mod1.gate * self.img_attn.proj(img_attn) + img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift) # calculate the txt bloks txt += txt_mod1.gate * self.txt_attn.proj(txt_attn) From a562c17e8ac52c6a3cb14902af43dee5a6f1adf4 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 12 Aug 2024 23:18:54 -0400 Subject: [PATCH 13/14] load_unet -> load_diffusion_model with a model_options argument. --- comfy/diffusers_load.py | 2 +- comfy/sd.py | 17 ++++++++++++++--- nodes.py | 8 ++++---- 3 files changed, 19 insertions(+), 8 deletions(-) diff --git a/comfy/diffusers_load.py b/comfy/diffusers_load.py index 98b888a19..56e63a756 100644 --- a/comfy/diffusers_load.py +++ b/comfy/diffusers_load.py @@ -22,7 +22,7 @@ def load_diffusers(model_path, output_vae=True, output_clip=True, embedding_dire if text_encoder2_path is not None: text_encoder_paths.append(text_encoder2_path) - unet = comfy.sd.load_unet(unet_path) + unet = comfy.sd.load_diffusion_model(unet_path) clip = None if output_clip: diff --git a/comfy/sd.py b/comfy/sd.py index c8a2f086c..13909d674 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -590,7 +590,8 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c return (model_patcher, clip, vae, clipvision) -def load_unet_state_dict(sd, dtype=None): #load unet in diffusers or regular format +def load_diffusion_model_state_dict(sd, model_options={}): #load unet in diffusers or regular format + dtype = model_options.get("dtype", None) #Allow loading unets from checkpoint files diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd) @@ -632,6 +633,7 @@ def load_unet_state_dict(sd, dtype=None): #load unet in diffusers or regular for manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes) model_config.set_inference_dtype(unet_dtype, manual_cast_dtype) + model_config.custom_operations = model_options.get("custom_operations", None) model = model_config.get_model(new_sd, "") model = model.to(offload_device) model.load_model_weights(new_sd, "") @@ -640,14 +642,23 @@ def load_unet_state_dict(sd, dtype=None): #load unet in diffusers or regular for logging.info("left over keys in unet: {}".format(left_over)) return comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=offload_device) -def load_unet(unet_path, dtype=None): + +def load_diffusion_model(unet_path, model_options={}): sd = comfy.utils.load_torch_file(unet_path) - model = load_unet_state_dict(sd, dtype=dtype) + model = load_diffusion_model_state_dict(sd, model_options=model_options) if model is None: logging.error("ERROR UNSUPPORTED UNET {}".format(unet_path)) raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path)) return model +def load_unet(unet_path, dtype=None): + print("WARNING: the load_unet function has been deprecated and will be removed please switch to: load_diffusion_model") + return load_diffusion_model(unet_path, model_options={"dtype": dtype}) + +def load_unet_state_dict(sd, dtype=None): + print("WARNING: the load_unet_state_dict function has been deprecated and will be removed please switch to: load_diffusion_model_state_dict") + return load_diffusion_model_state_dict(sd, model_options={"dtype": dtype}) + def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, metadata=None, extra_keys={}): clip_sd = None load_models = [model] diff --git a/nodes.py b/nodes.py index e296597c5..525b28d85 100644 --- a/nodes.py +++ b/nodes.py @@ -826,14 +826,14 @@ class UNETLoader: CATEGORY = "advanced/loaders" def load_unet(self, unet_name, weight_dtype): - dtype = None + model_options = {} if weight_dtype == "fp8_e4m3fn": - dtype = torch.float8_e4m3fn + model_options["dtype"] = torch.float8_e4m3fn elif weight_dtype == "fp8_e5m2": - dtype = torch.float8_e5m2 + model_options["dtype"] = torch.float8_e5m2 unet_path = folder_paths.get_full_path("unet", unet_name) - model = comfy.sd.load_unet(unet_path, dtype=dtype) + model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options) return (model,) class CLIPLoader: From 74e124f4d784b859465e751a7b361c20f192f0f9 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 12 Aug 2024 23:42:21 -0400 Subject: [PATCH 14/14] Fix some issues with TE being in lowvram mode. --- comfy/model_management.py | 4 ++-- comfy/model_patcher.py | 2 +- comfy/sd.py | 4 ++-- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index fdf3308ba..152a76f37 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -434,7 +434,7 @@ def free_memory(memory_required, device, keep_loaded=[]): soft_empty_cache() return unloaded_models -def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimum_memory_required=None): +def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimum_memory_required=None, force_full_load=False): global vram_state inference_memory = minimum_inference_memory() @@ -513,7 +513,7 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu else: vram_set_state = vram_state lowvram_model_memory = 0 - if lowvram_available and (vram_set_state == VRAMState.LOW_VRAM or vram_set_state == VRAMState.NORMAL_VRAM): + if lowvram_available and (vram_set_state == VRAMState.LOW_VRAM or vram_set_state == VRAMState.NORMAL_VRAM) and not force_full_load: model_size = loaded_model.model_memory_required(torch_dev) current_free_mem = get_free_memory(torch_dev) lowvram_model_memory = max(64 * (1024 * 1024), (current_free_mem - minimum_memory_required), min(current_free_mem * 0.4, current_free_mem - minimum_inference_memory())) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 39d00d82e..1edbf24ab 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -411,7 +411,7 @@ class ModelPatcher: logging.debug("lowvram: loaded module regularly {} {}".format(n, m)) if lowvram_counter > 0: - logging.info("loaded partially {} {}".format(lowvram_model_memory / (1024 * 1024), patch_counter)) + logging.info("loaded partially {} {} {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024), patch_counter)) self.model.model_lowvram = True else: logging.info("loaded completely {} {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024))) diff --git a/comfy/sd.py b/comfy/sd.py index 13909d674..edd0b51d8 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -86,7 +86,7 @@ class CLIP: self.tokenizer = tokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data) self.patcher = comfy.model_patcher.ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device) if params['device'] == load_device: - model_management.load_model_gpu(self.patcher) + model_management.load_models_gpu([self.patcher], force_full_load=True) self.layer_idx = None logging.debug("CLIP model load device: {}, offload device: {}, current: {}".format(load_device, offload_device, params['device'])) @@ -585,7 +585,7 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c model_patcher = comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=model_management.unet_offload_device()) if inital_load_device != torch.device("cpu"): logging.info("loaded straight to GPU") - model_management.load_model_gpu(model_patcher) + model_management.load_models_gpu([model_patcher], force_full_load=True) return (model_patcher, clip, vae, clipvision)