171 lines
7.7 KiB
Python
171 lines
7.7 KiB
Python
# Import necessary libraries
|
|
import json
|
|
import torch
|
|
import math # Import math for isnan
|
|
|
|
class TWanSigmaGraph:
|
|
"""
|
|
A ComfyUI node that generates a sigma schedule tensor based on user-defined
|
|
points edited via a custom graph widget in the UI. It also outputs the
|
|
number of steps used. The schedule typically decreases from high sigma
|
|
(noise) to low sigma.
|
|
"""
|
|
# Define a small epsilon for floating point comparisons
|
|
EPSILON = 1e-6
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
default_points = json.dumps([
|
|
{"x": 0.0, "y": 1.0},
|
|
{"x": 1.0, "y": 0.0}
|
|
])
|
|
return {
|
|
"required": {
|
|
"steps": ("INT", {
|
|
"default": 20, "min": 1, "max": 1000,
|
|
}),
|
|
# 'graph_data': JSON string representing points [{x, y}, ...].
|
|
# This is now primarily for storing the state set by the JS widget.
|
|
# It no longer has a 'widget' association in Python.
|
|
"graph_data": ("STRING", {
|
|
"default": default_points,
|
|
"multiline": True, # Keep multiline for easier debugging if needed
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("SIGMAS", "INT",)
|
|
RETURN_NAMES = ("sigmas", "steps",)
|
|
FUNCTION = "calculate_sigmas"
|
|
CATEGORY = "sampling/custom_sampling"
|
|
|
|
def _validate_and_clean_points(self, points_data_str):
|
|
""" Parses, validates, and cleans the points data. """
|
|
points = []
|
|
default_points_list = [{"x": 0.0, "y": 1.0}, {"x": 1.0, "y": 0.0}]
|
|
try:
|
|
points_data = json.loads(points_data_str)
|
|
if not isinstance(points_data, list):
|
|
raise ValueError("Graph data is not a list.")
|
|
|
|
# Filter for valid point structure and numeric types, check for NaN/Infinity
|
|
valid_points = []
|
|
for p in points_data:
|
|
if isinstance(p, dict) and 'x' in p and 'y' in p and \
|
|
isinstance(p['x'], (int, float)) and not math.isnan(p['x']) and not math.isinf(p['x']) and \
|
|
isinstance(p['y'], (int, float)) and not math.isnan(p['y']) and not math.isinf(p['y']):
|
|
valid_points.append({"x": float(p['x']), "y": float(p['y'])}) # Ensure float type
|
|
else:
|
|
print(f"[TWanSigmaGraph Warning] Ignoring invalid point data: {p}")
|
|
|
|
points = valid_points
|
|
if len(points) != len(points_data):
|
|
print("[TWanSigmaGraph Warning] Some points in graph_data were invalid and ignored.")
|
|
|
|
except (json.JSONDecodeError, ValueError, TypeError) as e:
|
|
print(f"[TWanSigmaGraph Warning] Invalid graph_data input: {e}. Using default points.")
|
|
return default_points_list # Return default on parse error
|
|
|
|
if not points:
|
|
print("[TWanSigmaGraph Warning] No valid points found after filtering. Using default points.")
|
|
return default_points_list
|
|
|
|
# Ensure boundary points exist using a tolerance
|
|
has_start = any(abs(p['x'] - 0.0) < self.EPSILON for p in points)
|
|
has_end = any(abs(p['x'] - 1.0) < self.EPSILON for p in points)
|
|
|
|
if not has_start:
|
|
# Find point closest to x=0 to determine y-value, or default to 1.0
|
|
start_y = min(points, key=lambda p: abs(p['x'] - 0.0))['y'] if points else 1.0
|
|
points.append({"x": 0.0, "y": start_y})
|
|
print("[TWanSigmaGraph Info] Added missing start point (x=0).")
|
|
if not has_end:
|
|
# Find point closest to x=1 to determine y-value, or default to 0.0
|
|
end_y = min(points, key=lambda p: abs(p['x'] - 1.0))['y'] if points else 0.0
|
|
points.append({"x": 1.0, "y": end_y})
|
|
print("[TWanSigmaGraph Info] Added missing end point (x=1).")
|
|
|
|
# Sort points by x-coordinate
|
|
points.sort(key=lambda p: p["x"])
|
|
|
|
# Remove duplicate points based on x-coordinate with tolerance
|
|
unique_points = []
|
|
if points:
|
|
unique_points.append(points[0])
|
|
last_x = points[0]['x']
|
|
for i in range(1, len(points)):
|
|
if abs(points[i]["x"] - last_x) > self.EPSILON:
|
|
unique_points.append(points[i])
|
|
last_x = points[i]['x']
|
|
else:
|
|
print(f"[TWanSigmaGraph Warning] Removing duplicate point near x={points[i]['x']}.")
|
|
|
|
|
|
# Ensure minimum number of points (e.g., 2 for interpolation)
|
|
if len(unique_points) < 2:
|
|
print("[TWanSigmaGraph Warning] Not enough unique points after cleanup. Using default points.")
|
|
return default_points_list
|
|
|
|
return unique_points
|
|
|
|
def calculate_sigmas(self, steps, graph_data):
|
|
"""
|
|
Calculates a sigma schedule tensor and returns it along with the steps.
|
|
"""
|
|
steps = max(1, int(steps))
|
|
points = self._validate_and_clean_points(graph_data)
|
|
num_sigmas_to_generate = steps
|
|
|
|
# --- Perform Interpolation ---
|
|
sigma_values = []
|
|
current_point_idx = 0
|
|
|
|
# Handle steps = 1 case separately for clarity
|
|
if steps == 1:
|
|
# Find point at x=0
|
|
start_point = next((p for p in points if abs(p['x'] - 0.0) < self.EPSILON), points[0])
|
|
sigma_values.append(max(0.001, start_point['y']))
|
|
else:
|
|
for i in range(num_sigmas_to_generate):
|
|
step_progress = i / (steps - 1)
|
|
step_progress = min(1.0, max(0.0, step_progress)) # Clamp progress [0, 1]
|
|
|
|
# Find the correct segment [p1, p2] for interpolation
|
|
# Advance index while the *next* point's x is less than current progress
|
|
while (current_point_idx < len(points) - 2 and
|
|
points[current_point_idx + 1]["x"] < step_progress - self.EPSILON):
|
|
current_point_idx += 1
|
|
|
|
p1 = points[current_point_idx]
|
|
# Ensure p2 is the next point, or clamp to the last point if already there
|
|
p2 = points[min(current_point_idx + 1, len(points) - 1)]
|
|
|
|
x_diff = p2["x"] - p1["x"]
|
|
sigma = 0.0
|
|
|
|
# Check for vertical segment or identical points
|
|
if x_diff <= self.EPSILON:
|
|
# If progress is closer to p2.x, use p2.y, otherwise use p1.y
|
|
sigma = p2["y"] if abs(step_progress - p2["x"]) < abs(step_progress - p1["x"]) else p1["y"]
|
|
else:
|
|
# Standard linear interpolation
|
|
# Clamp progress to the segment bounds before calculating ratio
|
|
clamped_progress = max(p1["x"], min(p2["x"], step_progress))
|
|
ratio = (clamped_progress - p1["x"]) / x_diff
|
|
# Ensure ratio is valid [0, 1] due to potential float errors
|
|
ratio = max(0.0, min(1.0, ratio))
|
|
sigma = p1["y"] + ratio * (p2["y"] - p1["y"])
|
|
|
|
# Ensure sigma is positive and non-zero
|
|
sigma_values.append(max(0.001, sigma))
|
|
|
|
# Convert the final list of sigma values to a PyTorch tensor
|
|
sigmas_tensor = torch.tensor(sigma_values, dtype=torch.float32)
|
|
|
|
# print(f"[TWanSigmaGraph Debug] Calculated Sigmas ({len(sigma_values)} values, first 5): {sigmas_tensor[:5]}...")
|
|
return (sigmas_tensor, steps,)
|
|
|
|
# --- Node Registration ---
|
|
NODE_CLASS_MAPPINGS = { "TWanSigmaGraph": TWanSigmaGraph }
|
|
NODE_DISPLAY_NAME_MAPPINGS = { "TWanSigmaGraph": "Sigma Schedule Graph" }
|