Aníbal Fuentes
+7
-9
@@ -1,18 +1,16 @@
|
||||
from .resolution_picker import EasyResolutionPicker
|
||||
from .step_calculator import StepCalculator
|
||||
from .denoise_calculator import DenoiseCalculator
|
||||
|
||||
from .nodes import EasyResolutionPicker, SmoothCurveMerger
|
||||
from .utils import UNETSave
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"EasyResolutionPicker": EasyResolutionPicker,
|
||||
"StepCalculator": StepCalculator,
|
||||
"DenoiseCalculator": DenoiseCalculator
|
||||
"UNETSave": UNETSave,
|
||||
"SmoothCurveMerger": SmoothCurveMerger
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"EasyResolutionPicker": "Easy Resolution Picker",
|
||||
"StepCalculator": "Step Calculator",
|
||||
"DenoiseCalculator": "Denoise Calculator"
|
||||
"UNETSave": "Save Diffusion Model",
|
||||
"SmoothCurveMerger": "Smooth Curve Merger"
|
||||
}
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
||||
@@ -1,36 +0,0 @@
|
||||
class DenoiseCalculator:
|
||||
NAME = "Denoise Calculator"
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"target_steps": ("INT", {
|
||||
"default": 20,
|
||||
"min": 1,
|
||||
"max": 500,
|
||||
"step": 1
|
||||
}),
|
||||
"ratio": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": 0.01,
|
||||
"max": 1.0,
|
||||
"step": 0.01
|
||||
})
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("INT", "FLOAT")
|
||||
RETURN_NAMES = ("total_steps", "denoise")
|
||||
FUNCTION = "calculate_denoise"
|
||||
OUTPUT_NODE = False
|
||||
CATEGORY = "UX Nodes"
|
||||
|
||||
def calculate_denoise(self, target_steps, ratio):
|
||||
total_steps = round(target_steps / ratio)
|
||||
denoise = ratio
|
||||
|
||||
return total_steps, denoise
|
||||
@@ -0,0 +1,213 @@
|
||||
import numpy as np
|
||||
from scipy.interpolate import interp1d
|
||||
|
||||
class EasyResolutionPicker:
|
||||
NAME = "Easy Resolution Picker"
|
||||
|
||||
def __init__(self):
|
||||
self.resolution_str = ""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"aspect_ratio": (
|
||||
[
|
||||
'1:1',
|
||||
'5:4',
|
||||
'4:3',
|
||||
'3:2',
|
||||
'7:4',
|
||||
'16:10',
|
||||
'16:9',
|
||||
'2:1'
|
||||
], { "default": "1:1" }),
|
||||
"orientation": (
|
||||
[
|
||||
'Horizontal',
|
||||
'Vertical'
|
||||
], { "default": "Horizontal" }),
|
||||
"length_type": (
|
||||
[
|
||||
'Long Side',
|
||||
'Short Side'
|
||||
], { "default": "Long Side" }),
|
||||
"side_length": ("INT", {
|
||||
"default": 1024,
|
||||
"min": 512,
|
||||
"max": 3072,
|
||||
"step": 2
|
||||
})
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("INT", "INT", "STRING")
|
||||
RETURN_NAMES = ("width", "height", "resolution_str")
|
||||
FUNCTION = "calculate_resolution"
|
||||
OUTPUT_NODE = False
|
||||
CATEGORY = "UX Nodes"
|
||||
|
||||
def calculate_resolution(self, aspect_ratio, orientation, length_type, side_length):
|
||||
aspect_ratios = {
|
||||
'1:1': (1, 1),
|
||||
'5:4': (5, 4),
|
||||
'4:3': (4, 3),
|
||||
'3:2': (3, 2),
|
||||
'7:4': (7, 4),
|
||||
'16:10': (16, 10),
|
||||
'16:9': (16, 9),
|
||||
'2:1': (2, 1)
|
||||
}
|
||||
|
||||
ar_width, ar_height = aspect_ratios[aspect_ratio]
|
||||
|
||||
if length_type == 'Long Side':
|
||||
long_side = side_length
|
||||
short_side = int(side_length * ar_height / ar_width)
|
||||
else:
|
||||
short_side = side_length
|
||||
long_side = int(side_length * ar_width / ar_height)
|
||||
|
||||
if orientation == 'Horizontal':
|
||||
width, height = long_side, short_side
|
||||
else:
|
||||
width, height = short_side, long_side
|
||||
|
||||
# Ensure dimensions are divisible by 2
|
||||
width = width - (width % 2)
|
||||
height = height - (height % 2)
|
||||
|
||||
# Generate resolution string
|
||||
resolution_str = f"{width}x{height}"
|
||||
self.resolution_str = resolution_str
|
||||
|
||||
# Update the node title with the resolution preview
|
||||
self.NAME = f"Easy Resolution Picker ({resolution_str})"
|
||||
|
||||
return width, height, resolution_str
|
||||
|
||||
@classmethod
|
||||
def WEB_DIRECTORY(cls):
|
||||
return {
|
||||
"html": """
|
||||
<div id="resolution_display" style="margin-top: 10px; font-weight: bold; color: #333;"></div>
|
||||
<script>
|
||||
function updateResolutionDisplay(resolution) {
|
||||
document.getElementById('resolution_display').innerText = 'Final Resolution: ' + resolution;
|
||||
}
|
||||
window.addEventListener('DOMContentLoaded', function() {
|
||||
window.updateResolutionDisplay = updateResolutionDisplay;
|
||||
});
|
||||
</script>
|
||||
"""
|
||||
}
|
||||
|
||||
def run(self, resolution_str):
|
||||
return {
|
||||
"ui": {"resolution_display": self.resolution_str},
|
||||
"result": (self.resolution_str,)
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
class SmoothCurveMerger:
|
||||
NAME = "Smooth Curve Merger"
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"curve1_points": ("STRING", {
|
||||
"default": "0.0,0.1,0.25,0.5,0.75,0.9,1.0",
|
||||
"multiline": False,
|
||||
}),
|
||||
"curve2_points": ("STRING", {
|
||||
"default": "1.0,0.9,0.75,0.5,0.25,0.1,0.0",
|
||||
"multiline": False,
|
||||
}),
|
||||
"adjustment_point": ("INT", {
|
||||
"default": 19,
|
||||
"min": 0,
|
||||
"max": 57,
|
||||
"step": 1
|
||||
}),
|
||||
"smoothness": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("FLOAT",) * 58 + ("STRING",)
|
||||
RETURN_NAMES = tuple(f"value_{i}" for i in range(58)) + ("values_preview",)
|
||||
FUNCTION = "generate_values"
|
||||
OUTPUT_NODE = False
|
||||
CATEGORY = "Utility"
|
||||
|
||||
def generate_values(self, curve1_points, curve2_points, adjustment_point, smoothness):
|
||||
# Parse the input strings into lists of floats
|
||||
curve1_points = list(map(float, curve1_points.split(',')))
|
||||
curve2_points = list(map(float, curve2_points.split(',')))
|
||||
|
||||
# Create x values based on the number of points
|
||||
x1 = np.linspace(0, 1, len(curve1_points))
|
||||
x2 = np.linspace(0, 1, len(curve2_points))
|
||||
|
||||
# Create the interpolation functions with adjustable smoothness
|
||||
interp_func1 = interp1d(x1, curve1_points, kind='quadratic' if smoothness > 0.5 else 'linear', fill_value="extrapolate")
|
||||
interp_func2 = interp1d(x2, curve2_points, kind='quadratic' if smoothness > 0.5 else 'linear', fill_value="extrapolate")
|
||||
|
||||
# Interpolate values
|
||||
curve1_values = interp_func1(np.linspace(0, 1, adjustment_point + 1))
|
||||
curve2_values = interp_func2(np.linspace(0, 1, 58 - adjustment_point - 1))
|
||||
|
||||
# Apply smoothness effect (smoothing out the linear interpolation)
|
||||
if smoothness < 1.0:
|
||||
curve1_values = self.apply_smoothness(curve1_values, smoothness)
|
||||
curve2_values = self.apply_smoothness(curve2_values, smoothness)
|
||||
|
||||
# Combine both curve segments
|
||||
values = np.concatenate((curve1_values, curve2_values))
|
||||
|
||||
# Generate a multiline string for the preview
|
||||
values_preview = "\n".join([f"{i}: {values[i]:.4f}" for i in range(58)])
|
||||
|
||||
# Generate ASCII art curve preview
|
||||
ascii_curve_preview = self.generate_ascii_art(values)
|
||||
|
||||
# Combine the numerical preview with the ASCII art preview
|
||||
preview = f"Numerical Preview:\n{values_preview}\n\nASCII Art Curve Preview:\n{ascii_curve_preview}"
|
||||
|
||||
return tuple(float(v) for v in values) + (preview,)
|
||||
|
||||
def apply_smoothness(self, values, smoothness):
|
||||
# Simple linear interpolation to smooth values
|
||||
smoothed_values = values * smoothness + np.linspace(values[0], values[-1], len(values)) * (1 - smoothness)
|
||||
return smoothed_values
|
||||
|
||||
def generate_ascii_art(self, values):
|
||||
# Determine the height of the ASCII art
|
||||
height = 10 # number of rows
|
||||
width = 58 # number of columns, one for each value
|
||||
|
||||
# Normalize the values to fit within the height
|
||||
normalized_values = np.interp(values, (0, 1), (height - 1, 0))
|
||||
|
||||
# Create a grid to hold the ASCII art
|
||||
grid = [[' ' for _ in range(width)] for _ in range(height)]
|
||||
|
||||
# Plot the values onto the grid
|
||||
for i, val in enumerate(normalized_values):
|
||||
grid[int(val)][i] = '*'
|
||||
|
||||
# Convert the grid to a string
|
||||
ascii_art = "\n".join(["".join(row) for row in grid])
|
||||
return ascii_art
|
||||
|
||||
|
||||
+2
-2
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "comfyui-ux-nodes"
|
||||
description = "Nodes: Easy Resolution Picker, Step Calculator, Denoise Calculator"
|
||||
version = "1.0.0"
|
||||
description = "Nodes: Easy Resolution Picker, Save Diffusion Model"
|
||||
version = "1.0.1"
|
||||
license = {file = "LICENSE"}
|
||||
|
||||
[project.urls]
|
||||
|
||||
@@ -1,99 +0,0 @@
|
||||
class EasyResolutionPicker:
|
||||
NAME = "Easy Resolution Picker"
|
||||
|
||||
def __init__(self):
|
||||
self.resolution_str = ""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"aspect_ratio": (
|
||||
[
|
||||
'1:1',
|
||||
'5:4',
|
||||
'4:3',
|
||||
'3:2',
|
||||
'16:10',
|
||||
'7:4',
|
||||
'16:9',
|
||||
'1.85:1',
|
||||
'2:1',
|
||||
'2.35:1'
|
||||
], { "default": "1:1" }),
|
||||
"orientation": (
|
||||
[
|
||||
'Horizontal',
|
||||
'Vertical'
|
||||
], { "default": "Horizontal" }),
|
||||
"length_type": (
|
||||
[
|
||||
'Long Side',
|
||||
'Short Side'
|
||||
], { "default": "Long Side" }),
|
||||
"side_length": ("INT", {
|
||||
"default": 1024,
|
||||
"min": 512,
|
||||
"max": 3072,
|
||||
"step": 2
|
||||
})
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("INT", "INT", "STRING")
|
||||
RETURN_NAMES = ("width", "height", "resolution_str")
|
||||
FUNCTION = "calculate_resolution"
|
||||
OUTPUT_NODE = False
|
||||
CATEGORY = "UX Nodes"
|
||||
|
||||
def calculate_resolution(self, aspect_ratio, orientation, length_type, side_length):
|
||||
aspect_ratios = {
|
||||
'1:1': (1, 1),
|
||||
'5:4': (5, 4),
|
||||
'4:3': (4, 3),
|
||||
'3:2': (3, 2),
|
||||
'16:10': (16, 10),
|
||||
'7:4': (7, 4),
|
||||
'16:9': (16, 9),
|
||||
'1.85:1': (1.85, 1),
|
||||
'2:1': (2, 1),
|
||||
'2.35:1': (2.35, 1)
|
||||
}
|
||||
|
||||
ar_width, ar_height = aspect_ratios[aspect_ratio]
|
||||
|
||||
if length_type == 'Long Side':
|
||||
long_side = side_length
|
||||
short_side = int(side_length * ar_height / ar_width)
|
||||
else:
|
||||
short_side = side_length
|
||||
long_side = int(side_length * ar_width / ar_height)
|
||||
|
||||
if orientation == 'Horizontal':
|
||||
width, height = long_side, short_side
|
||||
else:
|
||||
width, height = short_side, long_side
|
||||
|
||||
# Ensure dimensions are divisible by 2
|
||||
width = width - (width % 2)
|
||||
height = height - (height % 2)
|
||||
|
||||
# Generate resolution string
|
||||
resolution_str = f"{width}x{height}"
|
||||
self.resolution_str = resolution_str
|
||||
|
||||
# Update the node title with the resolution preview
|
||||
self.NAME = f"Easy Resolution Picker ({resolution_str})"
|
||||
|
||||
return width, height, resolution_str
|
||||
|
||||
@classmethod
|
||||
def OUTPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"resolution_text": ("STRING", {"default": "", "multiline": True})
|
||||
}
|
||||
}
|
||||
|
||||
def run(self, resolution_str):
|
||||
return {"ui": {"resolution_text": self.resolution_str}, "result": (self.resolution_str,)}
|
||||
@@ -1,36 +0,0 @@
|
||||
class StepCalculator:
|
||||
NAME = "Step Calculator"
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"steps": ("INT", {
|
||||
"default": 4,
|
||||
"min": 1,
|
||||
"max": 100,
|
||||
"step": 1
|
||||
}),
|
||||
"ratio": ("FLOAT", {
|
||||
"default": 0.5,
|
||||
"min": 0.01,
|
||||
"max": 1.0,
|
||||
"step": 0.01
|
||||
})
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("INT", "INT")
|
||||
RETURN_NAMES = ("calculated_steps", "start_at_step")
|
||||
FUNCTION = "calculate_steps"
|
||||
OUTPUT_NODE = False
|
||||
CATEGORY = "UX Nodes"
|
||||
|
||||
def calculate_steps(self, steps, ratio):
|
||||
calculated_steps = round(steps / ratio)
|
||||
start_at_step = round(calculated_steps - steps)
|
||||
|
||||
return calculated_steps, start_at_step
|
||||
@@ -0,0 +1,142 @@
|
||||
import torch
|
||||
from enum import Enum
|
||||
import logging
|
||||
|
||||
from comfy import model_management
|
||||
from comfy.ldm.models.autoencoder import AutoencoderKL, AutoencodingEngine
|
||||
from comfy.ldm.cascade.stage_a import StageA
|
||||
from comfy.ldm.cascade.stage_c_coder import StageC_coder
|
||||
from comfy.ldm.audio.autoencoder import AudioOobleckVAE
|
||||
import yaml
|
||||
|
||||
from comfy import clip_vision
|
||||
from comfy import gligen
|
||||
from comfy import diffusers_convert
|
||||
from comfy import model_detection
|
||||
|
||||
from comfy import sd1_clip
|
||||
from comfy import sdxl_clip
|
||||
import comfy.text_encoders.sd2_clip
|
||||
import comfy.text_encoders.sd3_clip
|
||||
import comfy.text_encoders.sa_t5
|
||||
import comfy.text_encoders.aura_t5
|
||||
import comfy.text_encoders.hydit
|
||||
import comfy.text_encoders.flux
|
||||
|
||||
import comfy.model_patcher
|
||||
import comfy.lora
|
||||
import comfy.t2i_adapter.adapter
|
||||
import comfy.supported_models_base
|
||||
import comfy.taesd.taesd
|
||||
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
import comfy.model_base
|
||||
import comfy.model_management
|
||||
import comfy.model_sampling
|
||||
|
||||
import torch
|
||||
import folder_paths
|
||||
import json
|
||||
import os
|
||||
|
||||
from comfy.cli_args import args
|
||||
|
||||
|
||||
def save_diffusion_model(output_path, model, metadata=None, extra_keys={}):
|
||||
# Ensure the model is loaded on the GPU for saving
|
||||
model_management.load_models_gpu([model], force_patch_weights=True)
|
||||
|
||||
# Get the state dictionary for saving, excluding CLIP and VAE
|
||||
sd = model.model.state_dict_for_saving()
|
||||
|
||||
# Add any extra keys to the state dictionary
|
||||
for k in extra_keys:
|
||||
sd[k] = extra_keys[k]
|
||||
|
||||
# Ensure all tensors are contiguous before saving
|
||||
for k in sd:
|
||||
t = sd[k]
|
||||
if not t.is_contiguous():
|
||||
sd[k] = t.contiguous()
|
||||
|
||||
# Save the state dictionary to the specified output path
|
||||
comfy.utils.save_torch_file(sd, output_path, metadata=metadata)
|
||||
|
||||
|
||||
def save_model(model, filename_prefix=None, output_dir=None, prompt=None, extra_pnginfo=None):
|
||||
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, output_dir)
|
||||
prompt_info = ""
|
||||
if prompt is not None:
|
||||
prompt_info = json.dumps(prompt)
|
||||
|
||||
metadata = {}
|
||||
|
||||
enable_modelspec = True
|
||||
if isinstance(model.model, comfy.model_base.SDXL):
|
||||
if isinstance(model.model, comfy.model_base.SDXL_instructpix2pix):
|
||||
metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-edit"
|
||||
else:
|
||||
metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-base"
|
||||
elif isinstance(model.model, comfy.model_base.SDXLRefiner):
|
||||
metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-refiner"
|
||||
elif isinstance(model.model, comfy.model_base.SVD_img2vid):
|
||||
metadata["modelspec.architecture"] = "stable-video-diffusion-img2vid-v1"
|
||||
elif isinstance(model.model, comfy.model_base.SD3):
|
||||
metadata["modelspec.architecture"] = "stable-diffusion-v3-medium" #TODO: other SD3 variants
|
||||
else:
|
||||
enable_modelspec = False
|
||||
|
||||
if enable_modelspec:
|
||||
metadata["modelspec.sai_model_spec"] = "1.0.0"
|
||||
metadata["modelspec.implementation"] = "sgm"
|
||||
metadata["modelspec.title"] = "{} {}".format(filename, counter)
|
||||
|
||||
#TODO:
|
||||
# "stable-diffusion-v1", "stable-diffusion-v1-inpainting", "stable-diffusion-v2-512",
|
||||
# "stable-diffusion-v2-768-v", "stable-diffusion-v2-unclip-l", "stable-diffusion-v2-unclip-h",
|
||||
# "v2-inpainting"
|
||||
|
||||
extra_keys = {}
|
||||
model_sampling = model.get_model_object("model_sampling")
|
||||
if isinstance(model_sampling, comfy.model_sampling.ModelSamplingContinuousEDM):
|
||||
if isinstance(model_sampling, comfy.model_sampling.V_PREDICTION):
|
||||
extra_keys["edm_vpred.sigma_max"] = torch.tensor(model_sampling.sigma_max).float()
|
||||
extra_keys["edm_vpred.sigma_min"] = torch.tensor(model_sampling.sigma_min).float()
|
||||
|
||||
if model.model.model_type == comfy.model_base.ModelType.EPS:
|
||||
metadata["modelspec.predict_key"] = "epsilon"
|
||||
elif model.model.model_type == comfy.model_base.ModelType.V_PREDICTION:
|
||||
metadata["modelspec.predict_key"] = "v"
|
||||
|
||||
if not args.disable_metadata:
|
||||
metadata["prompt"] = prompt_info
|
||||
if extra_pnginfo is not None:
|
||||
for x in extra_pnginfo:
|
||||
metadata[x] = json.dumps(extra_pnginfo[x])
|
||||
|
||||
output_checkpoint = f"{filename}_{counter:05}_.safetensors"
|
||||
output_checkpoint = os.path.join(full_output_folder, output_checkpoint)
|
||||
|
||||
save_diffusion_model(output_checkpoint, model, metadata=metadata, extra_keys=extra_keys)
|
||||
|
||||
|
||||
class UNETSave:
|
||||
NAME = "Save Diffusion Model"
|
||||
def __init__(self):
|
||||
self.output_dir = folder_paths.get_output_directory()
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "model": ("MODEL",),
|
||||
"filename_prefix": ("STRING", {"default": "checkpoints/ComfyUI"}),},
|
||||
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},}
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "save"
|
||||
OUTPUT_NODE = True
|
||||
|
||||
CATEGORY = "advanced/model_merging"
|
||||
|
||||
def save(self, model, filename_prefix, prompt=None, extra_pnginfo=None):
|
||||
save_model(model, filename_prefix=filename_prefix, output_dir=self.output_dir, prompt=prompt, extra_pnginfo=extra_pnginfo)
|
||||
return {}
|
||||
Reference in New Issue
Block a user