Version 2
Lots of new stuff
This commit is contained in:
kantsche
@@ -2,17 +2,12 @@
|
||||
|
||||
A custom node extension for ComfyUI that allows mixing multiple models during the sampling process for enhanced image generation.
|
||||
|
||||
## Please share workflows or good combinations here!
|
||||
|
||||
## Description
|
||||
|
||||
VRAM requirement: Fitting two or multiple models. I think around 12gb is minimum for two sdxl.
|
||||
|
||||
ComfyUI-MixMod provides a powerful way to combine multiple models during stable diffusion sampling. This extension introduces two primary mixing modes:
|
||||
|
||||
- **Team Mode**: Combines multiple models by weighted averaging of their predictions and applying their respective guidance scales.
|
||||
- If the models disagree on tags and prompts it will get bad, you have to prompt so both model are roughly on one page.
|
||||
- **2Model FFT Mode**: (EXPERIMENTAL) Splits the frequency domain between two models, using one model for low frequencies and another for high frequencies. Last model in the chain decides the split weight.
|
||||
- **2Model FFT Mode**: Splits the frequency domain between two models, using one model for low frequencies and another for high frequencies.
|
||||
|
||||
## Features
|
||||
|
||||
@@ -22,13 +17,6 @@ ComfyUI-MixMod provides a powerful way to combine multiple models during stable
|
||||
- Frequency domain splitting for controlling different image characteristics
|
||||
- Chain multiple components together to create complex model combinations
|
||||
|
||||
## Roadmap
|
||||
|
||||
- Add more modes
|
||||
- Better lora support
|
||||
- Make it upscale better/ fix resolution
|
||||
- Make this work with SD1.5, Flux and other models
|
||||
|
||||
## Installation
|
||||
|
||||
1. Clone this repository into your ComfyUI's `custom_nodes` directory:
|
||||
@@ -90,31 +78,46 @@ Here's a simple workflow to mix two models:
|
||||
│ │
|
||||
▼ ▼
|
||||
┌───────────────────────┐ ┌───────────────────────┐
|
||||
│ MixMod Component A │----│ MixMod Component B │
|
||||
│ MixMod Component A │ │ MixMod Component B │
|
||||
│ weight: 0.6, cfg: 7.5 │ │ weight: 0.4, cfg: 8.0 │
|
||||
└────────┬──────────────┘ └───────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌──────────────────┐
|
||||
│ MixMod Guider │
|
||||
│ mode: team │
|
||||
└────────┬─────────┘
|
||||
│
|
||||
▼
|
||||
┌──────────────────┐
|
||||
│ KSampler │
|
||||
│ (use as guider) │
|
||||
└────────┬─────────┘
|
||||
│
|
||||
▼
|
||||
┌──────────────────┐
|
||||
│ VAE Decode │
|
||||
└────────┬─────────┘
|
||||
│
|
||||
▼
|
||||
Image
|
||||
└───────────┬───────────┘ └───────────┬───────────┘
|
||||
│ │
|
||||
└────────────┬───────────────┘
|
||||
│
|
||||
▼
|
||||
┌──────────────────┐
|
||||
│ MixMod Guider │
|
||||
│ mode: team │
|
||||
└────────┬─────────┘
|
||||
│
|
||||
▼
|
||||
┌──────────────────┐
|
||||
│ KSampler │
|
||||
│ (use as guider) │
|
||||
└────────┬─────────┘
|
||||
│
|
||||
▼
|
||||
┌──────────────────┐
|
||||
│ VAE Decode │
|
||||
└────────┬─────────┘
|
||||
│
|
||||
▼
|
||||
Image
|
||||
```
|
||||
|
||||
### Frequency Split Mixing (2model_fft)
|
||||
|
||||
For frequency domain splitting:
|
||||
1. Create two "MixMod Guider Component" nodes
|
||||
2. Connect different models to each component
|
||||
3. Use the weight of the first model to control the frequency split point
|
||||
4. Set the mode to "2model_fft" in the "MixMod Guider" node
|
||||
|
||||
## How It Works
|
||||
|
||||
- **Team Mode**: Combines multiple models by computing a weighted average of their unconditioned predictions, then adds the guided differences according to each model's CFG value.
|
||||
- **2model_fft**: Transforms predictions to the frequency domain using FFT, applies the first model to low frequencies and the second model to high frequencies based on a split ratio derived from the first model's weight.
|
||||
|
||||
## Compatibility
|
||||
|
||||
- This extension is compatible with ComfyUI version 1.0.0 and above
|
||||
@@ -128,12 +131,10 @@ Here's a simple workflow to mix two models:
|
||||
- Experiment with different CFG values for each model to control its influence
|
||||
- For models with very different styles, start with lower weights for specialty models
|
||||
|
||||
Example with Ponyv6 and NoobAI:
|
||||
|
||||
Only Pony, Pony+Noob, Only Noob
|
||||
|
||||
|
||||
|
||||
## Credits
|
||||
|
||||
Created by Kantsche
|
||||
|
||||
## License
|
||||
|
||||
[Add license information here]
|
||||
@@ -0,0 +1,293 @@
|
||||
import torch
|
||||
import comfy.samplers
|
||||
import comfy.utils
|
||||
import logging
|
||||
import torch.nn.functional as F
|
||||
from .utils import sumweights
|
||||
|
||||
# Initialize other necessary attributes
|
||||
|
||||
import comfy.model_management as model_management
|
||||
from .utils import load_taesd, ensure_model, create_adaptive_low_res_model_wrapper, create_noise_injection_wrapper, load_vae
|
||||
|
||||
|
||||
# Import mixing functions
|
||||
from .mixes.fft import mix_teamfft, mix_teamfft4freq, mix_2model_fft, mix_2m2f, mix_fft_overlap, mix_fft_full, mix_bandfft
|
||||
from .mixes.standard import mix_standard, mix_masked, mix_dynamic_mask, mix_foreground_background, mix_dynamic_mask_alternative
|
||||
|
||||
|
||||
class MixModGuider(comfy.samplers.CFGGuider):
|
||||
def __init__(self, component, **kwargs):
|
||||
|
||||
self.component = component
|
||||
self.mode = kwargs.get("mode", "team")
|
||||
self.ratio = kwargs.get("ratio", 1.0)
|
||||
self.original_conds = {}
|
||||
self.prepared_models = []
|
||||
self.highres_mode = kwargs.get("highres_mode", False)
|
||||
self.bands = kwargs.get("bands", None)
|
||||
self.kwargs = kwargs
|
||||
self.last_latent = None
|
||||
|
||||
# Load the depth model
|
||||
if(self.mode == "foreground_background"):
|
||||
#check if controlnet_aux is installed
|
||||
try:
|
||||
from custom_nodes.comfyui_controlnet_aux.src.custom_controlnet_aux.depth_anything import DepthAnythingDetector
|
||||
self.depth_model = DepthAnythingDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
self.vae = comfy.sd.VAE(sd=load_taesd("taesdxl"))
|
||||
except Exception as e:
|
||||
self.mode = "team"
|
||||
raise Exception("Please install Comfyui ControlNet Aux custom node. Falling back to team mode.")
|
||||
|
||||
if(self.mode == "styleinjection" and not self.kwargs.get("image", None) is None):
|
||||
self.vae = comfy.sd.VAE(sd=load_vae("fixFP16ErrorsSDXLLowerMemoryUse_v10.safetensors"))
|
||||
|
||||
self.noiseinjection_latent = self.vae.encode(self.kwargs.get("image", None))
|
||||
|
||||
|
||||
self.models = []
|
||||
self.weights = []
|
||||
self.cfgs = []
|
||||
self.masks = [] # Initialize empty masks list
|
||||
self.step = 0
|
||||
self.total_steps = 0
|
||||
self.start_steps = []
|
||||
self.end_steps = []
|
||||
self.scale_factors = []
|
||||
self.options = []
|
||||
self.mask_history = None # For tracking dynamic masks between steps
|
||||
self.decay_factor = kwargs.get("decay_factor", 0.9) # Dynamic mask decay factor
|
||||
self.threshold = kwargs.get("threshold", 0.1) # Dynamic mask threshold
|
||||
self.active_masks = []
|
||||
self.active_weights = []
|
||||
self.active_cfgs = []
|
||||
self.uncond_decays = []
|
||||
self.foreground_background_masks = []
|
||||
self.seed = -1
|
||||
# Process component chain
|
||||
current = component
|
||||
conds = {}
|
||||
while current is not None:
|
||||
self.models.append(current.get("model"))
|
||||
if len(self.models) == 1:
|
||||
self.model_patcher = self.models[0]
|
||||
self.model_options = self.models[0].model_options
|
||||
self.weights.append(current.get("weight", 1.0))
|
||||
self.cfgs.append(current.get("cfg", 7.5))
|
||||
opt = current.get("options", {})
|
||||
options = {}
|
||||
while opt is not None:
|
||||
options.update(opt)
|
||||
opt = opt.get("prev")
|
||||
self.scale_factors.append(options.get("scale_factor", 0.0))
|
||||
self.start_steps.append(options.get("start_step", 0))
|
||||
self.end_steps.append(options.get("end_step", -1))
|
||||
self.masks.append(options.get("mask", None))
|
||||
self.uncond_decays.append(options.get("uncond_decay", 1.0))
|
||||
# Add to original_conds with indexed keys
|
||||
idx = len(self.models)
|
||||
conds.update({f"positive_{idx}": current["positive"], f"negative_{idx}": current["negative"]})
|
||||
|
||||
current = current.get("prev")
|
||||
|
||||
|
||||
self.inner_set_conds(conds)
|
||||
|
||||
|
||||
|
||||
|
||||
def set_mode(self, mode):
|
||||
self.mode = mode
|
||||
|
||||
def predict_noise(self, x, timestep, model_options={}, seed=None):
|
||||
pred = []
|
||||
self.active_masks = []
|
||||
self.active_weights = []
|
||||
self.active_cfgs = []
|
||||
self.active_uncond_decays = []
|
||||
def latent_callback(latent):
|
||||
self.noiseinjection_latent = latent
|
||||
for i in range(len(self.models)):
|
||||
cond = [self.conds.get(f"positive_{i+1}", None), self.conds.get(f"negative_{i+1}", None)]
|
||||
if(self.start_steps[i] <= self.step <= self.end_steps[i] or self.end_steps[i] == -1):
|
||||
if(self.scale_factors[i] > 0):
|
||||
model_options["model_function_wrapper"] = create_adaptive_low_res_model_wrapper(
|
||||
min_scale_factor=self.scale_factors[i],
|
||||
max_scale_factor=self.scale_factors[i]
|
||||
)
|
||||
|
||||
ensure_model(self.models[i],"normal")
|
||||
pred.append(comfy.samplers.calc_cond_batch(self.prepared_models[i][0], cond, x, timestep, model_options))
|
||||
del model_options["model_function_wrapper"]
|
||||
elif(self.mode == "styleinjection" and self.kwargs.get("image", None) is not None and i == 1 and self.kwargs.get("start_step", 4) <= self.step <= self.kwargs.get("end_step", 7)):
|
||||
#calculate noise weight, decreasing from self.kwargs.get("noise_weight", 1.0) to 0 between start_step and end_step
|
||||
noise_weight = self.kwargs.get("noise_weight", 1.0) - (self.step - self.kwargs.get("start_step", 4)) * (self.kwargs.get("noise_weight", 1.0) - 0) / (self.kwargs.get("end_step", 7) - self.kwargs.get("start_step", 4))
|
||||
model_options["model_function_wrapper"] = create_noise_injection_wrapper(self.noiseinjection_latent, noise_weight, latent_callback)
|
||||
ensure_model(self.models[i],"normal")
|
||||
pred.append(comfy.samplers.calc_cond_batch(self.prepared_models[i][0], cond, x, timestep, model_options))
|
||||
del model_options["model_function_wrapper"]
|
||||
else:
|
||||
ensure_model(self.models[i],"normal")
|
||||
pred.append(comfy.samplers.calc_cond_batch(self.prepared_models[i][0], cond, x, timestep, model_options))
|
||||
|
||||
self.active_masks.append(self.masks[i])
|
||||
self.active_weights.append(self.weights[i])
|
||||
self.active_cfgs.append(self.cfgs[i])
|
||||
self.active_uncond_decays.append(self.uncond_decays[i])
|
||||
|
||||
if not pred:
|
||||
logging.warning(f"No models were active at step {self.step}, using default model 0")
|
||||
cond = [self.conds.get(f"positive_1", None), self.conds.get(f"negative_1", None)]
|
||||
pred.append(comfy.samplers.calc_cond_batch(self.prepared_models[0][0], cond, x, timestep, model_options))
|
||||
self.active_masks.append(self.masks[0])
|
||||
self.active_weights.append(self.weights[0])
|
||||
self.active_cfgs.append(self.cfgs[0])
|
||||
self.active_uncond_decays.append(self.uncond_decays[0])
|
||||
return self.mix_function(self.prepared_models[0][0], pred, x, timestep, model_options, seed)
|
||||
|
||||
def mix_function(self, model, pred, x, timestep, model_options, seed):
|
||||
haspostfunctions = False
|
||||
|
||||
postuncond = torch.zeros_like(pred[0][1])
|
||||
postcond = torch.zeros_like(pred[0][0])
|
||||
|
||||
self.active_weights = sumweights(self.active_weights)
|
||||
""" if(len(pred) == 1 and not self.mode == "foreground_background"):
|
||||
mix = pred[0][1] + (pred[0][0] - pred[0][1]) * self.cfgs[0]
|
||||
for fn in model_options.get("sampler_post_cfg_function", []):
|
||||
args = {"denoised": mix, "cond": self.conds.get(f"positive_1", None), "uncond": self.conds.get(f"negative_1", None), "cond_scale": self.cfgs[0], "model": model, "uncond_denoised": pred[0][1], "cond_denoised": pred[0][0],
|
||||
"sigma": timestep, "model_options": model_options, "input": x}
|
||||
mix = fn(args)
|
||||
self.step += 1
|
||||
return mix """
|
||||
if(self.mode == "team"):
|
||||
if any(mask is not None for mask in self.active_masks):
|
||||
mix, postuncond = mix_masked(pred, self.active_weights, self.active_cfgs, self.active_masks,)
|
||||
else:
|
||||
mix, postuncond = mix_standard(pred, self.active_weights, self.active_cfgs, self.active_uncond_decays)
|
||||
|
||||
elif self.mode == "styleinjection":
|
||||
mix, postuncond = mix_standard(pred, self.active_weights, self.active_cfgs)
|
||||
|
||||
elif self.mode == "dynamic_mask":
|
||||
mix, postuncond, new_masks = mix_dynamic_mask(pred, self.active_weights, self.active_cfgs, self.step, self.mask_history, self.decay_factor, self.threshold, self.kwargs.get("blur_sigma", 0.5))
|
||||
self.mask_history = new_masks # Store the updated masks for next step
|
||||
|
||||
elif self.mode == "dynamic_mask_alternative":
|
||||
mix, postuncond = mix_dynamic_mask_alternative(pred, self.active_weights, self.active_cfgs, self.depth_model, self.threshold, self.kwargs.get("invert", False), self.kwargs.get("blur_sigma", 1.0))
|
||||
|
||||
elif self.mode == "foreground_background":
|
||||
if(self.step == 0):
|
||||
self.foreground_background_masks = []
|
||||
|
||||
masks, mix, postuncond = mix_foreground_background(pred, self.active_weights, self.last_latent, self.active_cfgs, vae=self.vae, depth_model=self.depth_model, sharpness=self.kwargs.get("sharpness", 0.0), invert=self.kwargs.get("invert", False), step=self.step, masks=self.foreground_background_masks, startmask=self.kwargs.get("startmask", 5))
|
||||
|
||||
if any(mask is not None for mask in masks):
|
||||
self.foreground_background_masks = masks
|
||||
|
||||
elif self.mode == "bandfft":
|
||||
mix, postuncond = mix_bandfft(pred, self.active_weights, self.active_cfgs, self.bands)
|
||||
|
||||
elif self.mode == "2m2f":
|
||||
mix, postuncond = mix_2m2f(pred, self.active_weights, self.active_cfgs, self.ratio)
|
||||
|
||||
elif self.mode == "teamfft":
|
||||
mix, postuncond = mix_teamfft(pred, self.active_weights, self.active_cfgs)
|
||||
|
||||
elif self.mode == "teamfft4freq":
|
||||
mix, postuncond = mix_teamfft4freq(pred, self.active_weights, self.active_cfgs)
|
||||
|
||||
elif self.mode == "2model_fft":
|
||||
mix, postuncond = mix_2model_fft(pred, self.active_weights, self.active_cfgs, self.ratio)
|
||||
|
||||
elif self.mode == "fft_overlap":
|
||||
mix, postuncond = mix_fft_overlap(pred, self.active_weights, self.active_cfgs, self.ratio)
|
||||
|
||||
elif self.mode == "fft_full":
|
||||
mix, postuncond = mix_fft_full(pred, self.active_weights, self.active_cfgs)
|
||||
else:
|
||||
# Fallback to standard mixing if mode is not recognized
|
||||
mix, postuncond = mix_standard(pred, self.active_weights, self.active_cfgs)
|
||||
|
||||
for fn in model_options.get("sampler_post_cfg_function", []):
|
||||
args = {"denoised": mix, "cond": self.conds.get(f"positive_1", None), "uncond": self.conds.get(f"negative_1", None), "cond_scale": self.active_cfgs[0], "model": model, "uncond_denoised": postuncond, "cond_denoised": postcond,
|
||||
"sigma": timestep, "model_options": model_options, "input": x}
|
||||
mix = fn(args)
|
||||
self.step += 1
|
||||
self.seed = seed
|
||||
return mix
|
||||
|
||||
def outer_sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
|
||||
# Get the device from the first model
|
||||
device = self.models[0].load_device
|
||||
|
||||
# Prepare the models, conditions, and loaded models for each model
|
||||
prepared_models = []
|
||||
all_loaded_models = []
|
||||
|
||||
for i, model in enumerate(self.models):
|
||||
# Get relevant conditions for this model
|
||||
model_conds = {k: v for k, v in self.conds.items() if k.endswith(f"_{i+1}")}
|
||||
|
||||
|
||||
# Prepare sampling for this model
|
||||
inner_model, prepared_conds, loaded_models = comfy.sampler_helpers.prepare_sampling(model, noise.shape, model_conds, model.model_options)
|
||||
|
||||
# Save the prepared data
|
||||
self.prepared_models.append((inner_model, prepared_conds))
|
||||
|
||||
all_loaded_models.extend(loaded_models)
|
||||
|
||||
# Update the conditions for this model
|
||||
for k, v in prepared_conds.items():
|
||||
self.conds[k] = v
|
||||
|
||||
self.inner_model = self.prepared_models[0][0]
|
||||
# Prepare the device-related tensors
|
||||
if denoise_mask is not None:
|
||||
denoise_mask = comfy.sampler_helpers.prepare_mask(denoise_mask, noise.shape, device)
|
||||
|
||||
noise = noise.to(device)
|
||||
latent_image = latent_image.to(device)
|
||||
sigmas = sigmas.to(device)
|
||||
|
||||
|
||||
# Set model options for the main model patcher
|
||||
comfy.samplers.cast_to_load_options(self.model_options, device=device, dtype=self.model_patcher.model_dtype())
|
||||
|
||||
try:
|
||||
# Pre-run all models
|
||||
for model in self.models:
|
||||
model.pre_run()
|
||||
|
||||
# Set the inner_model to be a list of prepared models
|
||||
|
||||
# Run the inner sample with prepared models
|
||||
output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
|
||||
finally:
|
||||
# Clean up all models
|
||||
for model in self.models:
|
||||
model.cleanup()
|
||||
del self.inner_model
|
||||
# Clean up models and loaded models
|
||||
comfy.sampler_helpers.cleanup_models(self.conds, all_loaded_models)
|
||||
|
||||
|
||||
return output
|
||||
|
||||
def inner_sample(self, noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed):
|
||||
# Create a wrapper callback that captures intermediate denoised results (x0)
|
||||
original_callback = callback
|
||||
|
||||
def wrapped_callback(step, x0, x, total_steps):
|
||||
# Store the current x0 (denoised latent)
|
||||
self.step = step
|
||||
self.total_steps = total_steps
|
||||
self.last_latent = x0.detach().clone()
|
||||
# Call the original callback if provided
|
||||
if original_callback is not None:
|
||||
original_callback(step, x0, x, total_steps)
|
||||
|
||||
# Pass our wrapped callback to the parent method
|
||||
return comfy.samplers.CFGGuider.inner_sample(self, noise, latent_image, device, sampler, sigmas, denoise_mask, wrapped_callback, disable_pbar, seed)
|
||||
@@ -0,0 +1 @@
|
||||
# FFT mixing functions for ComfyUI MixMod
|
||||
+649
@@ -0,0 +1,649 @@
|
||||
import torch
|
||||
|
||||
def mix_teamfft(pred, weights, cfgs):
|
||||
# First, get shape information from the first model's predictions
|
||||
_, _, h, w = pred[0][0].shape
|
||||
|
||||
# Initialize an empty tensor in frequency domain for the result
|
||||
fft_mix = torch.zeros_like(torch.fft.rfft2(pred[0][1]))
|
||||
|
||||
postuncond = None
|
||||
|
||||
# First loop: mix the unconditioned predictions in frequency domain
|
||||
for i in range(len(pred)):
|
||||
# Convert unconditional prediction to frequency domain
|
||||
fft_u_pred = torch.fft.rfft2(pred[i][1])
|
||||
# Add weighted contribution
|
||||
fft_mix = fft_mix + fft_u_pred * weights[i]
|
||||
|
||||
|
||||
postuncond = torch.fft.irfft2(fft_mix, s=(h, w))
|
||||
|
||||
# Second loop: apply the guidance in frequency domain
|
||||
for i in range(len(pred)):
|
||||
# Convert conditional and unconditional predictions to frequency domain
|
||||
fft_c_pred = torch.fft.rfft2(pred[i][0])
|
||||
fft_u_pred = torch.fft.rfft2(pred[i][1])
|
||||
|
||||
# Apply the CFG formula in frequency domain
|
||||
fft_mix = fft_mix + (fft_c_pred - fft_u_pred) * cfgs[i]
|
||||
|
||||
# Convert the final result back to spatial domain
|
||||
mix = torch.fft.irfft2(fft_mix, s=(h, w))
|
||||
|
||||
return mix, postuncond
|
||||
|
||||
def mix_teamfft4freq(pred, weights, cfgs):
|
||||
# Get shape information from the first model's predictions
|
||||
_, _, h, w = pred[0][0].shape
|
||||
|
||||
# Initialize an empty tensor in frequency domain for the result
|
||||
fft_mix = torch.zeros_like(torch.fft.rfft2(pred[0][1]))
|
||||
|
||||
# Get frequency domain dimensions
|
||||
freq_h, freq_w = fft_mix.shape[2], fft_mix.shape[3]
|
||||
|
||||
# Create frequency band masks - dividing into 4 frequency bands
|
||||
# Low, low-mid, high-mid, and high frequencies
|
||||
masks = []
|
||||
|
||||
# Calculate coordinates for frequency space
|
||||
h_coords = torch.linspace(0, 1, freq_h, device=fft_mix.device)
|
||||
w_coords = torch.linspace(0, 1, freq_w, device=fft_mix.device)
|
||||
h_grid, w_grid = torch.meshgrid(h_coords, w_coords, indexing='ij')
|
||||
|
||||
# Calculate frequency magnitude (distance from DC component)
|
||||
freq_magnitude = torch.sqrt(h_grid**2 + w_grid**2)
|
||||
|
||||
# Normalize to [0,1] range
|
||||
max_freq = torch.sqrt(torch.tensor(2.0, device=fft_mix.device))
|
||||
normalized_freq = freq_magnitude / max_freq
|
||||
|
||||
# Define frequency band boundaries
|
||||
band_boundaries = [0.0, 0.25, 0.5, 0.75, 1.0]
|
||||
|
||||
# Create masks for each frequency band with smooth transitions
|
||||
transition_width = 0.05 # 5% transition width
|
||||
|
||||
for i in range(4): # 4 frequency bands
|
||||
band_start = band_boundaries[i]
|
||||
band_end = band_boundaries[i+1]
|
||||
|
||||
# Calculate transition regions
|
||||
trans_start = max(0.0, band_start - transition_width/2)
|
||||
trans_end = min(1.0, band_end + transition_width/2)
|
||||
|
||||
# Create mask
|
||||
mask = torch.zeros_like(normalized_freq)
|
||||
|
||||
# Full strength in the main band region
|
||||
main_region = (normalized_freq >= band_start) & (normalized_freq <= band_end)
|
||||
mask = torch.where(main_region, torch.ones_like(mask), mask)
|
||||
|
||||
# Smooth transition at the start of the band
|
||||
if band_start > 0:
|
||||
start_trans = (normalized_freq >= trans_start) & (normalized_freq < band_start)
|
||||
mask = torch.where(
|
||||
start_trans,
|
||||
(normalized_freq - trans_start) / (band_start - trans_start),
|
||||
mask
|
||||
)
|
||||
|
||||
# Smooth transition at the end of the band
|
||||
if band_end < 1.0:
|
||||
end_trans = (normalized_freq > band_end) & (normalized_freq <= trans_end)
|
||||
mask = torch.where(
|
||||
end_trans,
|
||||
1.0 - (normalized_freq - band_end) / (trans_end - band_end),
|
||||
mask
|
||||
)
|
||||
|
||||
masks.append(mask)
|
||||
|
||||
# Ensure masks sum to 1.0 at each point for perfect reconstruction
|
||||
mask_sum = sum(masks)
|
||||
masks = [m / (mask_sum + 1e-8) for m in masks] # Add small epsilon to avoid division by zero
|
||||
|
||||
uncondresult = torch.zeros_like(fft_mix)
|
||||
postuncond = None
|
||||
|
||||
# Process each frequency band separately
|
||||
for band_idx in range(4):
|
||||
band_mix = torch.zeros_like(fft_mix)
|
||||
|
||||
# First loop: mix the unconditioned predictions in this frequency band
|
||||
for i in range(len(pred)):
|
||||
# Convert unconditional prediction to frequency domain
|
||||
fft_u_pred = torch.fft.rfft2(pred[i][1])
|
||||
# Add weighted contribution for this frequency band
|
||||
band_mix = band_mix + fft_u_pred * weights[i] * masks[band_idx]
|
||||
|
||||
uncondresult = uncondresult + band_mix
|
||||
|
||||
# Second loop: apply the guidance in this frequency band
|
||||
for i in range(len(pred)):
|
||||
# Convert conditional and unconditional predictions to frequency domain
|
||||
fft_c_pred = torch.fft.rfft2(pred[i][0])
|
||||
fft_u_pred = torch.fft.rfft2(pred[i][1])
|
||||
|
||||
# Apply the CFG formula in this frequency band
|
||||
band_mix = band_mix + (fft_c_pred - fft_u_pred) * cfgs[i] * masks[band_idx]
|
||||
|
||||
# Add this band's contribution to the final mix
|
||||
fft_mix = fft_mix + band_mix
|
||||
|
||||
|
||||
postuncond = torch.fft.irfft2(uncondresult, s=(h, w))
|
||||
|
||||
# Convert the final result back to spatial domain
|
||||
mix = torch.fft.irfft2(fft_mix, s=(h, w))
|
||||
|
||||
return mix, postuncond
|
||||
|
||||
def mix_bandfft(pred, weights, cfgs, bands=2):
|
||||
# Get shape information from the first model's predictions
|
||||
_, _, h, w = pred[0][0].shape
|
||||
|
||||
# Initialize an empty tensor in frequency domain for the result
|
||||
fft_mix = torch.zeros_like(torch.fft.rfft2(pred[0][1]))
|
||||
|
||||
# Get frequency domain dimensions
|
||||
freq_h, freq_w = fft_mix.shape[2], fft_mix.shape[3]
|
||||
|
||||
# Create frequency band masks - dividing into 4 frequency bands
|
||||
# Low, low-mid, high-mid, and high frequencies
|
||||
masks = []
|
||||
|
||||
# Calculate coordinates for frequency space
|
||||
h_coords = torch.linspace(0, 1, freq_h, device=fft_mix.device)
|
||||
w_coords = torch.linspace(0, 1, freq_w, device=fft_mix.device)
|
||||
h_grid, w_grid = torch.meshgrid(h_coords, w_coords, indexing='ij')
|
||||
|
||||
# Calculate frequency magnitude (distance from DC component)
|
||||
freq_magnitude = torch.sqrt(h_grid**2 + w_grid**2)
|
||||
|
||||
# Normalize to [0,1] range
|
||||
max_freq = torch.sqrt(torch.tensor(2.0, device=fft_mix.device))
|
||||
normalized_freq = freq_magnitude / max_freq
|
||||
|
||||
#calculate the number of bands boundaries from bands size
|
||||
band_boundaries = [0.0]
|
||||
for i in range(bands-1):
|
||||
band_boundaries.append(band_boundaries[-1] + 1.0 / bands)
|
||||
band_boundaries.append(1.0)
|
||||
|
||||
# Create masks for each frequency band with smooth transitions
|
||||
transition_width = 0.01 # 1% transition width
|
||||
|
||||
for i in range(bands): # 4 frequency bands
|
||||
band_start = band_boundaries[i]
|
||||
band_end = band_boundaries[i+1]
|
||||
|
||||
# Calculate transition regions
|
||||
trans_start = max(0.0, band_start - transition_width/2)
|
||||
trans_end = min(1.0, band_end + transition_width/2)
|
||||
|
||||
# Create mask
|
||||
mask = torch.zeros_like(normalized_freq)
|
||||
|
||||
# Full strength in the main band region
|
||||
main_region = (normalized_freq >= band_start) & (normalized_freq <= band_end)
|
||||
mask = torch.where(main_region, torch.ones_like(mask), mask)
|
||||
|
||||
# Smooth transition at the start of the band
|
||||
if band_start > 0:
|
||||
start_trans = (normalized_freq >= trans_start) & (normalized_freq < band_start)
|
||||
mask = torch.where(
|
||||
start_trans,
|
||||
(normalized_freq - trans_start) / (band_start - trans_start),
|
||||
mask
|
||||
)
|
||||
|
||||
# Smooth transition at the end of the band
|
||||
if band_end < 1.0:
|
||||
end_trans = (normalized_freq > band_end) & (normalized_freq <= trans_end)
|
||||
mask = torch.where(
|
||||
end_trans,
|
||||
1.0 - (normalized_freq - band_end) / (trans_end - band_end),
|
||||
mask
|
||||
)
|
||||
|
||||
masks.append(mask)
|
||||
|
||||
# Ensure masks sum to 1.0 at each point for perfect reconstruction
|
||||
mask_sum = sum(masks)
|
||||
masks = [m / (mask_sum + 1e-8) for m in masks] # Add small epsilon to avoid division by zero
|
||||
|
||||
uncondresult = torch.zeros_like(fft_mix)
|
||||
postuncond = None
|
||||
|
||||
# Process each frequency band separately
|
||||
for band_idx in range(bands):
|
||||
band_mix = torch.zeros_like(fft_mix)
|
||||
|
||||
# First loop: mix the unconditioned predictions in this frequency band
|
||||
for i in range(len(pred)):
|
||||
# Convert unconditional prediction to frequency domain
|
||||
fft_u_pred = torch.fft.rfft2(pred[i][1])
|
||||
# Add weighted contribution for this frequency band
|
||||
band_mix = band_mix + fft_u_pred * weights[i] * masks[band_idx]
|
||||
|
||||
uncondresult = uncondresult + band_mix
|
||||
|
||||
# Second loop: apply the guidance in this frequency band
|
||||
for i in range(len(pred)):
|
||||
# Convert conditional and unconditional predictions to frequency domain
|
||||
fft_c_pred = torch.fft.rfft2(pred[i][0])
|
||||
fft_u_pred = torch.fft.rfft2(pred[i][1])
|
||||
|
||||
# Apply the CFG formula in this frequency band
|
||||
band_mix = band_mix + (fft_c_pred - fft_u_pred) * cfgs[i] * masks[band_idx]
|
||||
|
||||
# Add this band's contribution to the final mix
|
||||
fft_mix = fft_mix + band_mix
|
||||
|
||||
|
||||
postuncond = torch.fft.irfft2(uncondresult, s=(h, w))
|
||||
|
||||
# Convert the final result back to spatial domain
|
||||
mix = torch.fft.irfft2(fft_mix, s=(h, w))
|
||||
|
||||
return mix, postuncond
|
||||
|
||||
def mix_2model_fft(pred, weights, cfgs, ratio):
|
||||
# Apply FFT-based frequency splitting
|
||||
# Need at least 2 models for FFT mode
|
||||
if len(pred) < 2:
|
||||
mix = pred[0][0] * cfgs[0] + pred[0][1] * (1 - cfgs[0])
|
||||
return mix, None # Return properly conditioned prediction
|
||||
|
||||
# Get shapes from first prediction
|
||||
_, _, h, w = pred[0][0].shape
|
||||
|
||||
# Use the first model for low frequencies and second model for high frequencies
|
||||
# Extract predictions from first model
|
||||
c_pred_1 = pred[0][0] # Conditioned prediction
|
||||
u_pred_1 = pred[0][1] # Unconditioned prediction
|
||||
|
||||
# Extract predictions from second model
|
||||
c_pred_2 = pred[1][0] # Conditioned prediction
|
||||
u_pred_2 = pred[1][1] # Unconditioned prediction
|
||||
|
||||
# Convert to frequency domain
|
||||
fft_u_pred_1 = torch.fft.rfft2(u_pred_1)
|
||||
fft_c_pred_1 = torch.fft.rfft2(c_pred_1)
|
||||
fft_u_pred_2 = torch.fft.rfft2(u_pred_2)
|
||||
fft_c_pred_2 = torch.fft.rfft2(c_pred_2)
|
||||
|
||||
# Use weight of first model as frequency split ratio
|
||||
# Ensure ratio is between 0.05 and 0.95 to guarantee some frequency coverage for both models
|
||||
split_ratio = max(0.05, min(0.95, abs(ratio)))
|
||||
|
||||
# Calculate frequency split point
|
||||
split_h = int(h * split_ratio)
|
||||
split_w = int(w * split_ratio / 2) # rfft2 returns half the width in frequency domain
|
||||
|
||||
# Create masks for low and high frequencies
|
||||
mask_low = torch.ones_like(fft_u_pred_1, dtype=torch.float32)
|
||||
mask_high = torch.ones_like(fft_u_pred_1, dtype=torch.float32)
|
||||
|
||||
# Create transition area for smooth blending
|
||||
transition_h = max(1, int(h * 0.05)) # 5% transition zone
|
||||
transition_w = max(1, int(w * 0.05 / 2)) # 5% transition zone
|
||||
|
||||
# Set high frequencies to zero in low mask with smooth transition
|
||||
for i in range(mask_low.shape[2]):
|
||||
for j in range(mask_low.shape[3]):
|
||||
if i >= split_h + transition_h or j >= split_w + transition_w:
|
||||
mask_low[:, :, i, j] = 0.0
|
||||
elif i >= split_h or j >= split_w:
|
||||
# Calculate smooth transition value
|
||||
fade_h = 1.0 - min(1.0, (i - split_h) / transition_h) if i >= split_h else 1.0
|
||||
fade_w = 1.0 - min(1.0, (j - split_w) / transition_w) if j >= split_w else 1.0
|
||||
mask_low[:, :, i, j] = fade_h * fade_w
|
||||
|
||||
# Set low frequencies to zero in high mask with inverse smooth transition
|
||||
for i in range(mask_high.shape[2]):
|
||||
for j in range(mask_high.shape[3]):
|
||||
if i <= split_h - transition_h and j <= split_w - transition_w:
|
||||
mask_high[:, :, i, j] = 0.0
|
||||
elif i <= split_h or j <= split_w:
|
||||
# Calculate smooth transition value - inverse of the low mask
|
||||
fade_h = min(1.0, (split_h - i) / transition_h) if i <= split_h else 0.0
|
||||
fade_w = min(1.0, (split_w - j) / transition_w) if j <= split_w else 0.0
|
||||
mask_high[:, :, i, j] = 1.0 - (fade_h * fade_w)
|
||||
|
||||
# Ensure masks sum to 1.0 at each point for full spectrum coverage
|
||||
mask_sum = mask_low + mask_high
|
||||
mask_low = mask_low / mask_sum
|
||||
mask_high = mask_high / mask_sum
|
||||
|
||||
|
||||
result = fft_u_pred_1 * mask_low + fft_u_pred_2 * mask_high
|
||||
postuncond = torch.fft.irfft2(result, s=(h, w))
|
||||
|
||||
# Apply masks and compute CFG in frequency domain using the corresponding CFG values
|
||||
fft_result_low = (fft_u_pred_1 * mask_low) + ((fft_c_pred_1 - fft_u_pred_1) * cfgs[0] * mask_low)
|
||||
fft_result_high = (fft_u_pred_2 * mask_high) + ((fft_c_pred_2 - fft_u_pred_2) * cfgs[1] * mask_high)
|
||||
|
||||
# Combine low and high frequency results
|
||||
fft_result = fft_result_low + fft_result_high
|
||||
|
||||
# Convert back to spatial domain
|
||||
mix = torch.fft.irfft2(fft_result, s=(h, w))
|
||||
|
||||
return mix, postuncond
|
||||
|
||||
def create_low_freq_mask(shape, radius):
|
||||
rows, cols = shape
|
||||
crow, ccol = rows // 2, cols // 2 # center
|
||||
|
||||
# Create coordinate grids
|
||||
y, x = torch.meshgrid(torch.arange(rows), torch.arange(cols), indexing='ij')
|
||||
|
||||
# Distance from center
|
||||
dist_from_center = torch.sqrt((x - ccol)**2 + (y - crow)**2)
|
||||
|
||||
# Create mask: 1 inside radius, 0 outside
|
||||
mask = torch.zeros((rows, cols))
|
||||
mask[dist_from_center <= radius] = 1
|
||||
|
||||
return mask
|
||||
|
||||
|
||||
def mix_2m2f(pred, weights, cfgs, ratio):
|
||||
# Split the prediction into low frequency and high frequency components
|
||||
# ratio controls the cutoff frequency between low and high
|
||||
|
||||
# Need at least 2 models for this mode
|
||||
if len(pred) < 2:
|
||||
mix = pred[0][0] * cfgs[0] + pred[0][1] * (1 - cfgs[0])
|
||||
return mix, None # Return properly conditioned prediction
|
||||
|
||||
# Get shapes from first prediction
|
||||
_, _, h, w = pred[0][0].shape
|
||||
|
||||
# Extract predictions from first model (handles low frequencies)
|
||||
c_pred_low = pred[0][0] # Conditioned prediction for low frequencies
|
||||
u_pred_low = pred[0][1] # Unconditioned prediction for low frequencies
|
||||
|
||||
# Extract predictions from second model (handles high frequencies)
|
||||
c_pred_high = pred[1][0] # Conditioned prediction for high frequencies
|
||||
u_pred_high = pred[1][1] # Unconditioned prediction for high frequencies
|
||||
|
||||
# Convert to frequency domain
|
||||
fft_c_low = torch.fft.rfft2(c_pred_low)
|
||||
fft_u_low = torch.fft.rfft2(u_pred_low)
|
||||
fft_c_high = torch.fft.rfft2(c_pred_high)
|
||||
fft_u_high = torch.fft.rfft2(u_pred_high)
|
||||
|
||||
# Get frequency domain dimensions
|
||||
freq_h, freq_w = fft_c_low.shape[2], fft_c_low.shape[3]
|
||||
|
||||
# Create normalized coordinate grids
|
||||
h_coords = torch.linspace(0, 1, freq_h, device=fft_c_low.device)
|
||||
w_coords = torch.linspace(0, 1, freq_w, device=fft_c_low.device)
|
||||
h_grid, w_grid = torch.meshgrid(h_coords, w_coords, indexing='ij')
|
||||
|
||||
# Calculate frequency magnitude (normalized distance from DC component)
|
||||
freq_magnitude = torch.sqrt(h_grid**2 + w_grid**2)
|
||||
max_freq = torch.sqrt(torch.tensor(2.0, device=fft_c_low.device))
|
||||
normalized_freq = freq_magnitude / max_freq
|
||||
|
||||
# Create frequency masks with smooth transition around cutoff
|
||||
cutoff = ratio # Frequency cutoff point (normalized)
|
||||
transition_width = 0.05 # Width of transition band for smooth blending
|
||||
|
||||
# Create low frequency mask (1 for low frequencies, 0 for high)
|
||||
low_mask = torch.ones_like(normalized_freq)
|
||||
# Create smooth transition where freq > (cutoff - transition_width/2)
|
||||
transition_start = cutoff - transition_width/2
|
||||
transition_end = cutoff + transition_width/2
|
||||
|
||||
# Apply smooth transition
|
||||
transition_region = (normalized_freq >= transition_start) & (normalized_freq <= transition_end)
|
||||
low_mask[transition_region] = 0.5 - 0.5 * torch.sin(
|
||||
torch.pi * (normalized_freq[transition_region] - cutoff) / transition_width
|
||||
)
|
||||
low_mask[normalized_freq > transition_end] = 0.0
|
||||
|
||||
# High frequency mask is complement of low frequency mask
|
||||
high_mask = 1.0 - low_mask
|
||||
|
||||
# Add batch and channel dimensions to masks
|
||||
low_mask = low_mask.unsqueeze(0).unsqueeze(0)
|
||||
high_mask = high_mask.unsqueeze(0).unsqueeze(0)
|
||||
|
||||
# Separate frequency components
|
||||
# Low frequency unconditioned prediction
|
||||
fft_u_result = fft_u_low * low_mask + fft_u_high * high_mask
|
||||
|
||||
# Convert unconditioned result back to spatial domain for returning
|
||||
postuncond = torch.fft.irfft2(fft_u_result, s=(h, w))
|
||||
|
||||
# Apply CFG separately to low and high frequency components
|
||||
fft_low_result = fft_u_low * low_mask + (fft_c_low - fft_u_low) * cfgs[0] * low_mask
|
||||
fft_high_result = fft_u_high * high_mask + (fft_c_high - fft_u_high) * cfgs[1] * high_mask
|
||||
|
||||
# Combine frequency components
|
||||
fft_result = fft_low_result + fft_high_result
|
||||
|
||||
# Convert back to spatial domain
|
||||
mix = torch.fft.irfft2(fft_result, s=(h, w))
|
||||
|
||||
return mix, postuncond
|
||||
|
||||
def mix_fft_overlap(pred, weights, cfgs, ratio):
|
||||
# Split each model into its own frequency band with overlap - optimized version
|
||||
_, _, h, w = pred[0][0].shape
|
||||
|
||||
# Initialize result in frequency domain
|
||||
fft_result = torch.zeros_like(torch.fft.rfft2(pred[0][0]))
|
||||
|
||||
# Convert all predictions to frequency domain
|
||||
fft_c_preds = []
|
||||
fft_u_preds = []
|
||||
valid_weights = []
|
||||
valid_cfgs = []
|
||||
|
||||
for i in range(len(pred)):
|
||||
if weights[i] <= 0:
|
||||
continue
|
||||
|
||||
# Extract predictions for this model
|
||||
c_pred = pred[i][0]
|
||||
u_pred = pred[i][1]
|
||||
|
||||
# Convert to frequency domain
|
||||
fft_c_preds.append(torch.fft.rfft2(c_pred))
|
||||
fft_u_preds.append(torch.fft.rfft2(u_pred))
|
||||
valid_weights.append(weights[i])
|
||||
valid_cfgs.append(cfgs[i])
|
||||
|
||||
# Get number of valid models
|
||||
num_models = len(valid_weights)
|
||||
if num_models == 0:
|
||||
# Handle edge case with no valid models
|
||||
default_mix = pred[0][0] * cfgs[0] + pred[0][1] * (1 - cfgs[0])
|
||||
return default_mix, None # Return first model's conditioned output
|
||||
elif num_models == 1:
|
||||
# With only one model, use it for the full spectrum
|
||||
mix = pred[0][0] * valid_cfgs[0] + pred[0][1] * (1 - valid_cfgs[0])
|
||||
return mix, None
|
||||
|
||||
# Calculate frequency coordinates
|
||||
freq_h, freq_w = fft_result.shape[2], fft_result.shape[3]
|
||||
|
||||
# Create normalized coordinate grids - vectorized
|
||||
h_coords = torch.linspace(0, 1, freq_h, device=fft_result.device)
|
||||
w_coords = torch.linspace(0, 1, freq_w, device=fft_result.device)
|
||||
|
||||
# Create meshgrid of coordinates
|
||||
h_grid, w_grid = torch.meshgrid(h_coords, w_coords, indexing='ij')
|
||||
|
||||
# Calculate frequency magnitude for all points
|
||||
freq_magnitude = torch.sqrt(h_grid**2 + w_grid**2)
|
||||
|
||||
# Normalize frequency magnitude to range [0, 1]
|
||||
max_freq = torch.sqrt(torch.tensor(2.0, device=fft_result.device)) # Maximum possible frequency magnitude
|
||||
normalized_freq = freq_magnitude / max_freq
|
||||
|
||||
# Define overlap ratio (0.2 means 20% overlap between bands)
|
||||
overlap_ratio = ratio
|
||||
|
||||
# Calculate band width for each model, accounting for overlap
|
||||
band_width = (1.0 + overlap_ratio) / num_models
|
||||
|
||||
# Calculate frequency masks for each model - vectorized
|
||||
model_masks = []
|
||||
for i in range(num_models):
|
||||
# Calculate band start and end points
|
||||
band_start = max(0.0, i * band_width - (overlap_ratio * band_width / 2))
|
||||
band_end = min(1.0, (i + 1) * band_width + (overlap_ratio * band_width / 2))
|
||||
|
||||
# Create initial mask where freq is within the band
|
||||
mask = torch.zeros_like(normalized_freq)
|
||||
in_band = (normalized_freq >= band_start) & (normalized_freq <= band_end)
|
||||
|
||||
# Define transition regions
|
||||
start_trans = band_start + (overlap_ratio * band_width / 2)
|
||||
end_trans = band_end - (overlap_ratio * band_width / 2)
|
||||
|
||||
# Create mask for different regions
|
||||
full_region = (normalized_freq >= start_trans) & (normalized_freq <= end_trans)
|
||||
start_region = (normalized_freq >= band_start) & (normalized_freq < start_trans)
|
||||
end_region = (normalized_freq > end_trans) & (normalized_freq <= band_end)
|
||||
|
||||
# Set full strength in middle region
|
||||
mask[full_region] = 1.0
|
||||
|
||||
# Calculate and set smooth transitions at edges
|
||||
if torch.any(start_region):
|
||||
trans_width = start_trans - band_start
|
||||
mask[start_region] = (normalized_freq[start_region] - band_start) / trans_width
|
||||
|
||||
if torch.any(end_region):
|
||||
trans_width = band_end - end_trans
|
||||
mask[end_region] = (band_end - normalized_freq[end_region]) / trans_width
|
||||
|
||||
model_masks.append(mask)
|
||||
|
||||
# Stack masks for easier normalization
|
||||
stacked_masks = torch.stack(model_masks)
|
||||
|
||||
# Normalize masks to ensure they sum to 1.0 across all models
|
||||
sum_masks = torch.sum(stacked_masks, dim=0)
|
||||
sum_masks = torch.clamp(sum_masks, min=1e-10)
|
||||
|
||||
# Normalize each mask
|
||||
normalized_stacked_masks = stacked_masks / sum_masks
|
||||
uncondresult = torch.zeros_like(fft_result)
|
||||
|
||||
# Apply CFG to each frequency band and combine - vectorized
|
||||
for i in range(num_models):
|
||||
# Get mask, add batch and channel dimensions
|
||||
mask = normalized_stacked_masks[i].unsqueeze(0).unsqueeze(0)
|
||||
|
||||
# Apply CFG formula with frequency band mask
|
||||
fft_u = fft_u_preds[i]
|
||||
uncondresult = uncondresult + fft_u
|
||||
fft_c = fft_c_preds[i]
|
||||
cfg_scale = valid_cfgs[i]
|
||||
|
||||
# Add weighted contribution to result
|
||||
fft_result += (fft_u + (fft_c - fft_u) * cfg_scale) * mask
|
||||
|
||||
|
||||
postuncond = torch.fft.irfft2(uncondresult, s=(h, w))
|
||||
|
||||
# Convert back to spatial domain
|
||||
mix = torch.fft.irfft2(fft_result, s=(h, w))
|
||||
|
||||
return mix, postuncond
|
||||
|
||||
def mix_fft_full(pred, weights, cfgs):
|
||||
# Full spectrum frequency mixing implementation - optimized version
|
||||
_, _, h, w = pred[0][0].shape
|
||||
|
||||
# Initialize result in frequency domain
|
||||
fft_result = torch.zeros_like(torch.fft.rfft2(pred[0][0]))
|
||||
|
||||
# First, convert all predictions to frequency domain
|
||||
fft_c_preds = []
|
||||
fft_u_preds = []
|
||||
valid_model_indices = []
|
||||
|
||||
for i in range(len(pred)):
|
||||
if weights[i] <= 0:
|
||||
# Add placeholders for zero-weight models
|
||||
fft_c_preds.append(None)
|
||||
fft_u_preds.append(None)
|
||||
continue
|
||||
|
||||
# Extract predictions for this model
|
||||
c_pred = pred[i][0]
|
||||
u_pred = pred[i][1]
|
||||
|
||||
# Convert to frequency domain
|
||||
fft_c_preds.append(torch.fft.rfft2(c_pred))
|
||||
fft_u_preds.append(torch.fft.rfft2(u_pred))
|
||||
valid_model_indices.append(i)
|
||||
|
||||
# Calculate frequency coordinates
|
||||
frequency_h, frequency_w = fft_result.shape[2], fft_result.shape[3]
|
||||
|
||||
# Create normalized coordinate grids (vectorized)
|
||||
h_coords = torch.arange(frequency_h, device=fft_result.device).float() / frequency_h
|
||||
w_coords = torch.arange(frequency_w, device=fft_result.device).float() / frequency_w
|
||||
|
||||
# Create meshgrid of coordinates
|
||||
h_grid, w_grid = torch.meshgrid(h_coords, w_coords, indexing='ij')
|
||||
|
||||
# Calculate frequency magnitude for all points at once
|
||||
freq_magnitude = torch.sqrt(h_grid**2 + w_grid**2)
|
||||
uncondresult = torch.zeros_like(fft_result)
|
||||
|
||||
# Process each valid model
|
||||
if valid_model_indices:
|
||||
# Pre-calculate model weights for all frequency points
|
||||
all_model_weights = []
|
||||
for i in valid_model_indices:
|
||||
# Calculate model influence: higher weights → lower frequencies, lower weights → higher frequencies
|
||||
model_weight = weights[i] * (1.0 - freq_magnitude) + (1.0 - weights[i]) * freq_magnitude
|
||||
all_model_weights.append(model_weight)
|
||||
|
||||
# Convert to tensor for easier operations
|
||||
all_model_weights = torch.stack(all_model_weights)
|
||||
|
||||
# Normalize weights along model dimension
|
||||
weight_sum = all_model_weights.sum(dim=0, keepdim=True)
|
||||
# Avoid division by zero
|
||||
weight_sum = torch.clamp(weight_sum, min=1e-10)
|
||||
normalized_weights = all_model_weights / weight_sum
|
||||
|
||||
# Apply weights and CFG to each frequency component
|
||||
for idx, i in enumerate(valid_model_indices):
|
||||
# Get model weights for this model
|
||||
model_weights = normalized_weights[idx]
|
||||
|
||||
# Apply weights to each frequency point (broadcasting)
|
||||
model_weights = model_weights.unsqueeze(0).unsqueeze(0) # Add batch and channel dimensions
|
||||
|
||||
# Get unconditioned and conditioned predictions
|
||||
fft_u = fft_u_preds[i]
|
||||
uncondresult = uncondresult + fft_u
|
||||
fft_c = fft_c_preds[i]
|
||||
|
||||
# Apply CFG with model-specific weight
|
||||
contrib = fft_u * model_weights + (fft_c - fft_u) * cfgs[i] * model_weights
|
||||
|
||||
# Add to result
|
||||
fft_result += contrib
|
||||
|
||||
|
||||
postuncond = torch.fft.irfft2(uncondresult, s=(h, w))
|
||||
|
||||
# Convert back to spatial domain
|
||||
mix = torch.fft.irfft2(fft_result, s=(h, w))
|
||||
|
||||
return mix, postuncond
|
||||
@@ -0,0 +1,494 @@
|
||||
import torch
|
||||
|
||||
from comfy import latent_formats
|
||||
from ..utils import create_depth_map, save_debug_image, process_mask
|
||||
import comfy.sd
|
||||
import numpy as np
|
||||
import comfy.samplers
|
||||
from ..utils import sumweights
|
||||
|
||||
|
||||
def mix_standard(pred, weights, cfgs, uncond_decays=[1.0]):
|
||||
"""
|
||||
Standard mixing function that applies CFG for each model based on weights.
|
||||
|
||||
Args:
|
||||
pred: List of model predictions, each containing [conditional, unconditional]
|
||||
weights: List of weights for each model
|
||||
cfgs: List of cfg values for each model
|
||||
|
||||
Returns:
|
||||
mix: The mixed prediction
|
||||
postuncond: The unconditioned prediction for post-processing
|
||||
"""
|
||||
print("mix_standard")
|
||||
# Initialize with the shape of the first model's prediction
|
||||
mix = torch.zeros_like(pred[0][1])
|
||||
|
||||
|
||||
# First, combine all unconditional predictions based on weights
|
||||
for i in range(len(pred)):
|
||||
mix = mix + pred[i][1] * weights[i]
|
||||
|
||||
# Save the combined unconditional prediction for post-processing
|
||||
postuncond = mix.clone()
|
||||
|
||||
# For each model, apply the guidance according to weights and cfg values
|
||||
for i in range(len(pred)):
|
||||
# Extract conditional and unconditional predictions for this model
|
||||
c_pred = pred[i][0]
|
||||
u_pred = pred[i][1]
|
||||
|
||||
# Apply the CFG formula with the corresponding weight and cfg
|
||||
# Add the weighted difference between conditional and unconditional
|
||||
mix = mix + (c_pred - u_pred) * cfgs[i]
|
||||
|
||||
return mix, postuncond
|
||||
|
||||
def mix_masked(pred, weights, cfgs, masks=None):
|
||||
"""
|
||||
Mask-based mixing function that applies different models to different regions.
|
||||
|
||||
Args:
|
||||
pred: List of model predictions, each containing [conditional, unconditional]
|
||||
weights: List of weights for each model
|
||||
cfgs: List of cfg values for each model
|
||||
masks: List of masks (one per model) - should already be properly formatted/feathered
|
||||
|
||||
Returns:
|
||||
mix: The mixed prediction
|
||||
postuncond: The unconditioned prediction for post-processing
|
||||
"""
|
||||
print("mix_masked")
|
||||
# Check if we have valid mask in the list
|
||||
#if not any(mask is not None for mask in masks):
|
||||
# # Fall back to standard mixing if no valid masks
|
||||
# return mix_standard(pred, weights, cfgs)
|
||||
|
||||
# Initialize with zeros
|
||||
B, C, H, W = pred[0][0].shape
|
||||
mix = torch.zeros_like(pred[0][0])
|
||||
postuncond = torch.zeros_like(pred[0][1])
|
||||
# Prepare masks
|
||||
processed_masks = []
|
||||
for i, mask in enumerate(masks):
|
||||
try:
|
||||
if(mask is not None):
|
||||
processed_mask = process_mask(mask, H, W, device=pred[0][0].device)
|
||||
processed_masks.append(processed_mask)
|
||||
else:
|
||||
processed_masks.append(torch.ones((H, W), device=pred[0][0].device))
|
||||
except Exception as e:
|
||||
print(f"Error processing mask {i}: {e}")
|
||||
processed_masks.append(torch.zeros((H, W), device=pred[0][0].device))
|
||||
|
||||
# If we couldn't process any masks, fall back to standard mixing
|
||||
if len(processed_masks) == 0:
|
||||
return mix_standard(pred, weights, cfgs)
|
||||
|
||||
# If we have fewer masks than models, pad with zeros
|
||||
while len(processed_masks) < len(pred):
|
||||
processed_masks.append(torch.ones_like(processed_masks[0], device=pred[0][0].device))
|
||||
|
||||
print(len(processed_masks))
|
||||
# Stack and normalize masks
|
||||
stacked_masks = torch.stack(processed_masks[:len(pred)])
|
||||
mask_sum = torch.sum(stacked_masks, dim=0, keepdim=True)
|
||||
mask_sum = torch.clamp(mask_sum, min=1e-6) # Avoid division by zero
|
||||
normalized_masks = stacked_masks / mask_sum
|
||||
|
||||
# Apply each model according to its normalized mask
|
||||
for i in range(len(pred)):
|
||||
if i >= len(normalized_masks):
|
||||
continue
|
||||
|
||||
# Get mask for this model and add batch and channel dimensions
|
||||
model_mask = normalized_masks[i].unsqueeze(0).unsqueeze(0).repeat(B, C, 1, 1)
|
||||
|
||||
# Extract predictions
|
||||
c_pred = pred[i][0]
|
||||
u_pred = pred[i][1]
|
||||
|
||||
# Apply CFG
|
||||
model_result = u_pred + (c_pred - u_pred) * cfgs[i]
|
||||
|
||||
# Add to result based on mask
|
||||
mix = mix + model_result * model_mask
|
||||
postuncond = postuncond + u_pred * model_mask
|
||||
|
||||
return mix, postuncond
|
||||
|
||||
def mix_dynamic_mask(pred, weights, cfgs, step=0, history=None, decay_factor=0.9, threshold=0.1, blur_sigma=0.5, max_update_step=15):
|
||||
"""
|
||||
Dynamic mask-based mixing that evolves masks based on prediction magnitudes over time.
|
||||
|
||||
Args:
|
||||
pred: List of model predictions, each containing [conditional, unconditional]
|
||||
weights: List of weights for each model
|
||||
cfgs: List of cfg values for each model
|
||||
step: Current generation step (used for mask evolution)
|
||||
history: Previous mask history to evolve from (None for first step)
|
||||
decay_factor: How much to retain from previous mask (0-1)
|
||||
threshold: Threshold for magnitude difference to affect mask
|
||||
blur_sigma: Standard deviation for Gaussian blur applied to masks (0 for no blur)
|
||||
max_update_step: If set, stops updating masks after this step and uses existing masks
|
||||
|
||||
Returns:
|
||||
mix: The mixed prediction
|
||||
postuncond: The unconditioned prediction for post-processing
|
||||
masks: The updated masks to pass to the next step
|
||||
"""
|
||||
# Get dimensions from the first prediction
|
||||
B, C, H, W = pred[0][0].shape
|
||||
device = pred[0][0].device
|
||||
# Initialize output tensors
|
||||
mix = torch.zeros_like(pred[0][0])
|
||||
postuncond = torch.zeros_like(pred[0][1])
|
||||
|
||||
# Initialize or update masks
|
||||
if history is None or step == 0:
|
||||
# First step: initialize with equal masks or based on weights
|
||||
masks = []
|
||||
for i in range(len(pred)):
|
||||
masks.append(torch.ones((H, W), device=device) * weights[i])
|
||||
else:
|
||||
# Use previous masks as starting point
|
||||
masks = history
|
||||
# Save masks to debug folder
|
||||
for i, mask in enumerate(masks):
|
||||
save_debug_image(mask, f"mask_{i}.png")
|
||||
|
||||
# Check if we should update masks based on max_update_step
|
||||
should_update_masks = max_update_step is None or step <= max_update_step
|
||||
|
||||
if should_update_masks:
|
||||
# Calculate magnitudes for each model's prediction
|
||||
magnitudes = []
|
||||
for i in range(len(pred)):
|
||||
# Use the difference between conditional and unconditional as a measure of "importance"
|
||||
model_contrib = torch.abs(pred[i][0] - pred[i][1])
|
||||
# Average across channels to get a 2D magnitude map
|
||||
magnitude = torch.mean(model_contrib, dim=1)[0] # Take first batch
|
||||
magnitudes.append(magnitude)
|
||||
|
||||
# Stack magnitudes for comparison
|
||||
stacked_magnitudes = torch.stack(magnitudes)
|
||||
|
||||
# Find which model has highest magnitude at each pixel
|
||||
max_indices = torch.argmax(stacked_magnitudes, dim=0)
|
||||
|
||||
# Update masks based on magnitudes
|
||||
updated_masks = []
|
||||
for i in range(len(pred)):
|
||||
# Create binary mask where this model has highest magnitude
|
||||
model_dominant = (max_indices == i).float()
|
||||
|
||||
# Apply threshold to only update where difference is significant
|
||||
magnitude_threshold = torch.max(stacked_magnitudes) * threshold
|
||||
significant_diff = (stacked_magnitudes[i] > magnitude_threshold).float()
|
||||
|
||||
# Combine with threshold
|
||||
update_mask = model_dominant * significant_diff
|
||||
|
||||
# Evolve mask gradually (blend old and new)
|
||||
if i < len(masks):
|
||||
# Decay old mask and add new information
|
||||
new_mask = masks[i] * decay_factor + update_mask * (1 - decay_factor)
|
||||
updated_masks.append(new_mask)
|
||||
else:
|
||||
updated_masks.append(update_mask)
|
||||
|
||||
# Use the updated masks
|
||||
mask_list = updated_masks
|
||||
else:
|
||||
# Just use the existing masks without updating
|
||||
if isinstance(masks, list):
|
||||
mask_list = masks
|
||||
else:
|
||||
# Convert tensor masks back to list if needed
|
||||
mask_list = [masks[i] for i in range(masks.shape[0])]
|
||||
|
||||
# Normalize masks to sum to 1 at each pixel
|
||||
stacked_masks = torch.stack(mask_list)
|
||||
mask_sum = torch.sum(stacked_masks, dim=0, keepdim=True)
|
||||
mask_sum = torch.clamp(mask_sum, min=1e-6) # Avoid division by zero
|
||||
normalized_masks = stacked_masks / mask_sum
|
||||
|
||||
# Apply Gaussian blur to the masks if blur_sigma > 0
|
||||
if blur_sigma > 0:
|
||||
import torch.nn.functional as F
|
||||
import math
|
||||
blurred_masks = []
|
||||
|
||||
# Create a custom Gaussian blur function since F.gaussian_blur isn't available
|
||||
def custom_gaussian_blur(input_tensor, kernel_size, sigma):
|
||||
# Create Gaussian kernel
|
||||
channels = input_tensor.shape[1]
|
||||
kernel_size = [kernel_size, kernel_size]
|
||||
sigma = [sigma, sigma]
|
||||
|
||||
# Create a grid of coordinates
|
||||
x_coord = torch.arange(kernel_size[0], device=input_tensor.device)
|
||||
x_grid = x_coord.repeat(kernel_size[1]).view(kernel_size[1], kernel_size[0])
|
||||
y_grid = x_grid.t()
|
||||
|
||||
xy_grid = torch.stack([x_grid, y_grid], dim=-1).float()
|
||||
|
||||
# Calculate the center of the filter
|
||||
mean = (torch.tensor(kernel_size, device=input_tensor.device) - 1) / 2.
|
||||
variance = torch.tensor(sigma, device=input_tensor.device) ** 2
|
||||
|
||||
# Calculate the 2d gaussian kernel
|
||||
gaussian_kernel = torch.exp(
|
||||
-torch.sum((xy_grid - mean) ** 2., dim=-1) / (2 * variance[0])
|
||||
)
|
||||
|
||||
# Make sure sum of values in gaussian kernel equals 1.
|
||||
gaussian_kernel = gaussian_kernel / torch.sum(gaussian_kernel)
|
||||
|
||||
# Reshape to 2d depth-wise convolutional weight
|
||||
gaussian_kernel = gaussian_kernel.view(1, 1, kernel_size[0], kernel_size[1])
|
||||
gaussian_kernel = gaussian_kernel.repeat(channels, 1, 1, 1)
|
||||
|
||||
# Apply padding
|
||||
pad_size = (kernel_size[0] // 2, kernel_size[1] // 2)
|
||||
padded_input = F.pad(input_tensor, (pad_size[0], pad_size[0], pad_size[1], pad_size[1]), mode='reflect')
|
||||
|
||||
# Apply Gaussian blur using depthwise convolution
|
||||
return F.conv2d(padded_input, gaussian_kernel, groups=channels)
|
||||
|
||||
for i in range(len(normalized_masks)):
|
||||
mask = normalized_masks[i].unsqueeze(0).unsqueeze(0) # Add batch and channel dimensions
|
||||
# Apply Gaussian blur (kernel size should be odd and >= 3)
|
||||
kernel_size = max(3, int(blur_sigma * 4 + 0.5) * 2 + 1)
|
||||
blurred_mask = custom_gaussian_blur(mask, kernel_size, blur_sigma)
|
||||
blurred_masks.append(blurred_mask.squeeze(0).squeeze(0))
|
||||
|
||||
# Re-normalize after blurring to ensure they still sum to 1
|
||||
stacked_blurred_masks = torch.stack(blurred_masks)
|
||||
blurred_sum = torch.sum(stacked_blurred_masks, dim=0, keepdim=True)
|
||||
blurred_sum = torch.clamp(blurred_sum, min=1e-6)
|
||||
normalized_masks = stacked_blurred_masks / blurred_sum
|
||||
|
||||
# Apply each model according to its normalized mask
|
||||
for i in range(len(pred)):
|
||||
# Get mask for this model and add batch and channel dimensions
|
||||
model_mask = normalized_masks[i].unsqueeze(0).unsqueeze(0).repeat(B, C, 1, 1)
|
||||
|
||||
# Extract predictions
|
||||
c_pred = pred[i][0]
|
||||
u_pred = pred[i][1]
|
||||
|
||||
# Apply CFG
|
||||
model_result = u_pred + (c_pred - u_pred) * cfgs[i]
|
||||
|
||||
# Add to result based on mask
|
||||
mix = mix + model_result * model_mask
|
||||
postuncond = postuncond + u_pred * model_mask
|
||||
|
||||
# Return the mixed result, unconditioned prediction, and updated masks for next step
|
||||
return mix, postuncond, normalized_masks.detach()
|
||||
|
||||
def mix_dynamic_mask_alternative(pred, weights, cfgs, depth_model=None, threshold=0.5, invert=False, blur_sigma=1.0):
|
||||
"""
|
||||
Mixing function that automatically detects foreground/background and applies different models to each.
|
||||
|
||||
Args:
|
||||
pred: List of model predictions, each containing [conditional, unconditional]
|
||||
weights: List of weights for each model
|
||||
cfgs: List of cfg values for each model
|
||||
depth_model: Optional depth estimation model. If None, uses simple heuristics
|
||||
threshold: Threshold for foreground/background separation (0-1)
|
||||
invert: If True, inverts the mask (foreground becomes background)
|
||||
blur_sigma: Blur amount to apply to mask edges
|
||||
|
||||
Returns:
|
||||
mix: The mixed prediction
|
||||
postuncond: The unconditioned prediction for post-processing
|
||||
"""
|
||||
# Need at least 2 models to do foreground/background mixing
|
||||
if len(pred) < 2:
|
||||
return mix_standard(pred, weights, cfgs)
|
||||
|
||||
# Get dimensions from the first prediction
|
||||
B, C, H, W = pred[0][0].shape
|
||||
device = pred[0][0].device
|
||||
|
||||
# Generate foreground/background mask using prediction features
|
||||
if depth_model is None:
|
||||
# Simple heuristic: use the magnitude of the conditional prediction
|
||||
# as an approximation of feature importance
|
||||
c_pred = pred[0][0] # Use first model's conditional prediction
|
||||
|
||||
# Calculate magnitude across channels (approximate depth/saliency)
|
||||
magnitude = torch.mean(torch.abs(c_pred), dim=1, keepdim=True)
|
||||
|
||||
# Normalize to 0-1 range
|
||||
min_val = torch.min(magnitude)
|
||||
max_val = torch.max(magnitude)
|
||||
normalized = (magnitude - min_val) / (max_val - min_val + 1e-8)
|
||||
|
||||
# Apply threshold
|
||||
mask = (normalized > threshold).float()
|
||||
|
||||
# Invert if requested
|
||||
if invert:
|
||||
mask = 1.0 - mask
|
||||
else:
|
||||
# Use provided depth model (not implemented here, would need integration)
|
||||
# This would replace the simple heuristic above with actual depth estimation
|
||||
mask = torch.ones((B, 1, H, W), device=device) * 0.5
|
||||
# mask = depth_model(c_pred)
|
||||
|
||||
# Apply Gaussian blur to smooth mask edges
|
||||
if blur_sigma > 0:
|
||||
import torch.nn.functional as F
|
||||
|
||||
# Create a Gaussian kernel
|
||||
kernel_size = max(3, int(blur_sigma * 4 + 0.5) * 2 + 1)
|
||||
|
||||
# Apply Gaussian blur
|
||||
mask = F.pad(mask, (kernel_size//2, kernel_size//2, kernel_size//2, kernel_size//2), mode='reflect')
|
||||
mask = F.avg_pool2d(mask, kernel_size, stride=1, padding=0)
|
||||
|
||||
# Create two masks: foreground and background
|
||||
mask_fg = mask
|
||||
mask_bg = 1.0 - mask
|
||||
|
||||
# Ensure masks have the correct shape (H, W) without extra dimensions
|
||||
if mask_fg.dim() > 2:
|
||||
# Extract a 2D mask, squeezing out extra dimensions
|
||||
if mask_fg.dim() == 3: # [C, H, W] or [1, H, W]
|
||||
mask_fg = mask_fg.squeeze(0)
|
||||
elif mask_fg.dim() == 4: # [B, C, H, W]
|
||||
mask_fg = mask_fg.squeeze(0).squeeze(0)
|
||||
|
||||
# If mask still has more than 2 dimensions, use only the first 2D slice
|
||||
if mask_fg.dim() > 2:
|
||||
mask_fg = mask_fg[0]
|
||||
|
||||
# Do the same for background mask
|
||||
if mask_bg.dim() > 2:
|
||||
if mask_bg.dim() == 3:
|
||||
mask_bg = mask_bg.squeeze(0)
|
||||
elif mask_bg.dim() == 4:
|
||||
mask_bg = mask_bg.squeeze(0).squeeze(0)
|
||||
|
||||
if mask_bg.dim() > 2:
|
||||
mask_bg = mask_bg[0]
|
||||
|
||||
# Create list of masks for each model
|
||||
masks = [mask_fg, mask_bg]
|
||||
|
||||
# Create directory if it doesn't exist
|
||||
for i, mask in enumerate(masks):
|
||||
save_debug_image(mask, f"mask_{i}.png")
|
||||
|
||||
# If more than 2 models, create zero masks for the rest
|
||||
for i in range(2, len(pred)):
|
||||
masks.append(torch.zeros_like(mask_fg))
|
||||
|
||||
# Use the existing masked mixing function with our generated masks
|
||||
return mix_masked(pred, weights, cfgs, masks)
|
||||
|
||||
def mix_foreground_background(pred, weights, last_latent, cfgs, step=0, invert=False, masks=None, vae=None, sharpness=0.0, depth_model=None, startmask=5):
|
||||
"""
|
||||
Mixing function that uses Depth Anything V2 to separate foreground and background.
|
||||
|
||||
Args:
|
||||
pred: List of model predictions, each containing [conditional, unconditional]
|
||||
weights: List of weights for each model
|
||||
cfgs: List of cfg values for each model
|
||||
invert: If True, inverts the mask (foreground becomes background)
|
||||
ckpt_name: The Depth Anything V2 checkpoint to use
|
||||
threshold: Threshold for foreground/background separation (0-1)
|
||||
blur_sigma: Blur amount to apply to mask edges
|
||||
|
||||
Returns:
|
||||
mix: The mixed prediction
|
||||
postuncond: The unconditioned prediction for post-processing
|
||||
"""
|
||||
|
||||
# Get dimensions from the first prediction
|
||||
B, C, H, W = pred[0][0].shape
|
||||
#get the higes resolution of the latent image make sure its a power of 2
|
||||
res = 2**(max(H, W).bit_length() - 1)
|
||||
device = pred[0][0].device
|
||||
|
||||
# Get the unconditioned prediction from the first model to generate a depth map
|
||||
u_pred = pred[0][1]
|
||||
|
||||
# Convert latent to image for depth estimation
|
||||
# This is a simplified approach - assumes we're working with latents
|
||||
# In a real implementation, you'd need proper latent-to-image conversion
|
||||
|
||||
if not any(mask is not None for mask in masks):
|
||||
if step == startmask:
|
||||
# Use Depth Anything V2 to generate a depth map
|
||||
latent_image = last_latent
|
||||
try:
|
||||
import comfy.model_management as model_management
|
||||
from custom_nodes.comfyui_controlnet_aux.utils import common_annotator_call
|
||||
|
||||
|
||||
|
||||
# Convert latent to proper image using TaeSDXL Mini VAE
|
||||
try:
|
||||
latent_image = latent_formats.SDXL().process_out(latent_image)
|
||||
latent_image = vae.decode(latent_image)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error using TAESD: {e}")
|
||||
|
||||
|
||||
out_tensor = create_depth_map(latent_image, depth_model, res)
|
||||
|
||||
depth_norm = out_tensor
|
||||
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error using Depth Anything V2: {e}")
|
||||
print("Falling back to simple heuristic for foreground/background separation")
|
||||
|
||||
# Fallback: use the magnitude of the conditional prediction as approximation
|
||||
c_pred = pred[0][0]
|
||||
depth_norm = torch.mean(torch.abs(c_pred), dim=1, keepdim=True)
|
||||
depth_norm = (depth_norm - depth_norm.min()) / (depth_norm.max() - depth_norm.min() + 1e-8)
|
||||
mask = depth_norm * (1.0 + sharpness) - (depth_norm * sharpness).mean()
|
||||
mask = torch.clamp(mask, 0.0, 1.0)
|
||||
|
||||
mask = mask
|
||||
# Invert if requested
|
||||
if invert:
|
||||
mask = 1.0 - mask
|
||||
|
||||
# Create two masks: foreground and background
|
||||
mask_fg = mask
|
||||
mask_bg = 1.0 - mask
|
||||
|
||||
|
||||
|
||||
# Create list of masks for each model
|
||||
masks = [mask_fg, mask_bg]
|
||||
|
||||
# If more than 2 models, create zero masks for the rest
|
||||
for i in range(2, len(pred)):
|
||||
masks.append(torch.zeros_like(mask_fg))
|
||||
|
||||
mix, postuncond = mix_masked(pred, weights, cfgs, masks)
|
||||
# Debug: Save masks
|
||||
try:
|
||||
# Save various debug images
|
||||
save_debug_image(latent_image, "last_latent.png")
|
||||
|
||||
save_debug_image(mask, "mask.png")
|
||||
except Exception as e:
|
||||
print(f"Failed to save debug masks: {e}")
|
||||
else:
|
||||
mix, postuncond = mix_standard(pred, weights, cfgs)
|
||||
else:
|
||||
mix, postuncond = mix_masked(pred, weights, cfgs, masks)
|
||||
torch.cuda.empty_cache()
|
||||
# Use the existing masked mixing function with our generated masks
|
||||
return masks, mix, postuncond
|
||||
|
||||
@@ -1,9 +1,7 @@
|
||||
import torch
|
||||
import math
|
||||
import comfy.samplers
|
||||
import comfy.utils
|
||||
import nodes
|
||||
import logging
|
||||
from .guider import MixModGuider
|
||||
|
||||
|
||||
# COMPONENTS
|
||||
|
||||
class MixModGuiderComponentNode:
|
||||
@classmethod
|
||||
@@ -17,7 +15,8 @@ class MixModGuiderComponentNode:
|
||||
},
|
||||
"optional":
|
||||
{
|
||||
"prev_component": ("COMPONENT",),
|
||||
"prev_component": ("COMPONENT",),
|
||||
"options": ("OPTIONS",),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -25,24 +24,80 @@ class MixModGuiderComponentNode:
|
||||
FUNCTION = "create_component"
|
||||
CATEGORY = "mixmod/conditioning"
|
||||
|
||||
def create_component(self, model, positive, negative, base_weight, cfg, mask=None, prev_component=None):
|
||||
def create_component(self, model, positive, negative, base_weight, cfg, options=None, prev_component=None):
|
||||
component = {
|
||||
"model": model,
|
||||
"positive": positive,
|
||||
"negative": negative,
|
||||
"weight": base_weight,
|
||||
"cfg": cfg,
|
||||
"mask": mask,
|
||||
"options": options,
|
||||
"prev": prev_component
|
||||
}
|
||||
return (component,)
|
||||
|
||||
|
||||
class MixModGuiderComponentPipelineNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"pipeline": ("PIPELINE",),
|
||||
"base_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"cfg": ("FLOAT", {"default": 7.5, "min": -100.0, "max": 100.0, "step": 0.1}),
|
||||
|
||||
},
|
||||
"optional":
|
||||
{
|
||||
"prev_component": ("COMPONENT",),
|
||||
"options": ("OPTIONS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("COMPONENT",)
|
||||
FUNCTION = "create_component"
|
||||
CATEGORY = "mixmod/conditioning"
|
||||
|
||||
def create_component(self, pipeline, base_weight, cfg, options=None, prev_component=None):
|
||||
component = {
|
||||
"model": pipeline["model"],
|
||||
"positive": pipeline["positive"],
|
||||
"negative": pipeline["negative"],
|
||||
"weight": base_weight,
|
||||
"cfg": cfg,
|
||||
"options": options,
|
||||
"prev": prev_component
|
||||
}
|
||||
return (component,)
|
||||
|
||||
|
||||
class MixModPipelineNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"model": ("MODEL",),
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("PIPELINE",)
|
||||
FUNCTION = "create_pipeline"
|
||||
CATEGORY = "mixmod/conditioning"
|
||||
|
||||
def create_pipeline(self, model, positive, negative):
|
||||
pipeline = {
|
||||
"model": model,
|
||||
"positive": positive,
|
||||
"negative": negative,
|
||||
}
|
||||
return (pipeline,)
|
||||
|
||||
|
||||
class MixModGuiderNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"component": ("COMPONENT",),
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
@@ -50,18 +105,162 @@ class MixModGuiderNode:
|
||||
FUNCTION = "create_guider"
|
||||
CATEGORY = "mixmod/conditioning"
|
||||
|
||||
def create_guider(self, component, image=None,):
|
||||
def create_guider(self, component):
|
||||
# This is a placeholder for the actual guider implementation
|
||||
# The actual implementation will be filled in by the user
|
||||
guider = MixModGuider(component, mode="team", image=image)
|
||||
return (guider,)
|
||||
guider = MixModGuider(component, mode="team")
|
||||
return (guider,)
|
||||
|
||||
|
||||
# OPTIONS
|
||||
|
||||
class MixModOptionsUncondDecayNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"uncond_decay": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
},
|
||||
"optional":{
|
||||
"prev_options": ("OPTIONS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("OPTIONS",)
|
||||
FUNCTION = "options"
|
||||
CATEGORY = "mixmod/conditioning"
|
||||
|
||||
def options(self, uncond_decay, prev_options=None):
|
||||
options = {
|
||||
"uncond_decay": uncond_decay,
|
||||
"prev": prev_options
|
||||
}
|
||||
return (options,)
|
||||
|
||||
|
||||
class MixModOptionsMaskNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"mask": ("MASK",),
|
||||
},
|
||||
"optional":{
|
||||
"prev_options": ("OPTIONS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("OPTIONS",)
|
||||
FUNCTION = "options"
|
||||
CATEGORY = "mixmod/conditioning"
|
||||
|
||||
def options(self, mask, prev_options=None):
|
||||
options = {
|
||||
"mask": mask,
|
||||
"prev": prev_options
|
||||
}
|
||||
return (options,)
|
||||
|
||||
class MixModOptionsSchedulerNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{
|
||||
"start_step": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
|
||||
"end_step": ("INT", {"default": -1, "min": -1, "max": 100, "step": 1}),
|
||||
},
|
||||
"optional":{
|
||||
"prev_options": ("OPTIONS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("OPTIONS",)
|
||||
FUNCTION = "options"
|
||||
CATEGORY = "mixmod/conditioning"
|
||||
|
||||
def options(self, start_step, end_step, prev_options=None):
|
||||
options = {
|
||||
"start_step": start_step,
|
||||
"end_step": end_step,
|
||||
"prev": prev_options
|
||||
}
|
||||
return (options,)
|
||||
|
||||
|
||||
class MixModOptionsScaleNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{
|
||||
"scale_factor": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
},
|
||||
"optional":{
|
||||
"prev_options": ("OPTIONS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("OPTIONS",)
|
||||
FUNCTION = "options"
|
||||
CATEGORY = "mixmod/conditioning"
|
||||
|
||||
def options(self, scale_factor, prev_options=None):
|
||||
options = {
|
||||
"scale_factor": scale_factor,
|
||||
"prev": prev_options
|
||||
}
|
||||
return (options,)
|
||||
|
||||
class MixModOptionsMaskNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{
|
||||
"mask": ("MASK",),
|
||||
},
|
||||
"optional":{
|
||||
"prev_options": ("OPTIONS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("OPTIONS",)
|
||||
FUNCTION = "options"
|
||||
CATEGORY = "mixmod/conditioning"
|
||||
|
||||
def options(self, mask, prev_options=None):
|
||||
options = {
|
||||
"mask": mask,
|
||||
"prev": prev_options
|
||||
}
|
||||
return (options,)
|
||||
|
||||
|
||||
# SPECIAL GUIDERS
|
||||
|
||||
|
||||
class MixModStyleInjectionGuiderNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"component": ("COMPONENT",),
|
||||
"image": ("IMAGE",),
|
||||
"noise_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"start_step": ("INT", {"default": 4, "min": 0, "max": 100, "step": 1}),
|
||||
"end_step": ("INT", {"default": -1, "min": -1, "max": 100, "step": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GUIDER",)
|
||||
FUNCTION = "create_guider"
|
||||
CATEGORY = "mixmod/conditioning"
|
||||
|
||||
def create_guider(self, component, image, noise_weight=1.0, start_step=4, end_step=-1):
|
||||
guider = MixModGuider(component, mode="styleinjection", image=image, noise_weight=noise_weight, start_step=start_step, end_step=end_step)
|
||||
return (guider,)
|
||||
|
||||
class MixModFFTGuiderNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"component": ("COMPONENT",),
|
||||
"mode": (["teamfft", "2model_fft"],),
|
||||
"mode": (["teamfft", "2model_fft", "2m2f", "teamfft4freq", "fft_overlap", "fft_full"],),
|
||||
"ratio": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
}
|
||||
}
|
||||
@@ -70,577 +269,150 @@ class MixModFFTGuiderNode:
|
||||
FUNCTION = "create_guider"
|
||||
CATEGORY = "mixmod/conditioning"
|
||||
|
||||
def create_guider(self, component, mode="team", image=None, ratio=1.0):
|
||||
def create_guider(self, component, mode="team", ratio=1.0):
|
||||
# This is a placeholder for the actual guider implementation
|
||||
# The actual implementation will be filled in by the user
|
||||
guider = MixModGuider(component, mode=mode, image=image, ratio=ratio)
|
||||
guider = MixModGuider(component, mode=mode, ratio=ratio)
|
||||
return (guider,)
|
||||
|
||||
|
||||
class MixModBandFFTGuiderNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"component": ("COMPONENT",),
|
||||
"bands": ("INT", {"default": 2, "min": 1, "max": 100, "step": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GUIDER",)
|
||||
FUNCTION = "create_guider"
|
||||
CATEGORY = "mixmod/conditioning"
|
||||
|
||||
def create_guider(self, component, bands=2):
|
||||
# This is a placeholder for the actual guider implementation
|
||||
# The actual implementation will be filled in by the user
|
||||
guider = MixModGuider(component, mode="bandfft", bands=bands)
|
||||
return (guider,)
|
||||
|
||||
|
||||
class MixModDynamicMaskGuiderNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"component": ("COMPONENT",),
|
||||
"decay_factor": ("FLOAT", {"default": 0.9, "min": 0.1, "max": 0.99, "step": 0.01}),
|
||||
"threshold": ("FLOAT", {"default": 0.1, "min": 0.01, "max": 0.5, "step": 0.01}),
|
||||
"blur_sigma": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 10.0, "step": 0.01}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GUIDER",)
|
||||
FUNCTION = "create_guider"
|
||||
CATEGORY = "mixmod/conditioning"
|
||||
|
||||
def create_guider(self, component, decay_factor=0.9, threshold=0.1, blur_sigma=0.5):
|
||||
guider = MixModGuider(component, mode="dynamic_mask", decay_factor=decay_factor, threshold=threshold, blur_sigma=blur_sigma)
|
||||
return (guider,)
|
||||
|
||||
class MixModDynamicMaskAlternativeGuiderNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"component": ("COMPONENT",),
|
||||
"threshold": ("FLOAT", {"default": 0.1, "min": 0.01, "max": 0.5, "step": 0.01}),
|
||||
"blur_sigma": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 10.0, "step": 0.01}),
|
||||
"invert": ("BOOLEAN", {"default": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GUIDER",)
|
||||
FUNCTION = "create_guider"
|
||||
CATEGORY = "mixmod/conditioning"
|
||||
|
||||
def create_guider(self, component, threshold=0.1, blur_sigma=0.5, invert=False):
|
||||
guider = MixModGuider(component, mode="dynamic_mask_alternative", threshold=threshold, blur_sigma=blur_sigma, invert=invert)
|
||||
return (guider,)
|
||||
|
||||
class MixModDepthGuiderNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"component": ("COMPONENT",),
|
||||
"sharpness": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"start": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
|
||||
},
|
||||
"optional":
|
||||
{
|
||||
"invert": ("BOOLEAN", {"default": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GUIDER",)
|
||||
FUNCTION = "create_guider"
|
||||
CATEGORY = "mixmod/conditioning"
|
||||
|
||||
def create_guider(self, component, sharpness=0.0, invert=False, start=0):
|
||||
guider = MixModGuider(component, mode="foreground_background", sharpness=sharpness, invert=invert, start=start)
|
||||
return (guider,)
|
||||
|
||||
|
||||
class MixModGuider(comfy.samplers.CFGGuider):
|
||||
def __init__(self, component, mode="team", vae=None, image=None, prev_guider=None, ratio=1.0):
|
||||
|
||||
self.component = component
|
||||
self.mode = mode
|
||||
self.vae = vae
|
||||
self.image = image
|
||||
self.prev_guider = prev_guider
|
||||
self.ratio = ratio
|
||||
self.original_conds = {}
|
||||
self.prepared_models = []
|
||||
# Initialize other necessary attributes
|
||||
self.models = []
|
||||
self.weights = []
|
||||
self.cfgs = []
|
||||
self.masks = []
|
||||
# Process component chain
|
||||
current = component
|
||||
conds = {}
|
||||
while current is not None:
|
||||
self.models.append(current["model"])
|
||||
if len(self.models) == 1:
|
||||
self.model_patcher = self.models[0]
|
||||
self.model_options = self.models[0].model_options
|
||||
self.weights.append(current["weight"])
|
||||
self.cfgs.append(current["cfg"])
|
||||
self.masks.append(current["mask"])
|
||||
|
||||
# Add to original_conds with indexed keys
|
||||
idx = len(self.models)
|
||||
conds.update({f"positive_{idx}": current["positive"], f"negative_{idx}": current["negative"]})
|
||||
|
||||
current = current.get("prev")
|
||||
|
||||
self.inner_set_conds(conds)
|
||||
|
||||
weight_sum = sum(self.weights)
|
||||
|
||||
# Normalize weights if sum is not zero
|
||||
if weight_sum > 0:
|
||||
self.weights = [w / weight_sum for w in self.weights]
|
||||
|
||||
def set_mode(self, mode):
|
||||
self.mode = mode
|
||||
|
||||
|
||||
def predict_noise(self, x, timestep, model_options={}, seed=None):
|
||||
pred = []
|
||||
for i in range(len(self.models)):
|
||||
cond = [self.conds.get(f"positive_{i+1}", None), self.conds.get(f"negative_{i+1}", None)]
|
||||
|
||||
pred.append(comfy.samplers.calc_cond_batch(self.prepared_models[i][0], cond, x, timestep, model_options))
|
||||
|
||||
for fn in model_options.get("sampler_pre_cfg_function", []):
|
||||
args = {"conds":cond, "conds_out": pred, "cond_scale": self.cfgs[i], "timestep": timestep,
|
||||
"input": x, "sigma": timestep, "model": self.prepared_models[i][0], "model_options": model_options}
|
||||
pred[i] = fn(args)
|
||||
|
||||
return self.mix_function(self.prepared_models[0][0], pred, x, timestep, model_options)
|
||||
|
||||
def mix_function(self, model, pred, x, timestep, model_options):
|
||||
|
||||
if(self.mode == "team"):
|
||||
# mix = (uncond_pred * ratio + baduncond * (1.0 - ratio)) + (cond_pred - uncond_pred) * cond_scale + (badcond - baduncond) * cond_scale_2
|
||||
# Use pred array with self.weights and self.cfgs
|
||||
# First, initialize with the unconditioned prediction from the first model
|
||||
mix = torch.zeros_like(pred[0][1])
|
||||
c_predf = torch.zeros_like(pred[0][0])
|
||||
u_predf = torch.zeros_like(pred[0][1])
|
||||
|
||||
for i in range(len(pred)):
|
||||
mix = mix + pred[i][1] * self.weights[i]
|
||||
|
||||
# For each model, apply the guidance according to weights and cfg values
|
||||
for i in range(len(pred)):
|
||||
# Extract conditional and unconditional predictions for this model
|
||||
c_pred = pred[i][0]
|
||||
u_pred = pred[i][1]
|
||||
|
||||
c_predf = c_predf + c_pred
|
||||
u_predf = u_predf + u_pred
|
||||
|
||||
# Apply the CFG formula with the corresponding weight and cfg
|
||||
# Add the weighted difference between conditional and unconditional
|
||||
mix = mix + (c_pred - u_pred) * self.cfgs[i]
|
||||
|
||||
|
||||
for fn in model_options.get("sampler_post_cfg_function", []):
|
||||
args = {"denoised": mix, "cond": self.conds.get(f"positive_1", None), "uncond": self.conds.get(f"negative_1", None), "cond_scale": self.cfgs[0], "model": model, "uncond_denoised": pred[0][1], "cond_denoised": pred[0][0],
|
||||
"sigma": timestep, "model_options": model_options, "input": x}
|
||||
mix = fn(args)
|
||||
|
||||
if(self.mode == "teamfft"):
|
||||
# mix = (uncond_pred * ratio + baduncond * (1.0 - ratio)) + (cond_pred - uncond_pred) * cond_scale + (badcond - baduncond) * cond_scale_2
|
||||
# Use pred array with self.weights and self.cfgs
|
||||
# First, get shape information from the first model's predictions
|
||||
_, _, h, w = pred[0][0].shape
|
||||
|
||||
# Initialize an empty tensor in frequency domain for the result
|
||||
fft_mix = torch.zeros_like(torch.fft.rfft2(pred[0][1]))
|
||||
|
||||
# First loop: mix the unconditioned predictions in frequency domain
|
||||
for i in range(len(pred)):
|
||||
# Convert unconditional prediction to frequency domain
|
||||
fft_u_pred = torch.fft.rfft2(pred[i][1])
|
||||
# Add weighted contribution
|
||||
fft_mix = fft_mix + fft_u_pred * self.weights[i]
|
||||
|
||||
# Second loop: apply the guidance in frequency domain
|
||||
for i in range(len(pred)):
|
||||
# Convert conditional and unconditional predictions to frequency domain
|
||||
fft_c_pred = torch.fft.rfft2(pred[i][0])
|
||||
fft_u_pred = torch.fft.rfft2(pred[i][1])
|
||||
|
||||
# Apply the CFG formula in frequency domain
|
||||
fft_mix = fft_mix + (fft_c_pred - fft_u_pred) * self.cfgs[i]
|
||||
|
||||
# Convert the final result back to spatial domain
|
||||
mix = torch.fft.irfft2(fft_mix, s=(h, w))
|
||||
|
||||
if(self.mode == "teamfft4freq"):
|
||||
# Get shape information from the first model's predictions
|
||||
_, _, h, w = pred[0][0].shape
|
||||
|
||||
# Initialize an empty tensor in frequency domain for the result
|
||||
fft_mix = torch.zeros_like(torch.fft.rfft2(pred[0][1]))
|
||||
|
||||
# Get frequency domain dimensions
|
||||
freq_h, freq_w = fft_mix.shape[2], fft_mix.shape[3]
|
||||
|
||||
# Create frequency band masks - dividing into 4 frequency bands
|
||||
# Low, low-mid, high-mid, and high frequencies
|
||||
masks = []
|
||||
|
||||
# Calculate coordinates for frequency space
|
||||
h_coords = torch.linspace(0, 1, freq_h, device=fft_mix.device)
|
||||
w_coords = torch.linspace(0, 1, freq_w, device=fft_mix.device)
|
||||
h_grid, w_grid = torch.meshgrid(h_coords, w_coords, indexing='ij')
|
||||
|
||||
# Calculate frequency magnitude (distance from DC component)
|
||||
freq_magnitude = torch.sqrt(h_grid**2 + w_grid**2)
|
||||
|
||||
# Normalize to [0,1] range
|
||||
max_freq = torch.sqrt(torch.tensor(2.0, device=fft_mix.device))
|
||||
normalized_freq = freq_magnitude / max_freq
|
||||
|
||||
# Define frequency band boundaries
|
||||
band_boundaries = [0.0, 0.25, 0.5, 0.75, 1.0]
|
||||
|
||||
# Create masks for each frequency band with smooth transitions
|
||||
transition_width = 0.05 # 5% transition width
|
||||
|
||||
for i in range(4): # 4 frequency bands
|
||||
band_start = band_boundaries[i]
|
||||
band_end = band_boundaries[i+1]
|
||||
|
||||
# Calculate transition regions
|
||||
trans_start = max(0.0, band_start - transition_width/2)
|
||||
trans_end = min(1.0, band_end + transition_width/2)
|
||||
|
||||
# Create mask
|
||||
mask = torch.zeros_like(normalized_freq)
|
||||
|
||||
# Full strength in the main band region
|
||||
main_region = (normalized_freq >= band_start) & (normalized_freq <= band_end)
|
||||
mask = torch.where(main_region, torch.ones_like(mask), mask)
|
||||
|
||||
# Smooth transition at the start of the band
|
||||
if band_start > 0:
|
||||
start_trans = (normalized_freq >= trans_start) & (normalized_freq < band_start)
|
||||
mask = torch.where(
|
||||
start_trans,
|
||||
(normalized_freq - trans_start) / (band_start - trans_start),
|
||||
mask
|
||||
)
|
||||
|
||||
# Smooth transition at the end of the band
|
||||
if band_end < 1.0:
|
||||
end_trans = (normalized_freq > band_end) & (normalized_freq <= trans_end)
|
||||
mask = torch.where(
|
||||
end_trans,
|
||||
1.0 - (normalized_freq - band_end) / (trans_end - band_end),
|
||||
mask
|
||||
)
|
||||
|
||||
masks.append(mask)
|
||||
|
||||
# Ensure masks sum to 1.0 at each point for perfect reconstruction
|
||||
mask_sum = sum(masks)
|
||||
masks = [m / (mask_sum + 1e-8) for m in masks] # Add small epsilon to avoid division by zero
|
||||
|
||||
# Process each frequency band separately
|
||||
for band_idx in range(4):
|
||||
band_mix = torch.zeros_like(fft_mix)
|
||||
|
||||
# First loop: mix the unconditioned predictions in this frequency band
|
||||
for i in range(len(pred)):
|
||||
# Convert unconditional prediction to frequency domain
|
||||
fft_u_pred = torch.fft.rfft2(pred[i][1])
|
||||
# Add weighted contribution for this frequency band
|
||||
band_mix = band_mix + fft_u_pred * self.weights[i] * masks[band_idx]
|
||||
|
||||
# Second loop: apply the guidance in this frequency band
|
||||
for i in range(len(pred)):
|
||||
# Convert conditional and unconditional predictions to frequency domain
|
||||
fft_c_pred = torch.fft.rfft2(pred[i][0])
|
||||
fft_u_pred = torch.fft.rfft2(pred[i][1])
|
||||
|
||||
# Apply the CFG formula in this frequency band
|
||||
band_mix = band_mix + (fft_c_pred - fft_u_pred) * self.cfgs[i] * masks[band_idx]
|
||||
|
||||
# Add this band's contribution to the final mix
|
||||
fft_mix = fft_mix + band_mix
|
||||
|
||||
# Convert the final result back to spatial domain
|
||||
mix = torch.fft.irfft2(fft_mix, s=(h, w))
|
||||
|
||||
elif(self.mode == "2model_fft"):
|
||||
# Apply FFT-based frequency splitting
|
||||
# Need at least 2 models for FFT mode
|
||||
if len(pred) < 2:
|
||||
return pred[0][0] * self.cfgs[0] + pred[0][1] * (1 - self.cfgs[0]) # Return properly conditioned prediction
|
||||
# Get shapes from first prediction
|
||||
_, _, h, w = pred[0][0].shape
|
||||
|
||||
# Use the first model for low frequencies and second model for high frequencies
|
||||
# Extract predictions from first model
|
||||
c_pred_1 = pred[0][0] # Conditioned prediction
|
||||
u_pred_1 = pred[0][1] # Unconditioned prediction
|
||||
|
||||
# Extract predictions from second model
|
||||
c_pred_2 = pred[1][0] # Conditioned prediction
|
||||
u_pred_2 = pred[1][1] # Unconditioned prediction
|
||||
|
||||
# Convert to frequency domain
|
||||
fft_u_pred_1 = torch.fft.rfft2(u_pred_1)
|
||||
fft_c_pred_1 = torch.fft.rfft2(c_pred_1)
|
||||
fft_u_pred_2 = torch.fft.rfft2(u_pred_2)
|
||||
fft_c_pred_2 = torch.fft.rfft2(c_pred_2)
|
||||
|
||||
# Use weight of first model as frequency split ratio
|
||||
# Ensure ratio is between 0.05 and 0.95 to guarantee some frequency coverage for both models
|
||||
ratio = max(0.05, min(0.95, abs(self.ratio)))
|
||||
|
||||
# Calculate frequency split point
|
||||
split_h = int(h * ratio)
|
||||
split_w = int(w * ratio / 2) # rfft2 returns half the width in frequency domain
|
||||
|
||||
# Create masks for low and high frequencies
|
||||
mask_low = torch.ones_like(fft_u_pred_1, dtype=torch.float32)
|
||||
mask_high = torch.ones_like(fft_u_pred_1, dtype=torch.float32)
|
||||
|
||||
# Create transition area for smooth blending
|
||||
transition_h = max(1, int(h * 0.05)) # 5% transition zone
|
||||
transition_w = max(1, int(w * 0.05 / 2)) # 5% transition zone
|
||||
|
||||
# Set high frequencies to zero in low mask with smooth transition
|
||||
for i in range(mask_low.shape[2]):
|
||||
for j in range(mask_low.shape[3]):
|
||||
if i >= split_h + transition_h or j >= split_w + transition_w:
|
||||
mask_low[:, :, i, j] = 0.0
|
||||
elif i >= split_h or j >= split_w:
|
||||
# Calculate smooth transition value
|
||||
fade_h = 1.0 - min(1.0, (i - split_h) / transition_h) if i >= split_h else 1.0
|
||||
fade_w = 1.0 - min(1.0, (j - split_w) / transition_w) if j >= split_w else 1.0
|
||||
mask_low[:, :, i, j] = fade_h * fade_w
|
||||
|
||||
# Set low frequencies to zero in high mask with inverse smooth transition
|
||||
for i in range(mask_high.shape[2]):
|
||||
for j in range(mask_high.shape[3]):
|
||||
if i <= split_h - transition_h and j <= split_w - transition_w:
|
||||
mask_high[:, :, i, j] = 0.0
|
||||
elif i <= split_h or j <= split_w:
|
||||
# Calculate smooth transition value - inverse of the low mask
|
||||
fade_h = min(1.0, (split_h - i) / transition_h) if i <= split_h else 0.0
|
||||
fade_w = min(1.0, (split_w - j) / transition_w) if j <= split_w else 0.0
|
||||
mask_high[:, :, i, j] = 1.0 - (fade_h * fade_w)
|
||||
|
||||
# Ensure masks sum to 1.0 at each point for full spectrum coverage
|
||||
mask_sum = mask_low + mask_high
|
||||
mask_low = mask_low / mask_sum
|
||||
mask_high = mask_high / mask_sum
|
||||
|
||||
# Apply masks and compute CFG in frequency domain using the corresponding CFG values
|
||||
fft_result_low = (fft_u_pred_1 * mask_low) + ((fft_c_pred_1 - fft_u_pred_1) * self.cfgs[0] * mask_low)
|
||||
fft_result_high = (fft_u_pred_2 * mask_high) + ((fft_c_pred_2 - fft_u_pred_2) * self.cfgs[1] * mask_high)
|
||||
|
||||
# Combine low and high frequency results
|
||||
fft_result = fft_result_low + fft_result_high
|
||||
|
||||
# Convert back to spatial domain
|
||||
mix = torch.fft.irfft2(fft_result, s=(h, w))
|
||||
|
||||
elif(self.mode == "fft_overlap"):
|
||||
# Split each model into its own frequency band with overlap - optimized version
|
||||
_, _, h, w = pred[0][0].shape
|
||||
|
||||
# Initialize result in frequency domain
|
||||
fft_result = torch.zeros_like(torch.fft.rfft2(pred[0][0]))
|
||||
|
||||
# Convert all predictions to frequency domain
|
||||
fft_c_preds = []
|
||||
fft_u_preds = []
|
||||
valid_weights = []
|
||||
valid_cfgs = []
|
||||
|
||||
for i in range(len(pred)):
|
||||
if self.weights[i] <= 0:
|
||||
continue
|
||||
|
||||
# Extract predictions for this model
|
||||
c_pred = pred[i][0]
|
||||
u_pred = pred[i][1]
|
||||
|
||||
# Convert to frequency domain
|
||||
fft_c_preds.append(torch.fft.rfft2(c_pred))
|
||||
fft_u_preds.append(torch.fft.rfft2(u_pred))
|
||||
valid_weights.append(self.weights[i])
|
||||
valid_cfgs.append(self.cfgs[i])
|
||||
|
||||
# Get number of valid models
|
||||
num_models = len(valid_weights)
|
||||
if num_models == 0:
|
||||
# Handle edge case with no valid models
|
||||
return pred[0][0] # Return first model's conditioned output
|
||||
elif num_models == 1:
|
||||
# With only one model, use it for the full spectrum
|
||||
mix = pred[0][0] * valid_cfgs[0] + pred[0][1] * (1 - valid_cfgs[0])
|
||||
return mix
|
||||
|
||||
# Calculate frequency coordinates
|
||||
freq_h, freq_w = fft_result.shape[2], fft_result.shape[3]
|
||||
|
||||
# Create normalized coordinate grids - vectorized
|
||||
h_coords = torch.linspace(0, 1, freq_h, device=fft_result.device)
|
||||
w_coords = torch.linspace(0, 1, freq_w, device=fft_result.device)
|
||||
|
||||
# Create meshgrid of coordinates
|
||||
h_grid, w_grid = torch.meshgrid(h_coords, w_coords, indexing='ij')
|
||||
|
||||
# Calculate frequency magnitude for all points
|
||||
freq_magnitude = torch.sqrt(h_grid**2 + w_grid**2)
|
||||
|
||||
# Normalize frequency magnitude to range [0, 1]
|
||||
max_freq = torch.sqrt(torch.tensor(2.0, device=fft_result.device)) # Maximum possible frequency magnitude
|
||||
normalized_freq = freq_magnitude / max_freq
|
||||
|
||||
# Define overlap ratio (0.2 means 20% overlap between bands)
|
||||
overlap_ratio = 0.2
|
||||
|
||||
# Calculate band width for each model, accounting for overlap
|
||||
band_width = (1.0 + overlap_ratio) / num_models
|
||||
|
||||
# Calculate frequency masks for each model - vectorized
|
||||
model_masks = []
|
||||
for i in range(num_models):
|
||||
# Calculate band start and end points
|
||||
band_start = max(0.0, i * band_width - (overlap_ratio * band_width / 2))
|
||||
band_end = min(1.0, (i + 1) * band_width + (overlap_ratio * band_width / 2))
|
||||
|
||||
# Create initial mask where freq is within the band
|
||||
mask = torch.zeros_like(normalized_freq)
|
||||
in_band = (normalized_freq >= band_start) & (normalized_freq <= band_end)
|
||||
|
||||
# Define transition regions
|
||||
start_trans = band_start + (overlap_ratio * band_width / 2)
|
||||
end_trans = band_end - (overlap_ratio * band_width / 2)
|
||||
|
||||
# Create mask for different regions
|
||||
full_region = (normalized_freq >= start_trans) & (normalized_freq <= end_trans)
|
||||
start_region = (normalized_freq >= band_start) & (normalized_freq < start_trans)
|
||||
end_region = (normalized_freq > end_trans) & (normalized_freq <= band_end)
|
||||
|
||||
# Set full strength in middle region
|
||||
mask[full_region] = 1.0
|
||||
|
||||
# Calculate and set smooth transitions at edges
|
||||
if torch.any(start_region):
|
||||
trans_width = start_trans - band_start
|
||||
mask[start_region] = (normalized_freq[start_region] - band_start) / trans_width
|
||||
|
||||
if torch.any(end_region):
|
||||
trans_width = band_end - end_trans
|
||||
mask[end_region] = (band_end - normalized_freq[end_region]) / trans_width
|
||||
|
||||
model_masks.append(mask)
|
||||
|
||||
# Stack masks for easier normalization
|
||||
stacked_masks = torch.stack(model_masks)
|
||||
|
||||
# Normalize masks to ensure they sum to 1.0 across all models
|
||||
sum_masks = torch.sum(stacked_masks, dim=0)
|
||||
sum_masks = torch.clamp(sum_masks, min=1e-10)
|
||||
|
||||
# Normalize each mask
|
||||
normalized_stacked_masks = stacked_masks / sum_masks
|
||||
|
||||
# Apply CFG to each frequency band and combine - vectorized
|
||||
for i in range(num_models):
|
||||
# Get mask, add batch and channel dimensions
|
||||
mask = normalized_stacked_masks[i].unsqueeze(0).unsqueeze(0)
|
||||
|
||||
# Apply CFG formula with frequency band mask
|
||||
fft_u = fft_u_preds[i]
|
||||
fft_c = fft_c_preds[i]
|
||||
cfg_scale = valid_cfgs[i]
|
||||
|
||||
# Add weighted contribution to result
|
||||
fft_result += (fft_u + (fft_c - fft_u) * cfg_scale) * mask
|
||||
|
||||
# Convert back to spatial domain
|
||||
mix = torch.fft.irfft2(fft_result, s=(h, w))
|
||||
|
||||
elif(self.mode == "fft_full"):
|
||||
# Full spectrum frequency mixing implementation - optimized version
|
||||
_, _, h, w = pred[0][0].shape
|
||||
|
||||
# Initialize result in frequency domain
|
||||
fft_result = torch.zeros_like(torch.fft.rfft2(pred[0][0]))
|
||||
|
||||
# First, convert all predictions to frequency domain
|
||||
fft_c_preds = []
|
||||
fft_u_preds = []
|
||||
valid_model_indices = []
|
||||
|
||||
for i in range(len(pred)):
|
||||
if self.weights[i] <= 0:
|
||||
# Add placeholders for zero-weight models
|
||||
fft_c_preds.append(None)
|
||||
fft_u_preds.append(None)
|
||||
continue
|
||||
|
||||
# Extract predictions for this model
|
||||
c_pred = pred[i][0]
|
||||
u_pred = pred[i][1]
|
||||
|
||||
# Convert to frequency domain
|
||||
fft_c_preds.append(torch.fft.rfft2(c_pred))
|
||||
fft_u_preds.append(torch.fft.rfft2(u_pred))
|
||||
valid_model_indices.append(i)
|
||||
|
||||
# Calculate frequency coordinates
|
||||
frequency_h, frequency_w = fft_result.shape[2], fft_result.shape[3]
|
||||
|
||||
# Create normalized coordinate grids (vectorized)
|
||||
h_coords = torch.arange(frequency_h, device=fft_result.device).float() / frequency_h
|
||||
w_coords = torch.arange(frequency_w, device=fft_result.device).float() / frequency_w
|
||||
|
||||
# Create meshgrid of coordinates
|
||||
h_grid, w_grid = torch.meshgrid(h_coords, w_coords, indexing='ij')
|
||||
|
||||
# Calculate frequency magnitude for all points at once
|
||||
freq_magnitude = torch.sqrt(h_grid**2 + w_grid**2)
|
||||
|
||||
# Process each valid model
|
||||
if valid_model_indices:
|
||||
# Pre-calculate model weights for all frequency points
|
||||
all_model_weights = []
|
||||
for i in valid_model_indices:
|
||||
# Calculate model influence: higher weights → lower frequencies, lower weights → higher frequencies
|
||||
model_weight = self.weights[i] * (1.0 - freq_magnitude) + (1.0 - self.weights[i]) * freq_magnitude
|
||||
all_model_weights.append(model_weight)
|
||||
|
||||
# Convert to tensor for easier operations
|
||||
all_model_weights = torch.stack(all_model_weights)
|
||||
|
||||
# Normalize weights along model dimension
|
||||
weight_sum = all_model_weights.sum(dim=0, keepdim=True)
|
||||
# Avoid division by zero
|
||||
weight_sum = torch.clamp(weight_sum, min=1e-10)
|
||||
normalized_weights = all_model_weights / weight_sum
|
||||
|
||||
# Apply weights and CFG to each frequency component
|
||||
for idx, i in enumerate(valid_model_indices):
|
||||
# Get model weights for this model
|
||||
model_weights = normalized_weights[idx]
|
||||
|
||||
# Apply weights to each frequency point (broadcasting)
|
||||
model_weights = model_weights.unsqueeze(0).unsqueeze(0) # Add batch and channel dimensions
|
||||
|
||||
# Get unconditioned and conditioned predictions
|
||||
fft_u = fft_u_preds[i]
|
||||
fft_c = fft_c_preds[i]
|
||||
|
||||
# Apply CFG with model-specific weight
|
||||
contrib = fft_u * model_weights + (fft_c - fft_u) * self.cfgs[i] * model_weights
|
||||
|
||||
# Add to result
|
||||
fft_result += contrib
|
||||
|
||||
# Convert back to spatial domain
|
||||
mix = torch.fft.irfft2(fft_result, s=(h, w))
|
||||
|
||||
|
||||
|
||||
return mix
|
||||
|
||||
def outer_sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
|
||||
# Get the device from the first model
|
||||
device = self.models[0].load_device
|
||||
|
||||
# Prepare the models, conditions, and loaded models for each model
|
||||
prepared_models = []
|
||||
all_loaded_models = []
|
||||
|
||||
for i, model in enumerate(self.models):
|
||||
# Get relevant conditions for this model
|
||||
model_conds = {k: v for k, v in self.conds.items() if k.endswith(f"_{i+1}")}
|
||||
|
||||
# Prepare sampling for this model
|
||||
inner_model, prepared_conds, loaded_models = comfy.sampler_helpers.prepare_sampling(model, noise.shape, model_conds, model.model_options)
|
||||
|
||||
# Save the prepared data
|
||||
self.prepared_models.append((inner_model, prepared_conds))
|
||||
|
||||
all_loaded_models.extend(loaded_models)
|
||||
|
||||
# Update the conditions for this model
|
||||
for k, v in prepared_conds.items():
|
||||
self.conds[k] = v
|
||||
|
||||
self.inner_model = self.prepared_models[0][0]
|
||||
# Prepare the device-related tensors
|
||||
if denoise_mask is not None:
|
||||
denoise_mask = comfy.sampler_helpers.prepare_mask(denoise_mask, noise.shape, device)
|
||||
|
||||
noise = noise.to(device)
|
||||
latent_image = latent_image.to(device)
|
||||
sigmas = sigmas.to(device)
|
||||
|
||||
# Set model options for the main model patcher
|
||||
comfy.samplers.cast_to_load_options(self.model_options, device=device, dtype=self.model_patcher.model_dtype())
|
||||
|
||||
try:
|
||||
# Pre-run all models
|
||||
for model in self.models:
|
||||
model.pre_run()
|
||||
|
||||
# Set the inner_model to be a list of prepared models
|
||||
|
||||
# Run the inner sample with prepared models
|
||||
output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
|
||||
finally:
|
||||
# Clean up all models
|
||||
for model in self.models:
|
||||
model.cleanup()
|
||||
del self.inner_model
|
||||
# Clean up models and loaded models
|
||||
comfy.sampler_helpers.cleanup_models(self.conds, all_loaded_models)
|
||||
|
||||
|
||||
return output
|
||||
class MixModHighResGuiderNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"component": ("COMPONENT",),
|
||||
"ratio": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"start_step": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
|
||||
"end_step": ("INT", {"default": -1, "min": -1, "max": 100, "step": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GUIDER",)
|
||||
FUNCTION = "create_guider"
|
||||
CATEGORY = "mixmod/conditioning"
|
||||
|
||||
def create_guider(self, component, ratio=0.5, start_step=0, end_step=-1, mode="team", image=None):
|
||||
# Convert end_step = -1 to None for the wrapper function
|
||||
actual_end_step = None if end_step == -1 else end_step
|
||||
# This is a placeholder for the actual guider implementation
|
||||
# The actual implementation will be filled in by the user
|
||||
guider = MixModGuider(component, mode=mode, ratio=ratio, highres_mode=True, start_step=start_step, end_step=actual_end_step)
|
||||
return (guider,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"MixModGuiderComponentNode": MixModGuiderComponentNode,
|
||||
"MixModGuiderComponentPipelineNode": MixModGuiderComponentPipelineNode,
|
||||
"MixModGuiderNode": MixModGuiderNode,
|
||||
"MixModStyleInjectionGuiderNode": MixModStyleInjectionGuiderNode,
|
||||
"MixModFFTGuiderNode": MixModFFTGuiderNode,
|
||||
"MixModBandFFTGuiderNode": MixModBandFFTGuiderNode,
|
||||
"MixModHighResGuiderNode": MixModHighResGuiderNode,
|
||||
"MixModOptionsMaskNode": MixModOptionsMaskNode,
|
||||
"MixModOptionsSchedulerNode": MixModOptionsSchedulerNode,
|
||||
"MixModOptionsScaleNode": MixModOptionsScaleNode,
|
||||
"MixModOptionsUncondDecayNode": MixModOptionsUncondDecayNode,
|
||||
"MixModPipelineNode": MixModPipelineNode,
|
||||
"MixModDynamicMaskGuiderNode": MixModDynamicMaskGuiderNode,
|
||||
"MixModDynamicMaskAlternativeGuiderNode": MixModDynamicMaskAlternativeGuiderNode,
|
||||
"MixModDepthGuiderNode": MixModDepthGuiderNode,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MixModGuiderComponentNode": "MixMod Guider Component",
|
||||
"MixModGuiderComponentPipelineNode": "MixMod Guider Component Pipeline",
|
||||
"MixModGuiderNode": "MixMod Guider",
|
||||
"MixModStyleInjectionGuiderNode": "MixMod Style Injection Guider",
|
||||
"MixModFFTGuiderNode": "MixMod FFT Guider",
|
||||
}
|
||||
"MixModBandFFTGuiderNode": "MixMod Band FFT Guider",
|
||||
"MixModHighResGuiderNode": "MixMod High Res Guider",
|
||||
"MixModOptionsMaskNode": "MixMod Options Mask",
|
||||
"MixModOptionsSchedulerNode": "MixMod Options Scheduler",
|
||||
"MixModOptionsScaleNode": "MixMod Options Scale",
|
||||
"MixModOptionsUncondDecayNode": "MixMod Options Uncond Decay",
|
||||
"MixModPipelineNode": "MixMod Pipeline",
|
||||
"MixModDynamicMaskGuiderNode": "MixMod Dynamic Mask Guider",
|
||||
"MixModDynamicMaskAlternativeGuiderNode": "MixMod Dynamic Mask Alternative Guider",
|
||||
"MixModDepthGuiderNode": "MixMod Depth Guider",
|
||||
}
|
||||
@@ -0,0 +1,392 @@
|
||||
from __future__ import annotations
|
||||
import folder_paths
|
||||
import comfy.utils
|
||||
import torch
|
||||
import numpy as np
|
||||
import os
|
||||
from PIL import Image
|
||||
import torchvision.transforms.functional as F
|
||||
import math
|
||||
from comfy import model_management
|
||||
from comfy import latent_formats
|
||||
|
||||
|
||||
def ensure_model(model_patcher, mode = None):
|
||||
if not mode or mode == "disable":
|
||||
# Ensure model mode disabled.
|
||||
return
|
||||
mp = model_patcher
|
||||
lm = model_management.LoadedModel(mp)
|
||||
found_lm = None
|
||||
for list_lm in model_management.current_loaded_models:
|
||||
if lm == list_lm:
|
||||
found_lm = list_lm
|
||||
break
|
||||
if getattr(found_lm, "currently_used", False):
|
||||
# Model already exists and appears to be loaded.
|
||||
return
|
||||
if found_lm is not None:
|
||||
lm = found_lm
|
||||
mp = lm.model
|
||||
if mode.startswith("normal"):
|
||||
if mode == "normal_unload":
|
||||
model_management.unload_all_models()
|
||||
model_management.load_models_gpu((mp,))
|
||||
return
|
||||
model_management.unload_all_models()
|
||||
if mode == "lowvram":
|
||||
# Logic from comfy.model_management.load_models_gpu
|
||||
min_inference_memory = model_management.minimum_inference_memory()
|
||||
minimum_memory_required = max(
|
||||
min_inference_memory,
|
||||
model_management.extra_reserved_memory(),
|
||||
)
|
||||
loaded_memory = lm.model_loaded_memory()
|
||||
current_free_mem = model_management.get_free_memory(lm.device) + loaded_memory
|
||||
|
||||
lowvram_model_memory = max(
|
||||
64 * 1024 * 1024,
|
||||
(current_free_mem - minimum_memory_required),
|
||||
min(
|
||||
current_free_mem * model_management.MIN_WEIGHT_MEMORY_RATIO,
|
||||
current_free_mem - model_management.minimum_inference_memory(),
|
||||
),
|
||||
)
|
||||
lowvram_model_memory = max(0.1, lowvram_model_memory - loaded_memory)
|
||||
elif mode == "novram":
|
||||
lowvram_model_memory = 0.1
|
||||
else:
|
||||
raise ValueError("Bad ensure_model_mode")
|
||||
lm.model_load(lowvram_model_memory)
|
||||
model_management.current_loaded_models.insert(0, lm)
|
||||
|
||||
|
||||
def fallback(val, default, *, exclude=None, default_is_fun=False):
|
||||
return val if val is not exclude else (default() if default_is_fun else default)
|
||||
|
||||
|
||||
def scale_dim(n, factor=1.0, *, increment=64) -> int:
|
||||
return math.ceil((n * factor) / increment) * increment
|
||||
|
||||
|
||||
def sigma_to_float(sigma):
|
||||
return sigma.detach().cpu().max().item()
|
||||
|
||||
def load_taesd(name):
|
||||
sd = {}
|
||||
approx_vaes = folder_paths.get_filename_list("vae_approx")
|
||||
|
||||
encoder = next(filter(lambda a: a.startswith("{}_encoder.".format(name)), approx_vaes))
|
||||
decoder = next(filter(lambda a: a.startswith("{}_decoder.".format(name)), approx_vaes))
|
||||
|
||||
enc = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", encoder))
|
||||
for k in enc:
|
||||
sd["taesd_encoder.{}".format(k)] = enc[k]
|
||||
|
||||
dec = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", decoder),device=model_management.get_torch_device())
|
||||
for k in dec:
|
||||
sd["taesd_decoder.{}".format(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
|
||||
|
||||
def load_vae(name):
|
||||
vae = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae", name))
|
||||
return vae
|
||||
|
||||
def create_depth_map(latent_image, depth_model, resolution, **kwargs):
|
||||
batch_size = latent_image.shape[0]
|
||||
out_tensor = None
|
||||
for i, image in enumerate(latent_image):
|
||||
np_image = np.asarray(image.cpu() * 255., dtype=np.uint8)
|
||||
np_result = depth_model(np_image, output_type="np", detect_resolution=resolution, **kwargs)
|
||||
out = torch.from_numpy(np_result.astype(np.float32) / 255.0)
|
||||
if out_tensor is None:
|
||||
out_tensor = torch.zeros(batch_size, *out.shape, dtype=torch.float32)
|
||||
out_tensor[i] = out
|
||||
|
||||
return out_tensor
|
||||
|
||||
def save_debug_image(tensor, filename, scale=True):
|
||||
try:
|
||||
# Ensure tensor is detached and on CPU
|
||||
tensor = tensor.detach().cpu().float()
|
||||
print(tensor.shape)
|
||||
# Remove batch dimension if present
|
||||
if tensor.ndim == 4 and tensor.shape[0] == 1:
|
||||
tensor = tensor.squeeze(0)
|
||||
|
||||
# Handle 2D tensors (grayscale masks) by converting to RGB
|
||||
if tensor.ndim == 2:
|
||||
# Convert 2D tensor to 3D by repeating across 3 channels
|
||||
tensor = tensor.unsqueeze(-1).repeat(1, 1, 3)
|
||||
|
||||
# If scale is True, scale to [0, 255]
|
||||
if scale:
|
||||
array = (tensor.numpy() * 255).clip(0, 255).astype('uint8')
|
||||
else:
|
||||
array = tensor.numpy().clip(0, 255).astype('uint8')
|
||||
# Ensure shape is (H, W, 3)
|
||||
if array.shape[-1] != 3:
|
||||
raise ValueError(f"Expected last dimension to be 3 (RGB), got {array.shape[-1]}")
|
||||
img = Image.fromarray(array)
|
||||
|
||||
# Ensure debug directory exists
|
||||
os.makedirs("debug", exist_ok=True)
|
||||
|
||||
img.save("debug/" + filename)
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error saving debug image {filename}: {e}")
|
||||
return False
|
||||
|
||||
def process_mask(mask, H, W, device):
|
||||
# Convert to tensor if not already
|
||||
if not isinstance(mask, torch.Tensor):
|
||||
mask = torch.tensor(mask, device=device, dtype=torch.float32)
|
||||
mask = mask.to(device).float()
|
||||
|
||||
# Squeeze singleton dimensions
|
||||
mask = mask.squeeze()
|
||||
|
||||
# If mask is RGB, convert to grayscale
|
||||
if mask.ndim == 3 and mask.shape[-1] == 3:
|
||||
# [H, W, 3] -> [3, H, W]
|
||||
mask = mask.permute(2, 0, 1)
|
||||
mask = mask.mean(dim=0) # [H, W]
|
||||
elif mask.ndim == 4 and mask.shape[-1] == 3:
|
||||
# [B, H, W, 3] -> [B, 3, H, W]
|
||||
mask = mask.permute(0, 3, 1, 2)
|
||||
mask = mask.mean(dim=1) # [B, H, W]
|
||||
mask = mask.squeeze(0) # Remove batch if present
|
||||
|
||||
# Now mask should be [H, W]
|
||||
if mask.shape != (H, W):
|
||||
mask = torch.nn.functional.interpolate(mask.unsqueeze(0).unsqueeze(0), size=(H, W), mode='bilinear')[0, 0]
|
||||
|
||||
return mask
|
||||
|
||||
|
||||
def create_adaptive_low_res_model_wrapper( min_scale_factor=0.5, max_scale_factor=1.0, start_step=0, end_step=None):
|
||||
"""
|
||||
Creates a model function wrapper that runs the model at progressively higher resolution as sampling proceeds.
|
||||
|
||||
Args:
|
||||
min_scale_factor: Minimum scale factor to use (for early steps)
|
||||
max_scale_factor: Maximum scale factor to use (for final steps)
|
||||
start_step: Step number to start increasing resolution (0 means immediately)
|
||||
end_step: Step number when scaling should reach max_scale_factor (None means the last step)
|
||||
|
||||
Returns:
|
||||
A function wrapper that can be used as model_options['model_function_wrapper']
|
||||
"""
|
||||
def wrapper(apply_model_func, args_dict):
|
||||
input_x = args_dict["input"]
|
||||
timestep = args_dict["timestep"]
|
||||
c = args_dict["c"]
|
||||
|
||||
# Get sampling progress from sigmas in transformer_options
|
||||
sample_sigmas = c['transformer_options'].get('sample_sigmas', None)
|
||||
current_sigma = c['transformer_options'].get('sigmas', None)
|
||||
|
||||
# Determine scale factor based on current step
|
||||
if sample_sigmas is not None and current_sigma is not None:
|
||||
# Get total steps from the length of sample_sigmas
|
||||
total_steps = len(sample_sigmas)
|
||||
|
||||
# Find current step by matching current_sigma to sample_sigmas
|
||||
current_sigma_val = float(current_sigma[0])
|
||||
|
||||
# Find the index of the closest sigma value in sample_sigmas
|
||||
# This gives us the current step number
|
||||
sigma_diffs = [abs(float(s) - current_sigma_val) for s in sample_sigmas]
|
||||
current_step = sigma_diffs.index(min(sigma_diffs))
|
||||
|
||||
# If end_step is not provided, use the last step
|
||||
actual_end_step = end_step if end_step is not None else (total_steps - 1)
|
||||
|
||||
# Clamp start and end steps to valid range
|
||||
start_step_clamped = max(0, min(start_step, total_steps - 1))
|
||||
end_step_clamped = max(start_step_clamped, min(actual_end_step, total_steps - 1))
|
||||
|
||||
# Calculate scale factor based on current step
|
||||
if current_step < start_step_clamped:
|
||||
# Before start step - use minimum scale
|
||||
scale_factor = min_scale_factor
|
||||
elif current_step > end_step_clamped:
|
||||
# After end step - use maximum scale
|
||||
scale_factor = max_scale_factor
|
||||
else:
|
||||
# In between - interpolate between min and max
|
||||
step_range = end_step_clamped - start_step_clamped
|
||||
if step_range > 0:
|
||||
progress = (current_step - start_step_clamped) / step_range
|
||||
scale_factor = min_scale_factor + progress * (max_scale_factor - min_scale_factor)
|
||||
else:
|
||||
# If start_step equals end_step
|
||||
scale_factor = min_scale_factor
|
||||
else:
|
||||
# Fallback if we can't determine steps
|
||||
scale_factor = max_scale_factor
|
||||
|
||||
# Original shape
|
||||
original_shape = input_x.shape
|
||||
|
||||
# Skip resizing if we're at full resolution
|
||||
if scale_factor >= 0.99:
|
||||
return apply_model_func(input_x, timestep, **c)
|
||||
print(f"Scaling down to {scale_factor}x for step {current_step} using bislerp")
|
||||
# Calculate new dimensions
|
||||
new_h = int(original_shape[2] * scale_factor)
|
||||
new_w = int(original_shape[3] * scale_factor)
|
||||
|
||||
# Ensure dimensions are at least 8 (minimum size for stable diffusion)
|
||||
new_h = max(8, new_h)
|
||||
new_w = max(8, new_w)
|
||||
|
||||
# Downscale using comfy.utils.common_upscale with bislerp mode
|
||||
downscaled_x = comfy.utils.common_upscale(
|
||||
input_x,
|
||||
new_w,
|
||||
new_h,
|
||||
'bislerp',
|
||||
'disabled'
|
||||
)
|
||||
|
||||
# Run model at lower resolution
|
||||
output = apply_model_func(downscaled_x, timestep, **c)
|
||||
|
||||
# Upscale result back to original resolution using bislerp
|
||||
upscaled_output = comfy.utils.common_upscale(
|
||||
output,
|
||||
original_shape[3], # width
|
||||
original_shape[2], # height
|
||||
'bislerp',
|
||||
'disabled'
|
||||
)
|
||||
|
||||
""" # #blurring the output
|
||||
# Apply a small amount of Gaussian blur to smooth the upscaled output
|
||||
# This helps reduce artifacts that can appear from the upscaling process
|
||||
kernel_size = 3 # Small kernel for subtle blurring
|
||||
sigma = 0.5 # Low sigma for gentle blur
|
||||
|
||||
# Create Gaussian kernel
|
||||
kernel = torch.tensor([
|
||||
[0.0625, 0.125, 0.0625],
|
||||
[0.125, 0.25, 0.125],
|
||||
[0.0625, 0.125, 0.0625]
|
||||
], device=upscaled_output.device).view(1, 1, kernel_size, kernel_size)
|
||||
|
||||
Expand kernel for all channels
|
||||
channels = upscaled_output.shape[1]
|
||||
kernel = kernel.expand(channels, 1, kernel_size, kernel_size)
|
||||
|
||||
# Group kernel for batch processing
|
||||
kernel = kernel.repeat(1, 1, 1, 1)
|
||||
|
||||
# Apply convolution for blurring
|
||||
# Use padding='same' to maintain dimensions
|
||||
padding = kernel_size // 2
|
||||
blurred_output = torch.nn.functional.conv2d(
|
||||
upscaled_output,
|
||||
kernel.to(upscaled_output.dtype),
|
||||
padding=padding,
|
||||
groups=channels
|
||||
)
|
||||
|
||||
# Use the blurred output instead of the direct upscaled output
|
||||
upscaled_output = blurred_output """
|
||||
return upscaled_output
|
||||
|
||||
return wrapper
|
||||
|
||||
def create_noise_injection_wrapper(latent, noise_weight=1.0, callback=None):
|
||||
"""
|
||||
Creates a wrapper that takes an input image, converts it to latent space,
|
||||
noises it to match the current timestep/sigma, and injects its noise prediction.
|
||||
|
||||
Args:
|
||||
image: Input image to inject (tensor in pixel space)
|
||||
vae: VAE model for encoding to latent space
|
||||
noise_weight: Weight of the injected noise prediction (0-1)
|
||||
|
||||
Returns:
|
||||
A function wrapper to use as model_options['model_function_wrapper']
|
||||
|
||||
|
||||
"""
|
||||
def wrapper(apply_model_func, args_dict):
|
||||
# Extract args
|
||||
input_x = args_dict["input"]
|
||||
timestep = args_dict["timestep"]
|
||||
c = args_dict["c"]
|
||||
|
||||
|
||||
# Get current sampling parameters
|
||||
sigmas = c['transformer_options'].get('sigmas', None)
|
||||
sigma_idx = c['transformer_options'].get('sigma_idx', 0)
|
||||
current_sigma = sigmas[sigma_idx] if sigmas is not None else None
|
||||
|
||||
# Noise the image latent to match current noise level
|
||||
if current_sigma is not None:
|
||||
# Move tensors to the same device as input_x
|
||||
device = input_x.device
|
||||
image_latent = latent.to(device)
|
||||
current_sigma = current_sigma.to(device)
|
||||
|
||||
image_latent = latent_formats.SDXL().process_in(image_latent)
|
||||
|
||||
# Resize image_latent to match input_x dimensions
|
||||
if image_latent.shape[2:] != input_x.shape[2:]:
|
||||
image_latent = torch.nn.functional.interpolate(
|
||||
image_latent,
|
||||
size=(input_x.shape[2], input_x.shape[3]),
|
||||
mode='bilinear',
|
||||
align_corners=False
|
||||
)
|
||||
|
||||
noise = torch.randn_like(image_latent)
|
||||
noised_image_latent = image_latent + noise * current_sigma
|
||||
|
||||
noised_image_latent = (1 - noise_weight) * input_x + noise_weight * noised_image_latent
|
||||
|
||||
# Get noise prediction on the image latent
|
||||
image_args = args_dict.copy()
|
||||
image_args["input"] = noised_image_latent
|
||||
combined_pred = apply_model_func(noised_image_latent, timestep, **c)
|
||||
|
||||
# Run original model to get base noise prediction
|
||||
#original_noise_pred = apply_model_func(input_x, timestep, **c)
|
||||
|
||||
|
||||
|
||||
# Combine predictions
|
||||
#combined_pred = (1 - noise_weight) * original_noise_pred + noise_weight * image_noise_pred
|
||||
#if callback is not None:
|
||||
# callback(combined_pred)
|
||||
return combined_pred
|
||||
|
||||
# Fall back to original model if we can't determine sigma
|
||||
return apply_model_func(input_x, timestep, **c)
|
||||
|
||||
return wrapper
|
||||
|
||||
def sumweights(weights):
|
||||
weight_sum = sum(weights)
|
||||
|
||||
# Normalize weights if sum is not zero
|
||||
if weight_sum > 0:
|
||||
weights = [w / weight_sum for w in weights]
|
||||
return weights
|
||||
Reference in New Issue
Block a user