added sdxl+sd1.5 support
This commit is contained in:
kantsche
@@ -1,293 +1,347 @@
|
||||
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)
|
||||
import torch
|
||||
import comfy.samplers
|
||||
import comfy.sampler_helpers
|
||||
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_model_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 = 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 = 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.types = []
|
||||
|
||||
self.step = 0
|
||||
self.total_steps = 0
|
||||
self.start_steps = []
|
||||
self.end_steps = []
|
||||
|
||||
self.options = []
|
||||
self.scale_factors = []
|
||||
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.mask_history = None # For tracking dynamic masks between steps
|
||||
|
||||
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 = []
|
||||
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"]
|
||||
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(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
|
||||
|
||||
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()
|
||||
# Clean up models and loaded models
|
||||
comfy.sampler_helpers.cleanup_models(self.conds, all_loaded_models)
|
||||
self.prepared_models = []
|
||||
del 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)
|
||||
|
||||
if latent_image is not None and torch.count_nonzero(latent_image) > 0: #Don't shift the empty latent image.
|
||||
latent_image = self.inner_model.process_latent_in(latent_image)
|
||||
|
||||
# Process each cond with each prepared inner model separately
|
||||
processed_conds = {}
|
||||
for i, (inner_model, _) in enumerate(self.prepared_models):
|
||||
# Get the conditions for this model (positive_i, negative_i)
|
||||
model_conds = {k: v for k, v in self.conds.items() if k.endswith(f"_{i+1}")}
|
||||
|
||||
# Process conditions with this model
|
||||
processed_model_conds = comfy.samplers.process_conds(inner_model, noise, model_conds, device, latent_image, denoise_mask, seed)
|
||||
|
||||
# Add to processed conditions
|
||||
processed_conds.update(processed_model_conds)
|
||||
|
||||
# Merge back into self.conds
|
||||
self.conds = processed_conds
|
||||
|
||||
extra_model_options = comfy.model_patcher.create_model_options_clone(self.model_options)
|
||||
extra_model_options.setdefault("transformer_options", {})["sample_sigmas"] = sigmas
|
||||
extra_args = {"model_options": extra_model_options, "seed": seed}
|
||||
|
||||
executor = comfy.patcher_extension.WrapperExecutor.new_class_executor(
|
||||
sampler.sample,
|
||||
sampler,
|
||||
comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.SAMPLER_SAMPLE, extra_args["model_options"], is_model_options=True)
|
||||
)
|
||||
samples = executor.execute(self, sigmas, extra_args, wrapped_callback, noise, latent_image, denoise_mask, disable_pbar)
|
||||
return samples
|
||||
|
||||
def sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
|
||||
if sigmas.shape[-1] == 0:
|
||||
return latent_image
|
||||
|
||||
self.conds = {}
|
||||
for k in self.original_conds:
|
||||
self.conds[k] = list(map(lambda a: a.copy(), self.original_conds[k]))
|
||||
comfy.samplers.preprocess_conds_hooks(self.conds)
|
||||
|
||||
try:
|
||||
orig_model_options = self.model_options
|
||||
self.model_options = comfy.model_patcher.create_model_options_clone(self.model_options)
|
||||
# if one hook type (or just None), then don't bother caching weights for hooks (will never change after first step)
|
||||
orig_hook_mode = self.model_patcher.hook_mode
|
||||
if comfy.samplers.get_total_hook_groups_in_conds(self.conds) <= 1:
|
||||
self.model_patcher.hook_mode = comfy.hooks.EnumHookMode.MinVram
|
||||
|
||||
comfy.sampler_helpers.prepare_model_patcher(self.model_patcher, self.conds, self.model_options)
|
||||
|
||||
comfy.samplers.filter_registered_hooks_on_conds(self.conds, self.model_options)
|
||||
executor = comfy.patcher_extension.WrapperExecutor.new_class_executor(
|
||||
self.outer_sample,
|
||||
self,
|
||||
comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.OUTER_SAMPLE, self.model_options, is_model_options=True)
|
||||
)
|
||||
output = executor.execute(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
|
||||
finally:
|
||||
comfy.samplers.cast_to_load_options(self.model_options, device=self.model_patcher.offload_device)
|
||||
self.model_options = orig_model_options
|
||||
self.model_patcher.hook_mode = orig_hook_mode
|
||||
self.model_patcher.restore_hook_patches()
|
||||
|
||||
del self.conds
|
||||
return output
|
||||
@@ -65,7 +65,7 @@ class MixModGuiderComponentPipelineNode:
|
||||
"weight": base_weight,
|
||||
"cfg": cfg,
|
||||
"options": options,
|
||||
"prev": prev_component
|
||||
"prev": prev_component,
|
||||
}
|
||||
return (component,)
|
||||
|
||||
@@ -76,7 +76,7 @@ class MixModPipelineNode:
|
||||
return {"required":
|
||||
{"model": ("MODEL",),
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -387,6 +387,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"MixModGuiderNode": MixModGuiderNode,
|
||||
"MixModFFTGuiderNode": MixModFFTGuiderNode,
|
||||
"MixModBandFFTGuiderNode": MixModBandFFTGuiderNode,
|
||||
"MixModHighResGuiderNode": MixModHighResGuiderNode,
|
||||
"MixModOptionsMaskNode": MixModOptionsMaskNode,
|
||||
"MixModOptionsSchedulerNode": MixModOptionsSchedulerNode,
|
||||
"MixModOptionsScaleNode": MixModOptionsScaleNode,
|
||||
@@ -402,6 +403,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MixModGuiderNode": "MixMod 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",
|
||||
@@ -409,4 +411,4 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MixModDynamicMaskGuiderNode": "MixMod Dynamic Mask Guider",
|
||||
"MixModDynamicMaskAlternativeGuiderNode": "MixMod Dynamic Mask Alternative Guider",
|
||||
"MixModDepthGuiderNode": "MixMod Depth Guider",
|
||||
}
|
||||
}
|
||||
@@ -1,392 +1,528 @@
|
||||
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
|
||||
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)
|
||||
|
||||
vae = comfy.sd.VAE(sd=sd)
|
||||
return vae
|
||||
|
||||
def load_vae(name):
|
||||
vae = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae", name),device=model_management.get_torch_device())
|
||||
vae = comfy.sd.VAE(sd=vae)
|
||||
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 create_model_wrapper(type, vae15, vaexl):
|
||||
"""
|
||||
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"]
|
||||
|
||||
# Noise the image latent to match current noise level
|
||||
if type == "sd15":
|
||||
# Get device from input_x
|
||||
device = input_x.device
|
||||
|
||||
image_latent = convert_latent_format(input_x, "sdxl", "sd15", device)
|
||||
|
||||
#
|
||||
model_output = apply_model_func(image_latent, timestep, **c)
|
||||
model_output = model_output.to(device)
|
||||
|
||||
final_output = convert_latent_format(model_output, "sd15", "sdxl", device)
|
||||
|
||||
return final_output
|
||||
|
||||
# Fall back to original model if type is not SD15
|
||||
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
|
||||
|
||||
def convert_latent_format(latent, source_format="sd15", target_format="sdxl", device=None):
|
||||
"""
|
||||
Converts latents between SD1.5 and SDXL formats using their RGB factors and scales.
|
||||
|
||||
Args:
|
||||
latent: Input latent tensor (B, 4, H, W)
|
||||
source_format: Source format ("sd15" or "sdxl")
|
||||
target_format: Target format ("sd15" or "sdxl")
|
||||
device: Device to use (if None, uses the latent's device)
|
||||
|
||||
Returns:
|
||||
Converted latent tensor (B, 4, H, W)
|
||||
"""
|
||||
if source_format == target_format:
|
||||
return latent
|
||||
|
||||
if device is None:
|
||||
device = latent.device
|
||||
|
||||
# Define format parameters
|
||||
sd15_scale = 0.18215
|
||||
sdxl_scale = 0.13025
|
||||
|
||||
sd15_factors = torch.tensor([
|
||||
[ 0.3512, 0.2297, 0.3227],
|
||||
[ 0.3250, 0.4974, 0.2350],
|
||||
[-0.2829, 0.1762, 0.2721],
|
||||
[-0.2120, -0.2616, -0.7177]
|
||||
], device=device)
|
||||
|
||||
sdxl_factors = torch.tensor([
|
||||
[ 0.3651, 0.4232, 0.4341],
|
||||
[-0.2533, -0.0042, 0.1068],
|
||||
[ 0.1076, 0.1111, -0.0362],
|
||||
[-0.3165, -0.2492, -0.2188]
|
||||
], device=device)
|
||||
|
||||
sdxl_bias = torch.tensor([0.1084, -0.0175, -0.0011], device=device)
|
||||
|
||||
# Unscale the latent based on source format
|
||||
if source_format == "sd15":
|
||||
unscaled_latent = latent * sd15_scale
|
||||
else: # sdxl
|
||||
unscaled_latent = latent * sdxl_scale
|
||||
|
||||
# Get dimensions
|
||||
B, C, H, W = unscaled_latent.shape
|
||||
|
||||
# For simplicity, we'll process each sample in the batch separately
|
||||
target_latent = torch.zeros_like(unscaled_latent)
|
||||
|
||||
for b in range(B):
|
||||
# Get the current sample
|
||||
sample = unscaled_latent[b] # [4, H, W]
|
||||
|
||||
# Reshape to [4, H*W]
|
||||
sample_flat = sample.reshape(4, -1)
|
||||
|
||||
# Convert to RGB representation
|
||||
if source_format == "sd15":
|
||||
# Calculate pseudo-inverse (4x3 matrix becomes 3x4)
|
||||
sd15_pinv = torch.pinverse(sd15_factors)
|
||||
# Transform from latent to RGB: [3x4] x [4, H*W] -> [3, H*W]
|
||||
rgb = torch.matmul(sd15_pinv, sample_flat)
|
||||
else: # sdxl
|
||||
sdxl_pinv = torch.pinverse(sdxl_factors)
|
||||
# Transform from latent to RGB: [3x4] x [4, H*W] -> [3, H*W]
|
||||
rgb = torch.matmul(sdxl_pinv, sample_flat)
|
||||
# Remove bias (for each channel)
|
||||
for i in range(3):
|
||||
rgb[i] = rgb[i] - sdxl_bias[i]
|
||||
|
||||
# Convert RGB to target latent space
|
||||
if target_format == "sd15":
|
||||
# RGB to SD1.5 latent: [4x3] x [3, H*W] -> [4, H*W]
|
||||
out = torch.matmul(sd15_factors, rgb)
|
||||
# Apply scale
|
||||
out = out / sd15_scale
|
||||
else: # sdxl
|
||||
# Add bias
|
||||
for i in range(3):
|
||||
rgb[i] = rgb[i] + sdxl_bias[i]
|
||||
# RGB to SDXL latent: [4x3] x [3, H*W] -> [4, H*W]
|
||||
out = torch.matmul(sdxl_factors, rgb)
|
||||
# Apply scale
|
||||
out = out / sdxl_scale
|
||||
|
||||
# Reshape back to [4, H, W] and store
|
||||
target_latent[b] = out.reshape(4, H, W)
|
||||
|
||||
return target_latent
|
||||
|
||||
Reference in New Issue
Block a user