Pirog17000
7821135b0f
- Refactors KSamplers into `node_modules` for better organization. - Implements a robust tiling system in KSamplerMultiSeedPlus that uses gradient-mask inpainting to eliminate seams. - Fixes critical bugs related to VAE dimension mismatches, incorrect tensor shapes, and uneven image splits. - Corrects state pollution bugs, allowing for stable batch processing and chaining of multiple KSampler nodes. - Updates README with details on new tiling functionality and refactoring.
645 lines
34 KiB
Python
645 lines
34 KiB
Python
|
|
"""
|
|
Enhanced KSampler nodes with tiling support
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import torch
|
|
import importlib.util
|
|
import numpy as np
|
|
|
|
# Import our advanced noise generation library
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
from noise_generator import create_spectral_diverse_noise, create_hierarchical_noise
|
|
|
|
# Add the ComfyUI root directory to Python path to access the main nodes
|
|
comfy_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
|
if comfy_dir not in sys.path:
|
|
sys.path.insert(0, comfy_dir)
|
|
|
|
# Import the common_ksampler function, MAX_RESOLUTION from the main nodes.py
|
|
try:
|
|
from nodes import common_ksampler, MAX_RESOLUTION
|
|
except ImportError:
|
|
# Fallback: try to import from nodes module in current directory
|
|
spec = importlib.util.spec_from_file_location("nodes", os.path.join(comfy_dir, "nodes.py"))
|
|
nodes_module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(nodes_module)
|
|
common_ksampler = nodes_module.common_ksampler
|
|
MAX_RESOLUTION = nodes_module.MAX_RESOLUTION
|
|
|
|
# Import required ComfyUI modules
|
|
import comfy.samplers
|
|
import comfy.utils
|
|
import comfy.model_management
|
|
|
|
|
|
class KSamplerMultiSeed:
|
|
"""
|
|
KSampler Multi-Seed Node
|
|
|
|
This node is extracted from the main ComfyUI nodes.py and provides
|
|
multi-seed sampling functionality. It automatically uses the latest
|
|
common_ksampler function from ComfyUI, ensuring compatibility with updates.
|
|
|
|
Features optional noise injection for generating variations.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL", {"tooltip": "The model used for denoising the input latent."}),
|
|
"positive": ("CONDITIONING", {
|
|
"tooltip": "The conditioning describing the attributes you want to include in the image."}),
|
|
"negative": ("CONDITIONING", {
|
|
"tooltip": "The conditioning describing the attributes you want to exclude from the image."}),
|
|
"latent_image": ("LATENT", {"tooltip": "The latent image to denoise."}),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01,
|
|
"tooltip": "The amount of denoising applied."}),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000,
|
|
"tooltip": "The number of steps used in the denoising process."}),
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01,
|
|
"tooltip": "Classifier-Free Guidance scale."}),
|
|
"seed_count": ("INT", {"default": 1, "min": 1, "max": 1000,
|
|
"tooltip": "The number of seeds to generate images with."}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True,
|
|
"tooltip": "The starting random seed. It will be incremented for each image in the batch."}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "The algorithm used when sampling."}),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS,
|
|
{"tooltip": "The scheduler controls how noise is gradually removed."}),
|
|
"injected_noise": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01,
|
|
"tooltip": "Strength of noise injection for variation generation. 0.0=disabled, >0.0=blend base and variation noise."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
OUTPUT_TOOLTIPS = ("A batch of denoised latents, one for each seed.",)
|
|
FUNCTION = "sample"
|
|
|
|
CATEGORY = "pirog/sampling"
|
|
DESCRIPTION = "Generates multiple images by incrementing the seed for each generation for each latent in the input batch."
|
|
|
|
def slerp(self, val, low, high):
|
|
"""Spherical linear interpolation for noise blending"""
|
|
dims = low.shape
|
|
low = low.reshape(dims[0], -1)
|
|
high = high.reshape(dims[0], -1)
|
|
|
|
low_norm = low / torch.norm(low, dim=1, keepdim=True)
|
|
high_norm = high / torch.norm(high, dim=1, keepdim=True)
|
|
|
|
low_norm[low_norm != low_norm] = 0.0
|
|
high_norm[high_norm != high_norm] = 0.0
|
|
|
|
omega = torch.acos((low_norm * high_norm).sum(1))
|
|
so = torch.sin(omega)
|
|
res = (torch.sin((1.0 - val) * omega) / so).unsqueeze(1) * low + (torch.sin(val * omega) / so).unsqueeze(1) * high
|
|
|
|
return res.reshape(dims)
|
|
|
|
def sample(self, model, positive, negative, latent_image, denoise, steps, cfg, seed_count, seed, sampler_name,
|
|
scheduler, injected_noise=0.0):
|
|
output_latents = []
|
|
input_latents = latent_image["samples"]
|
|
|
|
total_generations = input_latents.shape[0] * seed_count
|
|
pbar = comfy.utils.ProgressBar(total_generations)
|
|
|
|
for latent_sample in input_latents:
|
|
for i in range(seed_count):
|
|
current_seed = seed + i
|
|
|
|
# Reshape the single sample to a batch of 1 for the sampler
|
|
work_latent = latent_sample.unsqueeze(0)
|
|
|
|
# Apply noise injection if enabled
|
|
if injected_noise > 0.0:
|
|
# Generate base noise from main seed
|
|
generator = torch.manual_seed(seed)
|
|
base_noise = torch.randn(work_latent.shape, generator=generator, device="cpu", dtype=work_latent.dtype)
|
|
|
|
# Generate variation noise from current seed (naturally varies per iteration)
|
|
generator = torch.manual_seed(current_seed + 12345) # Offset to ensure different noise
|
|
variation_noise = torch.randn(work_latent.shape, generator=generator, device="cpu", dtype=work_latent.dtype)
|
|
|
|
# Blend noises using SLERP
|
|
blended_noise = self.slerp(injected_noise, base_noise, variation_noise)
|
|
|
|
# Calculate sigma scaling factor
|
|
device = comfy.model_management.get_torch_device()
|
|
comfy.model_management.load_model_gpu(model)
|
|
sampler = comfy.samplers.KSampler(model, steps=steps, device=device, sampler=sampler_name,
|
|
scheduler=scheduler, denoise=1.0, model_options=model.model_options)
|
|
sigmas = sampler.sigmas
|
|
start_at_step = round(steps - steps * denoise)
|
|
end_at_step = steps
|
|
sigma = sigmas[start_at_step] - sigmas[end_at_step]
|
|
sigma /= model.model.latent_format.scale_factor
|
|
sigma = sigma.detach().cpu().item()
|
|
|
|
# Apply scaled noise to latent
|
|
work_latent = work_latent + blended_noise.to(work_latent.device) * sigma
|
|
|
|
latent_for_sampler = {"samples": work_latent}
|
|
|
|
# Call the common ksampler function from the main ComfyUI nodes
|
|
# This ensures we always use the latest version
|
|
result_latent, = common_ksampler(model, current_seed, steps, cfg, sampler_name, scheduler, positive,
|
|
negative, latent_for_sampler, denoise=denoise)
|
|
|
|
output_latents.append(result_latent["samples"])
|
|
pbar.update(1)
|
|
|
|
# Combine all the generated latents into a single batch tensor
|
|
final_samples = torch.cat(output_latents, dim=0)
|
|
|
|
# Create the final output dictionary
|
|
final_latent = {"samples": final_samples}
|
|
|
|
return (final_latent,)
|
|
|
|
|
|
class KSamplerMultiSeedPlus:
|
|
"""
|
|
KSamplerMultiSeed+ Node - Integrated Sampling Pipeline
|
|
|
|
This node combines the KSamplerMultiSeed functionality with integrated VAE encode/decode
|
|
and empty latent image generation. It provides a complete pipeline in one node:
|
|
- For denoise=1.0: Creates empty latent → samples → decodes to image
|
|
- For denoise<1.0: Encodes input image → samples → decodes to image
|
|
|
|
Uses minimal approach by calling vanilla ComfyUI functions directly.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.device = comfy.model_management.intermediate_device()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL", {"tooltip": "The model used for denoising."}),
|
|
"vae": ("VAE", {"tooltip": "The VAE model used for encoding/decoding."}),
|
|
"positive": ("CONDITIONING", {
|
|
"tooltip": "The conditioning describing the attributes you want to include in the image."}),
|
|
"negative": ("CONDITIONING", {
|
|
"tooltip": "The conditioning describing the attributes you want to exclude from the image."}),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01,
|
|
"tooltip": "The amount of denoising applied. 1.0=new image, <1.0=img2img."}),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000,
|
|
"tooltip": "The number of steps used in the denoising process."}),
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01,
|
|
"tooltip": "Classifier-Free Guidance scale."}),
|
|
"seed_count": ("INT", {"default": 1, "min": 1, "max": 1000,
|
|
"tooltip": "The number of seeds to generate images with."}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True,
|
|
"tooltip": "The starting random seed. It will be incremented for each image in the batch."}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "The algorithm used when sampling."}),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS,
|
|
{"tooltip": "The scheduler controls how noise is gradually removed."}),
|
|
"width": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8,
|
|
"tooltip": "The width of the generated image in pixels (used when denoise=1.0)."}),
|
|
"height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8,
|
|
"tooltip": "The height of the generated image in pixels (used when denoise=1.0)."}),
|
|
"noise_type": (["vanilla", "spectral-diverse", "hierarchical"], {
|
|
"default": "vanilla",
|
|
"tooltip": "🎲 Noise Generation Method:\n\n"
|
|
"• vanilla: Standard ComfyUI noise (torch.randn) - reliable baseline\n"
|
|
"• spectral-diverse: Frequency-controlled noise with pink/blue/hybrid patterns - enhanced diversity\n"
|
|
"• hierarchical: Multi-scale latent-aware noise with statistical modeling - maximum quality\n\n"
|
|
"Advanced methods produce more diverse and potentially higher quality results."
|
|
}),
|
|
"injected_noise": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01,
|
|
"tooltip": "Strength of noise injection for variation generation. 0.0=disabled, >0.0=blend base and variation noise."}),
|
|
"vertical_splits": ("INT", {"default": 1, "min": 1, "max": 8, "step": 1,
|
|
"tooltip": "Number of vertical splits (1 = no splitting, 2+ = process image in tiles)"}),
|
|
"horizontal_splits": ("INT", {"default": 1, "min": 1, "max": 8, "step": 1,
|
|
"tooltip": "Number of horizontal splits (1 = no splitting, 2+ = process image in tiles)"}),
|
|
"expansion_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 50.0, "step": 1.0,
|
|
"tooltip": "Expansion percentage for tile borders (0-50%). Helps reduce seam artifacts between tiles."}),
|
|
},
|
|
"optional": {
|
|
"input_image": ("IMAGE", {"tooltip": "Input image for img2img (used when denoise<1.0)."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
OUTPUT_TOOLTIPS = ("Generated images, one for each seed.",)
|
|
FUNCTION = "sample_integrated"
|
|
|
|
CATEGORY = "pirog/sampling"
|
|
DESCRIPTION = "Integrated sampling pipeline: Creates or encodes latent → multi-seed sampling → decodes to images. For denoise=1.0 uses width/height to create empty latent. For denoise<1.0 uses input_image."
|
|
|
|
def slerp(self, val, low, high):
|
|
"""Spherical linear interpolation for noise blending"""
|
|
dims = low.shape
|
|
low = low.reshape(dims[0], -1)
|
|
high = high.reshape(dims[0], -1)
|
|
|
|
low_norm = low / torch.norm(low, dim=1, keepdim=True)
|
|
high_norm = high / torch.norm(high, dim=1, keepdim=True)
|
|
|
|
low_norm[low_norm != low_norm] = 0.0
|
|
high_norm[high_norm != high_norm] = 0.0
|
|
|
|
omega = torch.acos((low_norm * high_norm).sum(1))
|
|
so = torch.sin(omega)
|
|
res = (torch.sin((1.0 - val) * omega) / so).unsqueeze(1) * low + (torch.sin(val * omega) / so).unsqueeze(1) * high
|
|
|
|
return res.reshape(dims)
|
|
|
|
def split_image_into_tiles(self, image, vertical_splits, horizontal_splits, expansion_percent):
|
|
"""Split image into tiles with expansion borders and generate proper gradient masks for img2img"""
|
|
_, height, width, channels = image.shape
|
|
|
|
tiles = []
|
|
masks = []
|
|
bboxes = []
|
|
|
|
# Use linspace to create perfectly spaced coordinates, handling any image size
|
|
y_coords = np.linspace(0, height, vertical_splits + 1, dtype=int)
|
|
x_coords = np.linspace(0, width, horizontal_splits + 1, dtype=int)
|
|
|
|
def create_expansion_gradient_mask(tile_h, tile_w, orig_x_in_tile, orig_y_in_tile, orig_w, orig_h):
|
|
"""Create gradient mask spanning full expansion width - 1.0 at core edges, 0.0 at expansion boundaries"""
|
|
mask = torch.ones((tile_h, tile_w), dtype=torch.float32, device=image.device)
|
|
|
|
# Top expansion gradient
|
|
if orig_y_in_tile > 0:
|
|
expansion_top = orig_y_in_tile
|
|
alpha = torch.linspace(1 / (expansion_top + 1), 1.0, steps=expansion_top, device=image.device)
|
|
mask[:expansion_top, :] = alpha.view(-1, 1)
|
|
|
|
# Bottom expansion gradient
|
|
bottom_expansion_start = orig_y_in_tile + orig_h
|
|
if bottom_expansion_start < tile_h:
|
|
expansion_bottom = tile_h - bottom_expansion_start
|
|
alpha = torch.linspace(1.0, 1 / (expansion_bottom + 1), steps=expansion_bottom, device=image.device)
|
|
mask[bottom_expansion_start:, :] = alpha.view(-1, 1)
|
|
|
|
# Left expansion gradient
|
|
if orig_x_in_tile > 0:
|
|
expansion_left = orig_x_in_tile
|
|
alpha = torch.linspace(1 / (expansion_left + 1), 1.0, steps=expansion_left, device=image.device)
|
|
mask[:, :expansion_left] *= alpha.view(1, -1)
|
|
|
|
# Right expansion gradient
|
|
right_expansion_start = orig_x_in_tile + orig_w
|
|
if right_expansion_start < tile_w:
|
|
expansion_right = tile_w - right_expansion_start
|
|
alpha = torch.linspace(1.0, 1 / (expansion_right + 1), steps=expansion_right, device=image.device)
|
|
mask[:, right_expansion_start:] *= alpha.view(1, -1)
|
|
|
|
return mask
|
|
|
|
for i in range(vertical_splits):
|
|
for j in range(horizontal_splits):
|
|
y1, y2 = y_coords[i], y_coords[i+1]
|
|
x1, x2 = x_coords[j], x_coords[j+1]
|
|
|
|
core_h = y2 - y1
|
|
core_w = x2 - x1
|
|
|
|
# Expansion is now calculated per-tile, which is more accurate for uneven splits
|
|
expansion_h = int(core_h * expansion_percent / 100)
|
|
expansion_w = int(core_w * expansion_percent / 100)
|
|
|
|
y_start = max(0, y1 - expansion_h)
|
|
y_end = min(height, y2 + expansion_h)
|
|
x_start = max(0, x1 - expansion_w)
|
|
x_end = min(width, x2 + expansion_w)
|
|
|
|
# Skip tiles that have no area
|
|
if y_start >= y_end or x_start >= x_end:
|
|
continue
|
|
|
|
tile = image[:, y_start:y_end, x_start:x_end, :]
|
|
tiles.append(tile)
|
|
|
|
orig_x_in_tile = x1 - x_start
|
|
orig_y_in_tile = y1 - y_start
|
|
|
|
actual_tile_h = y_end - y_start
|
|
actual_tile_w = x_end - x_start
|
|
gradient_mask = create_expansion_gradient_mask(actual_tile_h, actual_tile_w, orig_x_in_tile, orig_y_in_tile, core_w, core_h)
|
|
masks.append(gradient_mask)
|
|
|
|
bbox_info = {
|
|
'orig_x': x1, 'orig_y': y1, 'orig_w': core_w, 'orig_h': core_h,
|
|
'exp_x_start': x_start, 'exp_y_start': y_start,
|
|
'exp_x_end': x_end, 'exp_y_end': y_end,
|
|
'tile_index': i * horizontal_splits + j
|
|
}
|
|
bboxes.append(bbox_info)
|
|
|
|
# Return lists of tensors, as they may have different sizes
|
|
return tiles, masks, bboxes
|
|
|
|
def reassemble_tiles(self, original_image, processed_tiles, tile_masks, bboxes, vertical_splits, horizontal_splits):
|
|
"""Reassemble tiles using a robust weighted accumulation method for seamless blending."""
|
|
_, height, width, channels = original_image.shape
|
|
device = original_image.device
|
|
dtype = original_image.dtype
|
|
|
|
# Initialize accumulation tensors
|
|
result_acc = torch.zeros_like(original_image)
|
|
weight_sum_acc = torch.zeros((1, height, width, 1), device=device, dtype=dtype)
|
|
|
|
for i, (tile, mask) in enumerate(zip(processed_tiles, tile_masks)):
|
|
bbox_info = bboxes[i]
|
|
exp_y_start, exp_x_start = bbox_info['exp_y_start'], bbox_info['exp_x_start']
|
|
|
|
tile_h, tile_w = tile.shape[1], tile.shape[2]
|
|
|
|
# VAE encode/decode can alter dimensions by a few pixels if not a multiple of 8.
|
|
# We must resize the mask to match the actual processed tile dimensions.
|
|
if mask.shape[0] != tile_h or mask.shape[1] != tile_w:
|
|
mask = torch.nn.functional.interpolate(
|
|
mask.unsqueeze(0).unsqueeze(0),
|
|
size=(tile_h, tile_w),
|
|
mode='bilinear'
|
|
).squeeze(0).squeeze(0)
|
|
|
|
# Expand mask to have a channel dimension for broadcasting
|
|
mask_expanded = mask.unsqueeze(-1)
|
|
|
|
# Accumulate weighted pixel data
|
|
result_acc[0, exp_y_start:exp_y_start+tile_h, exp_x_start:exp_x_start+tile_w, :] += tile[0] * mask_expanded
|
|
|
|
# Accumulate weights
|
|
weight_sum_acc[0, exp_y_start:exp_y_start+tile_h, exp_x_start:exp_x_start+tile_w, :] += mask_expanded
|
|
|
|
# Avoid division by zero by clamping weights
|
|
weight_sum_clamped = torch.clamp(weight_sum_acc, min=1e-6)
|
|
|
|
# Normalize the accumulated result
|
|
final_result = result_acc / weight_sum_clamped
|
|
|
|
# Where no tile contributed (weight is near zero), use the original image content
|
|
untouched_mask = weight_sum_acc <= 1e-6
|
|
final_result[untouched_mask.expand_as(final_result)] = original_image[untouched_mask.expand_as(original_image)]
|
|
|
|
return final_result
|
|
|
|
def sample_integrated(self, model, vae, positive, negative, denoise, steps, cfg, seed_count, seed,
|
|
sampler_name, scheduler, width, height, noise_type="vanilla", injected_noise=0.0,
|
|
vertical_splits=1, horizontal_splits=1, expansion_percent=0.0, input_image=None):
|
|
|
|
# Determine if we're doing txt2img (denoise=1.0) or img2img (denoise<1.0)
|
|
is_txt2img = abs(denoise - 1.0) < 0.001
|
|
|
|
# Check if we need tiling
|
|
use_tiling = vertical_splits > 1 or horizontal_splits > 1
|
|
|
|
if is_txt2img:
|
|
# Create latent image (txt2img pipeline)
|
|
if input_image is not None:
|
|
print("Warning: input_image provided but denoise=1.0, ignoring input_image and using width/height")
|
|
|
|
# Generate initial latent - detect required channel count from model
|
|
try:
|
|
# Try to get the required latent channels from the model
|
|
if hasattr(model, 'model') and hasattr(model.model, 'latent_format'):
|
|
latent_channels = model.model.latent_format.latent_channels
|
|
else:
|
|
# Fallback: use 4 channels for SD/SDXL models
|
|
latent_channels = 4
|
|
|
|
# Additional fallback based on model type detection
|
|
if latent_channels is None or latent_channels <= 0:
|
|
# Try to detect if it's a Flux model by checking model properties
|
|
try:
|
|
# Check if this might be a Flux model based on common indicators
|
|
model_config = getattr(model.model, 'model_config', {})
|
|
if hasattr(model.model, 'diffusion_model'):
|
|
# Check for Flux-specific architecture indicators
|
|
diff_model = model.model.diffusion_model
|
|
if (hasattr(diff_model, 'x_embedder') or
|
|
getattr(model_config, 'unet_config', {}).get('in_channels') == 16):
|
|
latent_channels = 16 # Flux models use 16 channels
|
|
else:
|
|
latent_channels = 4 # SD/SDXL models use 4 channels
|
|
else:
|
|
latent_channels = 4 # Default fallback
|
|
except:
|
|
latent_channels = 4 # Safe fallback
|
|
|
|
except:
|
|
# Final fallback
|
|
latent_channels = 4
|
|
|
|
latent_shape = [1, latent_channels, height // 8, width // 8]
|
|
latent = torch.zeros(latent_shape, device=self.device)
|
|
latent_image = {"samples": latent}
|
|
else:
|
|
# Encode input image (img2img pipeline)
|
|
if input_image is None:
|
|
raise ValueError("input_image is required when denoise < 1.0 (img2img mode)")
|
|
# Use vanilla VAE encode approach
|
|
encoded_tensor = vae.encode(input_image[:,:,:,:3])
|
|
latent_image = {"samples": encoded_tensor}
|
|
|
|
# Handle tiling or standard processing
|
|
if use_tiling and not is_txt2img:
|
|
# Tiled img2img processing
|
|
return self._process_tiled_img2img(model, vae, positive, negative, denoise, steps, cfg, seed_count, seed,
|
|
sampler_name, scheduler, noise_type, injected_noise, input_image,
|
|
vertical_splits, horizontal_splits, expansion_percent)
|
|
else:
|
|
# Standard processing (txt2img or non-tiled img2img)
|
|
return self._process_standard(model, vae, positive, negative, latent_image, denoise, steps, cfg, seed_count, seed,
|
|
sampler_name, scheduler, noise_type, injected_noise, is_txt2img)
|
|
|
|
def _process_standard(self, model, vae, positive, negative, latent_image, denoise, steps, cfg, seed_count, seed,
|
|
sampler_name, scheduler, noise_type, injected_noise, is_txt2img):
|
|
"""Standard processing without tiling"""
|
|
output_latents = []
|
|
input_latents = latent_image["samples"]
|
|
|
|
total_generations = input_latents.shape[0] * seed_count
|
|
pbar = comfy.utils.ProgressBar(total_generations)
|
|
|
|
for latent_sample in input_latents:
|
|
for i in range(seed_count):
|
|
current_seed = seed + i
|
|
|
|
# For advanced noise methods in txt2img mode, replace zeros with our noise
|
|
if is_txt2img and noise_type != "vanilla":
|
|
# Get the shape of the current latent sample
|
|
base_latent = latent_sample.unsqueeze(0)
|
|
|
|
try:
|
|
if noise_type == "spectral-diverse":
|
|
# Generate noise directly on the target device with exact shape
|
|
noise = create_spectral_diverse_noise(
|
|
base_latent.shape, current_seed, device=base_latent.device, noise_type="hybrid")
|
|
# Ensure exact compatibility
|
|
current_latent = noise.to(device=base_latent.device, dtype=base_latent.dtype)
|
|
elif noise_type == "hierarchical":
|
|
# Generate noise directly on the target device with exact shape
|
|
noise = create_hierarchical_noise(
|
|
base_latent.shape, current_seed, device=base_latent.device, diversity_strength=1.0)
|
|
# Ensure exact compatibility
|
|
current_latent = noise.to(device=base_latent.device, dtype=base_latent.dtype)
|
|
else:
|
|
current_latent = base_latent
|
|
except Exception as e:
|
|
# Fallback to vanilla noise if advanced noise generation fails
|
|
print(f"Warning: Advanced noise generation failed ({e}), falling back to vanilla noise")
|
|
current_latent = base_latent
|
|
|
|
work_latent = current_latent
|
|
else:
|
|
# Use the existing latent (vanilla mode or img2img)
|
|
work_latent = latent_sample.unsqueeze(0)
|
|
|
|
# Apply noise injection if enabled
|
|
if injected_noise > 0.0:
|
|
# Generate base noise from main seed
|
|
generator = torch.manual_seed(seed)
|
|
base_noise = torch.randn(work_latent.shape, generator=generator, device="cpu", dtype=work_latent.dtype)
|
|
|
|
# Generate variation noise from current seed (naturally varies per iteration)
|
|
generator = torch.manual_seed(current_seed + 12345) # Offset to ensure different noise
|
|
variation_noise = torch.randn(work_latent.shape, generator=generator, device="cpu", dtype=work_latent.dtype)
|
|
|
|
# Blend noises using SLERP
|
|
blended_noise = self.slerp(injected_noise, base_noise, variation_noise)
|
|
|
|
# Calculate sigma scaling factor
|
|
device = comfy.model_management.get_torch_device()
|
|
comfy.model_management.load_model_gpu(model)
|
|
sampler = comfy.samplers.KSampler(model, steps=steps, device=device, sampler=sampler_name,
|
|
scheduler=scheduler, denoise=1.0, model_options=model.model_options)
|
|
sigmas = sampler.sigmas
|
|
start_at_step = round(steps - steps * denoise)
|
|
end_at_step = steps
|
|
sigma = sigmas[start_at_step] - sigmas[end_at_step]
|
|
sigma /= model.model.latent_format.scale_factor
|
|
sigma = sigma.detach().cpu().item()
|
|
|
|
# Apply scaled noise to latent
|
|
work_latent = work_latent + blended_noise.to(work_latent.device) * sigma
|
|
|
|
latent_for_sampler = {"samples": work_latent}
|
|
|
|
# Call the common ksampler function from the main ComfyUI nodes
|
|
# This ensures we always use the latest version
|
|
result_latent, = common_ksampler(model, current_seed, steps, cfg, sampler_name, scheduler, positive,
|
|
negative, latent_for_sampler, denoise=denoise)
|
|
|
|
output_latents.append(result_latent["samples"])
|
|
pbar.update(1)
|
|
|
|
# Combine all the generated latents into a single batch tensor
|
|
final_samples = torch.cat(output_latents, dim=0)
|
|
|
|
# Decode latents to images using vanilla VAE decode approach
|
|
images = vae.decode(final_samples)
|
|
if len(images.shape) == 5: # Combine batches
|
|
images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1])
|
|
|
|
return (images,)
|
|
|
|
def _prepare_inpaint_conditioning(self, conditioning, concat_latent, latent_mask):
|
|
c = []
|
|
for t in conditioning:
|
|
# Create a deep enough copy to avoid modifying the original conditioning
|
|
n = [t[0], t[1].copy()]
|
|
n[1]["concat_latent_image"] = concat_latent
|
|
# The mask is inverted for ComfyUI's inpainting model
|
|
n[1]["concat_mask"] = 1.0 - latent_mask
|
|
c.append(n)
|
|
return c
|
|
|
|
def _process_tiled_img2img(self, model, vae, positive, negative, denoise, steps, cfg, seed_count, seed,
|
|
sampler_name, scheduler, noise_type, injected_noise, input_image,
|
|
vertical_splits, horizontal_splits, expansion_percent):
|
|
"""Tiled img2img processing with proper batch support."""
|
|
if input_image is None:
|
|
raise ValueError("input_image is required for tiled processing")
|
|
|
|
# This list will hold the final results for all images and all seeds
|
|
all_results = []
|
|
|
|
# The main loop to iterate over each image in the input batch
|
|
for i in range(input_image.shape[0]):
|
|
single_image = input_image[i:i+1] # Keep it as a batch of 1 [1, H, W, C]
|
|
|
|
# Now, run the multi-seed process on this single image
|
|
for seed_idx in range(seed_count):
|
|
current_seed = seed + seed_idx
|
|
|
|
# Split the single image into tiles
|
|
tiles, tile_masks, bboxes = self.split_image_into_tiles(single_image, vertical_splits, horizontal_splits, expansion_percent)
|
|
|
|
processed_tiles = []
|
|
|
|
# Process each tile for the current image and seed
|
|
for tile_idx, tile in enumerate(tiles):
|
|
# Get the corresponding mask for this tile
|
|
tile_mask = tile_masks[tile_idx]
|
|
|
|
# Encode the original tile content for inpainting conditioning
|
|
concat_latent = vae.encode(tile)
|
|
|
|
# Apply noise injection if enabled
|
|
if injected_noise > 0.0:
|
|
generator = torch.manual_seed(seed)
|
|
base_noise = torch.randn(concat_latent.shape, generator=generator, device="cpu", dtype=concat_latent.dtype)
|
|
|
|
generator = torch.manual_seed(current_seed + tile_idx + 12345)
|
|
variation_noise = torch.randn(concat_latent.shape, generator=generator, device="cpu", dtype=concat_latent.dtype)
|
|
|
|
blended_noise = self.slerp(injected_noise, base_noise, variation_noise)
|
|
|
|
device = comfy.model_management.get_torch_device()
|
|
comfy.model_management.load_model_gpu(model)
|
|
sampler = comfy.samplers.KSampler(model, steps=steps, device=device, sampler=sampler_name,
|
|
scheduler=scheduler, denoise=1.0, model_options=model.model_options)
|
|
sigmas = sampler.sigmas
|
|
start_at_step = round(steps - steps * denoise)
|
|
end_at_step = steps
|
|
sigma = sigmas[start_at_step] - sigmas[end_at_step]
|
|
sigma /= model.model.latent_format.scale_factor
|
|
sigma = sigma.detach().cpu().item()
|
|
|
|
concat_latent = concat_latent + blended_noise.to(concat_latent.device) * sigma
|
|
|
|
# Downscale the pixel mask to latent dimensions
|
|
latent_h, latent_w = concat_latent.shape[2], concat_latent.shape[3]
|
|
latent_mask = torch.nn.functional.interpolate(
|
|
tile_mask.reshape((-1, 1, tile_mask.shape[-2], tile_mask.shape[-1])),
|
|
size=(latent_h, latent_w),
|
|
mode="bilinear"
|
|
)
|
|
|
|
# Set up the latent structure for the sampler
|
|
latent_for_sampler = {
|
|
"samples": concat_latent,
|
|
"noise_mask": latent_mask
|
|
}
|
|
|
|
# Set inpainting-specific conditioning info
|
|
positive_inp = self._prepare_inpaint_conditioning(positive, concat_latent, latent_mask)
|
|
negative_inp = self._prepare_inpaint_conditioning(negative, concat_latent, latent_mask)
|
|
|
|
# Sample tile
|
|
result_latent, = common_ksampler(
|
|
model, current_seed + tile_idx, steps, cfg, sampler_name, scheduler,
|
|
positive_inp, negative_inp, latent_for_sampler, denoise=denoise
|
|
)
|
|
|
|
# Decode tile back to image
|
|
tile_image = vae.decode(result_latent["samples"])
|
|
processed_tiles.append(tile_image)
|
|
|
|
# Reassemble the tiles for the current image and seed
|
|
result_image = self.reassemble_tiles(single_image, processed_tiles, tile_masks, bboxes, vertical_splits, horizontal_splits)
|
|
all_results.append(result_image)
|
|
|
|
# Stack all final results from all images and seeds into a single batch
|
|
return (torch.cat(all_results, dim=0),)
|