From 0849c80e2ab50b2821856c3de1ae6ebc1f879260 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 17 Sep 2024 01:57:59 -0400 Subject: [PATCH 1/4] get_key_patches now works without unloading the model. --- comfy/model_patcher.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index a89ad4fd7..6ca124e6d 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -283,17 +283,21 @@ class ModelPatcher: return list(p) def get_key_patches(self, filter_prefix=None): - comfy.model_management.unload_model_clones(self) model_sd = self.model_state_dict() p = {} for k in model_sd: if filter_prefix is not None: if not k.startswith(filter_prefix): continue - if k in self.patches: - p[k] = [model_sd[k]] + self.patches[k] + bk = self.backup.get(k, None) + if bk is not None: + weight = bk.weight else: - p[k] = (model_sd[k],) + weight = model_sd[k] + if k in self.patches: + p[k] = [weight] + self.patches[k] + else: + p[k] = (weight,) return p def model_state_dict(self, filter_prefix=None): From d514bb38ee11477298ba2fb361c7d1479d1b6a82 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 17 Sep 2024 03:49:54 -0400 Subject: [PATCH 2/4] Add some option to model_options for the text encoder. load_device, offload_device and the initial_device can now be set. --- comfy/sd.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/comfy/sd.py b/comfy/sd.py index 07310b9d4..8c5b058ce 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -70,14 +70,14 @@ class CLIP: clip = target.clip tokenizer = target.tokenizer - load_device = model_management.text_encoder_device() - offload_device = model_management.text_encoder_offload_device() + load_device = model_options.get("load_device", model_management.text_encoder_device()) + offload_device = model_options.get("offload_device", model_management.text_encoder_offload_device()) dtype = model_options.get("dtype", None) if dtype is None: 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)) + params['device'] = model_options.get("initial_device", model_management.text_encoder_initial_device(load_device, offload_device, parameters * model_management.dtype_size(dtype))) params['model_options'] = model_options self.cond_stage_model = clip(**(params)) From 0b7dfa986dafb49261f6835f444aab74bc8d3322 Mon Sep 17 00:00:00 2001 From: pharmapsychotic <96542870+pharmapsychotic@users.noreply.github.com> Date: Tue, 17 Sep 2024 02:51:10 -0500 Subject: [PATCH 3/4] Improve tiling calculations to reduce number of tiles that need to be processed. (#4944) --- comfy/utils.py | 29 +++++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/comfy/utils.py b/comfy/utils.py index 1bc35df7a..78f8314fc 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -713,7 +713,9 @@ def common_upscale(samples, width, height, upscale_method, crop): return torch.nn.functional.interpolate(s, size=(height, width), mode=upscale_method) def get_tiled_scale_steps(width, height, tile_x, tile_y, overlap): - return math.ceil((height / (tile_y - overlap))) * math.ceil((width / (tile_x - overlap))) + rows = 1 if height <= tile_y else math.ceil((height - overlap) / (tile_y - overlap)) + cols = 1 if width <= tile_x else math.ceil((width - overlap) / (tile_x - overlap)) + return rows * cols @torch.inference_mode() def tiled_scale_multidim(samples, function, tile=(64, 64), overlap = 8, upscale_amount = 4, out_channels = 3, output_device="cpu", pbar = None): @@ -722,10 +724,20 @@ def tiled_scale_multidim(samples, function, tile=(64, 64), overlap = 8, upscale_ for b in range(samples.shape[0]): s = samples[b:b+1] + + # handle entire input fitting in a single tile + if all(s.shape[d+2] <= tile[d] for d in range(dims)): + output[b:b+1] = function(s).to(output_device) + if pbar is not None: + pbar.update(1) + continue + out = torch.zeros([s.shape[0], out_channels] + list(map(lambda a: round(a * upscale_amount), s.shape[2:])), device=output_device) out_div = torch.zeros([s.shape[0], out_channels] + list(map(lambda a: round(a * upscale_amount), s.shape[2:])), device=output_device) - for it in itertools.product(*map(lambda a: range(0, a[0], a[1] - overlap), zip(s.shape[2:], tile))): + positions = [range(0, s.shape[d+2], tile[d] - overlap) if s.shape[d+2] > tile[d] else [0] for d in range(dims)] + + for it in itertools.product(*positions): s_in = s upscaled = [] @@ -734,15 +746,16 @@ def tiled_scale_multidim(samples, function, tile=(64, 64), overlap = 8, upscale_ l = min(tile[d], s.shape[d + 2] - pos) s_in = s_in.narrow(d + 2, pos, l) upscaled.append(round(pos * upscale_amount)) + ps = function(s_in).to(output_device) mask = torch.ones_like(ps) feather = round(overlap * upscale_amount) + for t in range(feather): for d in range(2, dims + 2): - m = mask.narrow(d, t, 1) - m *= ((1.0/feather) * (t + 1)) - m = mask.narrow(d, mask.shape[d] -1 -t, 1) - m *= ((1.0/feather) * (t + 1)) + a = (t + 1) / feather + mask.narrow(d, t, 1).mul_(a) + mask.narrow(d, mask.shape[d] - 1 - t, 1).mul_(a) o = out o_d = out_div @@ -750,8 +763,8 @@ def tiled_scale_multidim(samples, function, tile=(64, 64), overlap = 8, upscale_ o = o.narrow(d + 2, upscaled[d], mask.shape[d + 2]) o_d = o_d.narrow(d + 2, upscaled[d], mask.shape[d + 2]) - o += ps * mask - o_d += mask + o.add_(ps * mask) + o_d.add_(mask) if pbar is not None: pbar.update(1) From 254838f23c5a656cde5711088253c2ef29e31665 Mon Sep 17 00:00:00 2001 From: "Alex \"mcmonkey\" Goodwin" <4000772+mcmonkey4eva@users.noreply.github.com> Date: Tue, 17 Sep 2024 16:57:17 +0900 Subject: [PATCH 4/4] add simple error check to model loading (#4950) --- comfy_extras/nodes_hypernetwork.py | 2 +- comfy_extras/nodes_photomaker.py | 2 +- comfy_extras/nodes_sd3.py | 6 +++--- comfy_extras/nodes_upscale_model.py | 2 +- comfy_extras/nodes_video_model.py | 2 +- folder_paths.py | 8 ++++++++ nodes.py | 32 ++++++++++++++--------------- 7 files changed, 31 insertions(+), 23 deletions(-) diff --git a/comfy_extras/nodes_hypernetwork.py b/comfy_extras/nodes_hypernetwork.py index cafafa6ab..665632292 100644 --- a/comfy_extras/nodes_hypernetwork.py +++ b/comfy_extras/nodes_hypernetwork.py @@ -107,7 +107,7 @@ class HypernetworkLoader: CATEGORY = "loaders" def load_hypernetwork(self, model, hypernetwork_name, strength): - hypernetwork_path = folder_paths.get_full_path("hypernetworks", hypernetwork_name) + hypernetwork_path = folder_paths.get_full_path_or_raise("hypernetworks", hypernetwork_name) model_hypernetwork = model.clone() patch = load_hypernetwork_patch(hypernetwork_path, strength) if patch is not None: diff --git a/comfy_extras/nodes_photomaker.py b/comfy_extras/nodes_photomaker.py index 29d127d74..95d24dd22 100644 --- a/comfy_extras/nodes_photomaker.py +++ b/comfy_extras/nodes_photomaker.py @@ -126,7 +126,7 @@ class PhotoMakerLoader: CATEGORY = "_for_testing/photomaker" def load_photomaker_model(self, photomaker_model_name): - photomaker_model_path = folder_paths.get_full_path("photomaker", photomaker_model_name) + photomaker_model_path = folder_paths.get_full_path_or_raise("photomaker", photomaker_model_name) photomaker_model = PhotoMakerIDEncoder() data = comfy.utils.load_torch_file(photomaker_model_path, safe_load=True) if "id_encoder" in data: diff --git a/comfy_extras/nodes_sd3.py b/comfy_extras/nodes_sd3.py index 046096cba..c4bccaa6f 100644 --- a/comfy_extras/nodes_sd3.py +++ b/comfy_extras/nodes_sd3.py @@ -15,9 +15,9 @@ class TripleCLIPLoader: CATEGORY = "advanced/loaders" def load_clip(self, clip_name1, clip_name2, clip_name3): - clip_path1 = folder_paths.get_full_path("clip", clip_name1) - clip_path2 = folder_paths.get_full_path("clip", clip_name2) - clip_path3 = folder_paths.get_full_path("clip", clip_name3) + clip_path1 = folder_paths.get_full_path_or_raise("clip", clip_name1) + clip_path2 = folder_paths.get_full_path_or_raise("clip", clip_name2) + clip_path3 = folder_paths.get_full_path_or_raise("clip", clip_name3) clip = comfy.sd.load_clip(ckpt_paths=[clip_path1, clip_path2, clip_path3], embedding_directory=folder_paths.get_folder_paths("embeddings")) return (clip,) diff --git a/comfy_extras/nodes_upscale_model.py b/comfy_extras/nodes_upscale_model.py index bca79ef2e..6ba3e404f 100644 --- a/comfy_extras/nodes_upscale_model.py +++ b/comfy_extras/nodes_upscale_model.py @@ -25,7 +25,7 @@ class UpscaleModelLoader: CATEGORY = "loaders" def load_model(self, model_name): - model_path = folder_paths.get_full_path("upscale_models", model_name) + model_path = folder_paths.get_full_path_or_raise("upscale_models", model_name) sd = comfy.utils.load_torch_file(model_path, safe_load=True) if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd: sd = comfy.utils.state_dict_prefix_replace(sd, {"module.":""}) diff --git a/comfy_extras/nodes_video_model.py b/comfy_extras/nodes_video_model.py index 1a0189ed4..7f2146d1e 100644 --- a/comfy_extras/nodes_video_model.py +++ b/comfy_extras/nodes_video_model.py @@ -17,7 +17,7 @@ class ImageOnlyCheckpointLoader: CATEGORY = "loaders/video_models" def load_checkpoint(self, ckpt_name, output_vae=True, output_clip=True): - ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) + ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name) out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=False, output_clipvision=True, embedding_directory=folder_paths.get_folder_paths("embeddings")) return (out[0], out[3], out[2]) diff --git a/folder_paths.py b/folder_paths.py index 263704a4c..38fad6238 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -235,6 +235,14 @@ def get_full_path(folder_name: str, filename: str) -> str | None: return None + +def get_full_path_or_raise(folder_name: str, filename: str) -> str: + full_path = get_full_path(folder_name, filename) + if full_path is None: + raise FileNotFoundError(f"Model in folder '{folder_name}' with filename '{filename}' not found.") + return full_path + + def get_filename_list_(folder_name: str) -> tuple[list[str], dict[str, float], float]: folder_name = map_legacy(folder_name) global folder_names_and_paths diff --git a/nodes.py b/nodes.py index 84524b6b9..292ff9cfa 100644 --- a/nodes.py +++ b/nodes.py @@ -515,7 +515,7 @@ class CheckpointLoader: def load_checkpoint(self, config_name, ckpt_name): config_path = folder_paths.get_full_path("configs", config_name) - ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) + ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name) return comfy.sd.load_checkpoint(config_path, ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings")) class CheckpointLoaderSimple: @@ -536,7 +536,7 @@ class CheckpointLoaderSimple: DESCRIPTION = "Loads a diffusion model checkpoint, diffusion models are used to denoise latents." def load_checkpoint(self, ckpt_name): - ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) + ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name) out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings")) return out[:3] @@ -578,7 +578,7 @@ class unCLIPCheckpointLoader: CATEGORY = "loaders" def load_checkpoint(self, ckpt_name, output_vae=True, output_clip=True): - ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) + ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name) out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=True, embedding_directory=folder_paths.get_folder_paths("embeddings")) return out @@ -625,7 +625,7 @@ class LoraLoader: if strength_model == 0 and strength_clip == 0: return (model, clip) - lora_path = folder_paths.get_full_path("loras", lora_name) + lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) lora = None if self.loaded_lora is not None: if self.loaded_lora[0] == lora_path: @@ -704,11 +704,11 @@ class VAELoader: encoder = next(filter(lambda a: a.startswith("{}_encoder.".format(name)), approx_vaes)) decoder = next(filter(lambda a: a.startswith("{}_decoder.".format(name)), approx_vaes)) - enc = comfy.utils.load_torch_file(folder_paths.get_full_path("vae_approx", encoder)) + enc = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", encoder)) for k in enc: sd["taesd_encoder.{}".format(k)] = enc[k] - dec = comfy.utils.load_torch_file(folder_paths.get_full_path("vae_approx", decoder)) + dec = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", decoder)) for k in dec: sd["taesd_decoder.{}".format(k)] = dec[k] @@ -739,7 +739,7 @@ class VAELoader: if vae_name in ["taesd", "taesdxl", "taesd3", "taef1"]: sd = self.load_taesd(vae_name) else: - vae_path = folder_paths.get_full_path("vae", vae_name) + vae_path = folder_paths.get_full_path_or_raise("vae", vae_name) sd = comfy.utils.load_torch_file(vae_path) vae = comfy.sd.VAE(sd=sd) return (vae,) @@ -755,7 +755,7 @@ class ControlNetLoader: CATEGORY = "loaders" def load_controlnet(self, control_net_name): - controlnet_path = folder_paths.get_full_path("controlnet", control_net_name) + controlnet_path = folder_paths.get_full_path_or_raise("controlnet", control_net_name) controlnet = comfy.controlnet.load_controlnet(controlnet_path) return (controlnet,) @@ -771,7 +771,7 @@ class DiffControlNetLoader: CATEGORY = "loaders" def load_controlnet(self, model, control_net_name): - controlnet_path = folder_paths.get_full_path("controlnet", control_net_name) + controlnet_path = folder_paths.get_full_path_or_raise("controlnet", control_net_name) controlnet = comfy.controlnet.load_controlnet(controlnet_path, model) return (controlnet,) @@ -871,7 +871,7 @@ class UNETLoader: elif weight_dtype == "fp8_e5m2": model_options["dtype"] = torch.float8_e5m2 - unet_path = folder_paths.get_full_path("diffusion_models", unet_name) + unet_path = folder_paths.get_full_path_or_raise("diffusion_models", unet_name) model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options) return (model,) @@ -896,7 +896,7 @@ class CLIPLoader: else: clip_type = comfy.sd.CLIPType.STABLE_DIFFUSION - clip_path = folder_paths.get_full_path("clip", clip_name) + clip_path = folder_paths.get_full_path_or_raise("clip", clip_name) clip = comfy.sd.load_clip(ckpt_paths=[clip_path], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=clip_type) return (clip,) @@ -913,8 +913,8 @@ class DualCLIPLoader: CATEGORY = "advanced/loaders" def load_clip(self, clip_name1, clip_name2, type): - clip_path1 = folder_paths.get_full_path("clip", clip_name1) - clip_path2 = folder_paths.get_full_path("clip", clip_name2) + clip_path1 = folder_paths.get_full_path_or_raise("clip", clip_name1) + clip_path2 = folder_paths.get_full_path_or_raise("clip", clip_name2) if type == "sdxl": clip_type = comfy.sd.CLIPType.STABLE_DIFFUSION elif type == "sd3": @@ -936,7 +936,7 @@ class CLIPVisionLoader: CATEGORY = "loaders" def load_clip(self, clip_name): - clip_path = folder_paths.get_full_path("clip_vision", clip_name) + clip_path = folder_paths.get_full_path_or_raise("clip_vision", clip_name) clip_vision = comfy.clip_vision.load(clip_path) return (clip_vision,) @@ -966,7 +966,7 @@ class StyleModelLoader: CATEGORY = "loaders" def load_style_model(self, style_model_name): - style_model_path = folder_paths.get_full_path("style_models", style_model_name) + style_model_path = folder_paths.get_full_path_or_raise("style_models", style_model_name) style_model = comfy.sd.load_style_model(style_model_path) return (style_model,) @@ -1031,7 +1031,7 @@ class GLIGENLoader: CATEGORY = "loaders" def load_gligen(self, gligen_name): - gligen_path = folder_paths.get_full_path("gligen", gligen_name) + gligen_path = folder_paths.get_full_path_or_raise("gligen", gligen_name) gligen = comfy.sd.load_gligen(gligen_path) return (gligen,)