From e89a8b00a5c8381834ffae486d63a27b152f8158 Mon Sep 17 00:00:00 2001 From: Adrien Toupet Date: Thu, 18 Sep 2025 23:19:35 -0400 Subject: [PATCH] feat: promote cond_noise_scale to UI for controlling conditioning noise level helping mitigate noise artifacts in high-res upscaling --- inference_cli.py | 6 +++++- src/core/generation.py | 6 +++--- src/interfaces/comfyui_node.py | 18 ++++++++++++++---- 3 files changed, 22 insertions(+), 8 deletions(-) diff --git a/inference_cli.py b/inference_cli.py index 16ac0bc..6429476 100644 --- a/inference_cli.py +++ b/inference_cli.py @@ -335,7 +335,8 @@ def _worker_process(proc_idx, device_id, frames_np, shared_args, return_queue): debug=worker_debug, progress_callback=None, cfg_scale=shared_args["cfg_scale"], - seed=shared_args["seed"] + seed=shared_args["seed"], + cond_noise_scale=shared_args["cond_noise_scale"] ) # Phase 3: Decode all batches @@ -388,6 +389,7 @@ def _gpu_processing(frames_tensor, device_list, args): "model_dir": args.model_dir if args.model_dir is not None else "./models/SEEDVR2", "preserve_vram": args.preserve_vram, "color_correction": args.color_correction, + "cond_noise_scale": args.cond_noise_scale, "debug": args.debug, "cfg_scale": 1.0, "seed": args.seed, @@ -512,6 +514,8 @@ def parse_arguments(): parser.add_argument("--color_correction", type=str, default="wavelet", choices=["wavelet", "adain", "none"], help="Color correction method: 'wavelet' (natural, recommended), 'adain' (stylistic), 'none' (no correction)") + parser.add_argument("--cond_noise_scale", type=float, default=0.0, + help="Conditional noise scale (0.0-1.0). Higher values add more noise, affecting sharpness (default: 0.0)") parser.add_argument("--preserve_vram", action="store_true", help="Enable VRAM preservation mode") parser.add_argument("--debug", action="store_true", diff --git a/src/core/generation.py b/src/core/generation.py index a1e66c7..579994f 100644 --- a/src/core/generation.py +++ b/src/core/generation.py @@ -528,7 +528,7 @@ def encode_all_batches(runner, ctx=None, images=None, batch_size=90, preserve_vr def upscale_all_batches(runner, ctx=None, preserve_vram=False, debug=None, - progress_callback=None, cfg_scale=1.0, seed=100): + progress_callback=None, cfg_scale=1.0, seed=100, cond_noise_scale=0.0): """ Phase 2: DiT Upscaling for all encoded batches. @@ -543,6 +543,8 @@ def upscale_all_batches(runner, ctx=None, preserve_vram=False, debug=None, progress_callback: Optional callback(current, total, frames, phase_name) cfg_scale: Classifier-free guidance scale (default: 1.0) seed: Random seed for noise generation + cond_noise_scale: Conditional noise scale for latent augmentation (0.0-1.0). + Controls the amount of noise added to the conditioning latent. Returns: dict: Updated context containing: @@ -584,8 +586,6 @@ def upscale_all_batches(runner, ctx=None, preserve_vram=False, debug=None, # Set seed for generation set_seed(seed) - cond_noise_scale = 0.0 - # Count valid latents num_valid_latents = len([l for l in ctx['all_latents'] if l is not None]) diff --git a/src/interfaces/comfyui_node.py b/src/interfaces/comfyui_node.py index 070aa6b..5838c1d 100644 --- a/src/interfaces/comfyui_node.py +++ b/src/interfaces/comfyui_node.py @@ -94,6 +94,13 @@ class SeedVR2: "default": "wavelet", "tooltip": "Color correction method. Wavelet: Frequency-based for natural results (recommended). AdaIN: Statistical matching for stylized effects. None: No color correction." }), + "cond_noise_scale": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "tooltip": "Conditional noise scale for latent augmentation. Higher values add more noise to the conditioning, affecting sharpness and detail. 0.0 = no noise (crisp), 1.0 = maximum noise (softer)." + }), }, "optional": { "block_swap_config": ("block_swap_config", { @@ -112,7 +119,8 @@ class SeedVR2: CATEGORY = "SEEDVR2" def execute(self, images: torch.Tensor, model: str, seed: int, new_resolution: int, - batch_size: int, color_correction=str, block_swap_config=None, extra_args=None) -> Tuple[torch.Tensor]: + batch_size: int, color_correction: str, cond_noise_scale: float, + block_swap_config=None, extra_args=None) -> Tuple[torch.Tensor]: """Execute SeedVR2 video upscaling with progress reporting""" temporal_overlap = 0 @@ -155,7 +163,8 @@ class SeedVR2: cfg_scale = 1.0 try: return self._internal_execute(images, model, seed, new_resolution, cfg_scale, - batch_size, color_correction, tiled_vae, vae_tile_size, vae_tile_overlap, + batch_size, color_correction, cond_noise_scale, + tiled_vae, vae_tile_size, vae_tile_overlap, preserve_vram, temporal_overlap, cache_model, device, block_swap_config) except Exception as e: @@ -198,7 +207,7 @@ class SeedVR2: def _internal_execute(self, images, model, seed, new_resolution, cfg_scale, batch_size, - color_correction, tiled_vae, vae_tile_size, vae_tile_overlap, + color_correction, cond_noise_scale, tiled_vae, vae_tile_size, vae_tile_overlap, preserve_vram, temporal_overlap, cache_model, device, block_swap_config): """Internal execution logic with progress tracking""" @@ -278,7 +287,8 @@ class SeedVR2: debug=debug, progress_callback=self._progress_callback, cfg_scale=cfg_scale, - seed=seed + seed=seed, + cond_noise_scale=cond_noise_scale ) # Phase 3: Decode all batches