diff --git a/nodes.py b/nodes.py index c53724c..8e44778 100644 --- a/nodes.py +++ b/nodes.py @@ -311,7 +311,7 @@ class STARVideoSuperResolution: }), "prompt": ("STRING", { "default": "", "multiline": True, - "tooltip": "Text prompt describing the desired output. Leave empty to use STAR's built-in quality prompt.", + "tooltip": "Optional text prompt. STAR's built-in quality prompt is always appended. Leave empty to use only the quality prompt.", }), "solver_mode": (["fast", "normal"], { "default": "fast", @@ -325,6 +325,10 @@ class STARVideoSuperResolution: "default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, "tooltip": "Random seed for reproducible results.", }), + "denoise": ("FLOAT", { + "default": 0.9, "min": 0.1, "max": 1.0, "step": 0.05, + "tooltip": "How much noise to add before denoising. 0.9 = STAR default (preserves input structure). 1.0 = full denoise (starts from pure noise). Lower values preserve more of the original.", + }), "color_fix": (["adain", "wavelet", "none"], { "default": "adain", "tooltip": "Post-processing color correction. adain: match color stats from input. wavelet: preserve input low-frequency color. none: no correction.", @@ -353,6 +357,7 @@ class STARVideoSuperResolution: solver_mode, max_chunk_len, seed, + denoise, color_fix, segment_size=0, ): @@ -368,6 +373,7 @@ class STARVideoSuperResolution: solver_mode=solver_mode, max_chunk_len=max_chunk_len, seed=seed, + denoise=denoise, color_fix=color_fix, ) diff --git a/star_pipeline.py b/star_pipeline.py index 3a05fe8..3ce44e3 100644 --- a/star_pipeline.py +++ b/star_pipeline.py @@ -141,6 +141,7 @@ def run_star_inference_segmented( solver_mode: str = "fast", max_chunk_len: int = 32, seed: int = 0, + denoise: float = 0.9, color_fix: str = "adain", ) -> torch.Tensor: """Run STAR inference in overlapping segments to bound peak RAM usage. @@ -156,7 +157,8 @@ def run_star_inference_segmented( return run_star_inference( star_model=star_model, images=images, upscale=upscale, steps=steps, guide_scale=guide_scale, prompt=prompt, solver_mode=solver_mode, - max_chunk_len=max_chunk_len, seed=seed, color_fix=color_fix, + max_chunk_len=max_chunk_len, seed=seed, denoise=denoise, + color_fix=color_fix, ) overlap = max(2, segment_size // 4) @@ -195,6 +197,7 @@ def run_star_inference_segmented( solver_mode=solver_mode, max_chunk_len=max_chunk_len, seed=seed, + denoise=denoise, color_fix=color_fix, ) # seg_result: [F_seg, H, W, 3] float32 on CPU @@ -246,6 +249,7 @@ def run_star_inference( solver_mode: str = "fast", max_chunk_len: int = 32, seed: int = 0, + denoise: float = 0.9, color_fix: str = "adain", ) -> torch.Tensor: """Run STAR video super-resolution and return ComfyUI IMAGE batch.""" @@ -265,7 +269,7 @@ def run_star_inference( if offload == "aggressive": vae_dec_chunk = 1 - total_noise_levels = 900 + total_noise_levels = int(round(denoise * 1000)) # -- Convert ComfyUI frames to STAR format -- video_data = comfyui_to_star_frames(images) # [F, 3, H, W]