From f7086236d2a1b9cd40aba4be703a84e2f6551e1e Mon Sep 17 00:00:00 2001 From: kantsche Date: Tue, 29 Apr 2025 10:05:35 +0200 Subject: [PATCH] Version 2 Lots of new stuff --- README.md | 83 ++--- guider.py | 293 +++++++++++++++ mixes/__init__.py | 1 + mixes/fft.py | 649 +++++++++++++++++++++++++++++++++ mixes/standard.py | 494 +++++++++++++++++++++++++ nodes.py | 910 +++++++++++++++++----------------------------- utils.py | 392 ++++++++++++++++++++ 7 files changed, 2212 insertions(+), 610 deletions(-) create mode 100644 guider.py create mode 100644 mixes/__init__.py create mode 100644 mixes/fft.py create mode 100644 mixes/standard.py create mode 100644 utils.py diff --git a/README.md b/README.md index 508237a..f0e634f 100644 --- a/README.md +++ b/README.md @@ -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 -![image](https://github.com/user-attachments/assets/0108c1e4-bf3c-4060-9860-47ae8a52b627) - - ## Credits Created by Kantsche + +## License + +[Add license information here] \ No newline at end of file diff --git a/guider.py b/guider.py new file mode 100644 index 0000000..3395198 --- /dev/null +++ b/guider.py @@ -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) diff --git a/mixes/__init__.py b/mixes/__init__.py new file mode 100644 index 0000000..9b3b6a7 --- /dev/null +++ b/mixes/__init__.py @@ -0,0 +1 @@ +# FFT mixing functions for ComfyUI MixMod \ No newline at end of file diff --git a/mixes/fft.py b/mixes/fft.py new file mode 100644 index 0000000..1f2fdcf --- /dev/null +++ b/mixes/fft.py @@ -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 \ No newline at end of file diff --git a/mixes/standard.py b/mixes/standard.py new file mode 100644 index 0000000..e3d6e89 --- /dev/null +++ b/mixes/standard.py @@ -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 + diff --git a/nodes.py b/nodes.py index 2dea335..fccc3e6 100644 --- a/nodes.py +++ b/nodes.py @@ -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", +} \ No newline at end of file diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..5f3dc7c --- /dev/null +++ b/utils.py @@ -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