From 4eb1a327ab596c43491c4ab1c45132847978e520 Mon Sep 17 00:00:00 2001 From: BVH <82035780+bvhari@users.noreply.github.com> Date: Wed, 6 Sep 2023 00:14:28 +0530 Subject: [PATCH 1/2] Add Tonemap Noise + Rescale CFG --- __init__.py | 46 +++++++++++++++-------------- sampler_tonemap_rescalecfg.py | 55 +++++++++++++++++++++++++++++++++++ 2 files changed, 79 insertions(+), 22 deletions(-) create mode 100644 sampler_tonemap_rescalecfg.py diff --git a/__init__.py b/__init__.py index c47df67..18dd91e 100644 --- a/__init__.py +++ b/__init__.py @@ -1,22 +1,24 @@ -import importlib -import os - -node_list = [ #Add list of .py files containing nodes here - "advanced_model_merging", - "reference_only", - "sampler_rescalecfg", - "sampler_tonemap", - "sdxl_model_merging" -] - -NODE_CLASS_MAPPINGS = {} -NODE_DISPLAY_NAME_MAPPINGS = {} - -for module_name in node_list: - imported_module = importlib.import_module(".{}".format(module_name), __name__) - - NODE_CLASS_MAPPINGS = {**NODE_CLASS_MAPPINGS, **imported_module.NODE_CLASS_MAPPINGS} - if hasattr(imported_module, "NODE_DISPLAY_NAME_MAPPINGS"): - NODE_DISPLAY_NAME_MAPPINGS = {**NODE_DISPLAY_NAME_MAPPINGS, **imported_module.NODE_DISPLAY_NAME_MAPPINGS} - -__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] +import importlib +import os + + +node_list = [ # Add list of .py files containing nodes here + "advanced_model_merging", + "reference_only", + "sampler_rescalecfg", + "sampler_tonemap", + "sampler_tonemap_rescalecfg", + "sdxl_model_merging" +] + +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +for module_name in node_list: + imported_module = importlib.import_module(".{}".format(module_name), __name__) + + NODE_CLASS_MAPPINGS = {**NODE_CLASS_MAPPINGS, **imported_module.NODE_CLASS_MAPPINGS} + if hasattr(imported_module, "NODE_DISPLAY_NAME_MAPPINGS"): + NODE_DISPLAY_NAME_MAPPINGS = {**NODE_DISPLAY_NAME_MAPPINGS, **imported_module.NODE_DISPLAY_NAME_MAPPINGS} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/sampler_tonemap_rescalecfg.py b/sampler_tonemap_rescalecfg.py new file mode 100644 index 0000000..0132330 --- /dev/null +++ b/sampler_tonemap_rescalecfg.py @@ -0,0 +1,55 @@ +import torch + + +class TonemapNoiseWithRescaleCFG: + @classmethod + def INPUT_TYPES(s): + return {"required": {"model": ("MODEL",), + "tonemap_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}), + "rescale_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + + CATEGORY = "custom_node_experiments" + + def patch(self, model, tonemap_multiplier, rescale_multiplier): + + def tonemap_noise_rescale_cfg(args): + cond = args["cond"] + uncond = args["uncond"] + cond_scale = args["cond_scale"] + + # Tonemap + noise_pred = (cond - uncond) + noise_pred_vector_magnitude = (torch.linalg.vector_norm(noise_pred, dim=(1)) + 0.0000000001)[:, None] + noise_pred /= noise_pred_vector_magnitude + + mean = torch.mean(noise_pred_vector_magnitude, dim=(1, 2, 3), keepdim=True) + std = torch.std(noise_pred_vector_magnitude, dim=(1, 2, 3), keepdim=True) + + top = (std * 3 + mean) * tonemap_multiplier + + # Reinhard + noise_pred_vector_magnitude *= (1.0 / top) + new_magnitude = noise_pred_vector_magnitude / (noise_pred_vector_magnitude + 1.0) + new_magnitude *= top + + # Rescale CFG + x_cfg = uncond + (noise_pred * new_magnitude * cond_scale) + ro_pos = torch.std(cond, dim=(1, 2, 3), keepdim=True) + ro_cfg = torch.std(x_cfg, dim=(1, 2, 3), keepdim=True) + + x_rescaled = x_cfg * (ro_pos / ro_cfg) + x_final = rescale_multiplier * x_rescaled + (1.0 - rescale_multiplier) * x_cfg + + return x_final + + m = model.clone() + m.set_model_sampler_cfg_function(tonemap_noise_rescale_cfg) + return (m, ) + + +NODE_CLASS_MAPPINGS = { + "TonemapNoiseWithRescaleCFG": TonemapNoiseWithRescaleCFG, +} From 285f7a8183c8c80d217e727fa47e31b3aa61389c Mon Sep 17 00:00:00 2001 From: BVH <82035780+bvhari@users.noreply.github.com> Date: Wed, 6 Sep 2023 18:59:36 +0530 Subject: [PATCH 2/2] Revert formatting changes --- __init__.py | 47 +++++++++++++++++++++++------------------------ 1 file changed, 23 insertions(+), 24 deletions(-) diff --git a/__init__.py b/__init__.py index 18dd91e..c6273d1 100644 --- a/__init__.py +++ b/__init__.py @@ -1,24 +1,23 @@ -import importlib -import os - - -node_list = [ # Add list of .py files containing nodes here - "advanced_model_merging", - "reference_only", - "sampler_rescalecfg", - "sampler_tonemap", - "sampler_tonemap_rescalecfg", - "sdxl_model_merging" -] - -NODE_CLASS_MAPPINGS = {} -NODE_DISPLAY_NAME_MAPPINGS = {} - -for module_name in node_list: - imported_module = importlib.import_module(".{}".format(module_name), __name__) - - NODE_CLASS_MAPPINGS = {**NODE_CLASS_MAPPINGS, **imported_module.NODE_CLASS_MAPPINGS} - if hasattr(imported_module, "NODE_DISPLAY_NAME_MAPPINGS"): - NODE_DISPLAY_NAME_MAPPINGS = {**NODE_DISPLAY_NAME_MAPPINGS, **imported_module.NODE_DISPLAY_NAME_MAPPINGS} - -__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] +import importlib +import os + +node_list = [ #Add list of .py files containing nodes here + "advanced_model_merging", + "reference_only", + "sampler_rescalecfg", + "sampler_tonemap", + "sampler_tonemap_rescalecfg", + "sdxl_model_merging" +] + +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +for module_name in node_list: + imported_module = importlib.import_module(".{}".format(module_name), __name__) + + NODE_CLASS_MAPPINGS = {**NODE_CLASS_MAPPINGS, **imported_module.NODE_CLASS_MAPPINGS} + if hasattr(imported_module, "NODE_DISPLAY_NAME_MAPPINGS"): + NODE_DISPLAY_NAME_MAPPINGS = {**NODE_DISPLAY_NAME_MAPPINGS, **imported_module.NODE_DISPLAY_NAME_MAPPINGS} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']