From d7430c529a586ba4005bc46ba10ce02f71dba0d8 Mon Sep 17 00:00:00 2001 From: filtered <176114999+webfiltered@users.noreply.github.com> Date: Tue, 17 Jun 2025 15:58:28 -0700 Subject: [PATCH 01/13] Update frontend to 1.22.2 (#8567) --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 910634d87..15fde2849 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -comfyui-frontend-package==1.21.7 +comfyui-frontend-package==1.22.2 comfyui-workflow-templates==0.1.29 comfyui-embedded-docs==0.2.2 torch From e9e9a031a88f9cc4845b3322d4bae771d2854472 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Wed, 18 Jun 2025 03:55:21 -0700 Subject: [PATCH 02/13] Show a better error when the workflow OOMs. (#8574) --- execution.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/execution.py b/execution.py index d0012afda..f6006fa12 100644 --- a/execution.py +++ b/execution.py @@ -429,17 +429,20 @@ def execute(server, dynprompt, caches, current_item, extra_data, executed, promp logging.error(f"!!! Exception during processing !!! {ex}") logging.error(traceback.format_exc()) + tips = "" + + if isinstance(ex, comfy.model_management.OOM_EXCEPTION): + tips = "This error means you ran out of memory on your GPU.\n\nTIPS: If the workflow worked before you might have accidentally set the batch_size to a large number." + logging.error("Got an OOM, unloading all loaded models.") + comfy.model_management.unload_all_models() error_details = { "node_id": real_node_id, - "exception_message": str(ex), + "exception_message": "{}\n{}".format(ex, tips), "exception_type": exception_type, "traceback": traceback.format_tb(tb), "current_inputs": input_data_formatted } - if isinstance(ex, comfy.model_management.OOM_EXCEPTION): - logging.error("Got an OOM, unloading all loaded models.") - comfy.model_management.unload_all_models() return (ExecutionResult.FAILURE, error_details, ex) From 5b12b55e32aa2aa8fd47d545265a974d3b01ac7c Mon Sep 17 00:00:00 2001 From: coderfromthenorth93 Date: Wed, 18 Jun 2025 15:12:29 -0400 Subject: [PATCH 03/13] Add new fields to the config types (#8507) --- comfy_config/config_parser.py | 55 +++++++++++++++++++++++++++++++++++ comfy_config/types.py | 6 +++- 2 files changed, 60 insertions(+), 1 deletion(-) diff --git a/comfy_config/config_parser.py b/comfy_config/config_parser.py index a9cbd94dd..8da7bd901 100644 --- a/comfy_config/config_parser.py +++ b/comfy_config/config_parser.py @@ -11,6 +11,43 @@ from comfy_config.types import ( PyProjectSettings ) +def validate_and_extract_os_classifiers(classifiers: list) -> list: + os_classifiers = [c for c in classifiers if c.startswith("Operating System :: ")] + if not os_classifiers: + return [] + + os_values = [c[len("Operating System :: ") :] for c in os_classifiers] + valid_os_prefixes = {"Microsoft", "POSIX", "MacOS", "OS Independent"} + + for os_value in os_values: + if not any(os_value.startswith(prefix) for prefix in valid_os_prefixes): + return [] + + return os_values + + +def validate_and_extract_accelerator_classifiers(classifiers: list) -> list: + accelerator_classifiers = [c for c in classifiers if c.startswith("Environment ::")] + if not accelerator_classifiers: + return [] + + accelerator_values = [c[len("Environment :: ") :] for c in accelerator_classifiers] + + valid_accelerators = { + "GPU :: NVIDIA CUDA", + "GPU :: AMD ROCm", + "GPU :: Intel Arc", + "NPU :: Huawei Ascend", + "GPU :: Apple Metal", + } + + for accelerator_value in accelerator_values: + if accelerator_value not in valid_accelerators: + return [] + + return accelerator_values + + """ Extract configuration from a custom node directory's pyproject.toml file or a Python file. @@ -78,6 +115,24 @@ def extract_node_configuration(path) -> Optional[PyProjectConfig]: tool_data = raw_settings.tool comfy_data = tool_data.get("comfy", {}) if tool_data else {} + dependencies = project_data.get("dependencies", []) + supported_comfyui_frontend_version = "" + for dep in dependencies: + if isinstance(dep, str) and dep.startswith("comfyui-frontend-package"): + supported_comfyui_frontend_version = dep.removeprefix("comfyui-frontend-package") + break + + supported_comfyui_version = comfy_data.get("requires-comfyui", "") + + classifiers = project_data.get('classifiers', []) + supported_os = validate_and_extract_os_classifiers(classifiers) + supported_accelerators = validate_and_extract_accelerator_classifiers(classifiers) + + project_data['supported_os'] = supported_os + project_data['supported_accelerators'] = supported_accelerators + project_data['supported_comfyui_frontend_version'] = supported_comfyui_frontend_version + project_data['supported_comfyui_version'] = supported_comfyui_version + return PyProjectConfig(project=project_data, tool_comfy=comfy_data) diff --git a/comfy_config/types.py b/comfy_config/types.py index 5222cc59b..59448466b 100644 --- a/comfy_config/types.py +++ b/comfy_config/types.py @@ -51,7 +51,7 @@ class ComfyConfig(BaseModel): models: List[Model] = Field(default_factory=list, alias="Models") includes: List[str] = Field(default_factory=list) web: Optional[str] = None - + banner_url: str = "" class License(BaseModel): file: str = "" @@ -66,6 +66,10 @@ class ProjectConfig(BaseModel): dependencies: List[str] = Field(default_factory=list) license: License = Field(default_factory=License) urls: URLs = Field(default_factory=URLs) + supported_os: List[str] = Field(default_factory=list) + supported_accelerators: List[str] = Field(default_factory=list) + supported_comfyui_version: str = "" + supported_comfyui_frontend_version: str = "" @field_validator('license', mode='before') @classmethod From 91d40086db7956aabedef5cfcae0f6821529a3d1 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Thu, 19 Jun 2025 08:04:52 -0700 Subject: [PATCH 04/13] Fix pytorch warning. (#8593) --- comfy/ldm/modules/sub_quadratic_attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/ldm/modules/sub_quadratic_attention.py b/comfy/ldm/modules/sub_quadratic_attention.py index 21c72373f..fab145f1c 100644 --- a/comfy/ldm/modules/sub_quadratic_attention.py +++ b/comfy/ldm/modules/sub_quadratic_attention.py @@ -31,7 +31,7 @@ def dynamic_slice( starts: List[int], sizes: List[int], ) -> Tensor: - slicing = [slice(start, start + size) for start, size in zip(starts, sizes)] + slicing = tuple(slice(start, start + size) for start, size in zip(starts, sizes)) return x[slicing] class AttnChunk(NamedTuple): From 7e9267fa77c93355dd3bdf05b6cb8f02d41af5ae Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Thu, 19 Jun 2025 15:50:05 -0700 Subject: [PATCH 05/13] Make flux controlnet work with sd3 text enc. (#8599) --- comfy/ldm/flux/controlnet.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/comfy/ldm/flux/controlnet.py b/comfy/ldm/flux/controlnet.py index dbd2a47c0..7dcf82bbf 100644 --- a/comfy/ldm/flux/controlnet.py +++ b/comfy/ldm/flux/controlnet.py @@ -123,6 +123,8 @@ class ControlNetFlux(Flux): if y is None: y = torch.zeros((img.shape[0], self.params.vec_in_dim), device=img.device, dtype=img.dtype) + else: + y = y[:, :self.params.vec_in_dim] # running on sequences img img = self.img_in(img) From f7fb1937127a8ed011b99424598c9ab1e8565112 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Fri, 20 Jun 2025 02:37:32 -0700 Subject: [PATCH 06/13] Small flux optimization. (#8611) --- comfy/ldm/flux/layers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/ldm/flux/layers.py b/comfy/ldm/flux/layers.py index 76af967e6..113eb2096 100644 --- a/comfy/ldm/flux/layers.py +++ b/comfy/ldm/flux/layers.py @@ -118,7 +118,7 @@ class Modulation(nn.Module): def apply_mod(tensor, m_mult, m_add=None, modulation_dims=None): if modulation_dims is None: if m_add is not None: - return tensor * m_mult + m_add + return torch.addcmul(m_add, tensor, m_mult) else: return tensor * m_mult else: From 31ca603ccbc45ab4db1279aa485c715b96b2aae8 Mon Sep 17 00:00:00 2001 From: Lucas - BLOCK33 <95554128+tonynoce@users.noreply.github.com> Date: Sat, 21 Jun 2025 00:04:55 -0300 Subject: [PATCH 07/13] Improve the log time function for 10 minute + renders (#6207) * modified: main.py * Update main.py --- main.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/main.py b/main.py index c8c4194d4..79a652578 100644 --- a/main.py +++ b/main.py @@ -185,7 +185,13 @@ def prompt_worker(q, server_instance): current_time = time.perf_counter() execution_time = current_time - execution_start_time - logging.info("Prompt executed in {:.2f} seconds".format(execution_time)) + + # Log Time in a more readable way after 10 minutes + if execution_time > 600: + execution_time = time.strftime("%H:%M:%S", time.gmtime(execution_time)) + logging.info(f"Prompt executed in {execution_time}") + else: + logging.info("Prompt executed in {:.2f} seconds".format(execution_time)) flags = q.get_flags() free_memory = flags.get("free_memory", False) From 1883e70b4374a3317e0463a0bef292fc21182bad Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Sat, 21 Jun 2025 00:30:39 -0700 Subject: [PATCH 08/13] Fix exception when using a noise mask with cosmos predict2. (#8621) * Fix exception when using a noise mask with cosmos predict2. * Fix ruff. --- comfy/model_base.py | 2 ++ main.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/comfy/model_base.py b/comfy/model_base.py index cb7689e84..75ec42699 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -1024,6 +1024,8 @@ class CosmosPredict2(BaseModel): def process_timestep(self, timestep, x, denoise_mask=None, **kwargs): if denoise_mask is None: return timestep + if denoise_mask.ndim <= 4: + return timestep condition_video_mask_B_1_T_1_1 = denoise_mask.mean(dim=[1, 3, 4], keepdim=True) c_noise_B_1_T_1_1 = 0.0 * (1.0 - condition_video_mask_B_1_T_1_1) + timestep.reshape(timestep.shape[0], 1, 1, 1, 1) * condition_video_mask_B_1_T_1_1 out = c_noise_B_1_T_1_1.squeeze(dim=[1, 3, 4]) diff --git a/main.py b/main.py index 79a652578..0d7c97dcb 100644 --- a/main.py +++ b/main.py @@ -185,7 +185,7 @@ def prompt_worker(q, server_instance): current_time = time.perf_counter() execution_time = current_time - execution_start_time - + # Log Time in a more readable way after 10 minutes if execution_time > 600: execution_time = time.strftime("%H:%M:%S", time.gmtime(execution_time)) From 78f79266a9be9d9316e7feb6d6a0ed6ac616547d Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Sat, 21 Jun 2025 21:19:41 -0700 Subject: [PATCH 09/13] Allow padding in ImageStitch node to be white. (#8631) --- comfy_extras/nodes_images.py | 30 +++++++++++++++++------------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py index b1e0d4666..8d5fcdb85 100644 --- a/comfy_extras/nodes_images.py +++ b/comfy_extras/nodes_images.py @@ -304,10 +304,23 @@ Optional spacing can be added between images. image2.movedim(-1, 1), target_w, target_h, "lanczos", "disabled" ).movedim(1, -1) + color_map = { + "white": 1.0, + "black": 0.0, + "red": (1.0, 0.0, 0.0), + "green": (0.0, 1.0, 0.0), + "blue": (0.0, 0.0, 1.0), + } + + color_val = color_map[spacing_color] + # When not matching sizes, pad to align non-concat dimensions if not match_image_size: h1, w1 = image1.shape[1:3] h2, w2 = image2.shape[1:3] + pad_value = 0.0 + if not isinstance(color_val, tuple): + pad_value = color_val if direction in ["left", "right"]: # For horizontal concat, pad heights to match @@ -316,11 +329,11 @@ Optional spacing can be added between images. if h1 < target_h: pad_h = target_h - h1 pad_top, pad_bottom = pad_h // 2, pad_h - pad_h // 2 - image1 = torch.nn.functional.pad(image1, (0, 0, 0, 0, pad_top, pad_bottom), mode='constant', value=0.0) + image1 = torch.nn.functional.pad(image1, (0, 0, 0, 0, pad_top, pad_bottom), mode='constant', value=pad_value) if h2 < target_h: pad_h = target_h - h2 pad_top, pad_bottom = pad_h // 2, pad_h - pad_h // 2 - image2 = torch.nn.functional.pad(image2, (0, 0, 0, 0, pad_top, pad_bottom), mode='constant', value=0.0) + image2 = torch.nn.functional.pad(image2, (0, 0, 0, 0, pad_top, pad_bottom), mode='constant', value=pad_value) else: # up, down # For vertical concat, pad widths to match if w1 != w2: @@ -328,11 +341,11 @@ Optional spacing can be added between images. if w1 < target_w: pad_w = target_w - w1 pad_left, pad_right = pad_w // 2, pad_w - pad_w // 2 - image1 = torch.nn.functional.pad(image1, (0, 0, pad_left, pad_right), mode='constant', value=0.0) + image1 = torch.nn.functional.pad(image1, (0, 0, pad_left, pad_right), mode='constant', value=pad_value) if w2 < target_w: pad_w = target_w - w2 pad_left, pad_right = pad_w // 2, pad_w - pad_w // 2 - image2 = torch.nn.functional.pad(image2, (0, 0, pad_left, pad_right), mode='constant', value=0.0) + image2 = torch.nn.functional.pad(image2, (0, 0, pad_left, pad_right), mode='constant', value=pad_value) # Ensure same number of channels if image1.shape[-1] != image2.shape[-1]: @@ -366,15 +379,6 @@ Optional spacing can be added between images. if spacing_width > 0: spacing_width = spacing_width + (spacing_width % 2) # Ensure even - color_map = { - "white": 1.0, - "black": 0.0, - "red": (1.0, 0.0, 0.0), - "green": (0.0, 1.0, 0.0), - "blue": (0.0, 0.0, 1.0), - } - color_val = color_map[spacing_color] - if direction in ["left", "right"]: spacing_shape = ( image1.shape[0], From ae0e7c4dff827dae57d5b4c1d21b135036f42d64 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Sun, 22 Jun 2025 14:59:31 -0700 Subject: [PATCH 10/13] Resize and pad image node. (#8636) --- comfy_extras/nodes_images.py | 57 ++++++++++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py index 8d5fcdb85..ed54ccc57 100644 --- a/comfy_extras/nodes_images.py +++ b/comfy_extras/nodes_images.py @@ -414,6 +414,62 @@ Optional spacing can be added between images. concat_dim = 2 if direction in ["left", "right"] else 1 return (torch.cat(images, dim=concat_dim),) +class ResizeAndPadImage: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "target_width": ("INT", { + "default": 512, + "min": 1, + "max": MAX_RESOLUTION, + "step": 1 + }), + "target_height": ("INT", { + "default": 512, + "min": 1, + "max": MAX_RESOLUTION, + "step": 1 + }), + "padding_color": (["white", "black"],), + "interpolation": (["area", "bicubic", "nearest-exact", "bilinear", "lanczos"],), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "resize_and_pad" + CATEGORY = "image/transform" + + def resize_and_pad(self, image, target_width, target_height, padding_color, interpolation): + batch_size, orig_height, orig_width, channels = image.shape + + scale_w = target_width / orig_width + scale_h = target_height / orig_height + scale = min(scale_w, scale_h) + + new_width = int(orig_width * scale) + new_height = int(orig_height * scale) + + image_permuted = image.permute(0, 3, 1, 2) + + resized = comfy.utils.common_upscale(image_permuted, new_width, new_height, interpolation, "disabled") + + pad_value = 0.0 if padding_color == "black" else 1.0 + padded = torch.full( + (batch_size, channels, target_height, target_width), + pad_value, + dtype=image.dtype, + device=image.device + ) + + y_offset = (target_height - new_height) // 2 + x_offset = (target_width - new_width) // 2 + + padded[:, :, y_offset:y_offset + new_height, x_offset:x_offset + new_width] = resized + + output = padded.permute(0, 2, 3, 1) + return (output,) class SaveSVGNode: """ @@ -536,5 +592,6 @@ NODE_CLASS_MAPPINGS = { "SaveAnimatedPNG": SaveAnimatedPNG, "SaveSVGNode": SaveSVGNode, "ImageStitch": ImageStitch, + "ResizeAndPadImage": ResizeAndPadImage, "GetImageSize": GetImageSize, } From dd94416db24b93f81e9f4bd2aa91d1588041b489 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Mon, 23 Jun 2025 11:04:49 -0700 Subject: [PATCH 11/13] Indicate that directml is not recommended in the README. (#8644) --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index 0de4a6bb5..6366280e7 100644 --- a/README.md +++ b/README.md @@ -273,6 +273,8 @@ You can install ComfyUI in Apple Mac silicon (M1 or M2) with any recent macOS ve #### DirectML (AMD Cards on Windows) +This is very badly supported and is not recommended. There are some unofficial builds of pytorch ROCm on windows that exist that will give you a much better experience than this. This readme will be updated once official pytorch ROCm builds for windows come out. + ```pip install torch-directml``` Then you can launch ComfyUI with: ```python main.py --directml``` #### Ascend NPUs From bd9f166c1281e9da622a322d046897bbeee82d1d Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Tue, 24 Jun 2025 02:17:16 -0700 Subject: [PATCH 12/13] Cosmos predict2 model merging nodes. (#8647) --- .../nodes_model_merging_model_specific.py | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/comfy_extras/nodes_model_merging_model_specific.py b/comfy_extras/nodes_model_merging_model_specific.py index dc3411947..2c93cd84f 100644 --- a/comfy_extras/nodes_model_merging_model_specific.py +++ b/comfy_extras/nodes_model_merging_model_specific.py @@ -268,6 +268,52 @@ class ModelMergeWAN2_1(comfy_extras.nodes_model_merging.ModelMergeBlocks): return {"required": arg_dict} +class ModelMergeCosmosPredict2_2B(comfy_extras.nodes_model_merging.ModelMergeBlocks): + CATEGORY = "advanced/model_merging/model_specific" + + @classmethod + def INPUT_TYPES(s): + arg_dict = { "model1": ("MODEL",), + "model2": ("MODEL",)} + + argument = ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}) + + arg_dict["pos_embedder."] = argument + arg_dict["x_embedder."] = argument + arg_dict["t_embedder."] = argument + arg_dict["t_embedding_norm."] = argument + + + for i in range(28): + arg_dict["blocks.{}.".format(i)] = argument + + arg_dict["final_layer."] = argument + + return {"required": arg_dict} + +class ModelMergeCosmosPredict2_14B(comfy_extras.nodes_model_merging.ModelMergeBlocks): + CATEGORY = "advanced/model_merging/model_specific" + + @classmethod + def INPUT_TYPES(s): + arg_dict = { "model1": ("MODEL",), + "model2": ("MODEL",)} + + argument = ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}) + + arg_dict["pos_embedder."] = argument + arg_dict["x_embedder."] = argument + arg_dict["t_embedder."] = argument + arg_dict["t_embedding_norm."] = argument + + + for i in range(36): + arg_dict["blocks.{}.".format(i)] = argument + + arg_dict["final_layer."] = argument + + return {"required": arg_dict} + NODE_CLASS_MAPPINGS = { "ModelMergeSD1": ModelMergeSD1, "ModelMergeSD2": ModelMergeSD1, #SD1 and SD2 have the same blocks @@ -281,4 +327,6 @@ NODE_CLASS_MAPPINGS = { "ModelMergeCosmos7B": ModelMergeCosmos7B, "ModelMergeCosmos14B": ModelMergeCosmos14B, "ModelMergeWAN2_1": ModelMergeWAN2_1, + "ModelMergeCosmosPredict2_2B": ModelMergeCosmosPredict2_2B, + "ModelMergeCosmosPredict2_14B": ModelMergeCosmosPredict2_14B, } From 8042eb20c6fbcff54921ede9e983e1e848a65794 Mon Sep 17 00:00:00 2001 From: chaObserv <154517000+chaObserv@users.noreply.github.com> Date: Wed, 25 Jun 2025 02:59:09 +0800 Subject: [PATCH 13/13] Singlestep DPM++ SDE for RF (#8627) Refactor the algorithm, and apply alpha scaling. --- comfy/k_diffusion/sampling.py | 48 ++++++++++++++++++++++------------- 1 file changed, 30 insertions(+), 18 deletions(-) diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py index 8030048fc..739468872 100644 --- a/comfy/k_diffusion/sampling.py +++ b/comfy/k_diffusion/sampling.py @@ -710,6 +710,7 @@ def sample_dpmpp_2s_ancestral_RF(model, x, sigmas, extra_args=None, callback=Non # logged_x = torch.cat((logged_x, x.unsqueeze(0)), dim=0) return x + @torch.no_grad() def sample_dpmpp_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=1 / 2): """DPM-Solver++ (stochastic).""" @@ -721,38 +722,49 @@ def sample_dpmpp_sde(model, x, sigmas, extra_args=None, callback=None, disable=N seed = extra_args.get("seed", None) noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler s_in = x.new_ones([x.shape[0]]) - sigma_fn = lambda t: t.neg().exp() - t_fn = lambda sigma: sigma.log().neg() + + model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling') + sigma_fn = partial(half_log_snr_to_sigma, model_sampling=model_sampling) + lambda_fn = partial(sigma_to_half_log_snr, model_sampling=model_sampling) + sigmas = offset_first_sigma_for_snr(sigmas, model_sampling) for i in trange(len(sigmas) - 1, disable=disable): denoised = model(x, sigmas[i] * s_in, **extra_args) if callback is not None: callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) if sigmas[i + 1] == 0: - # Euler method - d = to_d(x, sigmas[i], denoised) - dt = sigmas[i + 1] - sigmas[i] - x = x + d * dt + # Denoising step + x = denoised else: # DPM-Solver++ - t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1]) - h = t_next - t - s = t + h * r + lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1]) + h = lambda_t - lambda_s + lambda_s_1 = lambda_s + r * h fac = 1 / (2 * r) + sigma_s_1 = sigma_fn(lambda_s_1) + + alpha_s = sigmas[i] * lambda_s.exp() + alpha_s_1 = sigma_s_1 * lambda_s_1.exp() + alpha_t = sigmas[i + 1] * lambda_t.exp() + # Step 1 - sd, su = get_ancestral_step(sigma_fn(t), sigma_fn(s), eta) - s_ = t_fn(sd) - x_2 = (sigma_fn(s_) / sigma_fn(t)) * x - (t - s_).expm1() * denoised - x_2 = x_2 + noise_sampler(sigma_fn(t), sigma_fn(s)) * s_noise * su - denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args) + sd, su = get_ancestral_step(lambda_s.neg().exp(), lambda_s_1.neg().exp(), eta) + lambda_s_1_ = sd.log().neg() + h_ = lambda_s_1_ - lambda_s + x_2 = (alpha_s_1 / alpha_s) * (-h_).exp() * x - alpha_s_1 * (-h_).expm1() * denoised + if eta > 0 and s_noise > 0: + x_2 = x_2 + alpha_s_1 * noise_sampler(sigmas[i], sigma_s_1) * s_noise * su + denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args) # Step 2 - sd, su = get_ancestral_step(sigma_fn(t), sigma_fn(t_next), eta) - t_next_ = t_fn(sd) + sd, su = get_ancestral_step(lambda_s.neg().exp(), lambda_t.neg().exp(), eta) + lambda_t_ = sd.log().neg() + h_ = lambda_t_ - lambda_s denoised_d = (1 - fac) * denoised + fac * denoised_2 - x = (sigma_fn(t_next_) / sigma_fn(t)) * x - (t - t_next_).expm1() * denoised_d - x = x + noise_sampler(sigma_fn(t), sigma_fn(t_next)) * s_noise * su + x = (alpha_t / alpha_s) * (-h_).exp() * x - alpha_t * (-h_).expm1() * denoised_d + if eta > 0 and s_noise > 0: + x = x + alpha_t * noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * su return x