import torch import math import comfy.samplers import comfy.utils import nodes import logging class MixModGuiderComponentNode: @classmethod def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), "positive": ("CONDITIONING",), "negative": ("CONDITIONING",), "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",), } } RETURN_TYPES = ("COMPONENT",) FUNCTION = "create_component" CATEGORY = "mixmod/conditioning" def create_component(self, model, positive, negative, base_weight, cfg, mask=None, prev_component=None): component = { "model": model, "positive": positive, "negative": negative, "weight": base_weight, "cfg": cfg, "mask": mask, "prev": prev_component } return (component,) class MixModGuiderNode: @classmethod def INPUT_TYPES(s): return {"required": {"component": ("COMPONENT",), } } RETURN_TYPES = ("GUIDER",) FUNCTION = "create_guider" CATEGORY = "mixmod/conditioning" def create_guider(self, component, image=None,): # 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,) class MixModFFTGuiderNode: @classmethod def INPUT_TYPES(s): return {"required": {"component": ("COMPONENT",), "mode": (["teamfft", "2model_fft"],), "ratio": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), } } RETURN_TYPES = ("GUIDER",) FUNCTION = "create_guider" CATEGORY = "mixmod/conditioning" def create_guider(self, component, mode="team", image=None, 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) 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 NODE_CLASS_MAPPINGS = { "MixModGuiderComponentNode": MixModGuiderComponentNode, "MixModGuiderNode": MixModGuiderNode, "MixModFFTGuiderNode": MixModFFTGuiderNode, } NODE_DISPLAY_NAME_MAPPINGS = { "MixModGuiderComponentNode": "MixMod Guider Component", "MixModGuiderNode": "MixMod Guider", "MixModFFTGuiderNode": "MixMod FFT Guider", }