diff --git a/__init__.py b/__init__.py index 78c5b74..d398bba 100644 --- a/__init__.py +++ b/__init__.py @@ -31,6 +31,8 @@ from .nodes.math_expression import NODE_CLASS_MAPPINGS as MATH_EXPR_MAPPINGS from .nodes.math_expression import NODE_DISPLAY_NAME_MAPPINGS as MATH_EXPR_DISPLAY from .nodes.make_image_batch import NODE_CLASS_MAPPINGS as IMAGE_BATCH_MAPPINGS from .nodes.make_image_batch import NODE_DISPLAY_NAME_MAPPINGS as IMAGE_BATCH_DISPLAY +from .nodes.mask_denoise_remap import NODE_CLASS_MAPPINGS as MASK_DENOISE_REMAP_MAPPINGS +from .nodes.mask_denoise_remap import NODE_DISPLAY_NAME_MAPPINGS as MASK_DENOISE_REMAP_DISPLAY # Combine all mappings NODE_CLASS_MAPPINGS = { @@ -48,6 +50,7 @@ NODE_CLASS_MAPPINGS = { **MULTI_DETECT_MAPPINGS, **MATH_EXPR_MAPPINGS, **IMAGE_BATCH_MAPPINGS, + **MASK_DENOISE_REMAP_MAPPINGS, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -65,6 +68,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { **MULTI_DETECT_DISPLAY, **MATH_EXPR_DISPLAY, **IMAGE_BATCH_DISPLAY, + **MASK_DENOISE_REMAP_DISPLAY, } WEB_DIRECTORY = "./js" diff --git a/nodes/mask_denoise_remap.py b/nodes/mask_denoise_remap.py new file mode 100644 index 0000000..b465652 --- /dev/null +++ b/nodes/mask_denoise_remap.py @@ -0,0 +1,56 @@ +""" +RM Mask Denoise Remap - Remaps mask values to control per-region denoise strength. +White areas get white_denoise, black areas get black_denoise, grays interpolate linearly. +""" + +import torch + + +class RMMaskDenoiseRemap: + """ + Remaps a mask so that black and white regions map to custom denoise levels. + Feed the output into Set Latent Noise Mask for per-region denoise control. + + Output = black_denoise + mask * (white_denoise - black_denoise) + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask": ("MASK",), + "white_denoise": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "display": "slider", + }), + "black_denoise": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "display": "slider", + }), + } + } + + CATEGORY = "RMAutomation/Mask" + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("mask",) + FUNCTION = "remap" + + def remap(self, mask, white_denoise, black_denoise): + remapped = black_denoise + mask * (white_denoise - black_denoise) + remapped = torch.clamp(remapped, 0.0, 1.0) + return (remapped,) + + +NODE_CLASS_MAPPINGS = { + "RMMaskDenoiseRemap": RMMaskDenoiseRemap, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "RMMaskDenoiseRemap": "RM Mask Denoise Remap", +}