stage v0.1.5
This commit is contained in:
@@ -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
@@ -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"]
|
||||
|
||||
@@ -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
@@ -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
@@ -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 } }
|
||||
@@ -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,)
|
||||
Reference in New Issue
Block a user