diff --git a/__init__.py b/__init__.py
index b08b00d..702925e 100644
--- a/__init__.py
+++ b/__init__.py
@@ -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']
\ No newline at end of file
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
diff --git a/denoise_calculator.py b/denoise_calculator.py
deleted file mode 100644
index 7dc57c2..0000000
--- a/denoise_calculator.py
+++ /dev/null
@@ -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
\ No newline at end of file
diff --git a/nodes.py b/nodes.py
new file mode 100644
index 0000000..0411ec8
--- /dev/null
+++ b/nodes.py
@@ -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": """
+
+
+ """
+ }
+
+ 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
+
+
diff --git a/pyproject.toml b/pyproject.toml
index b0df329..d27118f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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]
diff --git a/resolution_picker.py b/resolution_picker.py
deleted file mode 100644
index 56c5eb9..0000000
--- a/resolution_picker.py
+++ /dev/null
@@ -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,)}
\ No newline at end of file
diff --git a/step_calculator.py b/step_calculator.py
deleted file mode 100644
index d37e22b..0000000
--- a/step_calculator.py
+++ /dev/null
@@ -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
diff --git a/utils.py b/utils.py
new file mode 100644
index 0000000..620cf89
--- /dev/null
+++ b/utils.py
@@ -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 {}
\ No newline at end of file