stage v0.1.5

This commit is contained in:
Sonny Box
2026-01-03 15:42:37 -08:00
parent f80920f2cf
commit 1830397400
6 changed files with 211 additions and 125 deletions
+1 -1
View File
@@ -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
+6 -5
View File
@@ -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"]
+3 -3
View File
@@ -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
+3 -2
View File
@@ -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]
+198
View File
@@ -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 } }
-114
View File
@@ -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,)