Add ComfyUI-TkNodes pack with custom nodes, workflow, documentation, and dependencies
This commit is contained in:
TensorKaze
@@ -0,0 +1,3 @@
|
||||
__pycache__/
|
||||
*.pyc
|
||||
.ruff_cache/
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2025 TensorKaze
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -0,0 +1,119 @@
|
||||
# ComfyUI-TkNodes
|
||||
|
||||
A collection of custom nodes for ComfyUI by TensorKaze, designed to enhance workflows with advanced sampling, latent manipulation, image processing, and model loading. Includes nodes optimized for Flux.1 as well as general-purpose utilities for ComfyUI.
|
||||
These nodes aim to streamline complex workflows, offering quality-of-life improvements for both Flux-based and general diffusion tasks. This is a work in progress, with node descriptions available in the ComfyUI interface (click the ? icon on each node for details).
|
||||
Installation
|
||||
Clone this repository into the custom_nodes folder of ComfyUI:
|
||||
bash
|
||||
|
||||
cd ComfyUI/custom_nodes
|
||||
git clone https://github.com/TensorKaze/ComfyUI-TkNodes
|
||||
|
||||
Install dependencies:
|
||||
bash
|
||||
|
||||
cd ComfyUI-TkNodes
|
||||
pip install -r requirements.txt
|
||||
|
||||
If you're using the portable version of ComfyUI, run:
|
||||
bash
|
||||
|
||||
ComfyUI_windows_portable/python_embeded/python.exe -m pip install -r ComfyUI/custom_nodes/ComfyUI-TkNodes/requirements.txt
|
||||
|
||||
Note: Ensure you have a PyTorch version compatible with your hardware. For example:
|
||||
bash
|
||||
|
||||
pip install torch --index-url https://download.pytorch.org/whl/cu121 # For CUDA 12.1
|
||||
pip install torch # For CPU
|
||||
|
||||
Recommended: PyTorch 2.0.0 or higher for Flux.1 nodes.
|
||||
|
||||
Restart ComfyUI to load the nodes.
|
||||
|
||||
Requirements
|
||||
ComfyUI updated to the latest version (recommended: as of May 24, 2025).
|
||||
|
||||
Flux.1 model (Dev or Schnell) for Flux-specific nodes (FluxAdvancedSampler, FluxLatentSampler).
|
||||
|
||||
Minimum VRAM: 6 GB (for Flux GGUF) or 8-12 GB (for Flux.1 Dev/Schnell).
|
||||
|
||||
Optional: Upscale models for LoadModelAndUpscaleImage (available via ComfyUI's model downloader or Hugging Face).
|
||||
|
||||
Nodes
|
||||
FluxAdvancedSampler
|
||||
Advanced sampling for Flux.1 with customizable guidance, denoising, and scheduler settings.
|
||||
Inputs: Model, conditioning, sampler name, noise seed, steps, denoise, scheduler, guidance, latent image, VAE.
|
||||
|
||||
Output: Decoded image.
|
||||
|
||||
Category: sampling/flux
|
||||
|
||||
Use case: Generate high-quality images with precise control over Flux sampling parameters.
|
||||
|
||||
FluxLatentSampler
|
||||
Generates empty latents for Flux.1 with adjustable shift parameters.
|
||||
Inputs: Model, width, height, batch size, max shift, base shift.
|
||||
|
||||
Outputs: Modified model and latent.
|
||||
|
||||
Category: advanced/model
|
||||
|
||||
Use case: Create latents for Flux workflows with custom resolutions and sampling adjustments.
|
||||
|
||||
LoadImageAndScaleToTotalPixels
|
||||
Loads an image from the input directory and scales it to a specified total pixel count (in megapixels).
|
||||
Inputs: Image file, upscale method, target megapixels.
|
||||
|
||||
Outputs: Scaled image and mask (from alpha channel, if present).
|
||||
|
||||
Category: image
|
||||
|
||||
Use case: Preprocess images to desired resolutions for diffusion or other workflows.
|
||||
|
||||
LoadModelAndUpscaleImage
|
||||
Loads an upscale model and applies it to enhance the resolution of an input image.
|
||||
Inputs: Image, upscale model name.
|
||||
|
||||
Output: Upscaled image.
|
||||
|
||||
Category: image/upscaling
|
||||
|
||||
Use case: Improve image quality using pretrained upscale models.
|
||||
|
||||
MultiLatentSelector
|
||||
Selects one of up to four input latents based on a selector, simplifying dynamic workflows.
|
||||
Inputs: Up to four latents, selector.
|
||||
|
||||
Output: Selected latent.
|
||||
|
||||
Category: utilities/latent
|
||||
|
||||
Use case: Switch between latents without complex rerouting in workflows.
|
||||
|
||||
MultiModelLoader
|
||||
Loads a diffusion model, dual CLIP models, VAE, and optionally applies a LoRA with customizable strengths.
|
||||
Inputs: UNET name, weight dtype, CLIP names (two), CLIP type, VAE name, LoRA name, model/CLIP strengths, bypass LoRA option, CLIP device.
|
||||
|
||||
Outputs: Diffusion model, CLIP, VAE.
|
||||
|
||||
Category: loaders
|
||||
|
||||
Use case: Efficiently load and configure models for diffusion workflows, with LoRA support.
|
||||
|
||||
Known Limitations
|
||||
Flux-specific nodes (FluxAdvancedSampler, FluxLatentSampler) require Flux.1 models, which may need 6-12 GB of VRAM depending on the model.
|
||||
|
||||
LoadModelAndUpscaleImage requires compatible upscale models, downloadable separately.
|
||||
|
||||
Ensure PyTorch is installed with the correct CUDA version for GPU users to avoid compatibility issues.
|
||||
|
||||
Some nodes may not work with outdated ComfyUI versions or incompatible model types.
|
||||
|
||||
Heavy workflows (e.g., large batches or high-resolution images) may require significant computational resources.
|
||||
|
||||
Support
|
||||
For issues or feature requests, open an issue on this repository: GitHub Issues.
|
||||
|
||||
Check the ? icon on each node in ComfyUI for detailed descriptions.
|
||||
|
||||
Join the ComfyUI community on Reddit or Discord for additional support.
|
||||
+40
@@ -0,0 +1,40 @@
|
||||
from .multi_latent_selector import NODE_CLASS_MAPPINGS as MultiLatentSelectorMappings
|
||||
from .multi_latent_selector import NODE_DISPLAY_NAME_MAPPINGS as MultiLatentSelectorDisplayNames
|
||||
from .flux_latent_sampler import NODE_CLASS_MAPPINGS as FluxLatentSamplerMappings
|
||||
from .flux_latent_sampler import NODE_DISPLAY_NAME_MAPPINGS as FluxLatentSamplerDisplayNames
|
||||
from .flux_advanced_sampler import NODE_CLASS_MAPPINGS as FluxAdvancedSamplerMappings
|
||||
from .flux_advanced_sampler import NODE_DISPLAY_NAME_MAPPINGS as FluxAdvancedSamplerDisplayNames
|
||||
from .multi_model_loader import NODE_CLASS_MAPPINGS as MultiModelLoaderMappings
|
||||
from .multi_model_loader import NODE_DISPLAY_NAME_MAPPINGS as MultiModelLoaderDisplayNames
|
||||
from .load_image_and_scale import NODE_CLASS_MAPPINGS as LoadImageAndScaleMappings
|
||||
from .load_image_and_scale import NODE_DISPLAY_NAME_MAPPINGS as LoadImageAndScaleDisplayNames
|
||||
from .load_model_and_upscale import NODE_CLASS_MAPPINGS as LoadModelAndUpscaleMappings
|
||||
from .load_model_and_upscale import NODE_DISPLAY_NAME_MAPPINGS as LoadModelAndUpscaleDisplayNames
|
||||
from .vae_encode_optional import NODE_CLASS_MAPPINGS as VAEEncodeOptionalMappings
|
||||
from .vae_encode_optional import NODE_DISPLAY_NAME_MAPPINGS as VAEEncodeOptionalDisplayNames
|
||||
from .repeat_latent_batch_optional import NODE_CLASS_MAPPINGS as RepeatLatentBatchOptionalMappings
|
||||
from .repeat_latent_batch_optional import NODE_DISPLAY_NAME_MAPPINGS as RepeatLatentBatchOptionalDisplayNames
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
**MultiLatentSelectorMappings,
|
||||
**FluxLatentSamplerMappings,
|
||||
**MultiModelLoaderMappings,
|
||||
**FluxAdvancedSamplerMappings,
|
||||
**LoadImageAndScaleMappings,
|
||||
**LoadModelAndUpscaleMappings,
|
||||
**VAEEncodeOptionalMappings,
|
||||
**RepeatLatentBatchOptionalMappings
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
**MultiLatentSelectorDisplayNames,
|
||||
**FluxLatentSamplerDisplayNames,
|
||||
**MultiModelLoaderDisplayNames,
|
||||
**FluxAdvancedSamplerDisplayNames,
|
||||
**LoadImageAndScaleDisplayNames,
|
||||
**LoadModelAndUpscaleDisplayNames,
|
||||
**VAEEncodeOptionalDisplayNames,
|
||||
**RepeatLatentBatchOptionalDisplayNames
|
||||
}
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,91 @@
|
||||
import comfy.samplers
|
||||
import comfy.sample
|
||||
import latent_preview
|
||||
import torch
|
||||
import node_helpers
|
||||
from comfy.comfy_types import IO
|
||||
|
||||
class Guider_Basic(comfy.samplers.CFGGuider):
|
||||
def set_conds(self, positive):
|
||||
self.inner_set_conds({"positive": positive})
|
||||
|
||||
class FluxAdvancedSampler:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL", {"tooltip": "The diffusion model used for sampling."}),
|
||||
"conditioning": ("CONDITIONING", {"tooltip": "The conditioning to guide the sampling process."}),
|
||||
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "The name of the sampler to use for Flux sampling."}),
|
||||
"noise_seed": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xffffffffffffffff,
|
||||
"control_after_generate": True,
|
||||
"tooltip": "The random seed for noise generation."
|
||||
}),
|
||||
"steps": ("INT", {"default": 20, "min": 1, "max": 10000, "tooltip": "Number of sampling steps."}),
|
||||
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Denoising strength."}),
|
||||
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"tooltip": "The scheduler for sigma calculation."}),
|
||||
"guidance": ("FLOAT", {"default": 3.5, "min": 0.0, "max": 100.0, "step": 0.1, "tooltip": "Guidance scale for conditioning."}),
|
||||
"latent_image": ("LATENT", {"tooltip": "The input latent image to sample."}),
|
||||
"vae": (IO.VAE, {"tooltip": "The VAE model used to decode the latent image."}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = (IO.IMAGE,)
|
||||
RETURN_NAMES = ("IMAGE",)
|
||||
OUTPUT_TOOLTIPS = ("The decoded image from the sampled latent.",)
|
||||
FUNCTION = "sample_and_decode"
|
||||
|
||||
CATEGORY = "sampling/flux"
|
||||
DESCRIPTION = "Samples a latent image using Flux and decodes it to an image using a VAE."
|
||||
|
||||
def sample_and_decode(self, model, conditioning, sampler_name, noise_seed, steps, denoise, scheduler, guidance, latent_image, vae):
|
||||
# Modificar condicionamiento con guidance
|
||||
conditioning = node_helpers.conditioning_set_values(conditioning, {"guidance": guidance})
|
||||
|
||||
# Crear guider con Guider_Basic
|
||||
guider = Guider_Basic(model)
|
||||
guider.set_conds(conditioning)
|
||||
|
||||
# Obtener sampler
|
||||
sampler = comfy.samplers.sampler_object(sampler_name)
|
||||
|
||||
# Generar ruido
|
||||
noise = comfy.sample.prepare_noise(latent_image["samples"], noise_seed, latent_image.get("batch_index", None))
|
||||
|
||||
# Calcular sigmas
|
||||
total_steps = steps
|
||||
if denoise < 1.0:
|
||||
if denoise <= 0.0:
|
||||
sigmas = torch.FloatTensor([])
|
||||
else:
|
||||
total_steps = int(steps / denoise)
|
||||
sigmas = comfy.samplers.calculate_sigmas(model.get_model_object("model_sampling"), scheduler, total_steps).cpu()
|
||||
sigmas = sigmas[-(steps + 1):]
|
||||
|
||||
# Realizar muestreo
|
||||
latent = latent_image.copy()
|
||||
latent_image_samples = comfy.sample.fix_empty_latent_channels(model, latent_image["samples"])
|
||||
latent["samples"] = latent_image_samples
|
||||
noise_mask = latent.get("noise_mask", None)
|
||||
x0_output = {}
|
||||
callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output)
|
||||
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
|
||||
samples = guider.sample(noise, latent_image_samples, sampler, sigmas, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed)
|
||||
samples = samples.to(comfy.model_management.intermediate_device())
|
||||
|
||||
# Decodificar el latente con VAE
|
||||
decoded_samples = vae.decode(samples.to(vae.vae_dtype))
|
||||
|
||||
return (decoded_samples,)
|
||||
|
||||
# Mapeo de nodos
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FluxAdvancedSampler": FluxAdvancedSampler
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FluxAdvancedSampler": "Flux Advanced Sampler"
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
import torch
|
||||
import comfy.model_management
|
||||
import comfy.model_sampling
|
||||
import nodes
|
||||
|
||||
class FluxLatentSampler:
|
||||
def __init__(self):
|
||||
self.device = comfy.model_management.intermediate_device()
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"width": ("INT", {"default": 1024, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}),
|
||||
"height": ("INT", {"default": 1024, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}),
|
||||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
|
||||
"max_shift": ("FLOAT", {"default": 1.15, "min": 0.0, "max": 100.0, "step": 0.01}),
|
||||
"base_shift": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 100.0, "step": 0.01}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "LATENT")
|
||||
RETURN_NAMES = ("model", "latent")
|
||||
FUNCTION = "sample"
|
||||
CATEGORY = "advanced/model"
|
||||
|
||||
def sample(self, model, max_shift, base_shift, width, height, batch_size):
|
||||
# Generar el latente vacío
|
||||
latent = torch.zeros([batch_size, 16, height // 8, width // 8], device=self.device)
|
||||
latent_output = {"samples": latent}
|
||||
|
||||
# Ajustar parámetros de muestreo para Flux
|
||||
m = model.clone()
|
||||
|
||||
# Calcular shift
|
||||
x1 = 256
|
||||
x2 = 4096
|
||||
mm = (max_shift - base_shift) / (x2 - x1)
|
||||
b = base_shift - mm * x1
|
||||
shift = (width * height / (8 * 8 * 2 * 2)) * mm + b
|
||||
|
||||
# Crear la clase de muestreo
|
||||
sampling_base = comfy.model_sampling.ModelSamplingFlux
|
||||
sampling_type = comfy.model_sampling.CONST
|
||||
|
||||
class ModelSamplingAdvanced(sampling_base, sampling_type):
|
||||
pass
|
||||
|
||||
model_sampling = ModelSamplingAdvanced(model.model.model_config)
|
||||
model_sampling.set_parameters(shift=shift)
|
||||
m.add_object_patch("model_sampling", model_sampling)
|
||||
|
||||
return (m, latent_output)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FluxLatentSampler": FluxLatentSampler
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FluxLatentSampler": "Flux Latent Sampler"
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import torch
|
||||
import math
|
||||
from PIL import Image, ImageSequence
|
||||
import folder_paths
|
||||
import comfy
|
||||
from comfy.comfy_types import IO
|
||||
|
||||
class LoadImageAndScaleToTotalPixels:
|
||||
upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
input_dir = folder_paths.get_input_directory()
|
||||
files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
|
||||
files = folder_paths.filter_files_content_types(files, ["image"])
|
||||
return {
|
||||
"required": {
|
||||
"image": (["none"] + sorted(files), {"image_upload": True, "tooltip": "The image file to load from the input directory. Select 'none' to skip loading an image."}),
|
||||
"upscale_method": (cls.upscale_methods, {"tooltip": "The method used for scaling the image."}),
|
||||
"megapixels": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": 0.01,
|
||||
"max": 16.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "Target total pixels in megapixels (e.g., 1.0 for 1,048,576 pixels)."
|
||||
}),
|
||||
"bypass_sttp": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "If true, skips scaling and returns the original image and mask."
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = (IO.IMAGE, IO.MASK)
|
||||
RETURN_NAMES = ("IMAGE", "MASK")
|
||||
OUTPUT_TOOLTIPS = (
|
||||
"The loaded and scaled image (or original if bypassed), or None if no image is provided.",
|
||||
"The mask generated from the alpha channel (scaled to match the image or original if bypassed), or None if no image is provided."
|
||||
)
|
||||
FUNCTION = "load_and_scale"
|
||||
|
||||
CATEGORY = "image"
|
||||
DESCRIPTION = "Loads an image from the input directory, scales it to a specified total number of pixels (unless bypassed), and returns the scaled mask. If 'none' is selected, returns None for both outputs."
|
||||
|
||||
def load_and_scale(self, image, upscale_method, megapixels, bypass_sttp):
|
||||
# Si se selecciona "none", devolver None para ambas salidas
|
||||
if image == "none":
|
||||
return (None, None)
|
||||
|
||||
# Cargar la imagen (exactamente como LoadImage en nodes.py)
|
||||
image_path = os.path.join(folder_paths.get_input_directory(), image)
|
||||
i = None
|
||||
try:
|
||||
for img in ImageSequence.Iterator(Image.open(image_path)):
|
||||
i = img.convert("RGBA")
|
||||
break
|
||||
if i is None:
|
||||
raise ValueError(f"Failed to load image: {image_path}")
|
||||
except Exception:
|
||||
return (None, None)
|
||||
|
||||
img = np.array(i).astype(np.float32) / 255.0
|
||||
img = torch.from_numpy(img)[None,]
|
||||
|
||||
if img.shape[-1] == 4:
|
||||
mask = img[0, :, :, 3].clone() # (H, W), float32
|
||||
else:
|
||||
mask = torch.ones((img.shape[1], img.shape[2]), dtype=torch.float32, device=img.device)
|
||||
image_tensor = img[:, :, :, :3] # (1, H, W, 3)
|
||||
|
||||
# Si bypass_sttp está activado, devolver imagen y máscara originales
|
||||
if bypass_sttp:
|
||||
return (image_tensor, mask)
|
||||
|
||||
# Escalar la imagen
|
||||
samples = image_tensor.movedim(-1, 1) # (1, H, W, 3) -> (1, 3, H, W)
|
||||
total = int(megapixels * 1024 * 1024)
|
||||
scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2]))
|
||||
width = round(samples.shape[3] * scale_by)
|
||||
height = round(samples.shape[2] * scale_by)
|
||||
|
||||
scaled_samples = comfy.utils.common_upscale(samples, width, height, upscale_method, "disabled")
|
||||
scaled_samples = scaled_samples.movedim(1, -1) # (1, 3, H', W') -> (1, H', W', 3)
|
||||
|
||||
# Escalar la máscara
|
||||
mask_samples = mask[None, None, :, :] # (H, W) -> (1, 1, H, W)
|
||||
# Usar bicubic para la máscara si upscale_method es lanczos, si no, usar el mismo método
|
||||
mask_upscale_method = "bicubic" if upscale_method == "lanczos" else upscale_method
|
||||
scaled_mask = comfy.utils.common_upscale(mask_samples, width, height, mask_upscale_method, "disabled")
|
||||
scaled_mask = scaled_mask[0, 0, :, :] # (1, 1, H', W') -> (H', W')
|
||||
|
||||
return (scaled_samples, scaled_mask)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, image, upscale_method, megapixels, bypass_sttp, **kwargs):
|
||||
if image == "none":
|
||||
return None
|
||||
image_path = os.path.join(folder_paths.get_input_directory(), image)
|
||||
try:
|
||||
hash_value = comfy.utils.calculate_file_hash(image_path)
|
||||
return f"{hash_value}_{bypass_sttp}"
|
||||
except AttributeError:
|
||||
return f"{image_path}_{bypass_sttp}"
|
||||
|
||||
# Mapeo de nodos
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"LoadImageAndScaleToTotalPixels": LoadImageAndScaleToTotalPixels
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"LoadImageAndScaleToTotalPixels": "Load Image & Scale to Total Pixels"
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
import os
|
||||
import torch
|
||||
import folder_paths
|
||||
import comfy
|
||||
import comfy.utils
|
||||
from comfy.comfy_types import IO
|
||||
from spandrel import ModelLoader
|
||||
|
||||
class LoadModelAndUpscaleImage:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE", {"tooltip": "The image to upscale."}),
|
||||
"model_name": (folder_paths.get_filename_list("upscale_models"), {"tooltip": "The upscale model to use."}),
|
||||
"bypass_upscaler": ("BOOLEAN", {"default": False, "toggle": True, "label_on": "yes", "label_off": "no", "tooltip": "Bypass upscaling and return the original image if yes."}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "load_and_upscale"
|
||||
CATEGORY = "image/upscaling"
|
||||
DESCRIPTION = "Loads an upscale model and upscales the input image using the model, unless bypass_upscaler is yes."
|
||||
|
||||
def load_and_upscale(self, image, model_name, bypass_upscaler):
|
||||
# Si bypass_upscaler es True ("yes"), devolver la imagen original
|
||||
if bypass_upscaler:
|
||||
return (image,)
|
||||
|
||||
# Cargar el modelo de escalado (exactamente como UpscaleModelLoader en nodes_upscale_model.py)
|
||||
model_path = folder_paths.get_full_path_or_raise("upscale_models", model_name)
|
||||
sd = comfy.utils.load_torch_file(model_path, safe_load=True)
|
||||
if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd:
|
||||
sd = comfy.utils.state_dict_prefix_replace(sd, {"module.":""})
|
||||
upscale_model = ModelLoader().load_from_state_dict(sd).eval()
|
||||
|
||||
# Escalar la imagen con el modelo (exactamente como ImageUpscaleWithModel en nodes_upscale_model.py)
|
||||
device = comfy.model_management.get_torch_device()
|
||||
memory_required = comfy.model_management.module_size(upscale_model.model)
|
||||
memory_required += (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0
|
||||
memory_required += image.nelement() * image.element_size()
|
||||
comfy.model_management.free_memory(memory_required, device)
|
||||
|
||||
upscale_model.to(device)
|
||||
in_img = image.movedim(-1, -3).to(device) # (B, H, W, C) -> (B, C, H, W)
|
||||
|
||||
tile = 512
|
||||
overlap = 32
|
||||
oom = True
|
||||
while oom:
|
||||
try:
|
||||
steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap)
|
||||
pbar = comfy.utils.ProgressBar(steps)
|
||||
s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar)
|
||||
oom = False
|
||||
except comfy.model_management.OOM_EXCEPTION as e:
|
||||
tile //= 2
|
||||
if tile < 128:
|
||||
raise e
|
||||
|
||||
upscale_model.to("cpu")
|
||||
scaled_image = torch.clamp(s.movedim(-3, -1), min=0, max=1.0) # (B, C, H', W') -> (B, H', W', C)
|
||||
|
||||
# Devolver la imagen escalada
|
||||
return (scaled_image,)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, image, model_name, bypass_upscaler, **kwargs):
|
||||
model_path = folder_paths.get_full_path_or_raise("upscale_models", model_name)
|
||||
return comfy.utils.calculate_file_hash(model_path)
|
||||
|
||||
# Mapeo de nodos
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"LoadModelAndUpscaleImage": LoadModelAndUpscaleImage
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"LoadModelAndUpscaleImage": "Load Model & Upscale Image"
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
class MultiLatentSelector:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {},
|
||||
"optional": {
|
||||
"latent_1": ("LATENT",),
|
||||
"latent_2": ("LATENT",),
|
||||
"latent_3": ("LATENT",),
|
||||
"latent_4": ("LATENT",),
|
||||
"selector": (["latent_1", "latent_2", "latent_3", "latent_4"], {"default": "latent_1"}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "select_latent"
|
||||
CATEGORY = "utilities/latent"
|
||||
DISPLAY_NAME = "Multi-Latent Selector"
|
||||
|
||||
def select_latent(self, latent_1=None, latent_2=None, latent_3=None, latent_4=None, selector="latent_1"):
|
||||
if selector == "latent_1":
|
||||
if latent_1 is None:
|
||||
raise ValueError("latent_1 is not connected but was selected in the selector")
|
||||
return (latent_1,)
|
||||
elif selector == "latent_2":
|
||||
if latent_2 is None:
|
||||
raise ValueError("latent_2 is not connected but was selected in the selector")
|
||||
return (latent_2,)
|
||||
elif selector == "latent_3":
|
||||
if latent_3 is None:
|
||||
raise ValueError("latent_3 is not connected but was selected in the selector")
|
||||
return (latent_3,)
|
||||
elif selector == "latent_4":
|
||||
if latent_4 is None:
|
||||
raise ValueError("latent_4 is not connected but was selected in the selector")
|
||||
return (latent_4,)
|
||||
else:
|
||||
raise ValueError(f"Invalid selector value: {selector}")
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"MultiLatentSelector": MultiLatentSelector
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MultiLatentSelector": "Multi-Latent Selector"
|
||||
}
|
||||
@@ -0,0 +1,242 @@
|
||||
import torch
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
from comfy.comfy_types import IO, InputTypeDict
|
||||
|
||||
class MultiModelLoader:
|
||||
def __init__(self):
|
||||
# Cache para almacenar modelos y LoRA
|
||||
self.cached_model = None
|
||||
self.cached_clip = None
|
||||
self.cached_vae = None
|
||||
self.cached_lora = None
|
||||
self.last_unet_name = None
|
||||
self.last_weight_dtype = None
|
||||
self.last_clip_name1 = None
|
||||
self.last_clip_name2 = None
|
||||
self.last_clip_type = None
|
||||
self.last_clip_device = None
|
||||
self.last_vae_name = None
|
||||
self.last_lora_name = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> InputTypeDict:
|
||||
loras = ["none"] + folder_paths.get_filename_list("loras")
|
||||
return {
|
||||
"required": {
|
||||
"unet_name": (folder_paths.get_filename_list("diffusion_models"), {
|
||||
"tooltip": "The name of the diffusion model (UNET) to load."
|
||||
}),
|
||||
"weight_dtype": (["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"], {
|
||||
"tooltip": "The weight dtype for the diffusion model."
|
||||
}),
|
||||
"clip_name1": (folder_paths.get_filename_list("text_encoders"), {
|
||||
"tooltip": "The name of the first CLIP text encoder to load."
|
||||
}),
|
||||
"clip_name2": (folder_paths.get_filename_list("text_encoders"), {
|
||||
"tooltip": "The name of the second CLIP text encoder to load."
|
||||
}),
|
||||
"clip_type": (["sdxl", "sd3", "flux", "hunyuan_video", "hidream"], {
|
||||
"tooltip": "The type of CLIP configuration (e.g., flux: clip-l, t5)."
|
||||
}),
|
||||
"clip_device": (["default", "cpu"], {
|
||||
"advanced": True,
|
||||
"tooltip": "Device for loading CLIP models."
|
||||
}),
|
||||
"vae_name": (cls.vae_list(), {
|
||||
"tooltip": "The name of the VAE model to load."
|
||||
}),
|
||||
"lora_name": (loras, {
|
||||
"tooltip": "The name of the LoRA to load (select 'none' to skip loading)."
|
||||
}),
|
||||
"strength_model": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": -100.0,
|
||||
"max": 100.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "Strength of the LoRA applied to the diffusion model."
|
||||
}),
|
||||
"strength_clip": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": -100.0,
|
||||
"max": 100.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "Strength of the LoRA applied to the CLIP model."
|
||||
}),
|
||||
"bypass_lora": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"toggle": True,
|
||||
"label_on": "yes",
|
||||
"label_off": "no",
|
||||
"tooltip": "Bypass LoRA application to MODEL and CLIP if yes, but keep it loaded."
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def vae_list():
|
||||
vaes = folder_paths.get_filename_list("vae")
|
||||
approx_vaes = folder_paths.get_filename_list("vae_approx")
|
||||
sdxl_taesd_enc = False
|
||||
sdxl_taesd_dec = False
|
||||
sd1_taesd_enc = False
|
||||
sd1_taesd_dec = False
|
||||
sd3_taesd_enc = False
|
||||
sd3_taesd_dec = False
|
||||
f1_taesd_enc = False
|
||||
f1_taesd_dec = False
|
||||
|
||||
for v in approx_vaes:
|
||||
if v.startswith("taesd_decoder."):
|
||||
sd1_taesd_dec = True
|
||||
elif v.startswith("taesd_encoder."):
|
||||
sd1_taesd_enc = True
|
||||
elif v.startswith("taesdxl_decoder."):
|
||||
sdxl_taesd_dec = True
|
||||
elif v.startswith("taesdxl_encoder."):
|
||||
sdxl_taesd_enc = True
|
||||
elif v.startswith("taesd3_decoder."):
|
||||
sd3_taesd_dec = True
|
||||
elif v.startswith("taesd3_encoder."):
|
||||
sd3_taesd_enc = True
|
||||
elif v.startswith("taef1_encoder."):
|
||||
f1_taesd_dec = True
|
||||
elif v.startswith("taef1_decoder."):
|
||||
f1_taesd_enc = True
|
||||
if sd1_taesd_dec and sd1_taesd_enc:
|
||||
vaes.append("taesd")
|
||||
if sdxl_taesd_dec and sdxl_taesd_enc:
|
||||
vaes.append("taesdxl")
|
||||
if sd3_taesd_dec and sd3_taesd_enc:
|
||||
vaes.append("taesd3")
|
||||
if f1_taesd_dec and f1_taesd_enc:
|
||||
vaes.append("taef1")
|
||||
return vaes
|
||||
|
||||
@staticmethod
|
||||
def load_taesd(name):
|
||||
sd = {}
|
||||
approx_vaes = folder_paths.get_filename_list("vae_approx")
|
||||
encoder = next(filter(lambda a: a.startswith(f"{name}_encoder."), approx_vaes))
|
||||
decoder = next(filter(lambda a: a.startswith(f"{name}_decoder."), approx_vaes))
|
||||
enc = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", encoder))
|
||||
for k in enc:
|
||||
sd[f"taesd_encoder.{k}"] = enc[k]
|
||||
dec = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", decoder))
|
||||
for k in dec:
|
||||
sd[f"taesd_decoder.{k}"] = dec[k]
|
||||
if name == "taesd":
|
||||
sd["vae_scale"] = torch.tensor(0.18215)
|
||||
sd["vae_shift"] = torch.tensor(0.0)
|
||||
elif name == "taesdxl":
|
||||
sd["vae_scale"] = torch.tensor(0.13025)
|
||||
sd["vae_shift"] = torch.tensor(0.0)
|
||||
elif name == "taesd3":
|
||||
sd["vae_scale"] = torch.tensor(1.5305)
|
||||
sd["vae_shift"] = torch.tensor(0.0609)
|
||||
elif name == "taef1":
|
||||
sd["vae_scale"] = torch.tensor(0.3611)
|
||||
sd["vae_shift"] = torch.tensor(0.1159)
|
||||
return sd
|
||||
|
||||
RETURN_TYPES = (IO.MODEL, IO.CLIP, IO.VAE)
|
||||
RETURN_NAMES = ("MODEL", "CLIP", "VAE")
|
||||
OUTPUT_TOOLTIPS = (
|
||||
"The diffusion model (modified by LoRA if loaded and not bypassed).",
|
||||
"The CLIP model (modified by LoRA if loaded and not bypassed).",
|
||||
"The VAE model used for encoding/decoding latents."
|
||||
)
|
||||
FUNCTION = "load_unified"
|
||||
|
||||
CATEGORY = "loaders"
|
||||
DESCRIPTION = "Loads a diffusion model, dual CLIP models, VAE, and optionally applies a loaded LoRA to the model and CLIP with bypass option."
|
||||
|
||||
def load_unified(self, unet_name, weight_dtype, clip_name1, clip_name2, clip_type, clip_device, vae_name, lora_name, strength_model, strength_clip, bypass_lora):
|
||||
# Cargar modelo de difusión si no está cacheado o cambió
|
||||
if (self.cached_model is None or
|
||||
unet_name != self.last_unet_name or
|
||||
weight_dtype != self.last_weight_dtype):
|
||||
model_options = {}
|
||||
if weight_dtype == "fp8_e4m3fn":
|
||||
model_options["dtype"] = torch.float8_e4m3fn
|
||||
elif weight_dtype == "fp8_e4m3fn_fast":
|
||||
model_options["dtype"] = torch.float8_e4m3fn
|
||||
model_options["fp8_optimizations"] = True
|
||||
elif weight_dtype == "fp8_e5m2":
|
||||
model_options["dtype"] = torch.float8_e5m2
|
||||
|
||||
unet_path = folder_paths.get_full_path_or_raise("diffusion_models", unet_name)
|
||||
self.cached_model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options)
|
||||
self.last_unet_name = unet_name
|
||||
self.last_weight_dtype = weight_dtype
|
||||
|
||||
model = self.cached_model
|
||||
|
||||
# Cargar CLIPs si no están cacheados o cambiaron
|
||||
if (self.cached_clip is None or
|
||||
clip_name1 != self.last_clip_name1 or
|
||||
clip_name2 != self.last_clip_name2 or
|
||||
clip_type != self.last_clip_type or
|
||||
clip_device != self.last_clip_device):
|
||||
clip_type_enum = getattr(comfy.sd.CLIPType, clip_type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
|
||||
clip_path1 = folder_paths.get_full_path_or_raise("text_encoders", clip_name1)
|
||||
clip_path2 = folder_paths.get_full_path_or_raise("text_encoders", clip_name2)
|
||||
clip_options = {}
|
||||
if clip_device == "cpu":
|
||||
clip_options["load_device"] = clip_options["offload_device"] = torch.device("cpu")
|
||||
self.cached_clip = comfy.sd.load_clip(
|
||||
ckpt_paths=[clip_path1, clip_path2],
|
||||
embedding_directory=folder_paths.get_folder_paths("embeddings"),
|
||||
clip_type=clip_type_enum,
|
||||
model_options=clip_options
|
||||
)
|
||||
self.last_clip_name1 = clip_name1
|
||||
self.last_clip_name2 = clip_name2
|
||||
self.last_clip_type = clip_type
|
||||
self.last_clip_device = clip_device
|
||||
|
||||
clip = self.cached_clip
|
||||
|
||||
# Cargar VAE si no está cacheado o cambió
|
||||
if self.cached_vae is None or vae_name != self.last_vae_name:
|
||||
if vae_name in ["taesd", "taesdxl", "taesd3", "taef1"]:
|
||||
sd = self.load_taesd(vae_name)
|
||||
else:
|
||||
vae_path = folder_paths.get_full_path_or_raise("vae", vae_name)
|
||||
sd = comfy.utils.load_torch_file(vae_path)
|
||||
self.cached_vae = comfy.sd.VAE(sd=sd)
|
||||
self.cached_vae.throw_exception_if_invalid()
|
||||
self.last_vae_name = vae_name
|
||||
|
||||
vae = self.cached_vae
|
||||
|
||||
# Cargar LoRA solo si lora_name no es "none" y no está cacheado o cambió
|
||||
if lora_name != "none" and (self.cached_lora is None or lora_name != self.last_lora_name):
|
||||
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
||||
self.cached_lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
||||
self.last_lora_name = lora_name
|
||||
elif lora_name == "none":
|
||||
self.cached_lora = None
|
||||
self.last_lora_name = None
|
||||
|
||||
# Aplicar LoRA solo si está cargado, bypass_lora es False ("no"), y las intensidades no son 0
|
||||
if self.cached_lora is not None and not bypass_lora and (strength_model != 0 or strength_clip != 0):
|
||||
# Crear copias para no modificar el cache
|
||||
model = comfy.sd.load_lora_for_models(self.cached_model, self.cached_clip, self.cached_lora, strength_model, strength_clip)[0]
|
||||
clip = comfy.sd.load_lora_for_models(self.cached_model, self.cached_clip, self.cached_lora, strength_model, strength_clip)[1]
|
||||
else:
|
||||
# Usar cache directamente si bypass_lora es True ("yes") o no hay LoRA
|
||||
model = self.cached_model
|
||||
clip = self.cached_clip
|
||||
|
||||
return (model, clip, vae)
|
||||
|
||||
# Mapeo de nodos
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"MultiModelLoader": MultiModelLoader
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MultiModelLoader": "Multi-Model Loader"
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
import torch
|
||||
from comfy.comfy_types import IO
|
||||
|
||||
class RepeatLatentBatchOptional:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"samples": ("LATENT", {"tooltip": "The latent to repeat. If None, returns None."}),
|
||||
"amount": ("INT", {"default": 1, "min": 1, "max": 64, "tooltip": "Number of times to repeat the latent."}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
RETURN_NAMES = ("LATENT",)
|
||||
FUNCTION = "repeat"
|
||||
CATEGORY = "latent/batch"
|
||||
|
||||
def repeat(self, samples, amount):
|
||||
if samples is None:
|
||||
return (None,)
|
||||
|
||||
s = samples.copy()
|
||||
s_in = samples["samples"]
|
||||
s["samples"] = s_in.repeat((amount, 1, 1, 1))
|
||||
if "noise_mask" in samples and samples["noise_mask"].shape[0] > 1:
|
||||
masks = samples["noise_mask"]
|
||||
if masks.shape[0] < s_in.shape[0]:
|
||||
masks = masks.repeat(math.ceil(s_in.shape[0] / masks.shape[0]), 1, 1, 1)[:s_in.shape[0]]
|
||||
s["noise_mask"] = samples["noise_mask"].repeat((amount, 1, 1, 1))
|
||||
if "batch_index" in s:
|
||||
offset = max(s["batch_index"]) - min(s["batch_index"]) + 1
|
||||
s["batch_index"] = s["batch_index"] + [x + (i * offset) for i in range(1, amount) for x in s["batch_index"]]
|
||||
return (s,)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"RepeatLatentBatchOptional": RepeatLatentBatchOptional
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"RepeatLatentBatchOptional": "Repeat Latent Batch (Optional)"
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
torch
|
||||
numpy
|
||||
Pillow
|
||||
spandrel
|
||||
@@ -0,0 +1,45 @@
|
||||
import torch
|
||||
from comfy.comfy_types import IO
|
||||
|
||||
class VAEEncodeOptional:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"vae": ("VAE", {"tooltip": "The VAE model used for encoding the image to latent space."}),
|
||||
},
|
||||
"optional": {
|
||||
"image": ("IMAGE", {"tooltip": "The image to encode to latent space. If not provided, returns None."}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
RETURN_NAMES = ("LATENT",)
|
||||
OUTPUT_TOOLTIPS = ("The encoded latent image, or None if no image is provided.",)
|
||||
FUNCTION = "encode"
|
||||
|
||||
CATEGORY = "latent"
|
||||
DESCRIPTION = "Encodes an image to latent space using a VAE model. If no image is provided, acts as a bypass and returns None."
|
||||
|
||||
def encode(self, vae, image=None):
|
||||
# Modo bypass: si no hay imagen, devolver None
|
||||
if image is None:
|
||||
return (None,)
|
||||
|
||||
# Codificar la imagen con el VAE
|
||||
try:
|
||||
# Asegurarse de que la imagen solo use los canales RGB (ignorar alfa si existe)
|
||||
latent = vae.encode(image[:,:,:,:3])
|
||||
return ({"samples": latent},)
|
||||
except Exception as e:
|
||||
# En caso de error (por ejemplo, dimensiones inválidas), devolver None
|
||||
return (None,)
|
||||
|
||||
# Mapeo de nodos
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"VAEEncodeOptional": VAEEncodeOptional
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"VAEEncodeOptional": "VAE Encode (Optional)"
|
||||
}
|
||||
Reference in New Issue
Block a user