From 51714141435d050c255bba0cf588349f77bbc634 Mon Sep 17 00:00:00 2001 From: shiimizu Date: Sat, 17 Feb 2024 22:03:34 -0800 Subject: [PATCH 01/19] Support additional PNG info. --- web/scripts/pnginfo.js | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/web/scripts/pnginfo.js b/web/scripts/pnginfo.js index 83a4ebc86..169609209 100644 --- a/web/scripts/pnginfo.js +++ b/web/scripts/pnginfo.js @@ -24,7 +24,7 @@ export function getPngMetadata(file) { const length = dataView.getUint32(offset); // Get the chunk type const type = String.fromCharCode(...pngData.slice(offset + 4, offset + 8)); - if (type === "tEXt" || type == "comf") { + if (type === "tEXt" || type == "comf" || type === "iTXt") { // Get the keyword let keyword_end = offset + 8; while (pngData[keyword_end] !== 0) { @@ -33,7 +33,7 @@ export function getPngMetadata(file) { const keyword = String.fromCharCode(...pngData.slice(offset + 8, keyword_end)); // Get the text const contentArraySegment = pngData.slice(keyword_end + 1, offset + 8 + length); - const contentJson = Array.from(contentArraySegment).map(s=>String.fromCharCode(s)).join('') + const contentJson = new TextDecoder("utf-8").decode(contentArraySegment); txt_chunks[keyword] = contentJson; } From 3b2e579926d5cf8231de0e68e79096d1ee8091f2 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 19 Feb 2024 04:06:49 -0500 Subject: [PATCH 02/19] Support loading the Stable Cascade effnet and previewer as a VAE. The effnet can be used to encode images for img2img with Stage C. --- comfy/sd.py | 42 ++++++++++++++++++++++++++++++++++++++---- nodes.py | 20 ++++---------------- 2 files changed, 42 insertions(+), 20 deletions(-) diff --git a/comfy/sd.py b/comfy/sd.py index d8c0bfa7c..00633e107 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -4,6 +4,7 @@ from enum import Enum from comfy import model_management from .ldm.models.autoencoder import AutoencoderKL, AutoencodingEngine from .ldm.cascade.stage_a import StageA +from .ldm.cascade.stage_c_coder import StageC_coder import yaml @@ -158,6 +159,7 @@ class VAE: self.memory_used_encode = lambda shape, dtype: (1767 * shape[2] * shape[3]) * model_management.dtype_size(dtype) #These are for AutoencoderKL and need tweaking (should be lower) self.memory_used_decode = lambda shape, dtype: (2178 * shape[2] * shape[3] * 64) * model_management.dtype_size(dtype) self.downscale_ratio = 8 + self.upscale_ratio = 8 self.latent_channels = 4 self.process_input = lambda image: image * 2.0 - 1.0 self.process_output = lambda image: torch.clamp((image + 1.0) / 2.0, min=0.0, max=1.0) @@ -176,11 +178,31 @@ class VAE: elif "vquantizer.codebook.weight" in sd: #VQGan: stage a of stable cascade self.first_stage_model = StageA() self.downscale_ratio = 4 + self.upscale_ratio = 4 #TODO #self.memory_used_encode #self.memory_used_decode self.process_input = lambda image: image self.process_output = lambda image: image + elif "backbone.1.0.block.0.1.num_batches_tracked" in sd: #effnet: encoder for stage c latent of stable cascade + self.first_stage_model = StageC_coder() + self.downscale_ratio = 32 + self.latent_channels = 16 + new_sd = {} + for k in sd: + new_sd["encoder.{}".format(k)] = sd[k] + sd = new_sd + elif "blocks.11.num_batches_tracked" in sd: #previewer: decoder for stage c latent of stable cascade + self.first_stage_model = StageC_coder() + self.latent_channels = 16 + new_sd = {} + for k in sd: + new_sd["previewer.{}".format(k)] = sd[k] + sd = new_sd + elif "encoder.backbone.1.0.block.0.1.num_batches_tracked" in sd: #combined effnet and previewer for stable cascade + self.first_stage_model = StageC_coder() + self.downscale_ratio = 32 + self.latent_channels = 16 else: #default SD1.x/SD2.x VAE parameters ddconfig = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0} @@ -188,6 +210,7 @@ class VAE: if 'encoder.down.2.downsample.conv.weight' not in sd: #Stable diffusion x4 upscaler VAE ddconfig['ch_mult'] = [1, 2, 4] self.downscale_ratio = 4 + self.upscale_ratio = 4 self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=4) else: @@ -213,6 +236,15 @@ class VAE: self.patcher = comfy.model_patcher.ModelPatcher(self.first_stage_model, load_device=self.device, offload_device=offload_device) + def vae_encode_crop_pixels(self, pixels): + x = (pixels.shape[1] // self.downscale_ratio) * self.downscale_ratio + y = (pixels.shape[2] // self.downscale_ratio) * self.downscale_ratio + if pixels.shape[1] != x or pixels.shape[2] != y: + x_offset = (pixels.shape[1] % self.downscale_ratio) // 2 + y_offset = (pixels.shape[2] % self.downscale_ratio) // 2 + pixels = pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :] + return pixels + def decode_tiled_(self, samples, tile_x=64, tile_y=64, overlap = 16): steps = samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x, tile_y, overlap) steps += samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x // 2, tile_y * 2, overlap) @@ -221,9 +253,9 @@ class VAE: decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).float() output = self.process_output( - (comfy.utils.tiled_scale(samples, decode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = self.downscale_ratio, output_device=self.output_device, pbar = pbar) + - comfy.utils.tiled_scale(samples, decode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = self.downscale_ratio, output_device=self.output_device, pbar = pbar) + - comfy.utils.tiled_scale(samples, decode_fn, tile_x, tile_y, overlap, upscale_amount = self.downscale_ratio, output_device=self.output_device, pbar = pbar)) + (comfy.utils.tiled_scale(samples, decode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar) + + comfy.utils.tiled_scale(samples, decode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar) + + comfy.utils.tiled_scale(samples, decode_fn, tile_x, tile_y, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar)) / 3.0) return output @@ -248,7 +280,7 @@ class VAE: batch_number = int(free_memory / memory_used) batch_number = max(1, batch_number) - pixel_samples = torch.empty((samples_in.shape[0], 3, round(samples_in.shape[2] * self.downscale_ratio), round(samples_in.shape[3] * self.downscale_ratio)), device=self.output_device) + pixel_samples = torch.empty((samples_in.shape[0], 3, round(samples_in.shape[2] * self.upscale_ratio), round(samples_in.shape[3] * self.upscale_ratio)), device=self.output_device) for x in range(0, samples_in.shape[0], batch_number): samples = samples_in[x:x+batch_number].to(self.vae_dtype).to(self.device) pixel_samples[x:x+batch_number] = self.process_output(self.first_stage_model.decode(samples).to(self.output_device).float()) @@ -265,6 +297,7 @@ class VAE: return output.movedim(1,-1) def encode(self, pixel_samples): + pixel_samples = self.vae_encode_crop_pixels(pixel_samples) pixel_samples = pixel_samples.movedim(-1,1) try: memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype) @@ -284,6 +317,7 @@ class VAE: return samples def encode_tiled(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64): + pixel_samples = self.vae_encode_crop_pixels(pixel_samples) model_management.load_model_gpu(self.patcher) pixel_samples = pixel_samples.movedim(-1,1) samples = self.encode_tiled_(pixel_samples, tile_x=tile_x, tile_y=tile_y, overlap=overlap) diff --git a/nodes.py b/nodes.py index 47203f417..a577c2126 100644 --- a/nodes.py +++ b/nodes.py @@ -309,18 +309,7 @@ class VAEEncode: CATEGORY = "latent" - @staticmethod - def vae_encode_crop_pixels(pixels): - x = (pixels.shape[1] // 8) * 8 - y = (pixels.shape[2] // 8) * 8 - if pixels.shape[1] != x or pixels.shape[2] != y: - x_offset = (pixels.shape[1] % 8) // 2 - y_offset = (pixels.shape[2] % 8) // 2 - pixels = pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :] - return pixels - def encode(self, vae, pixels): - pixels = self.vae_encode_crop_pixels(pixels) t = vae.encode(pixels[:,:,:,:3]) return ({"samples":t}, ) @@ -336,7 +325,6 @@ class VAEEncodeTiled: CATEGORY = "_for_testing" def encode(self, vae, pixels, tile_size): - pixels = VAEEncode.vae_encode_crop_pixels(pixels) t = vae.encode_tiled(pixels[:,:,:,:3], tile_x=tile_size, tile_y=tile_size, ) return ({"samples":t}, ) @@ -350,14 +338,14 @@ class VAEEncodeForInpaint: CATEGORY = "latent/inpaint" def encode(self, vae, pixels, mask, grow_mask_by=6): - x = (pixels.shape[1] // 8) * 8 - y = (pixels.shape[2] // 8) * 8 + x = (pixels.shape[1] // vae.downscale_ratio) * vae.downscale_ratio + y = (pixels.shape[2] // vae.downscale_ratio) * vae.downscale_ratio mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear") pixels = pixels.clone() if pixels.shape[1] != x or pixels.shape[2] != y: - x_offset = (pixels.shape[1] % 8) // 2 - y_offset = (pixels.shape[2] % 8) // 2 + x_offset = (pixels.shape[1] % vae.downscale_ratio) // 2 + y_offset = (pixels.shape[2] % vae.downscale_ratio) // 2 pixels = pixels[:,x_offset:x + x_offset, y_offset:y + y_offset,:] mask = mask[:,:,x_offset:x + x_offset, y_offset:y + y_offset] From a7b5eaa7e3d5ecc1ae1d395e50d781124bb4e611 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 19 Feb 2024 04:25:46 -0500 Subject: [PATCH 03/19] Forgot to commit this. --- comfy/ldm/cascade/stage_c_coder.py | 96 ++++++++++++++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 comfy/ldm/cascade/stage_c_coder.py diff --git a/comfy/ldm/cascade/stage_c_coder.py b/comfy/ldm/cascade/stage_c_coder.py new file mode 100644 index 000000000..98c9a0b61 --- /dev/null +++ b/comfy/ldm/cascade/stage_c_coder.py @@ -0,0 +1,96 @@ +""" + This file is part of ComfyUI. + Copyright (C) 2024 Stability AI + + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License + along with this program. If not, see . +""" +import torch +import torchvision +from torch import nn + + +# EfficientNet +class EfficientNetEncoder(nn.Module): + def __init__(self, c_latent=16): + super().__init__() + self.backbone = torchvision.models.efficientnet_v2_s().features.eval() + self.mapper = nn.Sequential( + nn.Conv2d(1280, c_latent, kernel_size=1, bias=False), + nn.BatchNorm2d(c_latent, affine=False), # then normalize them to have mean 0 and std 1 + ) + self.mean = nn.Parameter(torch.tensor([0.485, 0.456, 0.406])) + self.std = nn.Parameter(torch.tensor([0.229, 0.224, 0.225])) + + def forward(self, x): + x = x * 0.5 + 0.5 + x = (x - self.mean.view([3,1,1])) / self.std.view([3,1,1]) + o = self.mapper(self.backbone(x)) + print(o.shape) + return o + + +# Fast Decoder for Stage C latents. E.g. 16 x 24 x 24 -> 3 x 192 x 192 +class Previewer(nn.Module): + def __init__(self, c_in=16, c_hidden=512, c_out=3): + super().__init__() + self.blocks = nn.Sequential( + nn.Conv2d(c_in, c_hidden, kernel_size=1), # 16 channels to 512 channels + nn.GELU(), + nn.BatchNorm2d(c_hidden), + + nn.Conv2d(c_hidden, c_hidden, kernel_size=3, padding=1), + nn.GELU(), + nn.BatchNorm2d(c_hidden), + + nn.ConvTranspose2d(c_hidden, c_hidden // 2, kernel_size=2, stride=2), # 16 -> 32 + nn.GELU(), + nn.BatchNorm2d(c_hidden // 2), + + nn.Conv2d(c_hidden // 2, c_hidden // 2, kernel_size=3, padding=1), + nn.GELU(), + nn.BatchNorm2d(c_hidden // 2), + + nn.ConvTranspose2d(c_hidden // 2, c_hidden // 4, kernel_size=2, stride=2), # 32 -> 64 + nn.GELU(), + nn.BatchNorm2d(c_hidden // 4), + + nn.Conv2d(c_hidden // 4, c_hidden // 4, kernel_size=3, padding=1), + nn.GELU(), + nn.BatchNorm2d(c_hidden // 4), + + nn.ConvTranspose2d(c_hidden // 4, c_hidden // 4, kernel_size=2, stride=2), # 64 -> 128 + nn.GELU(), + nn.BatchNorm2d(c_hidden // 4), + + nn.Conv2d(c_hidden // 4, c_hidden // 4, kernel_size=3, padding=1), + nn.GELU(), + nn.BatchNorm2d(c_hidden // 4), + + nn.Conv2d(c_hidden // 4, c_out, kernel_size=1), + ) + + def forward(self, x): + return (self.blocks(x) - 0.5) * 2.0 + +class StageC_coder(nn.Module): + def __init__(self): + super().__init__() + self.previewer = Previewer() + self.encoder = EfficientNetEncoder() + + def encode(self, x): + return self.encoder(x) + + def decode(self, x): + return self.previewer(x) From dbe0979b3f8ee4215e55012700f6d0afb0fec5b2 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 19 Feb 2024 08:59:34 -0500 Subject: [PATCH 04/19] Larger range for min/max compression for StableCascade_EmptyLatentImage. --- comfy_extras/nodes_stable_cascade.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_stable_cascade.py b/comfy_extras/nodes_stable_cascade.py index 5d31c1e59..efe3586d2 100644 --- a/comfy_extras/nodes_stable_cascade.py +++ b/comfy_extras/nodes_stable_cascade.py @@ -29,8 +29,8 @@ class StableCascade_EmptyLatentImage: return {"required": { "width": ("INT", {"default": 1024, "min": 256, "max": nodes.MAX_RESOLUTION, "step": 8}), "height": ("INT", {"default": 1024, "min": 256, "max": nodes.MAX_RESOLUTION, "step": 8}), - "compression": ("INT", {"default": 42, "min": 32, "max": 64, "step": 1}), - "batch_size": ("INT", {"default": 1, "min": 1, "max": 64}) + "compression": ("INT", {"default": 42, "min": 4, "max": 128, "step": 1}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}) }} RETURN_TYPES = ("LATENT", "LATENT") RETURN_NAMES = ("stage_c", "stage_b") From d91f45ef280a5acbdc22f3cc757f8fdbb254261b Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 19 Feb 2024 10:29:18 -0500 Subject: [PATCH 05/19] Some cleanups to how the text encoders are loaded. --- comfy/sd.py | 23 +++++++++++++---------- comfy/supported_models.py | 31 +++++++++++++++---------------- comfy/supported_models_base.py | 6 ++++-- 3 files changed, 32 insertions(+), 28 deletions(-) diff --git a/comfy/sd.py b/comfy/sd.py index 00633e107..7a77bb177 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -138,8 +138,11 @@ class CLIP: tokens = self.tokenize(text) return self.encode_from_tokens(tokens) - def load_sd(self, sd): - return self.cond_stage_model.load_sd(sd) + def load_sd(self, sd, full_model=False): + if full_model: + return self.cond_stage_model.load_state_dict(sd, strict=False) + else: + return self.cond_stage_model.load_sd(sd) def get_sd(self): return self.cond_stage_model.state_dict() @@ -494,9 +497,6 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o parameters = comfy.utils.calculate_parameters(sd, "model.diffusion_model.") load_device = model_management.get_torch_device() - class WeightsLoader(torch.nn.Module): - pass - model_config = model_detection.model_config_from_unet(sd, "model.diffusion_model.") unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=model_config.supported_inference_dtypes) manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes) @@ -521,14 +521,17 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o vae = VAE(sd=vae_sd) if output_clip: - w = WeightsLoader() clip_target = model_config.clip_target() if clip_target is not None: - sd = model_config.process_clip_state_dict(sd) - if any(k.startswith('cond_stage_model.') for k in sd): + clip_sd = model_config.process_clip_state_dict(sd) + if len(clip_sd) > 0: clip = CLIP(clip_target, embedding_directory=embedding_directory) - w.cond_stage_model = clip.cond_stage_model - load_model_weights(w, sd) + m, u = clip.load_sd(clip_sd, full_model=True) + if len(m) > 0: + print("clip missing:", m) + + if len(u) > 0: + print("clip unexpected:", u) else: print("no CLIP/text encoder weights in checkpoint, the text encoder model will not be loaded.") diff --git a/comfy/supported_models.py b/comfy/supported_models.py index 1a673646e..f29f7f3d9 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -40,8 +40,8 @@ class SD15(supported_models_base.BASE): state_dict['cond_stage_model.transformer.text_model.embeddings.position_ids'] = ids.round() replace_prefix = {} - replace_prefix["cond_stage_model."] = "cond_stage_model.clip_l." - state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix) + replace_prefix["cond_stage_model."] = "clip_l." + state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=True) return state_dict def process_clip_state_dict_for_saving(self, state_dict): @@ -72,10 +72,10 @@ class SD20(supported_models_base.BASE): def process_clip_state_dict(self, state_dict): replace_prefix = {} - replace_prefix["conditioner.embedders.0.model."] = "cond_stage_model.model." #SD2 in sgm format - state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix) - - state_dict = utils.transformers_convert(state_dict, "cond_stage_model.model.", "cond_stage_model.clip_h.transformer.text_model.", 24) + replace_prefix["conditioner.embedders.0.model."] = "clip_h." #SD2 in sgm format + replace_prefix["cond_stage_model.model."] = "clip_h." + state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=True) + state_dict = utils.transformers_convert(state_dict, "clip_h.", "clip_h.transformer.text_model.", 24) return state_dict def process_clip_state_dict_for_saving(self, state_dict): @@ -131,11 +131,10 @@ class SDXLRefiner(supported_models_base.BASE): def process_clip_state_dict(self, state_dict): keys_to_replace = {} replace_prefix = {} + replace_prefix["conditioner.embedders.0.model."] = "clip_g." + state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=True) - state_dict = utils.transformers_convert(state_dict, "conditioner.embedders.0.model.", "cond_stage_model.clip_g.transformer.text_model.", 32) - keys_to_replace["conditioner.embedders.0.model.text_projection"] = "cond_stage_model.clip_g.text_projection" - keys_to_replace["conditioner.embedders.0.model.logit_scale"] = "cond_stage_model.clip_g.logit_scale" - + state_dict = utils.transformers_convert(state_dict, "clip_g.", "clip_g.transformer.text_model.", 32) state_dict = utils.state_dict_key_replace(state_dict, keys_to_replace) return state_dict @@ -179,13 +178,13 @@ class SDXL(supported_models_base.BASE): keys_to_replace = {} replace_prefix = {} - replace_prefix["conditioner.embedders.0.transformer.text_model"] = "cond_stage_model.clip_l.transformer.text_model" - state_dict = utils.transformers_convert(state_dict, "conditioner.embedders.1.model.", "cond_stage_model.clip_g.transformer.text_model.", 32) - keys_to_replace["conditioner.embedders.1.model.text_projection"] = "cond_stage_model.clip_g.text_projection" - keys_to_replace["conditioner.embedders.1.model.text_projection.weight"] = "cond_stage_model.clip_g.text_projection" - keys_to_replace["conditioner.embedders.1.model.logit_scale"] = "cond_stage_model.clip_g.logit_scale" + replace_prefix["conditioner.embedders.0.transformer.text_model"] = "clip_l.transformer.text_model" + replace_prefix["conditioner.embedders.1.model."] = "clip_g." + state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=True) + + state_dict = utils.transformers_convert(state_dict, "clip_g.", "clip_g.transformer.text_model.", 32) + keys_to_replace["clip_g.text_projection.weight"] = "clip_g.text_projection" - state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix) state_dict = utils.state_dict_key_replace(state_dict, keys_to_replace) return state_dict diff --git a/comfy/supported_models_base.py b/comfy/supported_models_base.py index 3bd4f9c65..4d7e25936 100644 --- a/comfy/supported_models_base.py +++ b/comfy/supported_models_base.py @@ -22,6 +22,7 @@ class BASE: sampling_settings = {} latent_format = latent_formats.LatentFormat vae_key_prefix = ["first_stage_model."] + text_encoder_key_prefix = ["cond_stage_model."] supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32] manual_cast_dtype = None @@ -55,6 +56,7 @@ class BASE: return out def process_clip_state_dict(self, state_dict): + state_dict = utils.state_dict_prefix_replace(state_dict, {k: "" for k in self.text_encoder_key_prefix}, filter_keys=True) return state_dict def process_unet_state_dict(self, state_dict): @@ -64,7 +66,7 @@ class BASE: return state_dict def process_clip_state_dict_for_saving(self, state_dict): - replace_prefix = {"": "cond_stage_model."} + replace_prefix = {"": self.text_encoder_key_prefix[0]} return utils.state_dict_prefix_replace(state_dict, replace_prefix) def process_clip_vision_state_dict_for_saving(self, state_dict): @@ -78,7 +80,7 @@ class BASE: return utils.state_dict_prefix_replace(state_dict, replace_prefix) def process_vae_state_dict_for_saving(self, state_dict): - replace_prefix = {"": "first_stage_model."} + replace_prefix = {"": self.vae_key_prefix[0]} return utils.state_dict_prefix_replace(state_dict, replace_prefix) def set_inference_dtype(self, dtype, manual_cast_dtype): From 3711b31dff3530ba584e9a30f1bb32feaf2b5886 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 19 Feb 2024 11:20:48 -0500 Subject: [PATCH 06/19] Support Stable Cascade in checkpoint format. --- comfy/supported_models.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/comfy/supported_models.py b/comfy/supported_models.py index f29f7f3d9..5bb98d88a 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -319,6 +319,10 @@ class Stable_Cascade_C(supported_models_base.BASE): "shift": 2.0, } + vae_key_prefix = ["vae."] + text_encoder_key_prefix = ["text_encoder."] + clip_vision_prefix = "clip_l_vision." + def process_unet_state_dict(self, state_dict): key_list = list(state_dict.keys()) for y in ["weight", "bias"]: @@ -355,6 +359,8 @@ class Stable_Cascade_B(Stable_Cascade_C): "shift": 1.0, } + clip_vision_prefix = None + def get_model(self, state_dict, prefix="", device=None): out = model_base.StableCascade_B(self, device=device) return out From e93cdd0ad043a637a81c9d398498b1e8a8eca3b0 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 19 Feb 2024 11:47:26 -0500 Subject: [PATCH 07/19] Remove print. --- comfy/ldm/cascade/stage_c_coder.py | 1 - 1 file changed, 1 deletion(-) diff --git a/comfy/ldm/cascade/stage_c_coder.py b/comfy/ldm/cascade/stage_c_coder.py index 98c9a0b61..0cb7c49fc 100644 --- a/comfy/ldm/cascade/stage_c_coder.py +++ b/comfy/ldm/cascade/stage_c_coder.py @@ -36,7 +36,6 @@ class EfficientNetEncoder(nn.Module): x = x * 0.5 + 0.5 x = (x - self.mean.view([3,1,1])) / self.std.view([3,1,1]) o = self.mapper(self.backbone(x)) - print(o.shape) return o From 88f300401c0815eb5185683cb69ecf8b52cb6e7b Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 19 Feb 2024 12:00:48 -0500 Subject: [PATCH 08/19] Enable fp16 by default on mps. --- comfy/model_management.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 681208ea0..adcc0e8ac 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -720,9 +720,9 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma if FORCE_FP16: return True - if device is not None: #TODO + if device is not None: if is_device_mps(device): - return False + return True if FORCE_FP32: return False @@ -730,8 +730,11 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma if directml_enabled: return False - if cpu_mode() or mps_mode(): - return False #TODO ? + if mps_mode(): + return True + + if cpu_mode(): + return False if is_intel_xpu(): return True From a31152496990913211c6deb3267144bd3095c1ee Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 19 Feb 2024 13:36:20 -0500 Subject: [PATCH 09/19] Node to make stable cascade image to image easier. --- comfy_extras/nodes_stable_cascade.py | 35 ++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/comfy_extras/nodes_stable_cascade.py b/comfy_extras/nodes_stable_cascade.py index efe3586d2..b795d0083 100644 --- a/comfy_extras/nodes_stable_cascade.py +++ b/comfy_extras/nodes_stable_cascade.py @@ -18,6 +18,7 @@ import torch import nodes +import comfy.utils class StableCascade_EmptyLatentImage: @@ -47,6 +48,39 @@ class StableCascade_EmptyLatentImage: "samples": b_latent, }) +class StableCascade_StageC_VAEEncode: + def __init__(self, device="cpu"): + self.device = device + + @classmethod + def INPUT_TYPES(s): + return {"required": { + "image": ("IMAGE",), + "vae": ("VAE", ), + "compression": ("INT", {"default": 42, "min": 4, "max": 128, "step": 1}), + }} + RETURN_TYPES = ("LATENT", "LATENT") + RETURN_NAMES = ("stage_c", "stage_b") + FUNCTION = "generate" + + CATEGORY = "_for_testing/stable_cascade" + + def generate(self, image, vae, compression): + width = image.shape[-2] + height = image.shape[-3] + out_width = (width // compression) * vae.downscale_ratio + out_height = (height // compression) * vae.downscale_ratio + + s = comfy.utils.common_upscale(image.movedim(-1,1), out_width, out_height, "bicubic", "center").movedim(1,-1) + + c_latent = vae.encode(s[:,:,:,:3]) + b_latent = torch.zeros([c_latent.shape[0], 4, height // 4, width // 4]) + return ({ + "samples": c_latent, + }, { + "samples": b_latent, + }) + class StableCascade_StageB_Conditioning: @classmethod def INPUT_TYPES(s): @@ -71,4 +105,5 @@ class StableCascade_StageB_Conditioning: NODE_CLASS_MAPPINGS = { "StableCascade_EmptyLatentImage": StableCascade_EmptyLatentImage, "StableCascade_StageB_Conditioning": StableCascade_StageB_Conditioning, + "StableCascade_StageC_VAEEncode": StableCascade_StageC_VAEEncode, } From ec4d89cee946faacff60681a7f16443305350260 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 19 Feb 2024 13:41:55 -0500 Subject: [PATCH 10/19] Add to Readme that stable cascade is supported. --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index ff3ab6420..a94a212ad 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ This ui will let you design and execute advanced stable diffusion pipelines usin ## Features - Nodes/graph/flowchart interface to experiment and create complex Stable Diffusion workflows without needing to code anything. -- Fully supports SD1.x, SD2.x, [SDXL](https://comfyanonymous.github.io/ComfyUI_examples/sdxl/) and [Stable Video Diffusion](https://comfyanonymous.github.io/ComfyUI_examples/video/) +- Fully supports SD1.x, SD2.x, [SDXL](https://comfyanonymous.github.io/ComfyUI_examples/sdxl/), [Stable Video Diffusion](https://comfyanonymous.github.io/ComfyUI_examples/video/) and [Stable Cascade](https://comfyanonymous.github.io/ComfyUI_examples/stable_cascade/) - Asynchronous Queue system - Many optimizations: Only re-executes the parts of the workflow that changes between executions. - Command line option: ```--lowvram``` to make it work on GPUs with less than 3GB vram (enabled automatically on GPUs with low vram) From c6b7a157ed30eb2dd59891ba465b7b5be97a687a Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 20 Feb 2024 04:05:39 -0500 Subject: [PATCH 11/19] Align simple scheduling closer to official stable cascade scheduler. --- comfy/model_sampling.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy/model_sampling.py b/comfy/model_sampling.py index ae42d81f2..97e91a01d 100644 --- a/comfy/model_sampling.py +++ b/comfy/model_sampling.py @@ -150,10 +150,10 @@ class StableCascadeSampling(ModelSamplingDiscrete): self._init_alpha_cumprod = torch.cos(self.cosine_s / (1 + self.cosine_s) * torch.pi * 0.5) ** 2 #This part is just for compatibility with some schedulers in the codebase - self.num_timesteps = 1000 + self.num_timesteps = 10000 sigmas = torch.empty((self.num_timesteps), dtype=torch.float32) for x in range(self.num_timesteps): - t = x / self.num_timesteps + t = (x + 1) / self.num_timesteps sigmas[x] = self.sigma(t) self.set_sigmas(sigmas) From 0d0fbabd1d153611a1c21aea3515d16339abc84f Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 20 Feb 2024 04:23:25 -0500 Subject: [PATCH 12/19] Pass pooled CLIP to stage b. --- comfy/model_base.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/comfy/model_base.py b/comfy/model_base.py index fefce7637..421f271b2 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -474,15 +474,11 @@ class StableCascade_B(BaseModel): clip_text_pooled = kwargs["pooled_output"] if clip_text_pooled is not None: - out['clip_text_pooled'] = comfy.conds.CONDRegular(clip_text_pooled) + out['clip'] = comfy.conds.CONDRegular(clip_text_pooled) #size of prior doesn't really matter if zeros because it gets resized but I still want it to get batched prior = kwargs.get("stable_cascade_prior", torch.zeros((1, 16, (noise.shape[2] * 4) // 42, (noise.shape[3] * 4) // 42), dtype=noise.dtype, layout=noise.layout, device=noise.device)) out["effnet"] = comfy.conds.CONDRegular(prior) out["sca"] = comfy.conds.CONDRegular(torch.zeros((1,))) - - cross_attn = kwargs.get("cross_attn", None) - if cross_attn is not None: - out['clip'] = comfy.conds.CONDCrossAttn(cross_attn) return out From 18c151b3e3f6838fab4028e7a8ba526e30e610d3 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 20 Feb 2024 10:57:24 -0500 Subject: [PATCH 13/19] Add some latent2rgb matrices for previews. --- comfy/latent_formats.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py index 68fd73d0b..03fd59e3d 100644 --- a/comfy/latent_formats.py +++ b/comfy/latent_formats.py @@ -37,11 +37,41 @@ class SDXL(LatentFormat): class SD_X4(LatentFormat): def __init__(self): self.scale_factor = 0.08333 + self.latent_rgb_factors = [ + [-0.2340, -0.3863, -0.3257], + [ 0.0994, 0.0885, -0.0908], + [-0.2833, -0.2349, -0.3741], + [ 0.2523, -0.0055, -0.1651] + ] class SC_Prior(LatentFormat): def __init__(self): self.scale_factor = 1.0 + self.latent_rgb_factors = [ + [-0.0326, -0.0204, -0.0127], + [-0.1592, -0.0427, 0.0216], + [ 0.0873, 0.0638, -0.0020], + [-0.0602, 0.0442, 0.1304], + [ 0.0800, -0.0313, -0.1796], + [-0.0810, -0.0638, -0.1581], + [ 0.1791, 0.1180, 0.0967], + [ 0.0740, 0.1416, 0.0432], + [-0.1745, -0.1888, -0.1373], + [ 0.2412, 0.1577, 0.0928], + [ 0.1908, 0.0998, 0.0682], + [ 0.0209, 0.0365, -0.0092], + [ 0.0448, -0.0650, -0.1728], + [-0.1658, -0.1045, -0.1308], + [ 0.0542, 0.1545, 0.1325], + [-0.0352, -0.1672, -0.2541] + ] class SC_B(LatentFormat): def __init__(self): self.scale_factor = 1.0 + self.latent_rgb_factors = [ + [ 0.1121, 0.2006, 0.1023], + [-0.2093, -0.0222, -0.0195], + [-0.3087, -0.1535, 0.0366], + [ 0.0290, -0.1574, -0.4078] + ] From 820807c8ed2c2057a64539eccb4474563684254d Mon Sep 17 00:00:00 2001 From: feffy380 <114889020+feffy380@users.noreply.github.com> Date: Wed, 21 Feb 2024 10:33:03 +0100 Subject: [PATCH 14/19] Fix Perp-Neg math adjust perp-neg implementation to match the paper --- comfy_extras/nodes_perpneg.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy_extras/nodes_perpneg.py b/comfy_extras/nodes_perpneg.py index 45e4d418f..64bbc1dcd 100644 --- a/comfy_extras/nodes_perpneg.py +++ b/comfy_extras/nodes_perpneg.py @@ -35,7 +35,7 @@ class PerpNeg: pos = noise_pred_pos - noise_pred_nocond neg = noise_pred_neg - noise_pred_nocond - perp = ((torch.mul(pos, neg).sum())/(torch.norm(neg)**2)) * neg + perp = neg - ((torch.mul(neg, pos).sum())/(torch.norm(pos)**2)) * pos perp_neg = perp * neg_scale cfg_result = noise_pred_nocond + cond_scale*(pos - perp_neg) cfg_result = x - cfg_result From 7faa4507ecbd2ad67afcdea44b46ecdceec75232 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 21 Feb 2024 08:05:43 -0500 Subject: [PATCH 15/19] ModelSamplingDiscrete: x0 model support that predict a denoised image. --- comfy_extras/nodes_model_advanced.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/comfy_extras/nodes_model_advanced.py b/comfy_extras/nodes_model_advanced.py index ac7c1c17a..1b3f3945e 100644 --- a/comfy_extras/nodes_model_advanced.py +++ b/comfy_extras/nodes_model_advanced.py @@ -17,6 +17,10 @@ class LCM(comfy.model_sampling.EPS): return c_out * x0 + c_skip * model_input +class X0(comfy.model_sampling.EPS): + def calculate_denoised(self, sigma, model_output, model_input): + return model_output + class ModelSamplingDiscreteDistilled(comfy.model_sampling.ModelSamplingDiscrete): original_timesteps = 50 @@ -68,7 +72,7 @@ class ModelSamplingDiscrete: @classmethod def INPUT_TYPES(s): return {"required": { "model": ("MODEL",), - "sampling": (["eps", "v_prediction", "lcm"],), + "sampling": (["eps", "v_prediction", "lcm", "x0"],), "zsnr": ("BOOLEAN", {"default": False}), }} @@ -88,6 +92,8 @@ class ModelSamplingDiscrete: elif sampling == "lcm": sampling_type = LCM sampling_base = ModelSamplingDiscreteDistilled + elif sampling == "x0": + sampling_type = X0 class ModelSamplingAdvanced(sampling_base, sampling_type): pass From f81dbe26e2e363c28ad043db67b59c11bb33f446 Mon Sep 17 00:00:00 2001 From: Rick Love Date: Wed, 21 Feb 2024 19:21:24 -0600 Subject: [PATCH 16/19] FIX recursive_will_execute performance (simple ~300x performance increase} (#2852) * FIX recursive_will_execute performance * Minimize code changes * memo must be created outside lambda --- execution.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/execution.py b/execution.py index 00908eadd..3e9d53b0e 100644 --- a/execution.py +++ b/execution.py @@ -194,8 +194,12 @@ def recursive_execute(server, prompt, outputs, current_item, extra_data, execute return (True, None, None) -def recursive_will_execute(prompt, outputs, current_item): +def recursive_will_execute(prompt, outputs, current_item, memo={}): unique_id = current_item + + if unique_id in memo: + return memo[unique_id] + inputs = prompt[unique_id]['inputs'] will_execute = [] if unique_id in outputs: @@ -207,9 +211,10 @@ def recursive_will_execute(prompt, outputs, current_item): input_unique_id = input_data[0] output_index = input_data[1] if input_unique_id not in outputs: - will_execute += recursive_will_execute(prompt, outputs, input_unique_id) + will_execute += recursive_will_execute(prompt, outputs, input_unique_id, memo) - return will_execute + [unique_id] + memo[unique_id] = will_execute + [unique_id] + return memo[unique_id] def recursive_output_delete_if_changed(prompt, old_prompt, outputs, current_item): unique_id = current_item @@ -377,7 +382,8 @@ class PromptExecutor: while len(to_execute) > 0: #always execute the output that depends on the least amount of unexecuted nodes first - to_execute = sorted(list(map(lambda a: (len(recursive_will_execute(prompt, self.outputs, a[-1])), a[-1]), to_execute))) + memo = {} + to_execute = sorted(list(map(lambda a: (len(recursive_will_execute(prompt, self.outputs, a[-1], memo)), a[-1]), to_execute))) output_node_id = to_execute.pop(0)[-1] # This call shouldn't raise anything if there's an error deep in From 10847dfafeeb05152e72a9517beaa327ea510b8e Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 23 Feb 2024 02:39:35 -0500 Subject: [PATCH 17/19] Cleanup uni_pc inpainting. This causes some small changes to the uni pc inpainting behavior but it seems to improve results slightly. --- comfy/extra_samplers/uni_pc.py | 37 +++++++++------------------------- comfy/samplers.py | 12 ++--------- 2 files changed, 11 insertions(+), 38 deletions(-) diff --git a/comfy/extra_samplers/uni_pc.py b/comfy/extra_samplers/uni_pc.py index 08bf0fc9e..a30d1d03f 100644 --- a/comfy/extra_samplers/uni_pc.py +++ b/comfy/extra_samplers/uni_pc.py @@ -358,9 +358,6 @@ class UniPC: thresholding=False, max_val=1., variant='bh1', - noise_mask=None, - masked_image=None, - noise=None, ): """Construct a UniPC. @@ -372,9 +369,6 @@ class UniPC: self.predict_x0 = predict_x0 self.thresholding = thresholding self.max_val = max_val - self.noise_mask = noise_mask - self.masked_image = masked_image - self.noise = noise def dynamic_thresholding_fn(self, x0, t=None): """ @@ -391,10 +385,7 @@ class UniPC: """ Return the noise prediction model. """ - if self.noise_mask is not None: - return self.model(x, t) * self.noise_mask - else: - return self.model(x, t) + return self.model(x, t) def data_prediction_fn(self, x, t): """ @@ -409,8 +400,6 @@ class UniPC: s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1) s = expand_dims(torch.maximum(s, self.max_val * torch.ones_like(s).to(s.device)), dims) x0 = torch.clamp(x0, -s, s) / s - if self.noise_mask is not None: - x0 = x0 * self.noise_mask + (1. - self.noise_mask) * self.masked_image return x0 def model_fn(self, x, t): @@ -723,8 +712,6 @@ class UniPC: assert timesteps.shape[0] - 1 == steps # with torch.no_grad(): for step_index in trange(steps, disable=disable_pbar): - if self.noise_mask is not None: - x = x * self.noise_mask + (1. - self.noise_mask) * (self.masked_image * self.noise_schedule.marginal_alpha(timesteps[step_index]) + self.noise * self.noise_schedule.marginal_std(timesteps[step_index])) if step_index == 0: vec_t = timesteps[0].expand((x.shape[0])) model_prev_list = [self.model_fn(x, vec_t)] @@ -766,7 +753,7 @@ class UniPC: model_x = self.model_fn(x, vec_t) model_prev_list[-1] = model_x if callback is not None: - callback(step_index, model_prev_list[-1], x, steps) + callback({'x': x, 'i': step_index, 'denoised': model_prev_list[-1]}) else: raise NotImplementedError() # if denoise_to_zero: @@ -858,7 +845,7 @@ def predict_eps_sigma(model, input, sigma_in, **kwargs): return (input - model(input, sigma_in, **kwargs)) / sigma -def sample_unipc(model, noise, image, sigmas, max_denoise, extra_args=None, callback=None, disable=False, noise_mask=None, variant='bh1'): +def sample_unipc(model, noise, sigmas, extra_args=None, callback=None, disable=False, variant='bh1'): timesteps = sigmas.clone() if sigmas[-1] == 0: timesteps = sigmas[:] @@ -867,16 +854,7 @@ def sample_unipc(model, noise, image, sigmas, max_denoise, extra_args=None, call timesteps = sigmas.clone() ns = SigmaConvert() - if image is not None: - img = image * ns.marginal_alpha(timesteps[0]) - if max_denoise: - noise_mult = 1.0 - else: - noise_mult = ns.marginal_std(timesteps[0]) - img += noise * noise_mult - else: - img = noise - + noise = noise / torch.sqrt(1.0 + timesteps[0] ** 2.0) model_type = "noise" model_fn = model_wrapper( @@ -888,7 +866,10 @@ def sample_unipc(model, noise, image, sigmas, max_denoise, extra_args=None, call ) order = min(3, len(timesteps) - 2) - uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False, noise_mask=noise_mask, masked_image=image, noise=noise, variant=variant) - x = uni_pc.sample(img, timesteps=timesteps, skip_type="time_uniform", method="multistep", order=order, lower_order_final=True, callback=callback, disable_pbar=disable) + uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False, variant=variant) + x = uni_pc.sample(noise, timesteps=timesteps, skip_type="time_uniform", method="multistep", order=order, lower_order_final=True, callback=callback, disable_pbar=disable) x /= ns.marginal_alpha(timesteps[-1]) return x + +def sample_unipc_bh2(model, noise, sigmas, extra_args=None, callback=None, disable=False): + return sample_unipc(model, noise, sigmas, extra_args, callback, disable, variant='bh2') \ No newline at end of file diff --git a/comfy/samplers.py b/comfy/samplers.py index c795f208d..9c45eb605 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -513,14 +513,6 @@ class Sampler: sigma = float(sigmas[0]) return math.isclose(max_sigma, sigma, rel_tol=1e-05) or sigma > max_sigma -class UNIPC(Sampler): - def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False): - return uni_pc.sample_unipc(model_wrap, noise, latent_image, sigmas, max_denoise=self.max_denoise(model_wrap, sigmas), extra_args=extra_args, noise_mask=denoise_mask, callback=callback, disable=disable_pbar) - -class UNIPCBH2(Sampler): - def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False): - return uni_pc.sample_unipc(model_wrap, noise, latent_image, sigmas, max_denoise=self.max_denoise(model_wrap, sigmas), extra_args=extra_args, noise_mask=denoise_mask, callback=callback, variant='bh2', disable=disable_pbar) - KSAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_ancestral", "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu", "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm"] @@ -640,9 +632,9 @@ def calculate_sigmas_scheduler(model, scheduler_name, steps): def sampler_object(name): if name == "uni_pc": - sampler = UNIPC() + sampler = KSAMPLER(uni_pc.sample_unipc) elif name == "uni_pc_bh2": - sampler = UNIPCBH2() + sampler = KSAMPLER(uni_pc.sample_unipc_bh2) elif name == "ddim": sampler = ksampler("euler", inpaint_options={"random": True}) else: From e1cb93c3835a4c1e975914b5ce898abf509b21fe Mon Sep 17 00:00:00 2001 From: logtd Date: Fri, 23 Feb 2024 01:19:43 -0700 Subject: [PATCH 18/19] Fix model and cond transformer options merge --- comfy/samplers.py | 1 + 1 file changed, 1 insertion(+) diff --git a/comfy/samplers.py b/comfy/samplers.py index 9c45eb605..491c95d39 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -208,6 +208,7 @@ def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options): cur_patches[p] = cur_patches[p] + patches[p] else: cur_patches[p] = patches[p] + transformer_options["patches"] = cur_patches else: transformer_options["patches"] = patches From 1e5f0f66be93182f30d5d453dc8df50f3f7c1826 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 23 Feb 2024 12:21:20 -0500 Subject: [PATCH 19/19] Support lora keys with lora_prior_unet_ and lora_prior_te_ --- comfy/lora.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/comfy/lora.py b/comfy/lora.py index 5e4009b47..82d57fe13 100644 --- a/comfy/lora.py +++ b/comfy/lora.py @@ -197,6 +197,8 @@ def model_lora_keys_clip(model, key_map={}): key_map[lora_key] = k lora_key = "text_encoder.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora key_map[lora_key] = k + lora_key = "lora_prior_te_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #cascade lora: TODO put lora key prefix in the model config + key_map[lora_key] = k return key_map @@ -207,6 +209,7 @@ def model_lora_keys_unet(model, key_map={}): if k.startswith("diffusion_model.") and k.endswith(".weight"): key_lora = k[len("diffusion_model."):-len(".weight")].replace(".", "_") key_map["lora_unet_{}".format(key_lora)] = k + key_map["lora_prior_unet_{}".format(key_lora)] = k #cascade lora: TODO put lora key prefix in the model config diffusers_keys = comfy.utils.unet_to_diffusers(model.model_config.unet_config) for k in diffusers_keys: