From 18303974008d4f77b7f5dcc3cccd6a042cebb7b5 Mon Sep 17 00:00:00 2001 From: Sonny Box Date: Sat, 3 Jan 2026 15:42:37 -0800 Subject: [PATCH] stage v0.1.5 --- LICENSE | 2 +- __init__.py | 11 +-- image.py | 6 +- pyproject.toml | 5 +- scheduler.py | 198 +++++++++++++++++++++++++++++++++++++++++++++++++ video.py | 114 ---------------------------- 6 files changed, 211 insertions(+), 125 deletions(-) delete mode 100644 video.py diff --git a/LICENSE b/LICENSE index 0daab8b..614a40b 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2025 Sonny Box +Copyright (c) 2026 Sonny Box Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/__init__.py b/__init__.py index 50498e5..45f98bf 100644 --- a/__init__.py +++ b/__init__.py @@ -1,8 +1,7 @@ from .image import ImageMaskCrop, RestoreMaskCrop from .qol import ImageSizeCalculator -from .scheduler import SigmaSmoother, SigmasRescale +from .scheduler import SigmaSmoother, SigmasGraph, SigmasRescale, LoadDiffusersScheduler from .tiling import CreateTiles, SeedVRCalculateTiles, StitchTiles -from .video import WanExtendI2VPlus NODE_CLASS_MAPPINGS = { "SuperCreateTiles": CreateTiles, @@ -13,7 +12,8 @@ NODE_CLASS_MAPPINGS = { "ImageSizeCalculator": ImageSizeCalculator, "ImageMaskCrop": ImageMaskCrop, "RestoreMaskCrop": RestoreMaskCrop, - "WanExtendI2VPlus": WanExtendI2VPlus, + "LoadDiffusersScheduler": LoadDiffusersScheduler, + "SigmasGraph": SigmasGraph, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -25,9 +25,10 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ImageSizeCalculator": "🐧 Image Size Calculator", "ImageMaskCrop": "🐧 Image Mask Crop", "RestoreMaskCrop": "🐧 Restore Mask Crop", - "WanExtendI2VPlus": "🐧 Wan Extend I2V+", + "LoadDiffusersScheduler": "🐧 Load Diffusers Scheduler", + "SigmasGraph": "🐧 Sigmas Graph", } -print("\033[34m[SuperNodes]\033[0m Loaded successfully.") +print("\033[34m[SuperNodes]\033[0m Classes initialized.") __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/image.py b/image.py index d4db4a1..d10444c 100644 --- a/image.py +++ b/image.py @@ -217,9 +217,9 @@ class RestoreMaskCrop: def resize_tensor(tensor, width, height, method, is_mask=False): """ Optimized resizing using torch.interpolate where possible, - falling back to PIL for unsupported methods (e.g. Lanczos). + falling back to PIL for unsupported methods (Lanczos). """ - # Fallback to PIL for Lanczos + # Use PIL for Lanczos if method == "lanczos": return pil_resize_fallback( tensor, width, height, method, is_mask @@ -335,7 +335,7 @@ class RestoreMaskCrop: elif strategy == "scale_original": # Calculate scale factor based on the difference between the NEW crop width and the OLD crop width - # This ensures the original image scales proportionally to how the crop was scaled + # This causes the original image to scale proportionally to how the crop was scaled scale_factor = cropped_image.shape[2] / orig_w # New dimensions for the full original image diff --git a/pyproject.toml b/pyproject.toml index 0cbfc64..74506f6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-supernodes" description = "A miscellaneous node pack of custom implementations needed for workflows released by my SuperComfy YouTube channel." -version = "0.1.4" +version = "0.1.5" license = { file = "LICENSE" } [project.urls] @@ -14,7 +14,8 @@ Repository = "https://github.com/sonnybox/ComfyUI-SuperNodes" PublisherId = "sonnybox" DisplayName = "ComfyUI-SuperNodes" Icon = "https://raw.githubusercontent.com/sonnybox/ComfyUI-SuperNodes/refs/heads/main/logo.png" -includes = [] +includes = [] +dependencies = ["numpy","pillow",] # "requires-comfyui" = ">=1.0.0" # ComfyUI version compatibility [tool.basedpyright] diff --git a/scheduler.py b/scheduler.py index 3857c14..2d09b2d 100644 --- a/scheduler.py +++ b/scheduler.py @@ -1,4 +1,14 @@ +import json +import math +import os +import random +import diffusers import torch +import folder_paths +import numpy as np +import matplotlib.pyplot as plt +import io +from PIL import Image class SigmaSmoother: @@ -153,3 +163,191 @@ class SigmasRescale: new_sigmas = normalized_curve * (max - min) + min return (new_sigmas,) + + +class LoadDiffusersScheduler: + @classmethod + def INPUT_TYPES(s): + scheduler_type = "scheduler" + default_scheduler_path = os.path.join(folder_paths.models_dir, scheduler_type) + folder_paths.add_model_folder_path(scheduler_type, default_scheduler_path) + paths, current_exts = folder_paths.folder_names_and_paths[scheduler_type] + if ".json" not in current_exts: + folder_paths.folder_names_and_paths[scheduler_type] = (paths, current_exts | {".json"}) + + return { + "required": { + # Now we use the standard get_filename_list with our new key + "scheduler_name": (folder_paths.get_filename_list(scheduler_type), {"tooltip": "Select a diffusers scheduler JSON config file from 'models/scheduler'."}), + "steps": ("INT", {"default": 20, "min": 1, "max": 10000, "tooltip": "The number of steps for the schedule."}), + "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The amount of denoising to apply."}), + }, + "optional": { + "width": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 8, "tooltip": "Used for 'dynamic shifting' calculations."}), + "height": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 8, "tooltip": "Used for 'dynamic shifting' calculations."}), + } + } + + RETURN_TYPES = ("SIGMAS",) + RETURN_NAMES = ("SIGMAS",) + CATEGORY = "advanced/loaders" + FUNCTION = "load_scheduler" + DESCRIPTION = "Loads a scheduler from a Diffusers JSON config and generates a SIGMAS schedule." + + def load_scheduler(self, scheduler_name, steps, denoise, width=1024, height=1024): + if diffusers is None: + raise ImportError("The 'diffusers' library is required to use this node. Please install it via pip.") + + scheduler_type = "scheduler" + + config_path = folder_paths.get_full_path(scheduler_type, scheduler_name) + + if config_path is None: + raise FileNotFoundError(f"Scheduler config '{scheduler_name}' not found. Please ensure it is in 'models/scheduler' or defined in extra_model_paths.yaml.") + + # Load JSON config + with open(config_path, 'r') as f: + config = json.load(f) + + class_name = config.get("_class_name", None) + if not class_name: + raise ValueError("Invalid scheduler JSON: missing '_class_name'.") + + scheduler_cls = getattr(diffusers, class_name, None) + if scheduler_cls is None: + raise ImportError(f"Scheduler class '{class_name}' not found in diffusers library.") + + # --- Dynamic Shifting Logic --- + if config.get("use_dynamic_shifting", False): + base_seq_len = config.get("base_image_seq_len", 256) + max_seq_len = config.get("max_image_seq_len", 4096) + base_shift = config.get("base_shift", 0.5) + max_shift = config.get("max_shift", 1.15) + + image_seq_len = (width // 16) * (height // 16) + + m = image_seq_len + m1 = base_seq_len ** 2 + m2 = max_seq_len ** 2 + + # Clamp mu between 0 and 1 + if m2 > m1: + mu = (m - m1) / (m2 - m1) + else: + mu = 0 + + mu = max(0.0, min(1.0, mu)) + shift = math.exp(math.log(base_shift) + mu * (math.log(max_shift) - math.log(base_shift))) + + config["shift"] = shift + + # Instantiate Scheduler + try: + scheduler = scheduler_cls.from_config(config) + except Exception as e: + raise RuntimeError(f"Failed to instantiate {class_name}: {e}") + + # --- Calculate Sigmas --- + total_steps = steps + if denoise < 1.0 and denoise > 0.0: + total_steps = int(steps / denoise) + if denoise == 0.0: + total_steps = steps + + scheduler.set_timesteps(total_steps) + + if hasattr(scheduler, "sigmas"): + sigmas = scheduler.sigmas + else: + raise AttributeError(f"Scheduler {class_name} does not expose 'sigmas'.") + + # Truncate if denoise was used + if denoise < 1.0 and denoise > 0.0: + sigmas = sigmas[-(steps + 1):] + + # Ensure it's a CPU tensor for ComfyUI + sigmas = sigmas.clone().detach().cpu() + + return (sigmas,) + + +class SigmasGraph: + def __init__(self): + # Based on PreviewImage logic here to save to the temp directory + self.output_dir = folder_paths.get_temp_directory() + self.type = "temp" + self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5)) + self.compress_level = 1 + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "sigmas": ("SIGMAS", {"tooltip": "The sigma schedule tensor to visualize."}), + } + } + + RETURN_TYPES = ("IMAGE",) + OUTPUT_NODE = True + FUNCTION = "plot_sigmas" + + CATEGORY = "SuperNodes" + DESCRIPTION = "Generates a visual graph of the sigma decay schedule and displays it in the node." + + def plot_sigmas(self, sigmas): + # 1. Prepare Data + if isinstance(sigmas, torch.Tensor): + s_data = sigmas.detach().cpu().numpy().flatten() + elif isinstance(sigmas, list): + s_data = np.array(sigmas).flatten() + else: + s_data = np.array(sigmas).flatten() + + # 2. Generate Plot + plt.figure(figsize=(8, 7)) + plt.plot(s_data, marker='o', linestyle='-', markersize=4, color='#1f77b4') + + steps = len(s_data) - 1 if len(s_data) > 0 else 0 + plt.title(f"Sigma Schedule ({steps} steps)") + plt.xlabel("Step") + plt.ylabel("Value") + plt.grid(True, which='both', linestyle='--', alpha=0.7) + plt.legend() + plt.tight_layout() + + # 3. Save Plot to Buffer + buf = io.BytesIO() + plt.savefig(buf, format='png', dpi=100) + buf.seek(0) + plt.close() + + # 4. Convert to Tensor (Batch, H, W, C) + image = Image.open(buf).convert("RGB") + image_np = np.array(image).astype(np.float32) / 255.0 + image_tensor = torch.from_numpy(image_np).unsqueeze(0) + + # 5. Call internal save method to display in UI + ui_output = self.save_images(image_tensor, filename_prefix="SigmasGraph") + + # Return both the UI dictionary and the image tensor + return {"ui": ui_output["ui"], "result": (image_tensor,)} + + def save_images(self, images, filename_prefix="ComfyUI"): + filename_prefix += self.prefix_append + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) + results = list() + for (batch_number, image) in enumerate(images): + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + + filename_with_batch_num = filename.replace("%batch_num%", str(batch_number)) + file = f"{filename_with_batch_num}_{counter:05}_.png" + img.save(os.path.join(full_output_folder, file), pnginfo=None, compress_level=self.compress_level) + results.append({ + "filename": file, + "subfolder": subfolder, + "type": self.type + }) + counter += 1 + + return { "ui": { "images": results } } \ No newline at end of file diff --git a/video.py b/video.py deleted file mode 100644 index 1aa73c7..0000000 --- a/video.py +++ /dev/null @@ -1,114 +0,0 @@ -import torch - - -class WanExtendI2VPlus: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "samples": ( - "LATENT", - { - "tooltip": "The output latent from the previous KSampler." - }, - ), - "context_frames": ( - "INT", - { - "default": 16, - "min": 1, - "max": 128, - "step": 1, - "tooltip": "Pixel frames to keep. Must match your VAE compression (e.g., 16 pixels = 4 latents).", - }, - ), - "temporal_compression": ( - "INT", - { - "default": 4, - "min": 1, - "max": 32, - "step": 1, - "tooltip": "Wan usually compresses 4 frames into 1 latent.", - }, - ), - } - } - - RETURN_TYPES = ("LATENT",) - OUTPUT_TOOLTIPS = ( - "The new latent batch with context frames 'locked' via a noise mask.", - ) - FUNCTION = "extend_latent" - - CATEGORY = "Wan/latent" - DESCRIPTION = "Extends a video latent batch by copying the last N frames to the start of a new batch and masking them so they don't change." - - def extend_latent(self, samples, context_frames, temporal_compression=4): - source_samples = samples["samples"] - - # Calculate how many LATENT frames correspond to the requested PIXEL frames - # e.g., 16 pixel frames / 4 compression = 4 latent frames - ctx_lat_len = context_frames // temporal_compression - - # --- SHAPE DETECTION --- - # 5D: [Batch, Channels, Time, Height, Width] (Native Video Format) - # 4D: [Batch, Channels, Height, Width] (Standard Comfy Format where Batch=Time) - - is_5d = len(source_samples.shape) == 5 - - if is_5d: - b, c, t, h, w = source_samples.shape - total_len = t - else: - b, c, h, w = source_samples.shape - total_len = b - - if ctx_lat_len >= total_len: - raise ValueError( - f"Context ({ctx_lat_len} latents) is larger than input video ({total_len} latents)." - ) - - # --- CREATE NEW TENSOR --- - new_samples = torch.zeros_like(source_samples) - - # --- SLICE & PASTE --- - if is_5d: - # Copy last N from source -> Paste to first N of new - # Shape: [Batch, Channel, TIME, Height, Width] - context_slice = source_samples[:, :, -ctx_lat_len:, :, :] - new_samples[:, :, :ctx_lat_len, :, :] = context_slice - - # --- CREATE MASK (5D) --- - # Mask shape must be [Batch, 1, Time, Height, Width] - mask = torch.ones( - (b, 1, t, h, w), - dtype=source_samples.dtype, - device=source_samples.device, - ) - mask[:, :, :ctx_lat_len, :, :] = 0.0 # Lock the context frames - - else: - # Shape: [BATCH(Time), Channel, Height, Width] - context_slice = source_samples[-ctx_lat_len:, :, :, :] - new_samples[:ctx_lat_len, :, :, :] = context_slice - - # --- CREATE MASK (4D) --- - # Mask shape must be [Batch, 1, Height, Width] - mask = torch.ones( - (b, 1, h, w), - dtype=source_samples.dtype, - device=source_samples.device, - ) - mask[:ctx_lat_len, :, :, :] = 0.0 # Lock the context frames - - # --- OUTPUT --- - out = samples.copy() - out["samples"] = new_samples - out["noise_mask"] = mask - - # Handle batch index for correct noise generation in some samplers - if "batch_index" in out: - out["batch_index"] = [x for x in range(source_samples.shape[0])] - - return (out,)