Upload 91 files

.gitattributes

@@ -34,3 +34,10 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 models/clip/qwen-4b-zimage-hereticV2-q8.gguf filter=lfs diff=lfs merge=lfs -text
+custom_nodes/comfyui-faded-branding/js/logo/logo01.png filter=lfs diff=lfs merge=lfs -text
+custom_nodes/comfyui-faded-branding/js/logo/logo40.png filter=lfs diff=lfs merge=lfs -text
+custom_nodes/comfyui-faded-branding/js/logo/logo44.png filter=lfs diff=lfs merge=lfs -text
+custom_nodes/srgb_noise/repo/doc/baseline.png filter=lfs diff=lfs merge=lfs -text
+custom_nodes/srgb_noise/repo/doc/visual_results.png filter=lfs diff=lfs merge=lfs -text
+custom_nodes/srgb_noise/repo/test/imgs/clean.png filter=lfs diff=lfs merge=lfs -text
+custom_nodes/srgb_noise/repo/test/imgs/noisy.png filter=lfs diff=lfs merge=lfs -text

custom_nodes/ComfyUI-DyPE/src/patch_utils.py

@@ -0,0 +1,272 @@
+import math
+import types
+import torch
+import torch.nn.functional as F
+import comfy
+from comfy.model_patcher import ModelPatcher
+from comfy import model_sampling
+
+from .models.flux import PosEmbedFlux
+from .models.nunchaku import PosEmbedNunchaku
+from .models.qwen import PosEmbedQwen
+from .models.zimage import PosEmbedZImage
+
+def apply_dype_to_model(model: ModelPatcher, model_type: str, width: int, height: int, method: str, yarn_alt_scaling: bool, enable_dype: bool, dype_scale: float, dype_exponent: float, base_shift: float, max_shift: float, base_resolution: int = 1024, dype_start_sigma: float = 1.0) -> ModelPatcher:
+    m = model.clone()
+
+    is_nunchaku = False
+    is_qwen = False
+    is_z_image = False
+
+    if model_type == "nunchaku":
+        is_nunchaku = True
+    elif model_type == "qwen":
+        is_qwen = True
+    elif model_type == "z_image":
+        is_z_image = True
+    elif model_type == "flux":
+        pass
+    else: # auto
+        if hasattr(m.model, "diffusion_model"):
+            dm = m.model.diffusion_model
+            model_class_name = dm.__class__.__name__
+            if "QwenImage" in model_class_name:
+                is_qwen = True
+            elif hasattr(dm, "rope_embedder"):
+                is_z_image = True
+            elif hasattr(dm, "model") and hasattr(dm.model, "pos_embed"):
+                is_nunchaku = True
+        else:
+            raise ValueError("The provided model is not a compatible model.")
+
+    new_dype_params = (width, height, base_shift, max_shift, method, yarn_alt_scaling, base_resolution, dype_start_sigma, is_nunchaku, is_qwen, is_z_image)
+
+    should_patch_schedule = True
+    if hasattr(m.model, "_dype_params"):
+        if m.model._dype_params == new_dype_params:
+            should_patch_schedule = False
+
+    base_patch_h_tokens = None
+    base_patch_w_tokens = None
+    if is_z_image:
+        axes_lens = getattr(m.model.diffusion_model, "axes_lens", None)
+        if isinstance(axes_lens, (list, tuple)) and len(axes_lens) >= 3:
+            base_patch_h_tokens = int(axes_lens[1])
+            base_patch_w_tokens = int(axes_lens[2])
+
+    patch_size = 2
+    try:
+        if is_nunchaku:
+            patch_size = m.model.diffusion_model.model.config.patch_size
+        else:
+            patch_size = m.model.diffusion_model.patch_size
+    except:
+        pass
+
+    if base_patch_h_tokens is not None and base_patch_w_tokens is not None:
+        derived_base_patches = max(base_patch_h_tokens, base_patch_w_tokens)
+        derived_base_seq_len = base_patch_h_tokens * base_patch_w_tokens
+    else:
+        derived_base_patches = (base_resolution // 8) // 2
+        derived_base_seq_len = derived_base_patches * derived_base_patches
+
+    if enable_dype and should_patch_schedule:
+        try:
+            if isinstance(m.model.model_sampling, model_sampling.ModelSamplingFlux) or is_qwen or is_z_image:
+                latent_h, latent_w = height // 8, width // 8
+                padded_h, padded_w = math.ceil(latent_h / patch_size) * patch_size, math.ceil(latent_w / patch_size) * patch_size
+                image_seq_len = (padded_h // patch_size) * (padded_w // patch_size)
+
+                base_seq_len = derived_base_seq_len
+                max_seq_len = image_seq_len
+
+                if max_seq_len <= base_seq_len:
+                    dype_shift = base_shift
+                else:
+                    slope = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+                    intercept = base_shift - slope * base_seq_len
+                    dype_shift = image_seq_len * slope + intercept
+
+                dype_shift = max(0.0, dype_shift)
+
+                class DypeModelSamplingFlux(model_sampling.ModelSamplingFlux, model_sampling.CONST):
+                    pass
+
+                new_model_sampler = DypeModelSamplingFlux(m.model.model_config)
+                new_model_sampler.set_parameters(shift=dype_shift)
+
+                m.add_object_patch("model_sampling", new_model_sampler)
+                m.model._dype_params = new_dype_params
+        except:
+            pass
+
+    elif not enable_dype:
+        if hasattr(m.model, "_dype_params"):
+            class DefaultModelSamplingFlux(model_sampling.ModelSamplingFlux, model_sampling.CONST): pass
+            default_sampler = DefaultModelSamplingFlux(m.model.model_config)
+            m.add_object_patch("model_sampling", default_sampler)
+            del m.model._dype_params
+
+    try:
+        if is_nunchaku:
+            orig_embedder = m.model.diffusion_model.model.pos_embed
+            target_patch_path = "diffusion_model.model.pos_embed"
+        elif is_z_image:
+            orig_embedder = m.model.diffusion_model.rope_embedder
+            target_patch_path = "diffusion_model.rope_embedder"
+        else:
+            orig_embedder = m.model.diffusion_model.pe_embedder
+            target_patch_path = "diffusion_model.pe_embedder"
+
+        theta, axes_dim = orig_embedder.theta, orig_embedder.axes_dim
+    except AttributeError:
+        raise ValueError("The provided model is not a compatible FLUX/Qwen model structure.")
+
+    embedder_cls = PosEmbedFlux
+    if is_nunchaku:
+        embedder_cls = PosEmbedNunchaku
+    elif is_qwen:
+        embedder_cls = PosEmbedQwen
+    elif is_z_image:
+        embedder_cls = PosEmbedZImage
+
+    embedder_base_patches = derived_base_patches if is_z_image else None
+
+    new_pe_embedder = embedder_cls(
+        theta, axes_dim, method, yarn_alt_scaling, enable_dype,
+        dype_scale, dype_exponent, base_resolution, dype_start_sigma, embedder_base_patches
+    )
+        
+    m.add_object_patch(target_patch_path, new_pe_embedder)
+
+    if is_z_image:
+        base_hw_override = None
+        if base_patch_h_tokens is not None and base_patch_w_tokens is not None:
+            base_hw_override = (base_patch_h_tokens, base_patch_w_tokens)
+        elif derived_base_patches is not None:
+            base_hw_override = (derived_base_patches, derived_base_patches)
+
+        if base_hw_override is not None:
+            m.model.diffusion_model._dype_base_hw = base_hw_override
+
+        def dype_patchify_and_embed(self, x, cap_feats, cap_mask, t, num_tokens, transformer_options={}, **kwargs):
+            bsz = len(x)
+            pH = pW = self.patch_size
+            device = x[0].device
+
+            if self.pad_tokens_multiple is not None:
+                pad_extra = (-cap_feats.shape[1]) % self.pad_tokens_multiple
+                if pad_extra:
+                    # FIX 2: Zero padding instead of cap_pad_token due to dimension mismatch in Turbo
+                    cap_pad = torch.zeros(1, 1, cap_feats.shape[-1], device=cap_feats.device, dtype=cap_feats.dtype)
+                    cap_feats = torch.cat((cap_feats, cap_pad.repeat(cap_feats.shape[0], pad_extra, 1)), dim=1)
+
+            cap_pos_ids = torch.zeros(bsz, cap_feats.shape[1], 3, dtype=torch.float32, device=device)
+            cap_pos_ids[:, :, 0] = torch.arange(cap_feats.shape[1], dtype=torch.float32, device=device) + 1.0
+
+            B, C, H, W = x.shape
+            x = self.x_embedder(x.view(B, C, H // pH, pH, W // pW, pW).permute(0, 2, 4, 3, 5, 1).flatten(3).flatten(1, 2))
+
+            requested_hw = transformer_options.get("dype_requested_hw", (height, width))
+            rope_base_resolution = transformer_options.get("dype_base_resolution", base_resolution)
+
+            raw_scale_y = float(rope_base_resolution) / max(1.0, float(requested_hw[0]))
+            raw_scale_x = float(rope_base_resolution) / max(1.0, float(requested_hw[1]))
+            
+            iso_scale = min(raw_scale_y, raw_scale_x)
+            rope_scale_y = iso_scale
+            rope_scale_x = iso_scale
+            
+            freq_scale_factor = 1.0 / iso_scale
+            new_pe_embedder.set_scale_hint(freq_scale_factor)
+
+            h_start = 0.0
+            w_start = 0.0
+
+            original_hw = transformer_options.get("dype_original_hw")
+            if original_hw is None:
+                original_hw = (H, W)
+
+            H_tokens = math.ceil(original_hw[0] / pH)
+            W_tokens = math.ceil(original_hw[1] / pW)
+            
+            token_stride_y = (original_hw[0] / max(1, H_tokens)) * rope_scale_y
+            token_stride_x = (original_hw[1] / max(1, W_tokens)) * rope_scale_x
+            
+            shift_y = h_start * (original_hw[0] / max(1, H_tokens))
+            shift_x = w_start * (original_hw[1] / max(1, W_tokens))
+            
+            def _build_spatial_pos_ids(batch: int, total_len: int, width_tokens: int, cap_len: int, stride_y: float, stride_x: float, h_start: float, w_start: float, device: torch.device):
+                base_pos = torch.arange(total_len, device=device, dtype=torch.float32)
+                y = torch.div(base_pos, width_tokens, rounding_mode='floor') * stride_y + h_start
+                x = torch.remainder(base_pos, width_tokens) * stride_x + w_start
+
+                pos = torch.stack([
+                    torch.full_like(base_pos, cap_len + 1),
+                    y,
+                    x
+                ], dim=-1)
+                return pos.unsqueeze(0).repeat(batch, 1, 1)
+
+            base_img_tokens = H_tokens * W_tokens
+            x_pos_ids = _build_spatial_pos_ids(bsz, base_img_tokens, W_tokens, cap_feats.shape[1], token_stride_y, token_stride_x, shift_y, shift_x, device)
+
+            if self.pad_tokens_multiple is not None:
+                pad_extra = (-x.shape[1]) % self.pad_tokens_multiple
+                if pad_extra:
+                    x = torch.cat((x, self.x_pad_token.to(device=x.device, dtype=x.dtype, copy=True).unsqueeze(0).repeat(x.shape[0], pad_extra, 1)), dim=1)
+
+            if x.shape[1] != x_pos_ids.shape[1]:
+                x_pos_ids = _build_spatial_pos_ids(bsz, x.shape[1], W_tokens, cap_feats.shape[1], token_stride_y, token_stride_x, shift_y, shift_x, device)
+
+            freqs_cis = self.rope_embedder(torch.cat((cap_pos_ids, x_pos_ids), dim=1)).movedim(1, 2)
+
+            # FIX 3: Missing cap_embedder projection to map text encoder dimension to model's hidden dimension
+            cap_feats = self.cap_embedder(cap_feats)
+            
+            for layer in self.context_refiner:
+                cap_feats = layer(cap_feats, cap_mask, freqs_cis[:, :cap_pos_ids.shape[1]], transformer_options=transformer_options)
+
+            padded_img_mask = None
+            for layer in self.noise_refiner:
+                x = layer(x, padded_img_mask, freqs_cis[:, cap_pos_ids.shape[1]:], t, transformer_options=transformer_options)
+
+            padded_full_embed = torch.cat((cap_feats, x), dim=1)
+            mask = None
+            img_sizes = [(H, W)] * bsz
+            l_effective_cap_len = [cap_feats.shape[1]] * bsz
+            
+            # FIX 4: Return tuple missing sixth value (None for non-omni path)
+            return padded_full_embed, mask, img_sizes, l_effective_cap_len, freqs_cis, None
+
+        m.add_object_patch(
+            "diffusion_model.patchify_and_embed",
+            types.MethodType(dype_patchify_and_embed, m.model.diffusion_model)
+        )
+
+    sigma_max = m.model.model_sampling.sigma_max.item()
+    
+    def dype_wrapper_function(model_function, args_dict):
+        timestep_tensor = args_dict.get("timestep")
+        if timestep_tensor is not None and timestep_tensor.numel() > 0:
+            current_sigma = timestep_tensor.flatten()[0].item()
+            
+            if sigma_max > 0:
+                normalized_timestep = min(max(current_sigma / sigma_max, 0.0), 1.0)
+                new_pe_embedder.set_timestep(normalized_timestep)
+        
+        input_x, c = args_dict.get("input"), args_dict.get("c", {})
+
+        if is_z_image and isinstance(input_x, torch.Tensor) and input_x.dim() >= 4:
+            c = dict(c)
+            transformer_options = dict(c.get("transformer_options", {}))
+            transformer_options["dype_original_hw"] = (input_x.shape[-2], input_x.shape[-1])
+            transformer_options["dype_requested_hw"] = (height, width)
+            transformer_options["dype_base_resolution"] = base_resolution
+            c["transformer_options"] = transformer_options
+
+        return model_function(input_x, args_dict.get("timestep"), **c)
+
+    m.set_model_unet_function_wrapper(dype_wrapper_function)
+
+    return m

custom_nodes/ComfyUI-EsesImageEffectLevels/js/image_effect_levels.js

@@ -0,0 +1,967 @@
+// ==========================================================================
+// Eses Image Effect Levels
+// ==========================================================================
+//
+// Description:
+// The 'Eses Image Effect Levels' node provides a versatile, image editor style
+// levels adjustment tool directly within the ComfyUI interface. It allows for
+// precise, interactive control over the tonal range of images and masks,
+// complete with a live histogram for visual feedback.
+//
+// Key Features:
+//
+// - Interactive Controls:
+//   - An interactive preview of the level adjustments in the node.
+//   - Features sliders for input Black, Mid, and White points, as well as
+//     sliders for the final Output Black and White points.
+//   - A live histogram visualizes the tonal distribution of the selected
+//     channel.
+//
+// - Multi-Channel Adjustments:
+//   - Apply level adjustments to the combined RGB channels for general
+//     corrections.
+//   - Isolate adjustments to the individual Red, Green, or Blue channels.
+//   - Apply a separate level adjustment directly to an input mask.
+//
+// - State Serialization:
+//   - All level adjustments are saved with the workflow and restored on
+//     reload.
+//   - The node's state persists after refreshing the browser page.
+//
+// - Live Preview:
+//   - The node displays a preview of the connected image with the level
+//     adjustments applied in real-time as you move the sliders.
+//
+// - Quality of Life Features:
+//   - "Set Auto Levels" button to automatically calculate optimal
+//      input levels.
+//   - A simple file-based preset system for saving and loading adjustment
+//     presets
+//   - "Reset" buttons to revert the current channel or all channels to
+//     default values.
+//   - Adjustable auto levels sensitivity for fine-tuned automatic adjustments.
+//
+// Version: 1.3.2
+//
+// License: See LICENSE.txt
+//
+// ==========================================================================
+
+
+import { app } from "/scripts/app.js";
+import { api } from "/scripts/api.js";
+
+const HEADER_HEIGHT = 445;
+const PADDING = 10;
+const HISTOGRAM_AREA_HEIGHT = 40;
+const AREA_SPACING = 8;
+const SLIDER_BOX_HEIGHT = 15;
+const OUTPUT_SLIDER_AREA_HEIGHT = 20;
+
+
+// --- PRESET HANDLING START ---
+
+// Global cache for presets to avoid
+// fetching them for every single node
+let allPresets = null;
+
+async function getPresets() {
+    if (allPresets === null) {
+        try {
+            const response = await api.fetchApi("/eses_levels/get_presets", { cache: "no-store" });
+            allPresets = await response.json();
+            //console.log("[EsesImageEffectLevels] Presets loaded:", allPresets);
+        }
+        catch (e) {
+            //console.error("[EsesImageEffectLevels] Failed to load presets:", e);
+            allPresets = {};
+        }
+    }
+
+    return allPresets;
+}
+
+// Pre-fetch presets when the script is loaded
+getPresets();
+
+// --- PRESET HANDLING END ---
+
+
+
+app.registerExtension({
+    name: "Eses.EsesImageEffectLevels",
+
+    // --- PRESET HANDLING START ---
+    async beforeRegisterNodeDef(nodeType, nodeData, app) {
+
+        // Runs once when the node definition is being registered.
+        // Dynamically add the preset names to the combo box definition.
+        if (nodeData.name === "EsesImageEffectLevels") {
+            const presets = await getPresets();
+            const presetNames = ["None", ...Object.keys(presets)];
+
+            // Find the 'preset' input in the node's
+            // definition and update its options
+            const presetInput = nodeData.input.required.preset;
+            if (presetInput) {
+                presetInput[0] = presetNames;
+            }
+        }
+    },
+    // --- PRESET HANDLING END ---
+
+
+    nodeCreated(node) {
+
+        if (node.comfyClass === "EsesImageEffectLevels") {
+            const originalConfigure = node.configure;
+            const originalSerialize = node.serialize;
+
+            node.size = [256, 575];
+            node.imagePreview = null;
+            node.originalImage = null;
+            node.originalMask = null;
+            node.isManuallyResized = false;
+            node.dragging = null;
+            node.resetButtonState = 'idle';
+
+            // Add state for auto color button
+            node.autoColorButtonState = 'idle';
+            node.onResize = function () { this.isManuallyResized = true; };
+
+            const defaultLevels = () => ({ black_point: 0.0, white_point: 1.0, mid_point: 1.0, output_black: 0, output_white: 255 });
+            const resetAllLevels = () => ({ rgb: defaultLevels(), r: defaultLevels(), g: defaultLevels(), b: defaultLevels(), mask: defaultLevels() });
+
+            node.allLevels = resetAllLevels();
+            node.histograms = {};
+
+            let blackWidget, midWidget, whiteWidget, outputBlackWidget, outputWhiteWidget;
+
+            function updateWidgetsFromState() {
+                const levels = node.allLevels[node.activeChannel];
+
+                if (!levels)
+                    return;
+
+                if (blackWidget)
+                    blackWidget.value = parseFloat(levels.black_point.toFixed(2));
+
+                if (midWidget)
+                    midWidget.value = parseFloat(levels.mid_point.toFixed(2));
+
+                if (whiteWidget)
+                    whiteWidget.value = parseFloat(levels.white_point.toFixed(2));
+
+                if (outputBlackWidget)
+                    outputBlackWidget.value = Math.round(levels.output_black);
+
+                if (outputWhiteWidget)
+                    outputWhiteWidget.value = Math.round(levels.output_white);
+            }
+
+            node.getResetButtonName = (isConfirm = false) => {
+                const channelWidget = node.widgets.find(w => w.name === "channel");
+
+                if (node.activeChannel === 'rgb')
+                    return isConfirm ? "CONFIRM RESET ALL" : "Reset All Levels";
+
+                const channelName = channelWidget.value;
+
+                return isConfirm ? `CONFIRM RESET ${channelName}` : `Reset ${channelName} Level`;
+            };
+
+            node.resetButtonCallback = function () {
+                const resetButton = node.widgets.find(w => w.name.startsWith("Reset") || w.name.startsWith("CONFIRM"));
+
+                if (node.resetButtonState === 'armed') {
+                    clearTimeout(node.resetTimeout);
+
+                    if (node.activeChannel === 'rgb') {
+                        node.allLevels = resetAllLevels();
+                    }
+                    else {
+                        node.allLevels[node.activeChannel] = defaultLevels();
+                    }
+
+                    updateWidgetsFromState();
+                    node.sendLevelUpdateToPython();
+                    resetButton.name = node.getResetButtonName();
+                    node.resetButtonState = 'idle';
+                }
+                else {
+                    node.resetButtonState = 'armed';
+                    resetButton.name = node.getResetButtonName(true);
+
+                    node.resetTimeout = setTimeout(() => {
+                        resetButton.name = node.getResetButtonName();
+                        node.resetButtonState = 'idle';
+                        app.graph.setDirtyCanvas(true);
+                    }, 3000);
+                }
+                app.graph.setDirtyCanvas(true);
+            };
+
+            function setupCustomWidgets() {
+                const levelsWidget = node.widgets.find(w => w.name === "all_levels_json");
+                const channelWidget = node.widgets.find(w => w.name === "channel");
+
+                // Hide the boolean widgets that are only used for triggering
+                const autoLevelsWidget = node.widgets.find(w => w.name === "auto_levels");
+                const autoColorWidget = node.widgets.find(w => w.name === "auto_color");
+
+                if (!levelsWidget || !channelWidget)
+                    return;
+
+                levelsWidget.type = "hidden";
+                if (levelsWidget.element) {
+                    levelsWidget.element.style.display = "none";
+                }
+
+                if (autoLevelsWidget) {
+                    autoLevelsWidget.type = "hidden";
+                    if (autoLevelsWidget.element) {
+                        autoLevelsWidget.element.style.display = "none";
+                    }
+                }
+
+                if (autoColorWidget) {
+                    autoColorWidget.type = "hidden";
+                    if (autoColorWidget.element) {
+                        autoColorWidget.element.style.display = "none";
+                    }
+                }
+                const channelMap = { "RGB": "rgb", "Red": "r", "Green": "g", "Blue": "b", "Mask": "mask" };
+                node.activeChannel = channelMap[channelWidget.value] || "rgb";
+
+                if (!node.widgets.find(w => w.name === "Black")) {
+                    blackWidget = node.addWidget("number", "Black", 0.0, (val) => {
+
+                        if (node.allLevels[node.activeChannel]) {
+                            node.allLevels[node.activeChannel].black_point = val;
+                            node.sendLevelUpdateToPython();
+                        }
+                    }, { min: 0.0, max: 1.0, step: 0.01, precision: 2 });
+
+                    midWidget = node.addWidget("number", "Mid", 1.0, (val) => {
+
+                        if (node.allLevels[node.activeChannel]) {
+                            node.allLevels[node.activeChannel].mid_point = val;
+                            node.sendLevelUpdateToPython();
+                        }
+                    }, { min: 0.01, max: 9.99, step: 0.01, precision: 2 });
+
+                    whiteWidget = node.addWidget("number", "White", 1.0, (val) => {
+
+                        if (node.allLevels[node.activeChannel]) {
+                            node.allLevels[node.activeChannel].white_point = val;
+                            node.sendLevelUpdateToPython();
+                        }
+                    }, { min: 0.0, max: 1.0, step: 0.01, precision: 2 });
+
+                    outputBlackWidget = node.addWidget("number", "Output Black", 0, (val) => {
+
+                        if (node.allLevels[node.activeChannel]) {
+                            node.allLevels[node.activeChannel].output_black = val;
+                            node.sendLevelUpdateToPython();
+                        }
+                    }, { min: 0, max: 255, step: 1, precision: 0 });
+
+                    outputWhiteWidget = node.addWidget("number", "Output White", 255, (val) => {
+
+                        if (node.allLevels[node.activeChannel]) {
+                            node.allLevels[node.activeChannel].output_white = val;
+                            node.sendLevelUpdateToPython();
+                        }
+                    }, { min: 0, max: 255, step: 1, precision: 0 });
+                }
+                else {
+                    blackWidget = node.widgets.find(w => w.name === "Black");
+                    midWidget = node.widgets.find(w => w.name === "Mid");
+                    whiteWidget = node.widgets.find(w => w.name === "White");
+                    outputBlackWidget = node.widgets.find(w => w.name === "Output Black");
+                    outputWhiteWidget = node.widgets.find(w => w.name === "Output White");
+                }
+
+                // --- PRESET HANDLING START ---
+
+                // Remove old "Copy Settings" button if it exists
+                const copySettingsButton = node.widgets.find(w => w.name === "Copy Settings");
+
+                if (copySettingsButton)
+                    node.widgets.splice(node.widgets.indexOf(copySettingsButton), 1);
+
+                // Add "Preset Name" text field and "Save Preset" button
+                if (!node.widgets.find(w => w.name === "Preset Name")) {
+                    const presetNameWidget = node.addWidget("string", "Preset Name", "", () => { }, {});
+
+                    node.addWidget("button", "Save Preset", null, async () => {
+                        const name = presetNameWidget.value;
+
+                        if (!name || !name.trim()) {
+                            alert("Please enter a preset name.");
+                            return;
+                        }
+
+                        try {
+                            const resp = await api.fetchApi("/eses_levels/save_preset", {
+                                method: "POST",
+                                headers: { "Content-Type": "application/json" },
+                                body: JSON.stringify({
+                                    preset_name: name,
+                                    preset_data: node.allLevels,
+                                }),
+                            });
+
+                            const res = await resp.json();
+
+                            if (res.success) {
+                                alert(`Preset '${name}' saved successfully!`);
+
+                                // Refresh the presets in the dropdown
+                                allPresets = res.presets;
+                                const presetWidget = node.widgets.find(w => w.name === "preset");
+
+                                if (presetWidget) {
+                                    presetWidget.options.values = ["None", ...Object.keys(allPresets)];
+                                    presetWidget.value = name;
+                                }
+                                presetNameWidget.value = "";
+                            }
+                            else {
+                                alert(`Error saving preset: ${res.message}`);
+                            }
+                        }
+                        catch (e) {
+                            console.error("[EsesImageEffectLevels] Error saving preset:", e);
+                            alert(`An error occurred: ${e}`);
+                        }
+                    });
+                }
+
+                const presetWidget = node.widgets.find(w => w.name === "preset");
+                if (presetWidget) {
+                    presetWidget.callback = async (value) => {
+
+                        if (value === "None")
+                            return;
+
+                        const presets = await getPresets();
+                        const selectedPreset = presets[value];
+
+                        if (selectedPreset) {
+
+                            // Create a fresh default state
+                            const newLevels = resetAllLevels();
+                            const loadedLevels = JSON.parse(JSON.stringify(selectedPreset));
+
+                            // Deep merge the loaded preset into the default state.
+                            // This ensures compatibility with older presets.
+                            for (const channelKey in newLevels) {
+
+                                if (loadedLevels[channelKey]) {
+                                    Object.assign(newLevels[channelKey], loadedLevels[channelKey]);
+                                }
+                            }
+
+                            node.allLevels = newLevels;
+                            updateWidgetsFromState();
+                            node.sendLevelUpdateToPython();
+                            node.setDirtyCanvas(true, true);
+                        }
+                    };
+                }
+
+                // --- PRESET HANDLING END ---
+
+                if (!node.widgets.find(w => w.name === "Set Auto Levels")) {
+                    node.addWidget("button", "Set Auto Levels", null, () => {
+                        const autoLevelsWidget = node.widgets.find(w => w.name === "auto_levels");
+
+                        if (autoLevelsWidget) {
+                            const originalValue = autoLevelsWidget.value;
+                            autoLevelsWidget.value = true;
+                            app.queuePrompt();
+                            setTimeout(() => { autoLevelsWidget.value = originalValue; }, 100);
+                        }
+                    });
+                }
+
+                // Add "Set Auto Color" button
+                if (!node.widgets.find(w => w.name === "Set Auto Color")) {
+
+                    node.addWidget("button", "Set Auto Color", null, () => {
+
+                        const autoColorWidget = node.widgets.find(w => w.name === "auto_color");
+
+                        if (autoColorWidget) {
+                            const originalValue = autoColorWidget.value;
+                            autoColorWidget.value = true;
+                            app.queuePrompt();
+
+                            // Reset the value back to false after a short delay
+                            setTimeout(() => { autoColorWidget.value = originalValue; }, 100);
+                        }
+                    });
+                }
+
+
+                const oldResetBtn = node.widgets.find(w => w && w.name && (w.name.startsWith("Reset") || w.name.startsWith("CONFIRM")));
+
+                if (oldResetBtn)
+                    node.widgets.splice(node.widgets.indexOf(oldResetBtn), 1);
+
+                const resetButton = node.addWidget("button", node.getResetButtonName(), null, node.resetButtonCallback.bind(node));
+
+                channelWidget.callback = (value) => {
+                    node.activeChannel = channelMap[value] || "rgb";
+                    resetButton.name = node.getResetButtonName();
+
+                    updateWidgetsFromState();
+                    node.updateHistogram();
+                    node.sendLevelUpdateToPython();
+                };
+            }
+
+            node.configure = function (data) {
+                originalConfigure.apply(this, arguments);
+
+                if (data.custom_all_levels) {
+                    this.allLevels = data.custom_all_levels;
+                }
+
+                setupCustomWidgets();
+                updateWidgetsFromState();
+            };
+
+            node.serialize = function () {
+                const data = originalSerialize.call(this);
+                data.custom_all_levels = node.allLevels;
+                return data;
+            };
+
+            setupCustomWidgets();
+
+            Object.assign(node, {
+
+                getImageArea: function () {
+                    const interactiveHeight = HISTOGRAM_AREA_HEIGHT + OUTPUT_SLIDER_AREA_HEIGHT + AREA_SPACING;
+
+                    // Moved image top edge down by the 
+                    // amount increased by HEADER_HEIGHT
+                    const totalContentHeight = this.size[1] - HEADER_HEIGHT - PADDING * 2;
+                    const imageAreaHeight = totalContentHeight - interactiveHeight;
+
+                    if (imageAreaHeight < 10)
+                        return null;
+
+                    return { x: PADDING, y: HEADER_HEIGHT + PADDING, width: this.size[0] - PADDING * 2, height: imageAreaHeight };
+                },
+
+                getBarArea: function () {
+                    const imageArea = this.getImageArea();
+
+                    if (!imageArea)
+                        return null;
+
+                    const barAreaHeight = HISTOGRAM_AREA_HEIGHT - SLIDER_BOX_HEIGHT - 4;
+                    return { x: PADDING, y: imageArea.y + imageArea.height + AREA_SPACING, width: this.size[0] - PADDING * 2, height: barAreaHeight };
+                },
+
+                getInputSliderArea: function () {
+                    const barArea = this.getBarArea();
+
+                    if (!barArea)
+                        return null;
+
+                    return { x: PADDING, y: barArea.y + barArea.height + 4, width: this.size[0] - PADDING * 2, height: SLIDER_BOX_HEIGHT };
+                },
+
+                getOutputSliderArea: function () {
+                    const inputSliderArea = this.getInputSliderArea();
+
+                    if (!inputSliderArea)
+                        return null;
+
+                    return { x: PADDING, y: inputSliderArea.y + inputSliderArea.height + AREA_SPACING, width: this.size[0] - PADDING * 2, height: SLIDER_BOX_HEIGHT };
+                },
+
+
+                updateHistogram: async function () {
+                    const channelWidget = node.widgets.find(w => w.name === "channel");
+                    let base_data_image = (this.activeChannel === 'mask') ? this.originalMask : this.originalImage;
+
+                    if (!base_data_image)
+                        return;
+
+                    this.histograms[this.activeChannel] = null;
+                    const canvas = document.createElement('canvas');
+                    canvas.width = base_data_image.naturalWidth;
+                    canvas.height = base_data_image.naturalHeight;
+                    canvas.getContext('2d').drawImage(base_data_image, 0, 0);
+                    const base_data_b64 = canvas.toDataURL('image/png').split(',')[1];
+
+                    try {
+                        const result = await api.fetchApi("/eses_levels/get_histogram", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ base_data_b64, channel_mode: channelWidget.value }) }).then(r => r.json());
+
+                        if (result.histogram) {
+                            this.histograms[this.activeChannel] = result.histogram;
+                            this.setDirtyCanvas(true, true);
+                        }
+                    }
+                    catch (err) {
+                        console.error("Failed to fetch histogram:", err);
+                    }
+                },
+
+                sendLevelUpdateToPython: async function () {
+                    this.setDirtyCanvas(true, true);
+
+                    const levelsWidget = this.widgets.find(w => w.name === "all_levels_json");
+
+                    if (!levelsWidget)
+                        return;
+
+                    levelsWidget.value = JSON.stringify(this.allLevels);
+                    let base_data_type, base_data_image;
+
+                    if (this.activeChannel === 'mask') {
+                        if (!this.originalMask) {
+                            this.imagePreview = null; this.setDirtyCanvas(true, true); return;
+                        }
+
+                        base_data_type = 'mask';
+                        base_data_image = this.originalMask;
+                    }
+                    else {
+                        if (!this.originalImage) {
+                            this.imagePreview = null; this.setDirtyCanvas(true, true); return;
+                        }
+
+                        base_data_type = 'image';
+                        base_data_image = this.originalImage;
+                    }
+
+                    const canvas = document.createElement('canvas');
+                    canvas.width = base_data_image.naturalWidth;
+                    canvas.height = base_data_image.naturalHeight;
+                    canvas.getContext('2d').drawImage(base_data_image, 0, 0);
+                    const base_data_b64 = canvas.toDataURL('image/png').split(',')[1];
+
+                    const body = {
+                        base_data_b64,
+                        all_levels_json: levelsWidget.value,
+                        base_data_type
+                    };
+
+                    try {
+                        const result = await api.fetchApi("/eses_levels/apply", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(body) }).then(r => r.json());
+
+                        if (result.adjusted_image_data) {
+                            const img = new Image();
+                            img.src = `data:image/png;base64,${result.adjusted_image_data}`;
+
+                            img.onload = () => {
+                                this.imagePreview = img; this.setDirtyCanvas(true, true);
+                            };
+                        }
+                    }
+                    catch (err) { }
+                },
+
+                drawHistogram: function (ctx, area, data, step = 1) {
+                    if (!data || data.length === 0)
+                        return;
+
+                    step = Math.max(1, Math.floor(step));
+                    ctx.fillStyle = "#999";
+                    const baseBarWidth = area.width / 256;
+
+                    for (let i = 0; i < 256; i += step) {
+                        const chunk = data.slice(i, i + step);
+
+                        if (chunk.length === 0)
+                            continue;
+
+                        const maxVal = Math.max(...chunk);
+                        const barHeight = maxVal * area.height;
+                        const barWidth = baseBarWidth * chunk.length;
+                        const x = area.x + i * baseBarWidth;
+                        ctx.fillRect(x, area.y + area.height - barHeight, barWidth, barHeight);
+                    }
+                },
+
+                drawHandle: function (ctx, x, y, size, color, strokeColor) {
+                    ctx.fillStyle = color;
+                    ctx.strokeStyle = strokeColor;
+                    ctx.lineWidth = 1;
+                    ctx.beginPath();
+                    ctx.moveTo(x, y);
+                    ctx.lineTo(x - size, y + size * 2);
+                    ctx.lineTo(x + size, y + size * 2);
+                    ctx.closePath();
+                    ctx.fill();
+                    ctx.stroke();
+                },
+
+
+                onDrawForeground: function (ctx) {
+
+                    if (this.flags?.collapsed)
+                        return;
+
+                    const imageArea = this.getImageArea();
+                    const barArea = this.getBarArea();
+                    const inputSliderArea = this.getInputSliderArea();
+                    const outputSliderArea = this.getOutputSliderArea();
+
+                    if (!imageArea || !barArea || !inputSliderArea || !outputSliderArea)
+                        return;
+
+                    if (this.imagePreview) {
+                        ctx.drawImage(this.imagePreview, imageArea.x, imageArea.y, imageArea.width, imageArea.height);
+                    }
+                    else {
+                        ctx.save();
+                        ctx.fillStyle = LiteGraph.NODE_TEXT_COLOR || "#CCC";
+                        ctx.textAlign = "center";
+                        ctx.textBaseline = "middle";
+                        ctx.font = "bold 14px Arial";
+                        const message = "Connect inputs and run workflow";
+                        ctx.fillText(message, imageArea.x + imageArea.width / 2, imageArea.y + imageArea.height / 2);
+                        ctx.restore();
+                    }
+
+                    // Draw Input Levels Area
+                    ctx.fillStyle = "#282828";
+                    ctx.fillRect(barArea.x, barArea.y, barArea.width, barArea.height);
+                    ctx.fillStyle = "#202020";
+                    ctx.fillRect(inputSliderArea.x, inputSliderArea.y, inputSliderArea.width, inputSliderArea.height);
+
+                    const histogramData = this.histograms[this.activeChannel];
+
+                    if (histogramData) {
+                        this.drawHistogram(ctx, barArea, histogramData, 1);
+                    }
+
+                    const levels = this.allLevels[this.activeChannel];
+                    if (!levels) return;
+
+                    // Draw Input Handles
+                    const blackX = inputSliderArea.x + levels.black_point * inputSliderArea.width;
+                    const whiteX = inputSliderArea.x + levels.white_point * inputSliderArea.width;
+                    const relativeRangeWidth = whiteX - blackX;
+
+                    const midPointToSlider = (gamma) => {
+                        if (gamma >= 1.0)
+                            return 1.0 - (0.5 + ((gamma - 1.0) / 8.99) * 0.5);
+                        else return 1.0 - (((gamma - 0.01) / 0.99) * 0.5);
+                    };
+
+                    const midHandleX = blackX + (midPointToSlider(levels.mid_point) * relativeRangeWidth);
+                    const handleTipY = inputSliderArea.y + 2.5;
+
+                    this.drawHandle(ctx, blackX, handleTipY, 5, "black", "white");
+                    this.drawHandle(ctx, whiteX, handleTipY, 5, "white", "black");
+                    this.drawHandle(ctx, midHandleX, handleTipY, 5, "#888", "white");
+
+                    // Draw Output Levels Area
+                    const gradient = ctx.createLinearGradient(outputSliderArea.x, 0, outputSliderArea.x + outputSliderArea.width, 0);
+                    const blackOutPos = levels.output_black / 255.0;
+                    const whiteOutPos = levels.output_white / 255.0;
+
+                    gradient.addColorStop(0, "black");
+                    gradient.addColorStop(blackOutPos, "black");
+                    gradient.addColorStop(whiteOutPos, "white");
+                    gradient.addColorStop(1, "white");
+                    ctx.fillStyle = gradient;
+                    ctx.fillRect(outputSliderArea.x, outputSliderArea.y, outputSliderArea.width, outputSliderArea.height);
+
+                    // Draw Output Handles
+                    const outBlackX = outputSliderArea.x + blackOutPos * outputSliderArea.width;
+                    const outWhiteX = outputSliderArea.x + whiteOutPos * outputSliderArea.width;
+                    const outHandleTipY = outputSliderArea.y + 2.5;
+
+                    this.drawHandle(ctx, outBlackX, outHandleTipY, 5, "black", "white");
+                    this.drawHandle(ctx, outWhiteX, outHandleTipY, 5, "white", "black");
+                },
+
+
+                onMouseDown: function (event) {
+
+                    if (event.button !== 0)
+                        return false;
+
+                    const localPos = app.canvas.convertEventToCanvasOffset(event);
+                    const mouseX = localPos[0] - this.pos[0];
+                    const mouseY = localPos[1] - this.pos[1];
+
+                    const levels = this.allLevels[this.activeChannel];
+
+                    if (!levels)
+                        return false;
+
+                    // Check output slider first
+                    const outputSliderArea = this.getOutputSliderArea();
+                    if (outputSliderArea && mouseX > outputSliderArea.x && mouseX < outputSliderArea.x + outputSliderArea.width &&
+                        mouseY > outputSliderArea.y && mouseY < outputSliderArea.y + outputSliderArea.height) {
+
+                        const outBlackX = outputSliderArea.x + (levels.output_black / 255.0) * outputSliderArea.width;
+
+                        if (Math.abs(mouseX - outBlackX) < 10) {
+                            this.dragging = 'output_black'; return true;
+                        }
+
+                        const outWhiteX = outputSliderArea.x + (levels.output_white / 255.0) * outputSliderArea.width;
+
+                        if (Math.abs(mouseX - outWhiteX) < 10) {
+                            this.dragging = 'output_white'; return true;
+                        }
+                    }
+
+
+                    // Check input slider
+                    const inputSliderArea = this.getInputSliderArea();
+                    if (inputSliderArea && mouseX > inputSliderArea.x && mouseX < inputSliderArea.x + inputSliderArea.width &&
+                        mouseY > inputSliderArea.y && mouseY < inputSliderArea.y + inputSliderArea.height) {
+
+                        const blackX = inputSliderArea.x + levels.black_point * inputSliderArea.width;
+                        const whiteX = inputSliderArea.x + levels.white_point * inputSliderArea.width;
+                        const relativeRangeWidth = whiteX - blackX;
+
+                        const midPointToSlider = (gamma) => {
+                            if (gamma >= 1.0) {
+                                return 1.0 - (0.5 + ((gamma - 1.0) / 8.99) * 0.5);
+                            }
+                            else {
+                                return 1.0 - (((gamma - 0.01) / 0.99) * 0.5);
+                            }
+                        };
+
+                        const midHandleX = blackX + (midPointToSlider(levels.mid_point) * relativeRangeWidth);
+
+                        if (Math.abs(mouseX - midHandleX) < 10) {
+                            this.dragging = 'mid';
+                        }
+                        else if (Math.abs(mouseX - blackX) < 10) {
+                            this.dragging = 'black';
+                        }
+                        else if (Math.abs(mouseX - whiteX) < 10) {
+                            this.dragging = 'white';
+                        }
+                    }
+
+                    return this.dragging !== null;
+                },
+
+                onMouseMove: function (event) {
+                    if (event.buttons !== 1) {
+
+                        if (this.dragging) {
+                            this.dragging = null; this.sendLevelUpdateToPython();
+                        }
+                        return;
+                    }
+
+                    if (!this.dragging)
+                        return;
+
+                    const localPos = app.canvas.convertEventToCanvasOffset(event);
+                    const mouseX = localPos[0] - this.pos[0];
+                    const levels = this.allLevels[this.activeChannel];
+
+                    if (!levels)
+                        return;
+
+                    let needsUpdate = false;
+
+                    // Handle input slider dragging
+                    if (this.dragging === 'black' || this.dragging === 'white' || this.dragging === 'mid') {
+                        const inputSliderArea = this.getInputSliderArea();
+
+                        if (!inputSliderArea)
+                            return;
+
+                        if (this.dragging === 'black' || this.dragging === 'white') {
+                            let newLevel = (mouseX - inputSliderArea.x) / inputSliderArea.width;
+                            newLevel = Math.max(0.0, Math.min(1.0, newLevel));
+
+                            if (this.dragging === 'black')
+                                levels.black_point = Math.min(newLevel, levels.white_point - 0.01);
+                            else
+                                levels.white_point = Math.max(newLevel, levels.black_point + 0.01);
+                        }
+                        else if (this.dragging === 'mid') {
+                            const blackX = inputSliderArea.x + levels.black_point * inputSliderArea.width;
+                            const whiteX = inputSliderArea.x + levels.white_point * inputSliderArea.width;
+                            const relativeRangeWidth = whiteX - blackX;
+
+                            if (relativeRangeWidth < 1)
+                                return;
+
+                            let relativePos = (mouseX - blackX) / relativeRangeWidth;
+                            relativePos = Math.max(0.0, Math.min(1.0, relativePos));
+                            let invertedPos = 1.0 - relativePos;
+
+                            if (invertedPos >= 0.5) {
+                                levels.mid_point = 1.0 + ((invertedPos - 0.5) / 0.5) * (9.99 - 1.0);
+                            }
+                            else {
+                                levels.mid_point = 0.01 + (invertedPos / 0.5) * (1.0 - 0.01);
+                            }
+
+                            levels.mid_point = parseFloat(levels.mid_point.toFixed(2));
+                        }
+                        needsUpdate = true;
+                    }
+                    // Handle output slider dragging
+                    else if (this.dragging === 'output_black' || this.dragging === 'output_white') {
+                        const outputSliderArea = this.getOutputSliderArea();
+
+                        if (!outputSliderArea)
+                            return;
+
+                        let newLevel = ((mouseX - outputSliderArea.x) / outputSliderArea.width) * 255;
+                        newLevel = Math.max(0, Math.min(255, newLevel));
+
+                        if (this.dragging === 'output_black') {
+                            levels.output_black = Math.min(newLevel, levels.output_white - 1);
+                        }
+                        else {
+                            levels.output_white = Math.max(newLevel, levels.output_black + 1);
+                        }
+                        needsUpdate = true;
+                    }
+
+
+                    if (needsUpdate) {
+                        updateWidgetsFromState();
+                        this.setDirtyCanvas(true, true);
+                        clearTimeout(this.updateTimeout);
+                        this.updateTimeout = setTimeout(() => this.sendLevelUpdateToPython(), 50);
+                    }
+                },
+
+                onMouseUp: function (event) {
+                    if (this.dragging) {
+                        this.dragging = null;
+                        this.sendLevelUpdateToPython();
+                    }
+                }
+            });
+
+            const methodsToBind = ['getImageArea', 'getBarArea', 'getInputSliderArea', 'getOutputSliderArea', 'updateHistogram', 'sendLevelUpdateToPython', 'onDrawForeground', 'onMouseDown', 'onMouseMove', 'onMouseUp', 'drawHistogram', 'drawHandle'];
+            methodsToBind.forEach(m => { if (node[m]) node[m] = node[m].bind(node); });
+        }
+    },
+});
+
+
+
+// Event listeners -----------
+
+api.addEventListener("eses.levels_preview_base", ({ detail }) => {
+    const node = app.graph.getNodeById(detail.node_id);
+
+    if (node) {
+        node.isManuallyResized = false;
+
+        // Check for both auto_levels_updated 
+        // and auto_color_updated
+        if (detail.updated_levels) {
+
+            // Runs for auto-levels or auto-color.
+            // Update all widgets.
+            node.allLevels = detail.updated_levels;
+
+            const channelWidget = node.widgets.find(w => w.name === "channel");
+            const channelMap = { "RGB": "rgb", "Red": "r", "Green": "g", "Blue": "b", "Mask": "mask" };
+
+            node.activeChannel = channelMap[channelWidget.value] || "rgb";
+
+            // Refresh the widgets for the currently
+            // selected channel to show the newly reset
+            // output levels.
+            const blackWidget = node.widgets.find(w => w.name === "Black");
+            const midWidget = node.widgets.find(w => w.name === "Mid");
+            const whiteWidget = node.widgets.find(w => w.name === "White");
+            const outputBlackWidget = node.widgets.find(w => w.name === "Output Black");
+            const outputWhiteWidget = node.widgets.find(w => w.name === "Output White");
+
+            const levels = node.allLevels[node.activeChannel];
+
+            if (levels) {
+                if (blackWidget)
+                    blackWidget.value = parseFloat(levels.black_point.toFixed(2));
+
+                if (midWidget)
+                    midWidget.value = parseFloat(levels.mid_point.toFixed(2));
+
+                if (whiteWidget)
+                    whiteWidget.value = parseFloat(levels.white_point.toFixed(2));
+
+                if (outputBlackWidget)
+                    outputBlackWidget.value = Math.round(levels.output_black);
+
+                if (outputWhiteWidget)
+                    outputWhiteWidget.value = Math.round(levels.output_white);
+            }
+        }
+
+        const onAssetLoad = async () => {
+            if (!node.isManuallyResized) {
+                const sizingObj = node.originalImage || node.originalMask;
+
+                if (sizingObj) {
+                    const aspectRatio = sizingObj.naturalWidth / sizingObj.naturalHeight;
+                    const imageDrawWidth = node.size[0] - PADDING * 2;
+                    const imageDrawHeight = imageDrawWidth / aspectRatio;
+                    const interactiveHeight = HISTOGRAM_AREA_HEIGHT + OUTPUT_SLIDER_AREA_HEIGHT + AREA_SPACING;
+
+                    node.size[1] = imageDrawHeight + interactiveHeight + HEADER_HEIGHT + PADDING * 2;
+                    node.isManuallyResized = true; // Set to true after auto-resize
+                }
+            }
+            await node.updateHistogram();
+            await node.sendLevelUpdateToPython();
+        };
+
+        let assetsToLoad = 0;
+
+        if (detail.image_data)
+            assetsToLoad++;
+
+        if (detail.mask_data)
+            assetsToLoad++;
+
+        if (assetsToLoad === 0) {
+            onAssetLoad(); return;
+        }
+
+        let loadedCount = 0;
+        const assetLoadedCallback = async () => {
+            if (++loadedCount === assetsToLoad) {
+                await onAssetLoad();
+            }
+        };
+
+        if (detail.image_data) {
+            const img = new Image();
+            img.src = "data:image/png;base64," + detail.image_data;
+            img.onload = assetLoadedCallback;
+            node.originalImage = img;
+        }
+        else {
+            node.originalImage = null;
+        }
+
+        if (detail.mask_data) {
+            const mask = new Image();
+            mask.src = "data:image/png;base64," + detail.mask_data;
+            mask.onload = assetLoadedCallback;
+            node.originalMask = mask;
+        }
+        else {
+            node.originalMask = null;
+        }
+    }
+});

custom_nodes/NeuralISP/__init__.py

@@ -0,0 +1,10 @@
+from .neural_isp_node import (
+    NODE_CLASS_MAPPINGS,
+    NODE_DISPLAY_NAME_MAPPINGS
+)
+
+
+__all__ = [
+    "NODE_CLASS_MAPPINGS",
+    "NODE_DISPLAY_NAME_MAPPINGS"
+]

custom_nodes/NeuralISP/mirnetv2_model.py

@@ -0,0 +1,301 @@
+
+## Learning Enriched Features for Fast Image Restoration and Enhancement
+## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan, Ming-Hsuan Yang, and Ling Shao
+## IEEE Transactions on Pattern Analysis and Machine Intelligence (TPAMI)
+## https://www.waqaszamir.com/publication/zamir-2022-mirnetv2/
+
+# --- Imports --- #
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from pdb import set_trace as stx
+import math
+
+##########################################################################
+##---------- Selective Kernel Feature Fusion (SKFF) ----------
+class SKFF(nn.Module):
+    def __init__(self, in_channels, height=3,reduction=8,bias=False):
+        super(SKFF, self).__init__()
+        
+        self.height = height
+        d = max(int(in_channels/reduction),4)
+        
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.conv_du = nn.Sequential(nn.Conv2d(in_channels, d, 1, padding=0, bias=bias), nn.LeakyReLU(0.2))
+
+        self.fcs = nn.ModuleList([])
+        for i in range(self.height):
+            self.fcs.append(nn.Conv2d(d, in_channels, kernel_size=1, stride=1,bias=bias))
+        
+        self.softmax = nn.Softmax(dim=1)
+
+    def forward(self, inp_feats):
+        batch_size = inp_feats[0].shape[0]
+        n_feats =  inp_feats[0].shape[1]
+        
+
+        inp_feats = torch.cat(inp_feats, dim=1)
+        inp_feats = inp_feats.view(batch_size, self.height, n_feats, inp_feats.shape[2], inp_feats.shape[3])
+        
+        feats_U = torch.sum(inp_feats, dim=1)
+        feats_S = self.avg_pool(feats_U)
+        feats_Z = self.conv_du(feats_S)
+
+        attention_vectors = [fc(feats_Z) for fc in self.fcs]
+        attention_vectors = torch.cat(attention_vectors, dim=1)
+        attention_vectors = attention_vectors.view(batch_size, self.height, n_feats, 1, 1)
+        # stx()
+        attention_vectors = self.softmax(attention_vectors)
+        
+        feats_V = torch.sum(inp_feats*attention_vectors, dim=1)
+        
+        return feats_V        
+
+
+
+class ContextBlock(nn.Module):
+
+    def __init__(self, n_feat, bias=False):
+        super(ContextBlock, self).__init__()
+
+        self.conv_mask = nn.Conv2d(n_feat, 1, kernel_size=1, bias=bias)
+        self.softmax = nn.Softmax(dim=2)
+
+        self.channel_add_conv = nn.Sequential(
+            nn.Conv2d(n_feat, n_feat, kernel_size=1, bias=bias),
+            nn.LeakyReLU(0.2),
+            nn.Conv2d(n_feat, n_feat, kernel_size=1, bias=bias)
+        )
+
+    def modeling(self, x):
+        batch, channel, height, width = x.size()
+        input_x = x
+        # [N, C, H * W]
+        input_x = input_x.view(batch, channel, height * width)
+        # [N, 1, C, H * W]
+        input_x = input_x.unsqueeze(1)
+        # [N, 1, H, W]
+        context_mask = self.conv_mask(x)
+        # [N, 1, H * W]
+        context_mask = context_mask.view(batch, 1, height * width)
+        # [N, 1, H * W]
+        context_mask = self.softmax(context_mask)
+        # [N, 1, H * W, 1]
+        context_mask = context_mask.unsqueeze(3)
+        # [N, 1, C, 1]
+        context = torch.matmul(input_x, context_mask)
+        # [N, C, 1, 1]
+        context = context.view(batch, channel, 1, 1)
+
+        return context
+
+    def forward(self, x):
+        # [N, C, 1, 1]
+        context = self.modeling(x)
+
+        # [N, C, 1, 1]
+        channel_add_term = self.channel_add_conv(context)
+        x = x + channel_add_term
+
+        return x
+
+##########################################################################
+### --------- Residual Context Block (RCB) ----------
+class RCB(nn.Module):
+    def __init__(self, n_feat, kernel_size=3, reduction=8, bias=False, groups=1):
+        super(RCB, self).__init__()
+        
+        act = nn.LeakyReLU(0.2)
+
+        self.body = nn.Sequential(
+            nn.Conv2d(n_feat, n_feat, kernel_size=3, stride=1, padding=1, bias=bias, groups=groups),
+            act, 
+            nn.Conv2d(n_feat, n_feat, kernel_size=3, stride=1, padding=1, bias=bias, groups=groups)
+        )
+
+        self.act = act
+        
+        self.gcnet = ContextBlock(n_feat, bias=bias)
+
+    def forward(self, x):
+        res = self.body(x)
+        res = self.act(self.gcnet(res))
+        res += x
+        return res
+
+
+##########################################################################
+##---------- Resizing Modules ----------    
+class Down(nn.Module):
+    def __init__(self, in_channels, chan_factor, bias=False):
+        super(Down, self).__init__()
+
+        self.bot = nn.Sequential(
+            nn.AvgPool2d(2, ceil_mode=True, count_include_pad=False),
+            nn.Conv2d(in_channels, int(in_channels*chan_factor), 1, stride=1, padding=0, bias=bias)
+            )
+
+    def forward(self, x):
+        return self.bot(x)
+
+class DownSample(nn.Module):
+    def __init__(self, in_channels, scale_factor, chan_factor=2, kernel_size=3):
+        super(DownSample, self).__init__()
+        self.scale_factor = int(np.log2(scale_factor))
+
+        modules_body = []
+        for i in range(self.scale_factor):
+            modules_body.append(Down(in_channels, chan_factor))
+            in_channels = int(in_channels * chan_factor)
+        
+        self.body = nn.Sequential(*modules_body)
+
+    def forward(self, x):
+        x = self.body(x)
+        return x
+
+class Up(nn.Module):
+    def __init__(self, in_channels, chan_factor, bias=False):
+        super(Up, self).__init__()
+
+        self.bot = nn.Sequential(
+            nn.Conv2d(in_channels, int(in_channels//chan_factor), 1, stride=1, padding=0, bias=bias),
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=bias)
+            )
+
+    def forward(self, x):
+        return self.bot(x)
+
+class UpSample(nn.Module):
+    def __init__(self, in_channels, scale_factor, chan_factor=2, kernel_size=3):
+        super(UpSample, self).__init__()
+        self.scale_factor = int(np.log2(scale_factor))
+
+        modules_body = []
+        for i in range(self.scale_factor):
+            modules_body.append(Up(in_channels, chan_factor))
+            in_channels = int(in_channels // chan_factor)
+        
+        self.body = nn.Sequential(*modules_body)
+
+    def forward(self, x):
+        x = self.body(x)
+        return x
+
+
+##########################################################################
+##---------- Multi-Scale Resiudal Block (MRB) ----------
+class MRB(nn.Module):
+    def __init__(self, n_feat, height, width, chan_factor, bias,groups):
+        super(MRB, self).__init__()
+
+        self.n_feat, self.height, self.width = n_feat, height, width
+
+        self.dau_top = RCB(int(n_feat*chan_factor**0), bias=bias, groups=groups)
+        self.dau_mid = RCB(int(n_feat*chan_factor**1), bias=bias, groups=groups)
+        self.dau_bot = RCB(int(n_feat*chan_factor**2), bias=bias, groups=groups)
+
+        self.down2 = DownSample(int((chan_factor**0)*n_feat),2,chan_factor)
+        self.down4 = nn.Sequential(
+            DownSample(int((chan_factor**0)*n_feat),2,chan_factor), 
+            DownSample(int((chan_factor**1)*n_feat),2,chan_factor)
+        )
+
+        self.up21_1 = UpSample(int((chan_factor**1)*n_feat),2,chan_factor)
+        self.up21_2 = UpSample(int((chan_factor**1)*n_feat),2,chan_factor)
+        self.up32_1 = UpSample(int((chan_factor**2)*n_feat),2,chan_factor)
+        self.up32_2 = UpSample(int((chan_factor**2)*n_feat),2,chan_factor)
+
+        self.conv_out = nn.Conv2d(n_feat, n_feat, kernel_size=1, padding=0, bias=bias)
+
+        # only two inputs for SKFF
+        self.skff_top = SKFF(int(n_feat*chan_factor**0), 2)
+        self.skff_mid = SKFF(int(n_feat*chan_factor**1), 2)
+
+    def forward(self, x):
+        x_top = x.clone()
+        x_mid = self.down2(x_top)
+        x_bot = self.down4(x_top)
+
+        x_top = self.dau_top(x_top)
+        x_mid = self.dau_mid(x_mid)
+        x_bot = self.dau_bot(x_bot)
+
+        x_mid = self.skff_mid([x_mid, self.up32_1(x_bot)])
+        x_top = self.skff_top([x_top, self.up21_1(x_mid)])
+
+        x_top = self.dau_top(x_top)
+        x_mid = self.dau_mid(x_mid)
+        x_bot = self.dau_bot(x_bot)
+
+        x_mid = self.skff_mid([x_mid, self.up32_2(x_bot)])
+        x_top = self.skff_top([x_top, self.up21_2(x_mid)])
+
+        out = self.conv_out(x_top)
+        out = out + x
+
+        return out
+
+##########################################################################
+##---------- Recursive Residual Group (RRG) ----------
+class RRG(nn.Module):
+    def __init__(self, n_feat, n_MRB, height, width, chan_factor, bias=False, groups=1):
+        super(RRG, self).__init__()
+        modules_body = [MRB(n_feat, height, width, chan_factor, bias, groups) for _ in range(n_MRB)]
+        modules_body.append(nn.Conv2d(n_feat, n_feat, kernel_size=3, stride=1, padding=1, bias=bias))
+        self.body = nn.Sequential(*modules_body)
+
+    def forward(self, x):
+        res = self.body(x)
+        res += x
+        return res
+
+
+##########################################################################
+##---------- MIRNet  -----------------------
+class MIRNet_v2(nn.Module):
+    def __init__(self,
+        inp_channels=3,
+        out_channels=3,
+        n_feat=80,
+        chan_factor=1.5,
+        n_RRG=4,
+        n_MRB=2,
+        height=3,
+        width=2,
+        scale=1,
+        bias=False,
+        task= None
+    ):
+        super(MIRNet_v2, self).__init__()
+
+        kernel_size=3
+        self.task = task
+
+        self.conv_in = nn.Conv2d(inp_channels, n_feat, kernel_size=3, padding=1, bias=bias)
+
+        modules_body = []
+        
+        modules_body.append(RRG(n_feat, n_MRB, height, width, chan_factor, bias, groups=1))
+        modules_body.append(RRG(n_feat, n_MRB, height, width, chan_factor, bias, groups=2))
+        modules_body.append(RRG(n_feat, n_MRB, height, width, chan_factor, bias, groups=4))
+        modules_body.append(RRG(n_feat, n_MRB, height, width, chan_factor, bias, groups=4))
+
+        self.body = nn.Sequential(*modules_body)
+        self.conv_out = nn.Conv2d(n_feat, out_channels, kernel_size=3, padding=1, bias=bias)
+        
+
+    def forward(self, inp_img):
+        shallow_feats = self.conv_in(inp_img)
+        deep_feats = self.body(shallow_feats)
+
+        if self.task == 'defocus_deblurring':
+            deep_feats += shallow_feats
+            out_img = self.conv_out(deep_feats)
+
+        else:
+            out_img = self.conv_out(deep_feats)
+            out_img += inp_img
+
+        return out_img

custom_nodes/NeuralISP/neural_isp_node.py

@@ -0,0 +1,407 @@
+import os
+import torch
+import torch.nn.functional as F
+import folder_paths
+
+from .mirnetv2_model import MIRNet_v2
+
+
+MODEL_DIR = os.path.join(
+    folder_paths.models_dir,
+    "mirnetv2"
+)
+
+
+class NeuralISPNode:
+
+
+    @classmethod
+    def INPUT_TYPES(cls):
+
+        models = []
+
+        if os.path.exists(MODEL_DIR):
+
+            for f in os.listdir(MODEL_DIR):
+                if f.endswith((".pth", ".pt")):
+                    models.append(f)
+
+
+        if not models:
+            models = [
+                "enhancement_fivek.pth",
+                "enhancement_lol.pth",
+                "real_denoising.pth"
+            ]
+
+
+        return {
+            "required": {
+
+                "image":
+                ("IMAGE",),
+
+                "model":
+                (models,),
+
+                "strength":
+                (
+                    "FLOAT",
+                    {
+                        "default":0.2,
+                        "min":0.0,
+                        "max":1.0,
+                        "step":0.05
+                    }
+                ),
+
+                "tile_size":
+                (
+                    "INT",
+                    {
+                        "default":512,
+                        "min":128,
+                        "max":2048,
+                        "step":64
+                    }
+                ),
+
+                "overlap":
+                (
+                    "INT",
+                    {
+                        "default":128,
+                        "min":32,
+                        "max":512,
+                        "step":32
+                    }
+                )
+            }
+        }
+
+
+    RETURN_TYPES = ("IMAGE",)
+
+    FUNCTION = "apply"
+
+    CATEGORY = "📸 shots by Faded/Neural Post FX"
+
+
+
+    def load_model(self, filename):
+
+        device = (
+            "cuda"
+            if torch.cuda.is_available()
+            else "cpu"
+        )
+
+
+        model = MIRNet_v2()
+
+
+        path = os.path.join(
+            MODEL_DIR,
+            filename
+        )
+
+
+        checkpoint = torch.load(
+            path,
+            map_location="cpu"
+        )
+
+
+        if "params" in checkpoint:
+            checkpoint = checkpoint["params"]
+
+        elif "state_dict" in checkpoint:
+            checkpoint = checkpoint["state_dict"]
+
+        elif "model" in checkpoint:
+            checkpoint = checkpoint["model"]
+
+
+
+        cleaned = {}
+
+        for k,v in checkpoint.items():
+
+            if k.startswith("module."):
+                k = k[7:]
+
+            cleaned[k] = v
+
+
+
+        model.load_state_dict(
+            cleaned,
+            strict=False
+        )
+
+
+        model.eval()
+
+        model.to(device)
+
+        # MIRNet лучше держать FP32
+        model.float()
+
+
+        return model, device
+
+
+
+    def create_mask(
+        self,
+        h,
+        w,
+        device
+    ):
+
+        mask = torch.ones(
+            1,
+            1,
+            h,
+            w,
+            device=device
+        )
+
+
+        return mask
+
+
+
+    def process_tile(
+        self,
+        model,
+        img,
+        tile_size,
+        overlap
+    ):
+
+
+        _,_,h,w = img.shape
+
+
+        stride = tile_size - overlap
+
+
+        output = torch.zeros_like(
+            img,
+            dtype=torch.float32
+        )
+
+
+        weight = torch.zeros_like(
+            img,
+            dtype=torch.float32
+        )
+
+
+
+        for y in range(0,h,stride):
+
+            for x in range(0,w,stride):
+
+
+                y1 = min(
+                    y+tile_size,
+                    h
+                )
+
+                x1 = min(
+                    x+tile_size,
+                    w
+                )
+
+
+                y0 = max(
+                    0,
+                    y1-tile_size
+                )
+
+                x0 = max(
+                    0,
+                    x1-tile_size
+                )
+
+
+                patch = img[
+                    :,
+                    :,
+                    y0:y1,
+                    x0:x1
+                ]
+
+
+                ph = patch.shape[-2]
+                pw = patch.shape[-1]
+
+
+                # MIRNet требует размеры кратные 32
+
+                pad_h = (
+                    32 - ph % 32
+                ) % 32
+
+
+                pad_w = (
+                    32 - pw % 32
+                ) % 32
+
+
+
+                patch_pad = F.pad(
+                    patch,
+                    (
+                        0,
+                        pad_w,
+                        0,
+                        pad_h
+                    ),
+                    mode="reflect"
+                )
+
+
+
+                with torch.no_grad():
+
+                    result = model(
+                        patch_pad.float()
+                    )
+
+
+
+                # возвращаем исходный размер
+
+                result = result[
+                    :,
+                    :,
+                    :ph,
+                    :pw
+                ]
+
+
+
+                mask = self.create_mask(
+                    ph,
+                    pw,
+                    img.device
+                )
+
+
+
+                output[
+                    :,
+                    :,
+                    y0:y1,
+                    x0:x1
+                ] += (
+                    result *
+                    mask
+                )
+
+
+
+                weight[
+                    :,
+                    :,
+                    y0:y1,
+                    x0:x1
+                ] += mask
+
+
+
+        return (
+            output /
+            weight.clamp(min=1e-6)
+        )
+
+
+
+    def apply(
+        self,
+        image,
+        model,
+        strength,
+        tile_size,
+        overlap
+    ):
+
+
+        net,device = self.load_model(
+            model
+        )
+
+
+        img = image[0]
+
+
+        img = img.permute(
+            2,
+            0,
+            1
+        ).unsqueeze(0)
+
+
+        img = img.to(device)
+
+
+
+        with torch.no_grad():
+
+            result = self.process_tile(
+                net,
+                img,
+                tile_size,
+                overlap
+            )
+
+
+
+        result = (
+            result * strength
+            +
+            img.float() *
+            (1-strength)
+        )
+
+
+        result = result.clamp(
+            0,
+            1
+        )
+
+
+        result = result.squeeze(0)
+
+        result = result.permute(
+            1,
+            2,
+            0
+        )
+
+
+        return (
+            result.unsqueeze(0)
+            .cpu(),
+        )
+
+
+
+NODE_CLASS_MAPPINGS = {
+
+    "NeuralISP":
+    NeuralISPNode
+
+}
+
+
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+
+    "NeuralISP":
+    "📸Neural ISP (x MIRNet)"
+
+}

custom_nodes/comfyui-faded-branding/__init__.py

@@ -0,0 +1,6 @@
+from .faded_logo_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+# КРИТИЧЕСКИ ВАЖНО: без этой строчки ComfyUI не загрузит ваш JavaScript!
+WEB_DIRECTORY = "./js"
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']

custom_nodes/comfyui-faded-branding/faded_logo_node.py

@@ -0,0 +1,30 @@
+class OceanLogoDisplay:
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        # Пустой словарь означает отсутствие входных пинов
+        return {
+            "required": {}, 
+        }
+
+    # Пустой кортеж гарантирует, что нода не выдает никаких данных в граф
+    RETURN_TYPES = ()
+    FUNCTION = "display_logo"
+    
+    CATEGORY = "📸 shots by Faded/Utils"
+    
+
+    def display_logo(self):
+        # Эта функция теперь никогда не будет вызвана при инференсе,
+        # так как нода не участвует в графе выполнения.
+        return {}
+
+NODE_CLASS_MAPPINGS = {
+    "OceanLogoDisplay": OceanLogoDisplay
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "OceanLogoDisplay": "🖼️ Faded Branding Banner"
+}

custom_nodes/comfyui-faded-branding/js/faded_logo.js

@@ -0,0 +1,117 @@
+import { app } from "../../../scripts/app.js";
+
+const original_drawNode = LGraphCanvas.prototype.drawNode;
+LGraphCanvas.prototype.drawNode = function(node, ctx) {
+    if (node && node.type === "OceanLogoDisplay") {
+        if (node.flags.collapsed) return;
+        node.title_height = 0;
+
+        ctx.save();
+        const w = node.size[0];
+        const h = node.size[1];
+
+        if (this.selected_nodes && this.selected_nodes[node.id]) {
+            ctx.lineWidth = 3;
+            ctx.strokeStyle = this.node_colors?.selected || "#ff9000";
+            ctx.beginPath();
+            const m = 3;
+            if (ctx.roundRect) ctx.roundRect(-m, -m, w + m * 2, h + m * 2, 12);
+            else ctx.rect(-m, -m, w + m * 2, h + m * 2);
+            ctx.stroke();
+        }
+
+        if (node.bgImage) {
+            const imgW = node.bgImage.width;
+            const imgH = node.bgImage.height;
+            const imgRatio = imgW / imgH;
+            const nodeRatio = w / h;
+            let drawW, drawH;
+            if (imgRatio > nodeRatio) { 
+                drawW = w; 
+                drawH = w / imgRatio; 
+            } else { 
+                drawH = h; 
+                drawW = h * imgRatio; 
+            }
+            const offsetX = (w - drawW) / 2;
+            const offsetY = (h - drawH) / 2;
+            ctx.drawImage(node.bgImage, offsetX, offsetY, drawW, drawH);
+        }
+        ctx.restore();
+        return;
+    }
+    return original_drawNode.apply(this, arguments);
+};
+
+app.registerExtension({
+    name: "Faded.BrandingBanner",
+    async beforeRegisterNodeDef(nodeType, nodeData, app) {
+        if (nodeData.name === "OceanLogoDisplay") {
+            
+            Object.defineProperty(nodeType.prototype, "title", {
+                get: function() { return ""; },
+                set: function(v) {},
+                enumerable: true,
+                configurable: true
+            });
+
+            nodeType.prototype.loadLocalImage = function(filename) {
+                if (!filename) return; 
+                
+                const imgUrl = `/extensions/comfyui-faded-branding/logo/${filename}?t=${Date.now()}`;
+                const img = new Image();
+                
+                img.onload = () => {
+                    this.bgImage = img;
+                    this.setDirtyCanvas(true, true);
+                    console.log(`[Faded] Загружена картинка: ${filename}`);
+                };
+                
+                img.onerror = (e) => {
+                    console.error(`[Faded] Ошибка загрузки: ${filename}`, e);
+                    // Fallback на logo.png
+                    if (filename !== "logo.png") {
+                        console.warn(`[Faded] Пробуем загрузить logo.png`);
+                        this.loadLocalImage("logo.png");
+                    }
+                };
+                
+                img.src = imgUrl;
+            };
+
+            const onNodeCreated = nodeType.prototype.onNodeCreated;
+            nodeType.prototype.onNodeCreated = function() {
+                const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined;
+                
+                this.addProperty("Image to show", "logo.png", "string");
+                this.bgImage = null;
+                this.size = [320, 180];
+                this.title_height = 0;
+
+                // Принудительно загружаем картинку при создании ноды
+                setTimeout(() => {
+                    this.loadLocalImage(this.properties["Image to show"] || "logo.png");
+                }, 100);
+
+                return r;
+            };
+
+            // Восстановление при загрузке воркфлоу
+            nodeType.prototype.onConfigure = function(info) {
+                if (info.properties && info.properties["Image to show"]) {
+                    // Небольшая задержка для надёжности
+                    setTimeout(() => {
+                        this.loadLocalImage(info.properties["Image to show"]);
+                    }, 200);
+                }
+            };
+
+            nodeType.prototype.onPropertyChanged = function(property, value) {
+                if (property === "Image to show") {
+                    this.loadLocalImage(value);
+                }
+                return true;
+            };
+        }
+    }
+});

custom_nodes/image_routing.py

@@ -0,0 +1,38 @@
+class ImageRouting:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {},
+            "optional": {
+                "image1": ("IMAGE",),
+                "image2": ("IMAGE",),
+                "image3": ("IMAGE",),
+                "image4": ("IMAGE",),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = "route"
+    CATEGORY = "📸 shots by Faded/Utils"
+
+    def route(self, image1=None, image2=None, image3=None, image4=None):
+        if image1 is not None:
+            return (image1,)
+        if image2 is not None:
+            return (image2,)
+        if image3 is not None:
+            return (image3,)
+        if image4 is not None:
+            return (image4,)
+
+        return (image1,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "ImageRouting": ImageRouting
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ImageRouting": "Image Routing"
+}

custom_nodes/latent_routing.py

@@ -0,0 +1,87 @@
+import torch
+
+# --- 1. АВТОМАТИЧЕСКАЯ НОДА (СТАРАЯ) ---
+class LatentRouting:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {},
+            "optional": {
+                "samples1": ("LATENT",),
+                "samples2": ("LATENT",),
+                "samples3": ("LATENT",),
+                "samples4": ("LATENT",),
+            }
+        }
+
+    RETURN_TYPES = ("LATENT",)
+    RETURN_NAMES = ("samples",)
+    FUNCTION = "route"
+    CATEGORY = "utils"
+
+    def route(self, samples1=None, samples2=None, samples3=None, samples4=None):
+        if samples1 is not None:
+            return (samples1,)
+        if samples2 is not None:
+            return (samples2,)
+        if samples3 is not None:
+            return (samples3,)
+        if samples4 is not None:
+            return (samples4,)
+
+        return (None,)
+
+
+# --- 2. НОДА С РУЧНЫМ ПЕРЕКЛЮЧАТЕЛЕМ (НОВАЯ) ---
+class LatentSwitchManual:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                # Выпадающий список для точного выбора входа
+                "select": (["1", "2", "3", "4"],),
+            },
+            "optional": {
+                "samples1": ("LATENT",),
+                "samples2": ("LATENT",),
+                "samples3": ("LATENT",),
+                "samples4": ("LATENT",),
+            }
+        }
+
+    RETURN_TYPES = ("LATENT",)
+    RETURN_NAMES = ("samples",)
+    FUNCTION = "switch"
+    CATEGORY = "📸 shots by Faded/Utils"
+
+    def switch(self, select, samples1=None, samples2=None, samples3=None, samples4=None):
+        # Строгая проверка по выбранному пункту меню
+        if select == "1" and samples1 is not None:
+            return (samples1,)
+        if select == "2" and samples2 is not None:
+            return (samples2,)
+        if select == "3" and samples3 is not None:
+            return (samples3,)
+        if select == "4" and samples4 is not None:
+            return (samples4,)
+
+        # Если выбранный вход пустой (например, отключен байпасом), 
+        # берем первый рабочий, чтобы не выдавать ошибку
+        if samples1 is not None: return (samples1,)
+        if samples2 is not None: return (samples2,)
+        if samples3 is not None: return (samples3,)
+        if samples4 is not None: return (samples4,)
+        
+        return (None,)
+
+
+# --- 3. РЕГИСТРАЦИЯ НОД В COMFYUI ---
+NODE_CLASS_MAPPINGS = {
+    "LatentRouting": LatentRouting,
+    "LatentSwitchManual": LatentSwitchManual
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "LatentRouting": "Latent Routing",
+    "LatentSwitchManual": "Latent Routing Manual"
+}

custom_nodes/latent_trick.py

@@ -0,0 +1,109 @@
+import torch
+import torch.fft as fft
+
+
+# ─────────────────────────────────────────────
+# Core (1:1 логика оригинала)
+# ─────────────────────────────────────────────
+
+def create_frequency_masks(height, width, device, low_cutoff, high_cutoff):
+    freq_y = torch.fft.fftfreq(height, device=device)
+    freq_x = torch.fft.fftfreq(width, device=device)
+    fy, fx = torch.meshgrid(freq_y, freq_x, indexing="ij")
+    freq_dist = torch.sqrt(fy ** 2 + fx ** 2)
+
+    max_freq = freq_dist.max()
+    if max_freq > 0:
+        freq_dist = freq_dist / max_freq
+
+    sharpness = 20.0
+
+    low_mask  = torch.sigmoid(sharpness * (low_cutoff  - freq_dist))
+    high_mask = torch.sigmoid(sharpness * (freq_dist   - high_cutoff))
+    mid_mask  = torch.clamp(1.0 - low_mask - high_mask, min=0.0)
+
+    return low_mask, mid_mask, high_mask
+
+
+def _expand_mask(mask, target):
+    while mask.dim() < target.dim():
+        mask = mask.unsqueeze(0)
+    return mask.expand_as(target)
+
+
+def spectral_latent_enhance(latent, detail_boost, denoise_high, low_cutoff, high_cutoff):
+    h, w = latent.shape[-2:]
+
+    if h <= 4 or w <= 4:
+        return latent
+
+    freq = fft.fft2(latent)
+
+    low_m, mid_m, high_m = create_frequency_masks(
+        h, w, latent.device, low_cutoff, high_cutoff
+    )
+
+    low_m  = _expand_mask(low_m,  freq)
+    mid_m  = _expand_mask(mid_m,  freq)
+    high_m = _expand_mask(high_m, freq)
+
+    gain = (
+        low_m * 1.0 +
+        mid_m * (1.0 + detail_boost) +
+        high_m * (1.0 - denoise_high)
+    )
+
+    return fft.ifft2(freq * gain).real
+
+
+# ─────────────────────────────────────────────
+# Node (как в оригинале, но без UI параметров)
+# ─────────────────────────────────────────────
+
+class LatentTrickNode:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "latent": ("LATENT",),
+            }
+        }
+
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "apply"
+    CATEGORY = "📸 shots by Faded/Utils"
+
+    def apply(self, latent):
+
+        # 🔒 фиксированные значения (как ты просил)
+        detail_boost = 0.10
+        denoise_high = 0.00
+        low_cutoff   = 0.40
+        high_cutoff  = 0.73
+
+        out = latent.copy()
+        samples = out["samples"]
+
+        out["samples"] = spectral_latent_enhance(
+            samples,
+            detail_boost,
+            denoise_high,
+            low_cutoff,
+            high_cutoff,
+        )
+
+        return (out,)
+
+
+# ─────────────────────────────────────────────
+# Registration
+# ─────────────────────────────────────────────
+
+NODE_CLASS_MAPPINGS = {
+    "latent_trick": LatentTrickNode,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "latent_trick": "Latent Trick",
+}

custom_nodes/save_image_jpeg.py

@@ -0,0 +1,124 @@
+import os
+import re
+import numpy as np
+from datetime import datetime
+from PIL import Image, ImageCms
+import folder_paths
+
+# Standard Apple ImageIO DQT Quantization Tables for JPEG (Luminance & Chrominance)
+APPLE_DQT_LUMA = [
+    2, 1, 1, 2, 2, 4, 5, 6,
+    1, 1, 1, 2, 3, 6, 6, 6,
+    1, 1, 2, 2, 4, 6, 7, 6,
+    1, 2, 2, 3, 5, 9, 8, 7,
+    2, 2, 4, 6, 7, 11, 10, 8,
+    2, 3, 5, 6, 8, 10, 11, 9,
+    5, 6, 8, 9, 10, 12, 12, 10,
+    7, 9, 9, 10, 11, 10, 10, 10
+]
+
+APPLE_DQT_CHROMA = [
+    2, 2, 2, 4, 8, 8, 8, 8,
+    2, 2, 3, 5, 8, 8, 8, 8,
+    2, 3, 5, 8, 8, 8, 8, 8,
+    4, 5, 8, 8, 8, 8, 8, 8,
+    8, 8, 8, 8, 8, 8, 8, 8,
+    8, 8, 8, 8, 8, 8, 8, 8,
+    8, 8, 8, 8, 8, 8, 8, 8,
+    8, 8, 8, 8, 8, 8, 8, 8
+]
+
+def process_filename_prefix(prefix):
+    def replace_date(match):
+        fmt_str = match.group(1)
+        replacements = [
+            ("yyyy", "%Y"), ("YY", "%y"), ("yy", "%y"),
+            ("MM", "%m"), ("dd", "%d"),
+            ("HH", "%H"), ("hh", "%H"),
+            ("mm", "%M"), ("ss", "%S")
+        ]
+        for old, new in replacements:
+            fmt_str = fmt_str.replace(old, new)
+        return datetime.now().strftime(fmt_str)
+
+    prefix = re.sub(r"%date:([^%]+)%", replace_date, prefix)
+    return prefix
+
+class SaveImageJPEG:
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type = "output"
+        self.prefix_append = ""
+        
+        # Generate binary sRGB IEC61966-2.1 ICC profile
+        srgb_profile = ImageCms.createProfile("sRGB")
+        self.srgb_icc_bytes = ImageCms.ImageCmsProfile(srgb_profile).tobytes()
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "Images": ("IMAGE", {"tooltip": "Input images tensor to save."}),
+                "Filename Prefix": ("STRING", {"default": "iPhone_JPEG"}),
+                "Quality": ("INT", {"default": 88, "min": 1, "max": 100, "step": 1, "tooltip": "JPEG Quality factor (Apple default is 85-90%)."}),
+                "Use Apple JPEG Tables": ("BOOLEAN", {"default": True, "tooltip": "Use custom Apple DQT quantization tables."}),
+                "Subsampling": (["4:2:0", "4:2:2", "4:4:4"], {"default": "4:2:0", "tooltip": "Chroma subsampling mode (Apple default is 4:2:0)."}),
+            },
+        }
+
+    RETURN_TYPES = ()
+    FUNCTION = "save_images"
+    OUTPUT_NODE = True
+    CATEGORY = "📸 shots by Faded/Utils"
+    DESCRIPTION = "Saves JPEG images matching Apple's compression profile with dynamic date formatting."
+
+    def save_images(self, Images, **kwargs):
+        filename_prefix = kwargs.get("Filename Prefix", "iPhone_JPEG")
+        quality = kwargs.get("Quality", 88)
+        use_apple_jpeg_tables = kwargs.get("Use Apple JPEG Tables", True)
+        subsampling = kwargs.get("Subsampling", "4:2:0")
+        
+        filename_prefix += self.prefix_append
+        filename_prefix = process_filename_prefix(filename_prefix)
+        
+        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
+            filename_prefix, self.output_dir, Images[0].shape[2], Images[0].shape[1]
+        )
+        
+        results = list()
+        qtables_param = [APPLE_DQT_LUMA, APPLE_DQT_CHROMA] if use_apple_jpeg_tables else None
+
+        for (batch_number, image) in enumerate(Images):
+            # Convert PyTorch tensor (float 0..1) to PIL Image (uint8 0..255 RGB)
+            i = 255. * image.cpu().numpy()
+            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+            
+            file = f"{filename}_{counter:05}_.jpg"
+            file_path = os.path.join(full_output_folder, file)
+            
+            img.save(
+                file_path,
+                format="JPEG",
+                quality=quality,
+                subsampling=subsampling,
+                progressive=False,
+                qtables=qtables_param,
+                icc_profile=self.srgb_icc_bytes,
+            )
+            
+            results.append({
+                "filename": file,
+                "subfolder": subfolder,
+                "type": self.type
+            })
+            counter += 1
+
+        return {"ui": {"images": results}}
+
+NODE_CLASS_MAPPINGS = {
+    "SaveImageJPEG": SaveImageJPEG
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "SaveImageJPEG": "Save Image (JPEG)"
+}

custom_nodes/srgb_noise/__init__.py

@@ -0,0 +1,6 @@
+from .srgb_noise_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = [
+    "NODE_CLASS_MAPPINGS",
+    "NODE_DISPLAY_NAME_MAPPINGS"
+]

custom_nodes/srgb_noise/repo/Crop_SIDD.py

@@ -0,0 +1,14 @@
+import os
+import utils.utils_image as util
+
+def main():
+    ori_dir = './Datasets/SIDD_Medium_Srgb_Parts/' # put original SIDD directory here
+    target_dir = './Datasets/SIDD_Medium_Crop/' # put target cropped SIDD directory here
+
+    if not os.path.exists(target_dir):
+        os.mkdir(target_dir)
+    
+    util.crop_sidd_dataset(ori_dir, target_dir, p_size=512, p_overlap=128, p_max=800)
+
+if __name__ == '__main__':
+    main()

custom_nodes/srgb_noise/repo/README.md

@@ -0,0 +1,94 @@
+# sRGB Real Noise Synthesizing with Neighboring Correlation-Aware Noise Model (CVPR'23)
+
+This is the official inplementation of CVPR 2023 paper [sRGB Real Noise Synthesizing with Neighboring Correlation-Aware Noise Model](https://openaccess.thecvf.com/content/CVPR2023/papers/Fu_sRGB_Real_Noise_Synthesizing_With_Neighboring_Correlation-Aware_Noise_Model_CVPR_2023_paper.pdf). 
+
+## Abstract 
+
+Modeling and synthesizing real noise in the standard RGB (sRGB) domain is challenging due to the complicated noise distribution. While most of the deep noise generators
+proposed to synthesize sRGB real noise using an end-to-end trained model, the lack of explicit noise modeling degrades the quality of their synthesized noise. In this work, we propose to model the real noise as not only dependent on the underlying clean image pixel intensity, but also highly correlated to its neighboring noise realization within the local region. Correspondingly, we propose a novel noise synthesizing framework by explicitly learning its neighboring correlation on top of the signal dependency. With the proposed noise model, our framework greatly bridges the distribution gap between synthetic noise and real noise. We show that our generated “real” sRGB noisy images can be used for training supervised deep denoisers, thus to improve their real denoising results with a large margin, comparing to the popular classic denoisers or the deep denoisers that are trained on other sRGB noise generators.
+
+For more details, please refer to our original [paper](https://openaccess.thecvf.com/content/CVPR2023/papers/Fu_sRGB_Real_Noise_Synthesizing_With_Neighboring_Correlation-Aware_Noise_Model_CVPR_2023_paper.pdf).
+
+
+![baseline](./doc/baseline.png)
+
+## Setup
+
+### Requirements
+
+- PyTorch 1.13.1
+- CUDA 11.7
+- Python 3.9
+- opencv 4.7.0
+- numpy 1.21.5
+
+### Directory
+
+Descriptions of code directory.
+
+```
+sRGB-Noise-Synthesis
+|-- data
+|-- data_preparation
+|-- Datasets
+    |-- SIDD_Medium_Srgb_Parts
+    |-- SIDD_Medium_Crop
+|-- models
+|-- options
+|-- pretrain
+|-- saves
+|-- test
+    |-- imgs
+    |-- saves
+|-- utils
+```
+
+### Pretrained models
+
+Download pretrained checkpoints here and place them into `/pretrain/` folder.
+
+Google drive: | Camera name | [IP](https://drive.google.com/drive/folders/1jouUbKfSz9vpvChuY5J-gvWN4JLkYs3H?usp=sharing) | [G4](https://drive.google.com/drive/folders/1oqhax1aESnqIx6q4jVqp_7XuoloGlJuk?usp=sharing) | [GP](https://drive.google.com/drive/folders/1L2lzAMmzUT1flYWdkmTtGnPAWBirGvAr?usp=sharing) | [N6](https://drive.google.com/drive/folders/1I7TEWZ2HnEVjAvE24wwYGBQvElqqdEH4?usp=sharing) | [S6](https://drive.google.com/drive/folders/1HgHNluBdvky5HHB-fvRDelfBRv9gebcT?usp=sharing) |
+
+## Quick test
+
+Test our noise synthesis pipeline with pre-trained models
+
+1. place noisy-clean image pairs into `/test/imgs/` folder. Name noisy image as `noisy.png`, clean image as `clean.png`. 
+2. Run the code as 
+  ```python main_test.py --cam_name IP --NeCA_type S```. You can also control detail test configurations (e.g. cam_name: IP/GP/G4/N6/S6, NecA_type: S/W).
+3. Images are saved in `/test/saves/` folder.
+
+## Train
+
+1. Download SIDD-Medium sRGB part at [here](https://www.eecs.yorku.ca/~kamel/sidd/dataset.php).
+2. Use `Crop_SIDD.py` to crop images into 512×512 patches.
+3. Run the code as
+
+```
+python main_train_jointly.py --dir_save './saves/noise_syn/' --train_img_dir './Datasets/SIDD_Medium_Crop/' --data_prep_dir './data_preparation/' --not_aug --patch_size 96 --cam_name GP --batch_size 16 --num_workers 8 --verbose --epochs 5000 --gpu_ids 0 --joint_start_epoch 50 --seperate_train --gan_mode wgangp --test_epoch 5 --lr_decay_2 80 
+```
+
+## Results
+
+### Visual results
+
+<p align=center><img width="50%" src="./doc/visual_results.png"/></p>
+
+### Denoising performance on SIDD/DND benchmark
+<p align=center><img width="70%" src="./doc/quantitative_results.png"/></p>
+
+## Reference
+
+Our implementation is based on [KAIR](https://github.com/cszn/KAIR), [CycleGAN](https://github.com/junyanz/CycleGAN). We would like to thank them.
+
+## Citation
+```
+@InProceedings{Fu_2023_CVPR,
+    author    = {Fu, Zixuan and Guo, Lanqing and Wen, Bihan},
+    title     = {sRGB Real Noise Synthesizing With Neighboring Correlation-Aware Noise Model},
+    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+    month     = {June},
+    year      = {2023},
+    pages     = {1683-1691}
+}
+```

custom_nodes/srgb_noise/repo/data/dataset_base.py

@@ -0,0 +1,62 @@
+from torch.utils.data import Dataset 
+import torch
+import os
+import numpy as np
+from abc import ABC, abstractmethod
+import cv2
+import glob
+
+"""
+# --------------------------------------------
+# define a base dataset
+# --------------------------------------------
+"""
+class dataset_base(Dataset, ABC):
+    """This is an abstract class for datesets"""
+
+    @abstractmethod
+    def __init__(self):
+        """Initialize the dataset, this should be overriden by the subclass"""
+        pass
+
+    @abstractmethod
+    def __len__(self):
+        """Return the number of images in the dataset, this should be overriden by the subclass"""
+
+        pass
+
+    @abstractmethod
+    def __getitem__(self, index):
+        """Return the image at index, this should be overriden by the subclass"""
+
+        pass
+
+    def _get_img_paths(self, root_dir, cam_name, flag='noisy'):
+        sub_dirs = os.listdir(root_dir)
+        img_lists = []
+
+        for sub_dir in sub_dirs:
+            if cam_name in sub_dir:
+                img_dir = os.path.join(root_dir, sub_dir, flag, '*.png')     
+                img_lists.extend(glob.glob(img_dir))
+
+        return img_lists
+    
+    def _padding(self, img, win_size=10, factor=8):
+        # pad the image if not_multiple of win_size * 8, win_size is the precision of padding stride
+        multiple_factor = win_size*factor
+        h, w, _ = img.shape
+        H, W = (h//multiple_factor+1)*multiple_factor, (w//multiple_factor+1)*multiple_factor
+        padh = H - h if h%multiple_factor != 0 else 0
+        padw = W - w if w%multiple_factor != 0 else 0
+        patch = cv2.copyMakeBorder(img, top=0, bottom=padh, left=0, right=padw, borderType=cv2.BORDER_REFLECT)
+        return patch, h, w
+    
+    def _PD_scheme(self, img, stride=2):
+        """input: img (tensor), output: img_mosaic (tensor)"""
+        num = stride*stride
+        idx = np.random.randint(0, num)
+        h_start = idx // stride
+        w_start = idx % stride
+        img_mosaic = img[:, h_start::stride, w_start::stride]
+        return img_mosaic

custom_nodes/srgb_noise/repo/data/dataset_noise_synthesis.py

@@ -0,0 +1,54 @@
+import os
+import random
+import utils.utils_image as util
+from data.dataset_base import dataset_base
+
+class dataset_noise_synthesis(dataset_base):
+    def __init__(self,
+                 img_lists,
+                 img_channels,
+                 not_train=False,
+                 not_aug=False,
+                 patch_size=96):
+        
+        self.noisy_img_lists = img_lists
+        self.img_channels = img_channels
+        self.not_train = not_train
+        self.not_aug = not_aug
+        self.patch_size = patch_size
+    
+    def __getitem__(self, index):
+        noisy_img_path = self.noisy_img_lists[index]
+        noisy_img_dir, noisy_img_name = os.path.split(noisy_img_path)
+        clean_img_dir = noisy_img_dir.replace('noisy', 'clean')
+        clean_img_name = noisy_img_name.replace('NOISY', 'GT')
+        clean_img_path = os.path.join(clean_img_dir, clean_img_name)
+        clean_img = util.imread_uint(clean_img_path, self.img_channels)
+        noisy_img = util.imread_uint(noisy_img_path, self.img_channels)
+        
+        if not self.not_train:
+            h, w, _ = clean_img.shape
+            rnd_h = random.randint(0, max(0, h - self.patch_size))
+            rnd_w = random.randint(0, max(0, w - self.patch_size))
+            clean_patch = clean_img[rnd_h:rnd_h + self.patch_size, rnd_w:rnd_w + self.patch_size, :]
+            noisy_patch = noisy_img[rnd_h:rnd_h + self.patch_size, rnd_w:rnd_w + self.patch_size, :]
+            
+            if not self.not_aug:
+                mode = random.randint(0, 7)
+                clean_patch = util.augment_img(clean_patch, mode)
+                noisy_patch = util.augment_img(noisy_patch, mode)
+
+                        
+            clean_patch = util.uint2tensor3(clean_patch)
+            noisy_patch = util.uint2tensor3(noisy_patch)
+
+        else:
+            clean_patch = util.uint2tensor3(clean_img)
+            noisy_patch = util.uint2tensor3(noisy_img)
+
+        noise_map = noisy_patch - clean_patch
+        return {'clean': clean_patch, 'noisy': noisy_patch, 'noise': noise_map}
+    
+    def __len__(self):
+        return len(self.noisy_img_lists) 
+            

custom_nodes/srgb_noise/repo/data_preparation/G4_test.txt

@@ -0,0 +1,4 @@
+0122_006_G4_00100_00050_3200_N
+0077_004_G4_00100_00025_3200_N
+0059_003_G4_00800_01000_5500_L
+0173_008_G4_00400_00400_4400_N

custom_nodes/srgb_noise/repo/data_preparation/G4_train.txt

@@ -0,0 +1,12 @@
+0172_008_G4_00100_00100_4400_N
+0147_007_G4_00100_00100_4400_L
+0045_002_G4_00100_00060_3200_L
+0025_001_G4_00100_00060_5500_L
+0098_005_G4_00100_00050_3200_N
+0057_003_G4_00100_00125_5500_L
+0149_007_G4_00800_00800_4400_L
+0027_001_G4_00800_00350_5500_L
+0047_002_G4_00800_00640_3200_L
+0078_004_G4_00200_00050_3200_N
+0123_006_G4_00400_00160_3200_N
+0099_005_G4_00400_00200_3200_N

custom_nodes/srgb_noise/repo/data_preparation/GP_test.txt

@@ -0,0 +1,5 @@
+0086_004_GP_00100_00100_5500_L
+0034_002_GP_00100_00160_3200_N
+0199_010_GP_00800_01600_5500_N
+0156_007_GP_00800_00800_5500_N
+0107_005_GP_01600_01600_4400_N

custom_nodes/srgb_noise/repo/data_preparation/GP_train.txt

@@ -0,0 +1,28 @@
+0132_006_GP_00100_00200_5500_L
+0063_003_GP_00100_00100_4400_N
+0066_003_GP_00100_00200_3200_L
+0180_008_GP_00100_00100_5500_N
+0129_006_GP_00100_00100_4400_N
+0017_001_GP_00100_00060_5500_N
+0105_005_GP_00100_00100_4400_N
+0155_007_GP_00100_00100_5500_N
+0198_010_GP_00100_00200_5500_N
+0018_001_GP_00100_00160_5500_L
+0181_008_GP_00800_00800_5500_N
+0035_002_GP_00800_00350_3200_N
+0020_001_GP_00800_00350_5500_N
+0133_006_GP_00800_01600_5500_L
+0019_001_GP_00800_00640_5500_L
+0038_002_GP_00800_00640_3200_L
+0087_004_GP_00800_00640_5500_L
+0108_005_GP_06400_06400_4400_N
+0036_002_GP_06400_03200_3200_N
+0064_003_GP_01600_01600_4400_N
+0157_007_GP_01600_01600_5500_N
+0200_010_GP_01600_03200_5500_N
+0065_003_GP_10000_08460_4400_N
+0083_004_GP_00050_00020_4400_N
+0084_004_GP_00200_00100_4400_N
+0130_006_GP_00400_00400_4400_N
+0106_005_GP_00400_00400_4400_N
+0182_008_GP_03200_03200_5500_N

custom_nodes/srgb_noise/repo/data_preparation/IP_test.txt

@@ -0,0 +1,10 @@
+0110_005_IP_00100_00100_5500_L
+0114_005_IP_00100_00200_5500_N
+0134_006_IP_00100_00100_5500_N
+0193_009_IP_00800_02000_3200_N
+0032_001_IP_00800_01000_3200_N
+0165_007_IP_00800_00800_3200_N
+0191_008_IP_01600_01600_3200_N
+0090_004_IP_01600_00750_5500_N
+0139_006_IP_00200_00200_3200_L
+0115_005_IP_00400_00750_5500_N

custom_nodes/srgb_noise/repo/data_preparation/IP_train.txt

@@ -0,0 +1,44 @@
+0163_007_IP_00100_00100_3200_N
+0197_009_IP_00100_00200_5500_L
+0192_009_IP_00100_00200_3200_N
+0044_002_IP_00100_00180_5500_N
+0138_006_IP_00100_00100_3200_L
+0028_001_IP_00100_00160_5500_N
+0159_007_IP_00100_00100_3200_L
+0184_008_IP_00100_00100_3200_L
+0033_001_IP_00100_00160_3200_N
+0188_008_IP_00100_00100_3200_N
+0039_002_IP_00100_00180_5500_L
+0088_004_IP_00100_00050_5500_N
+0029_001_IP_00800_01000_5500_N
+0043_002_IP_00800_01520_5500_N
+0040_002_IP_00800_02000_5500_L
+0136_006_IP_00800_00800_5500_N
+0116_005_IP_00800_01520_5500_N
+0140_006_IP_00800_00800_3200_L
+0196_009_IP_00800_02000_5500_L
+0186_008_IP_00800_00800_3200_L
+0161_007_IP_00800_00800_3200_L
+0190_008_IP_00800_00800_3200_N
+0137_006_IP_01600_01600_5500_N
+0042_002_IP_01600_03100_5500_N
+0194_009_IP_01600_04000_3200_N
+0195_009_IP_01600_04000_5500_L
+0030_001_IP_01600_02000_5500_N
+0113_005_IP_01600_01520_5500_L
+0166_007_IP_01600_01600_3200_N
+0117_005_IP_01600_04160_5500_N
+0068_003_IP_00200_00400_3200_N
+0073_003_IP_00200_01000_5500_L
+0069_003_IP_01000_02000_3200_N
+0072_003_IP_01000_02000_5500_L
+0070_003_IP_02000_04000_3200_N
+0089_004_IP_00500_00250_5500_N
+0091_004_IP_00320_00080_3200_L
+0092_004_IP_00640_00125_3200_L
+0160_007_IP_00400_00400_3200_L
+0111_005_IP_00400_00400_5500_L
+0135_006_IP_00400_00400_5500_N
+0164_007_IP_00400_00400_3200_N
+0189_008_IP_00400_00400_3200_N
+0185_008_IP_00400_00400_3200_L

custom_nodes/srgb_noise/repo/data_preparation/N6_test.txt

@@ -0,0 +1,4 @@
+0054_003_N6_00100_00160_5500_N
+0075_004_N6_00800_00080_3200_L
+0050_002_N6_03200_03200_5500_L
+0096_005_N6_01600_01000_3200_L

custom_nodes/srgb_noise/repo/data_preparation/N6_train.txt

@@ -0,0 +1,18 @@
+0022_001_N6_00100_00060_5500_N
+0142_007_N6_00100_00100_4400_N
+0048_002_N6_00100_00100_5500_L
+0167_008_N6_00100_00050_4400_L
+0094_005_N6_00100_00050_3200_L
+0118_006_N6_00100_00025_3200_L
+0055_003_N6_00800_01000_5500_N
+0169_008_N6_00800_00400_4400_L
+0023_001_N6_00800_00350_5500_N
+0097_005_N6_03200_02000_3200_L
+0145_007_N6_03200_03200_4400_N
+0076_004_N6_03200_00320_3200_L
+0121_006_N6_03200_01000_3200_L
+0120_006_N6_01600_00400_3200_L
+0144_007_N6_01600_01600_4400_N
+0170_008_N6_01600_00800_4400_L
+0146_007_N6_00400_00400_4400_N
+0168_008_N6_00400_00200_4400_L

custom_nodes/srgb_noise/repo/data_preparation/S6_test.txt

@@ -0,0 +1,7 @@
+0007_001_S6_00100_00100_5500_L
+0008_001_S6_00100_00100_5500_N
+0002_001_S6_00100_00020_3200_N
+0081_004_S6_00800_00160_4400_L
+0014_001_S6_03200_01250_3200_N
+0152_007_S6_01600_01600_5500_L
+0126_006_S6_00400_00200_4400_L

custom_nodes/srgb_noise/repo/data_preparation/S6_train.txt

@@ -0,0 +1,28 @@
+0051_002_S6_00100_00060_5500_N
+0125_006_S6_00100_00050_4400_L
+0101_005_S6_00100_00050_4400_L
+0001_001_S6_00100_00060_3200_L
+0006_001_S6_00100_00060_4400_H
+0175_008_S6_00100_00025_5500_L
+0150_007_S6_00100_00100_5500_L
+0060_003_S6_00100_00100_4400_L
+0003_001_S6_00100_00060_3200_H
+0004_001_S6_00100_00060_4400_L
+0005_001_S6_00100_00060_4400_N
+0177_008_S6_00800_00200_5500_L
+0010_001_S6_00800_00350_3200_N
+0012_001_S6_00800_00500_5500_N
+0011_001_S6_00800_00500_5500_L
+0151_007_S6_00800_00800_5500_L
+0179_008_S6_03200_00800_5500_L
+0013_001_S6_03200_01250_3200_L
+0016_001_S6_03200_01600_5500_N
+0062_003_S6_03200_02500_4400_L
+0015_001_S6_03200_01600_5500_L
+0104_005_S6_03200_01600_4400_L
+0127_006_S6_01600_00800_4400_L
+0178_008_S6_01600_00400_5500_L
+0052_002_S6_01600_01000_5500_N
+0080_004_S6_00200_00050_3200_N
+0102_005_S6_00400_00200_4400_L
+0154_007_S6_00400_00400_5500_L

custom_nodes/srgb_noise/repo/main_test.py

@@ -0,0 +1,78 @@
+import torch
+import torch.nn as nn
+import argparse
+import os
+import numpy as np
+
+from models.network_unet import UNet
+from models.network_plain import Noise_estimator, Noise_level_predictor, DnCNN
+import utils.utils_image as util
+
+def main():
+
+    opt = argparse.ArgumentParser() 
+    opt.add_argument('--cam_name', type=str, default='IP', help='camera name for training data, the dataset will load the images from the specified camera')
+    opt.add_argument('--dir_test_save', type=str, default='./test/saves/', help='save directory for test')
+    opt.add_argument('--dir_test_img', type=str, default='./test/imgs/', help='image directory for test data')
+    opt.add_argument('--NeCA_type', type=str, default='W', help='set mode of NeCA')
+    opt.add_argument('--noise_level', type=int, default=25, help='set Gaussian noise level for NeCA_S')
+    opt = opt.parse_args()
+
+    net_E = Noise_estimator(in_nc=3, out_nc=1, nc=96, nb=5, act_mode='R')
+    net_G1 = Noise_level_predictor(in_nc=3, out_nc=3, nc=96)
+    net_G2 = UNet(in_nc=3, out_nc=3, nc=64, act_mode='R', num_stages=4, downsample_mode='strideconv', upsample_mode='upsampling', bias=True, padding_mode='zeros', final_act='Tanh')
+
+    checkpoint_E = torch.load(os.path.join('./pretrain/', opt.cam_name, 'checkpoint_E.pth'), map_location='cpu')
+    checkpoint_G1 = torch.load(os.path.join('./pretrain/', opt.cam_name, 'checkpoint_G1.pth'), map_location='cpu')
+    checkpoint_G2 = torch.load(os.path.join('./pretrain/', opt.cam_name, 'checkpoint_G2.pth'), map_location='cpu')
+    
+    net_E.load_state_dict(checkpoint_E['model_state_dict'])
+    net_G1.load_state_dict(checkpoint_G1['model_state_dict'])
+    net_G2.load_state_dict(checkpoint_G2['model_state_dict'])
+
+    net_E.cuda()
+    net_G1.cuda()
+    net_G2.cuda()
+
+    net_E.eval()
+    net_G1.eval()
+    net_G2.eval()
+
+    clean = util.imread_uint(os.path.join(opt.dir_test_img, 'clean.png'), 3)
+    noisy = util.imread_uint(os.path.join(opt.dir_test_img, 'noisy.png'), 3)
+    h, w, _ = clean.shape
+    if h%8 != 0 or w%8 != 0:
+        clean = clean[:h//8*8, :w//8*8, :]
+        noisy = noisy[:h//8*8, :w//8*8, :]
+
+    clean = util.uint2tensor3(clean).unsqueeze(0).cuda()
+    noisy = util.uint2tensor3(noisy).unsqueeze(0).cuda()
+
+    with torch.no_grad():
+        if opt.NeCA_type == 'W':
+            z = torch.randn_like(clean).cuda()
+            gain_factor = net_E(noisy)
+            pred_noise_level_map = net_G1(clean, gain_factor)
+            sdnu_noise = pred_noise_level_map.mul(z)
+            sdnc_noise = net_G2(sdnu_noise) + sdnu_noise
+            fake_noisy = clean + sdnc_noise
+
+            util.imsave(util.tensor2uint(fake_noisy), os.path.join(opt.dir_test_save, 'fake_noisy.png'))
+            util.imsave(util.tensor2uint(sdnc_noise+0.5), os.path.join(opt.dir_test_save, 'sdnc_noise.png'))
+            util.imsave(util.tensor2uint(sdnu_noise+0.5), os.path.join(opt.dir_test_save, 'sdnu_noise.png'))
+
+        elif opt.NeCA_type == 'S':
+            z = torch.randn_like(clean).cuda()
+            z.mul_(opt.noise_level/255.0)
+            sinc_noise = net_G2(z) + z
+            fake_noisy = clean + sinc_noise
+
+            util.imsave(util.tensor2uint(fake_noisy), os.path.join(opt.dir_test_save, 'fake_noisy.png'))
+            util.imsave(util.tensor2uint(sinc_noise+0.5), os.path.join(opt.dir_test_save, 'sinc_noise.png'))
+
+
+        else:
+            raise NotImplementedError('NeCA_type [%s] is not found' % opt.NeCA_type)
+        
+if __name__ == '__main__':
+    main()

custom_nodes/srgb_noise/repo/main_train_jointly.py

@@ -0,0 +1,449 @@
+import torch
+import torch.nn as nn
+import logging
+import os
+from torch.utils.data import DataLoader
+import torch.optim as optim
+from torch.utils.tensorboard import SummaryWriter
+from collections import OrderedDict
+
+from options.option_noise_synthesis import BaseOptions
+from utils.utils_logger import logger_info
+from data.dataset_noise_synthesis import dataset_noise_synthesis
+from models.network_unet import UNet
+from models.network_plain import Noise_estimator, Noise_level_predictor, DnCNN 
+from models.network_discriminator import Discriminator_96, Discriminator_32 
+import utils.utils_network as net_util
+import utils.utils_image as util
+
+
+if __name__ == '__main__':
+    """
+    # -----------------------------------------------
+    # step 1 -- initiate logs and args
+    # -----------------------------------------------
+    """
+    # get training options
+    opt, opt_message = BaseOptions().parse() 
+
+    # initiate data folder
+    cam_name = opt.cam_name
+    dir_save_logs = os.path.join(opt.dir_save, cam_name, 'logs/')
+    dir_save_imgs = os.path.join(opt.dir_save, cam_name, 'imgs/')
+    dir_save_models = os.path.join(opt.dir_save, cam_name, 'models/')
+    dir_save_tb = os.path.join(opt.dir_save, cam_name, 'tb/')
+
+    if not os.path.exists(dir_save_logs):
+        os.makedirs(dir_save_logs)
+    if not os.path.exists(dir_save_imgs):
+        os.makedirs(dir_save_imgs)
+    if not os.path.exists(dir_save_models):
+        os.makedirs(dir_save_models)
+    if not os.path.exists(dir_save_tb):
+        os.makedirs(dir_save_tb)
+
+    # initiate logs
+    logger_name = 'train'
+    logger_info(logger_name, log_path=dir_save_logs+logger_name+'.log')
+    logger = logging.getLogger(logger_name)
+    logger.info(opt_message)
+    
+    """
+    # -----------------------------------------------
+    # step 2 -- prepare dataloader
+    # -----------------------------------------------
+    """
+    if opt.generate_data_lists:
+        util.train_test_split(opt.train_img_dir, opt.data_prep_dir, cam_name=cam_name, ratio=0.8)
+    
+    train_img_lists, val_img_lists = util.get_img_list(opt.train_img_dir, cam_name=cam_name, mode='train', ratio=1)
+    test_img_lists = util.get_img_list(opt.train_img_dir, cam_name=cam_name, mode='test')
+    data_message = 'train: {:d}, val: {:d}, test: {:d}'.format(len(train_img_lists), len(val_img_lists), len(test_img_lists))
+    print(data_message)
+    logger.info(data_message)
+
+    train_data = dataset_noise_synthesis(img_lists=train_img_lists,
+                                         img_channels=opt.img_channels,
+                                         not_train=opt.not_train,
+                                         not_aug=opt.not_aug,
+                                         patch_size=opt.patch_size)
+    val_data = dataset_noise_synthesis(img_lists=val_img_lists,
+                                       img_channels=opt.img_channels,
+                                       not_train=True)
+    test_data = dataset_noise_synthesis(img_lists=test_img_lists,
+                                        img_channels=opt.img_channels,
+                                        not_train=True)
+    
+    train_loader = DataLoader(dataset=train_data,
+                              batch_size=opt.batch_size,
+                              shuffle=not opt.not_train_shuffle,
+                              num_workers=opt.num_workers,
+                              pin_memory=True,
+                              drop_last=True)
+    val_loader = DataLoader(dataset=val_data,
+                            batch_size=1,
+                            shuffle=False,
+                            num_workers=opt.num_workers,
+                            pin_memory=True,
+                            drop_last=False)
+    test_loader = DataLoader(dataset=test_data,
+                             batch_size=1,
+                             shuffle=False,
+                             num_workers=opt.num_workers,
+                             pin_memory=True,
+                             drop_last=False)                            
+    
+    """
+    # ------------------------------------------------
+    # step 3 -- prepare networks, optims and loss
+    # ------------------------------------------------ 
+    """
+
+    """prepare networks
+    net_E -- Gain estimation network (GENet)
+    net_G1 -- Noise-level prediction network (NPNet)
+    net_G2 -- Neighboring correlation network (NCNet)
+    net_D1 -- Discriminator for neighboring uncorrelated noise
+    net_D2 -- Discriminator for neighboring correlated noise
+    net_G3 -- Denoiser for real noise, used for joint training
+    """
+    net_E = Noise_estimator(in_nc=opt.E_in_nc, 
+                            out_nc=opt.E_out_nc, 
+                            nc=opt.E_nc, 
+                            nb=opt.E_nb, 
+                            act_mode=opt.E_act_mode)
+    
+    net_G1 = Noise_level_predictor(in_nc=opt.G1_in_nc, out_nc=opt.G1_out_nc, nc=opt.G1_nc)
+
+    net_G2 = UNet(in_nc=opt.G2_in_nc,
+                 out_nc=opt.G2_out_nc,
+                 nc=opt.G2_nc,
+                 act_mode=opt.G2_act_mode,
+                 downsample_mode=opt.G2_downsample_mode,
+                 upsample_mode=opt.G2_upsample_mode,
+                 padding_mode=opt.G2_padding_mode,
+                 final_act=opt.G2_final_act)
+    
+    net_D1 = Discriminator_96(in_nc=opt.D1_in_nc, nc=opt.D1_nc, act_mode=opt.D1_act_mode)
+    
+    net_D2 = Discriminator_32(in_nc=opt.D2_in_nc, nc=opt.D2_nc, act_mode=opt.D2_act_mode)
+
+    net_G3 = DnCNN(in_nc=opt.G3_in_nc, 
+                   out_nc=opt.G3_out_nc,
+                   nc=opt.G3_nc,
+                   act_mode=opt.G3_act_mode)
+
+    # print networks
+    net_message = '------------ Network params -----------\n'
+    net_message += net_util.print_networks(network=net_E, network_name='GENet', verbose=opt.verbose)
+    net_message += net_util.print_networks(network=net_G1, network_name='NPNet', verbose=opt.verbose)
+    net_message += net_util.print_networks(network=net_G2, network_name='NCNet', verbose=opt.verbose)
+    net_message += net_util.print_networks(network=net_D1, network_name='D1', verbose=opt.verbose)
+    net_message += net_util.print_networks(network=net_D2, network_name='D2', verbose=opt.verbose)
+    net_message += net_util.print_networks(network=net_G3, network_name='Denoiser', verbose=opt.verbose)
+    net_message += '-----------------------------------------------\n'
+    logger.info(net_message)
+
+    # send networks to GPU
+    net_E, gpu_message, device = net_util.model_to_device(net_E, gpu_ids=opt.gpu_ids)
+    net_G1, _, _ = net_util.model_to_device(net_G1, gpu_ids=opt.gpu_ids)
+    net_G2, _, _ = net_util.model_to_device(net_G2, gpu_ids=opt.gpu_ids)
+    net_D1, _, _ = net_util.model_to_device(net_D1, gpu_ids=opt.gpu_ids)
+    net_D2, _, _ = net_util.model_to_device(net_D2, gpu_ids=opt.gpu_ids)
+    net_G3, _, _ = net_util.model_to_device(net_G3, gpu_ids=opt.gpu_ids)
+    logger.info(gpu_message)
+
+    # load pretrained networks 
+    if opt.use_pretrained:
+        checkpoint_E = torch.load(os.path.join(dir_save_models, opt.checkpoint_E_name), map_location=device)
+        checkpoint_G1 = torch.load(os.path.join(dir_save_models, opt.checkpoint_G1_name), map_location=device)
+        checkpoint_G2 = torch.load(os.path.join(dir_save_models, opt.checkpoint_G2_name), map_location=device)
+        checkpoint_D1 = torch.load(os.path.join(dir_save_models, opt.checkpoint_D1_name), map_location=device)
+        checkpoint_D2 = torch.load(os.path.join(dir_save_models, opt.checkpoint_D2_name), map_location=device)
+        # checkpoint_G3 = torch.load(os.path.join(dir_save_models, opt.checkpoint_G3_name), map_location=device)
+        net_E.load_state_dict(checkpoint_E['model_state_dict'])
+        net_G1.load_state_dict(checkpoint_G1['model_state_dict'])
+        net_G2.load_state_dict(checkpoint_G2['model_state_dict'])
+        net_D1.load_state_dict(checkpoint_D1['model_state_dict'])
+        net_D2.load_state_dict(checkpoint_D2['model_state_dict'])
+        # net_G3.load_state_dict(checkpoint_G3['model_state_dict'])
+
+    else:
+        net_util.init_weights(net_E, init_type=opt.init_type)
+        net_util.init_weights(net_G1, init_type=opt.init_type)
+        net_util.init_weights(net_G2, init_type=opt.init_type)
+        net_util.init_weights(net_D1, init_type=opt.init_type)
+        net_util.init_weights(net_D2, init_type=opt.init_type)
+        net_util.init_weights(net_G3, init_type=opt.init_type)
+    
+
+    # set optimizers
+    optim_E = optim.Adam(net_E.parameters(), lr=opt.lr_E, betas=(opt.beta1, opt.beta2))
+    optim_G1 = optim.Adam(net_G1.parameters(), lr=opt.lr_G1, betas=(opt.beta1, opt.beta2))
+    optim_G2 = optim.Adam(net_G2.parameters(), lr=opt.lr_G2, betas=(opt.beta1, opt.beta2))
+    optim_D1 = optim.Adam(net_D1.parameters(), lr=opt.lr_D1, betas=(opt.beta1, opt.beta2))
+    optim_D2 = optim.Adam(net_D2.parameters(), lr=opt.lr_D2, betas=(opt.beta1, opt.beta2))
+    optim_G3 = optim.Adam(net_G3.parameters(), lr=opt.lr_G3, betas=(opt.beta1, opt.beta2))
+
+    # set learning rate schedulers
+    lr_scheduler_E = optim.lr_scheduler.MultiStepLR(optimizer=optim_E, milestones=opt.lr_decay_1, gamma=0.1)
+    lr_scheduler_G1 = optim.lr_scheduler.MultiStepLR(optimizer=optim_G1, milestones=opt.lr_decay_1, gamma=0.1)
+    lr_scheduler_G2 = optim.lr_scheduler.MultiStepLR(optimizer=optim_G2, milestones=opt.lr_decay_1, gamma=0.1)
+    lr_scheduler_D1 = optim.lr_scheduler.MultiStepLR(optimizer=optim_D1, milestones=opt.lr_decay_1, gamma=0.1)
+    lr_scheduler_D2 = optim.lr_scheduler.MultiStepLR(optimizer=optim_D2, milestones=opt.lr_decay_1, gamma=0.1)
+    lr_scheduler_G3 = optim.lr_scheduler.MultiStepLR(optimizer=optim_G3, milestones=opt.lr_decay_2, gamma=0.1)
+
+    # set loss functions
+    Cri_rec = nn.MSELoss()
+    Cri_gan = net_util.GANLoss(gan_mode=opt.gan_mode).to(device)
+
+    # set tensorboard writer
+    if opt.use_tb:
+        tb = SummaryWriter(log_dir=dir_save_tb)
+    
+    """
+    # ------------------------------------------------
+    # step 4 -- training
+    # ------------------------------------------------
+    """
+    current_step = 0
+    start_epoch = 1 if not opt.use_pretrained else checkpoint_E['epoch']
+    iter_step = 5 if opt.gan_mode == 'wgangp' else 1
+    log_step = iter_step
+    for epoch in range(start_epoch, opt.epochs+1):
+        for i, data in enumerate(train_loader):
+            
+            current_step += 1
+            clean = data['clean'].to(device)
+            noisy = data['noisy'].to(device)
+            noise = data['noise'].to(device)
+            training_message = OrderedDict()
+            
+            # 1) feed data
+            z = torch.randn_like(clean).to(device)
+            gain_factor = net_E(noisy) # estimate gain factor from noisy image
+            pred_noise_level_map = net_G1(clean, gain_factor) 
+            sdnu_noise = pred_noise_level_map.mul(z) # synthesize signal-dependent neighboring uncorrelated noise
+            sdnc_noise = net_G2(sdnu_noise) + sdnu_noise # synthesize signal-dependent neighboring correlated noise
+
+
+            global_nl = torch.std(noise, dim=(1, 2, 3), keepdim=True) # global estimated noise level
+            mean_filter = nn.AvgPool2d((7, 7), stride=1, padding=3)
+            local_nl = torch.sqrt(mean_filter(noise**2) - (mean_filter(noise))**2) # local estimated noise level, var=E(x^2)-E(x)^2
+            pixel_shuffle = util.pixel_shuffle(stride=opt.pd_stride)
+            sub_noise = pixel_shuffle._crop(noise)
+            sub_sdnu_noise = pixel_shuffle._crop(sdnu_noise)
+
+            if opt.joint_start_epoch <= epoch:
+                if opt.seperate_train:
+                    fake_noisy = sdnc_noise.detach() + clean
+                else:
+                    fake_noisy = sdnc_noise + clean
+                rec_clean = net_G3(fake_noisy)
+
+            # 2) optimize generator
+            optim_E.zero_grad()
+            optim_G1.zero_grad()
+            optim_G2.zero_grad()
+            optim_G3.zero_grad()
+
+            std1_loss = Cri_rec(gain_factor, global_nl)
+            std2_loss = Cri_rec(pred_noise_level_map, local_nl)
+            reg_loss = Cri_rec(pred_noise_level_map, gain_factor.expand_as(pred_noise_level_map))
+            G_losses = opt.lambda_reg * reg_loss + opt.lambda_std1 * std1_loss + opt.lambda_std2 * std2_loss
+            if opt.joint_start_epoch <= epoch:
+                rec_loss = Cri_rec(rec_clean, clean)
+                G_losses += opt.lambda_rec * rec_loss
+
+            if current_step % iter_step == 0:
+
+                fake1 = net_D1(sdnc_noise)
+                fake2  =net_D2(sub_sdnu_noise)
+                adv1_loss = Cri_gan(fake1, True)
+                adv2_loss = Cri_gan(fake2, True)
+                G_losses += opt.lambda_adv1 * adv1_loss + opt.lambda_adv2 * adv2_loss
+            
+            G_losses.backward()
+            optim_E.step()
+            optim_G1.step()
+            optim_G2.step()
+            optim_G3.step() 
+
+            # 3) optimize discriminator
+            optim_D1.zero_grad()
+            optim_D2.zero_grad()
+
+            real1 = net_D1(noise)
+            fake1 = net_D1(sdnc_noise.detach())
+            real2 = net_D2(sub_noise)
+            fake2 = net_D2(sub_sdnu_noise.detach())
+            D_r_loss = Cri_gan(real1, True) + Cri_gan(real2, True)
+            D_f_loss = Cri_gan(fake1, False) + Cri_gan(fake2, False)
+            if opt.gan_mode == 'wgangp':
+                gp = net_util.gradient_penalty(net_D1, noise, sdnc_noise.detach(), device) + net_util.gradient_penalty(net_D2, sub_noise, sub_sdnu_noise.detach(), device)
+                D_losses = D_r_loss + D_f_loss+ opt.lambda_gp * gp
+            else:
+                D_losses = 0.5*(D_r_loss + D_f_loss)
+            D_losses.backward()
+            optim_D1.step()
+            optim_D2.step()
+
+            # 4) logs
+            if current_step % log_step == 0:
+                training_message['reg_loss'] = reg_loss.item()
+                training_message['std1_loss'] = std1_loss.item()
+                training_message['std2_loss'] = std2_loss.item()
+                training_message['adv1_loss'] = adv1_loss.item()
+                training_message['adv2_loss'] = adv2_loss.item()
+                training_message['D_r_loss'] = D_r_loss.item()
+                training_message['D_f_loss'] = D_f_loss.item()
+                loss_message = '<epoch:{:d}, step: {:08d}> '.format(epoch, current_step)
+                
+                for k, v in training_message.items():  # merge log information into message
+                    loss_message += '{:s}: {:.3e} '.format(k, v)
+                print(loss_message)
+                logger.info(loss_message)
+
+                if opt.use_tb:
+                    tb.add_scalar('reg_loss', training_message['reg_loss'], current_step)
+                    tb.add_scalar('std1_loss', training_message['std1_loss'], current_step)
+                    tb.add_scalar('std2_loss', training_message['std2_loss'], current_step)
+                    tb.add_scalar('adv1_loss', training_message['adv1_loss'], current_step)
+                    tb.add_scalar('adv2_loss', training_message['adv2_loss'], current_step)
+                    tb.add_scalar('D_r_loss', training_message['D_r_loss'], current_step)
+                    tb.add_scalar('D_f_loss', training_message['D_f_loss'], current_step)
+            
+            # if current_step % opt.val_step == 0:
+            #     # validation
+            #     net_E.eval()
+            #     net_G1.eval()
+            #     net_G2.eval()
+            #     net_G3.eval()
+            #     val_message = OrderedDict()
+            #     val_kl = 0
+            #     val_psnr = 0
+            #     val_num = 0
+
+            #     for data in val_loader: 
+            #         val_num += 1
+            #         clean = data['clean'].to(device)
+            #         noisy = data['noisy'].to(device)
+            #         noise = data['noise'].to(device)
+            #         with torch.no_grad():
+            #             z = torch.randn_like(clean).to(device)
+            #             gain_factor = net_E(noisy)
+            #             pred_noise_level_map = net_G1(clean, gain_factor)
+            #             sdnu_noise = pred_noise_level_map.mul(z)
+            #             sdnc_noise = sdnu_noise + net_G2(sdnu_noise)
+            #             fake_noisy = clean + sdnc_noise
+            #             quantized_noise = util.noise_quantization(fake_noisy, clean)
+            #             if opt.joint_start_epoch <= epoch:
+            #                 rec_clean = net_G3(noisy)
+
+            #             val_kl += util.cal_kld(util.tensor2data(noise), util.tensor2data(quantized_noise))
+            #             if opt.joint_start_epoch <= epoch:
+            #                 val_psnr += util.calculate_psnr(util.tensor2uint(rec_clean), util.tensor2uint(clean))
+
+            #     val_kl /= val_num
+            #     val_psnr /= val_num
+            #     val_message['val_kl'] = val_kl
+            #     val_message['val_psnr'] = val_psnr
+            #     message = '<epoch:{:d}, step: {:08d}> '.format(epoch, current_step)
+
+            #     for k, v in val_message.items():
+            #         message += '{:s}: {:.3e} '.format(k, v)
+            #     print(message)
+            #     logger.info(message)
+            #     if opt.use_tb:
+            #         tb.add_scalar('val_kl', val_kl, current_step)
+            #         tb.add_scalar('val_psnr', val_psnr, current_step)
+
+            #     net_E.train()
+            #     net_G1.train()
+            #     net_G2.train()
+            #     net_G3.train()
+        
+        # test
+        if epoch % opt.test_epoch == 0:
+            net_E.eval()
+            net_G1.eval()
+            net_G2.eval()
+            net_G3.eval()
+            test_message = OrderedDict()
+            test_kl = 0
+            test_psnr = 0
+            test_num = 0
+
+            for data in test_loader:
+                test_num += 1
+                clean = data['clean'].to(device)
+                noisy = data['noisy'].to(device)
+                noise = data['noise'].to(device)
+                with torch.no_grad():
+                    z = torch.randn_like(clean).to(device)
+                    gain_factor = net_E(noisy)
+                    pred_noise_level_map = net_G1(clean, gain_factor)
+                    sdnu_noise = pred_noise_level_map.mul(z)
+                    sdnc_noise = sdnu_noise + net_G2(sdnu_noise)
+                    fake_noisy = clean + sdnc_noise
+                    quantized_noise = util.noise_quantization(fake_noisy, clean)
+
+                    if opt.use_tb:
+                        tb.add_image('clean', clean.squeeze(), epoch)
+                        tb.add_image('noisy', noisy.squeeze(), epoch)
+                        tb.add_image('noise', noise.squeeze(), epoch)
+                        tb.add_image('sdnu_noise', sdnu_noise.squeeze()+0.5, epoch)
+                        tb.add_image('sdnc_noise', sdnc_noise.squeeze()+0.5, epoch)
+                        tb.add_image('fake_noisy', fake_noisy.squeeze(), epoch)
+                    
+                    if opt.joint_start_epoch <= epoch:
+                        rec_clean = net_G3(noisy)
+
+                    test_kl += util.cal_kld(util.tensor2data(noise), util.tensor2data(quantized_noise))
+                    if opt.joint_start_epoch <= epoch:
+                        test_psnr += util.calculate_psnr(util.tensor2uint(rec_clean), util.tensor2uint(clean))
+
+                if opt.joint_start_epoch <= epoch and opt.use_tb:
+                    tb.add_image('rec_clean', rec_clean.squeeze(), epoch)
+
+            test_kl /= test_num
+            test_psnr /= test_num
+            test_message['test_kl'] = test_kl
+            test_message['test_psnr'] = test_psnr
+            message = '<epoch:{:d}> '.format(epoch)
+
+            for k, v in test_message.items():
+                message += '{:s}: {:.3e} '.format(k, v)
+            print(message)
+            logger.info(message)
+            if opt.use_tb:
+                tb.add_scalar('test_kl', test_kl, current_step)
+                tb.add_scalar('test_psnr', test_psnr, current_step)
+            
+            net_E.train()
+            net_G1.train()
+            net_G2.train()
+            net_G3.train()
+
+            # save models
+            checkpoint_E = {'epoch': epoch, 'model_state_dict': net_util.get_bare_model(net_E).state_dict()}
+            checkpoint_G1 = {'epoch': epoch, 'model_state_dict': net_util.get_bare_model(net_G1).state_dict()}
+            checkpoint_G2 = {'epoch': epoch, 'model_state_dict': net_util.get_bare_model(net_G2).state_dict()}
+            checkpoint_D1 = {'epoch': epoch, 'model_state_dict': net_util.get_bare_model(net_D1).state_dict()}
+            checkpoint_D2 = {'epoch': epoch, 'model_state_dict': net_util.get_bare_model(net_D2).state_dict()}
+            checkpoint_G3 = {'epoch': epoch, 'model_state_dict': net_util.get_bare_model(net_G3).state_dict()}
+            torch.save(checkpoint_E, os.path.join(dir_save_models, 'checkpoint_E_{:08d}.pth'.format(epoch)))
+            torch.save(checkpoint_G1, os.path.join(dir_save_models, 'checkpoint_G1_{:08d}.pth'.format(epoch)))
+            torch.save(checkpoint_G2, os.path.join(dir_save_models, 'checkpoint_G2_{:08d}.pth'.format(epoch)))
+            torch.save(checkpoint_D1, os.path.join(dir_save_models, 'checkpoint_D1_{:08d}.pth'.format(epoch)))
+            torch.save(checkpoint_D2, os.path.join(dir_save_models, 'checkpoint_D2_{:08d}.pth'.format(epoch)))
+            torch.save(checkpoint_G3, os.path.join(dir_save_models, 'checkpoint_G3_{:08d}.pth'.format(epoch)))
+
+        # update learning rate
+        lr_scheduler_E.step()
+        lr_scheduler_G1.step()
+        lr_scheduler_G2.step()
+        lr_scheduler_D1.step()
+        lr_scheduler_D2.step()  
+        lr_scheduler_G3.step()
+