294 lines
15 KiB
Python
294 lines
15 KiB
Python
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)
|