diff --git a/guider.py b/guider.py index 3395198..382f565 100644 --- a/guider.py +++ b/guider.py @@ -1,293 +1,347 @@ -import torch -import comfy.samplers -import comfy.utils -import logging -import torch.nn.functional as F -from .utils import sumweights - -# Initialize other necessary attributes - -import comfy.model_management as model_management -from .utils import load_taesd, ensure_model, create_adaptive_low_res_model_wrapper, create_noise_injection_wrapper, load_vae - - -# Import mixing functions -from .mixes.fft import mix_teamfft, mix_teamfft4freq, mix_2model_fft, mix_2m2f, mix_fft_overlap, mix_fft_full, mix_bandfft -from .mixes.standard import mix_standard, mix_masked, mix_dynamic_mask, mix_foreground_background, mix_dynamic_mask_alternative - - -class MixModGuider(comfy.samplers.CFGGuider): - def __init__(self, component, **kwargs): - - self.component = component - self.mode = kwargs.get("mode", "team") - self.ratio = kwargs.get("ratio", 1.0) - self.original_conds = {} - self.prepared_models = [] - self.highres_mode = kwargs.get("highres_mode", False) - self.bands = kwargs.get("bands", None) - self.kwargs = kwargs - self.last_latent = None - - # Load the depth model - if(self.mode == "foreground_background"): - #check if controlnet_aux is installed - try: - from custom_nodes.comfyui_controlnet_aux.src.custom_controlnet_aux.depth_anything import DepthAnythingDetector - self.depth_model = DepthAnythingDetector.from_pretrained().to(model_management.get_torch_device()) - self.vae = comfy.sd.VAE(sd=load_taesd("taesdxl")) - except Exception as e: - self.mode = "team" - raise Exception("Please install Comfyui ControlNet Aux custom node. Falling back to team mode.") - - if(self.mode == "styleinjection" and not self.kwargs.get("image", None) is None): - self.vae = comfy.sd.VAE(sd=load_vae("fixFP16ErrorsSDXLLowerMemoryUse_v10.safetensors")) - - self.noiseinjection_latent = self.vae.encode(self.kwargs.get("image", None)) - - - self.models = [] - self.weights = [] - self.cfgs = [] - self.masks = [] # Initialize empty masks list - self.step = 0 - self.total_steps = 0 - self.start_steps = [] - self.end_steps = [] - self.scale_factors = [] - self.options = [] - self.mask_history = None # For tracking dynamic masks between steps - self.decay_factor = kwargs.get("decay_factor", 0.9) # Dynamic mask decay factor - self.threshold = kwargs.get("threshold", 0.1) # Dynamic mask threshold - self.active_masks = [] - self.active_weights = [] - self.active_cfgs = [] - self.uncond_decays = [] - self.foreground_background_masks = [] - self.seed = -1 - # Process component chain - current = component - conds = {} - while current is not None: - self.models.append(current.get("model")) - if len(self.models) == 1: - self.model_patcher = self.models[0] - self.model_options = self.models[0].model_options - self.weights.append(current.get("weight", 1.0)) - self.cfgs.append(current.get("cfg", 7.5)) - opt = current.get("options", {}) - options = {} - while opt is not None: - options.update(opt) - opt = opt.get("prev") - self.scale_factors.append(options.get("scale_factor", 0.0)) - self.start_steps.append(options.get("start_step", 0)) - self.end_steps.append(options.get("end_step", -1)) - self.masks.append(options.get("mask", None)) - self.uncond_decays.append(options.get("uncond_decay", 1.0)) - # Add to original_conds with indexed keys - idx = len(self.models) - conds.update({f"positive_{idx}": current["positive"], f"negative_{idx}": current["negative"]}) - - current = current.get("prev") - - - self.inner_set_conds(conds) - - - - - def set_mode(self, mode): - self.mode = mode - - def predict_noise(self, x, timestep, model_options={}, seed=None): - pred = [] - self.active_masks = [] - self.active_weights = [] - self.active_cfgs = [] - self.active_uncond_decays = [] - def latent_callback(latent): - self.noiseinjection_latent = latent - for i in range(len(self.models)): - cond = [self.conds.get(f"positive_{i+1}", None), self.conds.get(f"negative_{i+1}", None)] - if(self.start_steps[i] <= self.step <= self.end_steps[i] or self.end_steps[i] == -1): - if(self.scale_factors[i] > 0): - model_options["model_function_wrapper"] = create_adaptive_low_res_model_wrapper( - min_scale_factor=self.scale_factors[i], - max_scale_factor=self.scale_factors[i] - ) - - ensure_model(self.models[i],"normal") - pred.append(comfy.samplers.calc_cond_batch(self.prepared_models[i][0], cond, x, timestep, model_options)) - del model_options["model_function_wrapper"] - elif(self.mode == "styleinjection" and self.kwargs.get("image", None) is not None and i == 1 and self.kwargs.get("start_step", 4) <= self.step <= self.kwargs.get("end_step", 7)): - #calculate noise weight, decreasing from self.kwargs.get("noise_weight", 1.0) to 0 between start_step and end_step - noise_weight = self.kwargs.get("noise_weight", 1.0) - (self.step - self.kwargs.get("start_step", 4)) * (self.kwargs.get("noise_weight", 1.0) - 0) / (self.kwargs.get("end_step", 7) - self.kwargs.get("start_step", 4)) - model_options["model_function_wrapper"] = create_noise_injection_wrapper(self.noiseinjection_latent, noise_weight, latent_callback) - ensure_model(self.models[i],"normal") - pred.append(comfy.samplers.calc_cond_batch(self.prepared_models[i][0], cond, x, timestep, model_options)) - del model_options["model_function_wrapper"] - else: - ensure_model(self.models[i],"normal") - pred.append(comfy.samplers.calc_cond_batch(self.prepared_models[i][0], cond, x, timestep, model_options)) - - self.active_masks.append(self.masks[i]) - self.active_weights.append(self.weights[i]) - self.active_cfgs.append(self.cfgs[i]) - self.active_uncond_decays.append(self.uncond_decays[i]) - - if not pred: - logging.warning(f"No models were active at step {self.step}, using default model 0") - cond = [self.conds.get(f"positive_1", None), self.conds.get(f"negative_1", None)] - pred.append(comfy.samplers.calc_cond_batch(self.prepared_models[0][0], cond, x, timestep, model_options)) - self.active_masks.append(self.masks[0]) - self.active_weights.append(self.weights[0]) - self.active_cfgs.append(self.cfgs[0]) - self.active_uncond_decays.append(self.uncond_decays[0]) - return self.mix_function(self.prepared_models[0][0], pred, x, timestep, model_options, seed) - - def mix_function(self, model, pred, x, timestep, model_options, seed): - haspostfunctions = False - - postuncond = torch.zeros_like(pred[0][1]) - postcond = torch.zeros_like(pred[0][0]) - - self.active_weights = sumweights(self.active_weights) - """ if(len(pred) == 1 and not self.mode == "foreground_background"): - mix = pred[0][1] + (pred[0][0] - pred[0][1]) * self.cfgs[0] - for fn in model_options.get("sampler_post_cfg_function", []): - args = {"denoised": mix, "cond": self.conds.get(f"positive_1", None), "uncond": self.conds.get(f"negative_1", None), "cond_scale": self.cfgs[0], "model": model, "uncond_denoised": pred[0][1], "cond_denoised": pred[0][0], - "sigma": timestep, "model_options": model_options, "input": x} - mix = fn(args) - self.step += 1 - return mix """ - if(self.mode == "team"): - if any(mask is not None for mask in self.active_masks): - mix, postuncond = mix_masked(pred, self.active_weights, self.active_cfgs, self.active_masks,) - else: - mix, postuncond = mix_standard(pred, self.active_weights, self.active_cfgs, self.active_uncond_decays) - - elif self.mode == "styleinjection": - mix, postuncond = mix_standard(pred, self.active_weights, self.active_cfgs) - - elif self.mode == "dynamic_mask": - mix, postuncond, new_masks = mix_dynamic_mask(pred, self.active_weights, self.active_cfgs, self.step, self.mask_history, self.decay_factor, self.threshold, self.kwargs.get("blur_sigma", 0.5)) - self.mask_history = new_masks # Store the updated masks for next step - - elif self.mode == "dynamic_mask_alternative": - mix, postuncond = mix_dynamic_mask_alternative(pred, self.active_weights, self.active_cfgs, self.depth_model, self.threshold, self.kwargs.get("invert", False), self.kwargs.get("blur_sigma", 1.0)) - - elif self.mode == "foreground_background": - if(self.step == 0): - self.foreground_background_masks = [] - - masks, mix, postuncond = mix_foreground_background(pred, self.active_weights, self.last_latent, self.active_cfgs, vae=self.vae, depth_model=self.depth_model, sharpness=self.kwargs.get("sharpness", 0.0), invert=self.kwargs.get("invert", False), step=self.step, masks=self.foreground_background_masks, startmask=self.kwargs.get("startmask", 5)) - - if any(mask is not None for mask in masks): - self.foreground_background_masks = masks - - elif self.mode == "bandfft": - mix, postuncond = mix_bandfft(pred, self.active_weights, self.active_cfgs, self.bands) - - elif self.mode == "2m2f": - mix, postuncond = mix_2m2f(pred, self.active_weights, self.active_cfgs, self.ratio) - - elif self.mode == "teamfft": - mix, postuncond = mix_teamfft(pred, self.active_weights, self.active_cfgs) - - elif self.mode == "teamfft4freq": - mix, postuncond = mix_teamfft4freq(pred, self.active_weights, self.active_cfgs) - - elif self.mode == "2model_fft": - mix, postuncond = mix_2model_fft(pred, self.active_weights, self.active_cfgs, self.ratio) - - elif self.mode == "fft_overlap": - mix, postuncond = mix_fft_overlap(pred, self.active_weights, self.active_cfgs, self.ratio) - - elif self.mode == "fft_full": - mix, postuncond = mix_fft_full(pred, self.active_weights, self.active_cfgs) - else: - # Fallback to standard mixing if mode is not recognized - mix, postuncond = mix_standard(pred, self.active_weights, self.active_cfgs) - - for fn in model_options.get("sampler_post_cfg_function", []): - args = {"denoised": mix, "cond": self.conds.get(f"positive_1", None), "uncond": self.conds.get(f"negative_1", None), "cond_scale": self.active_cfgs[0], "model": model, "uncond_denoised": postuncond, "cond_denoised": postcond, - "sigma": timestep, "model_options": model_options, "input": x} - mix = fn(args) - self.step += 1 - self.seed = seed - return mix - - def outer_sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None): - # Get the device from the first model - device = self.models[0].load_device - - # Prepare the models, conditions, and loaded models for each model - prepared_models = [] - all_loaded_models = [] - - for i, model in enumerate(self.models): - # Get relevant conditions for this model - model_conds = {k: v for k, v in self.conds.items() if k.endswith(f"_{i+1}")} - - - # Prepare sampling for this model - inner_model, prepared_conds, loaded_models = comfy.sampler_helpers.prepare_sampling(model, noise.shape, model_conds, model.model_options) - - # Save the prepared data - self.prepared_models.append((inner_model, prepared_conds)) - - all_loaded_models.extend(loaded_models) - - # Update the conditions for this model - for k, v in prepared_conds.items(): - self.conds[k] = v - - self.inner_model = self.prepared_models[0][0] - # Prepare the device-related tensors - if denoise_mask is not None: - denoise_mask = comfy.sampler_helpers.prepare_mask(denoise_mask, noise.shape, device) - - noise = noise.to(device) - latent_image = latent_image.to(device) - sigmas = sigmas.to(device) - - - # Set model options for the main model patcher - comfy.samplers.cast_to_load_options(self.model_options, device=device, dtype=self.model_patcher.model_dtype()) - - try: - # Pre-run all models - for model in self.models: - model.pre_run() - - # Set the inner_model to be a list of prepared models - - # Run the inner sample with prepared models - output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed) - finally: - # Clean up all models - for model in self.models: - model.cleanup() - del self.inner_model - # Clean up models and loaded models - comfy.sampler_helpers.cleanup_models(self.conds, all_loaded_models) - - - return output - - def inner_sample(self, noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed): - # Create a wrapper callback that captures intermediate denoised results (x0) - original_callback = callback - - def wrapped_callback(step, x0, x, total_steps): - # Store the current x0 (denoised latent) - self.step = step - self.total_steps = total_steps - self.last_latent = x0.detach().clone() - # Call the original callback if provided - if original_callback is not None: - original_callback(step, x0, x, total_steps) - - # Pass our wrapped callback to the parent method - return comfy.samplers.CFGGuider.inner_sample(self, noise, latent_image, device, sampler, sigmas, denoise_mask, wrapped_callback, disable_pbar, seed) +import torch +import comfy.samplers +import comfy.sampler_helpers +import comfy.utils +import logging +import torch.nn.functional as F +from .utils import sumweights + +# Initialize other necessary attributes + +import comfy.model_management as model_management +from .utils import load_taesd, ensure_model, create_adaptive_low_res_model_wrapper, create_model_wrapper, load_vae + + +# Import mixing functions +from .mixes.fft import mix_teamfft, mix_teamfft4freq, mix_2model_fft, mix_2m2f, mix_fft_overlap, mix_fft_full, mix_bandfft +from .mixes.standard import mix_standard, mix_masked, mix_dynamic_mask, mix_foreground_background, mix_dynamic_mask_alternative + + +class MixModGuider(comfy.samplers.CFGGuider): + def __init__(self, component, **kwargs): + + self.component = component + self.mode = kwargs.get("mode", "team") + self.ratio = kwargs.get("ratio", 1.0) + self.original_conds = {} + self.prepared_models = [] + self.highres_mode = kwargs.get("highres_mode", False) + self.bands = kwargs.get("bands", None) + self.kwargs = kwargs + self.last_latent = None + + # Load the depth model + if(self.mode == "foreground_background"): + #check if controlnet_aux is installed + try: + from custom_nodes.comfyui_controlnet_aux.src.custom_controlnet_aux.depth_anything import DepthAnythingDetector + self.depth_model = DepthAnythingDetector.from_pretrained().to(model_management.get_torch_device()) + self.vae = load_taesd("taesdxl") + except Exception as e: + self.mode = "team" + raise Exception("Please install Comfyui ControlNet Aux custom node. Falling back to team mode.") + + if(self.mode == "styleinjection" and not self.kwargs.get("image", None) is None): + self.vae = load_vae("fixFP16ErrorsSDXLLowerMemoryUse_v10.safetensors") + + self.noiseinjection_latent = self.vae.encode(self.kwargs.get("image", None)) + + + self.models = [] + self.weights = [] + self.cfgs = [] + self.masks = [] # Initialize empty masks list + self.types = [] + + self.step = 0 + self.total_steps = 0 + self.start_steps = [] + self.end_steps = [] + + self.options = [] + self.scale_factors = [] + self.decay_factor = kwargs.get("decay_factor", 0.9) # Dynamic mask decay factor + self.threshold = kwargs.get("threshold", 0.1) # Dynamic mask threshold + + self.active_masks = [] + self.active_weights = [] + self.active_cfgs = [] + self.uncond_decays = [] + + + self.foreground_background_masks = [] + self.mask_history = None # For tracking dynamic masks between steps + + self.seed = -1 + # Process component chain + current = component + conds = {} + while current is not None: + self.models.append(current.get("model")) + if len(self.models) == 1: + self.model_patcher = self.models[0] + self.model_options = self.models[0].model_options + self.weights.append(current.get("weight", 1.0)) + self.cfgs.append(current.get("cfg", 7.5)) + opt = current.get("options", {}) + options = {} + while opt is not None: + options.update(opt) + opt = opt.get("prev") + self.scale_factors.append(options.get("scale_factor", 0.0)) + self.start_steps.append(options.get("start_step", 0)) + self.end_steps.append(options.get("end_step", -1)) + self.masks.append(options.get("mask", None)) + self.uncond_decays.append(options.get("uncond_decay", 1.0)) + # Add to original_conds with indexed keys + idx = len(self.models) + conds.update({f"positive_{idx}": current["positive"], f"negative_{idx}": current["negative"]}) + + current = current.get("prev") + + + self.inner_set_conds(conds) + + + + + def set_mode(self, mode): + self.mode = mode + + def predict_noise(self, x, timestep, model_options={}, seed=None): + pred = [] + self.active_masks = [] + self.active_weights = [] + self.active_cfgs = [] + self.active_uncond_decays = [] + for i in range(len(self.models)): + cond = [self.conds.get(f"positive_{i+1}", None), self.conds.get(f"negative_{i+1}", None)] + if(self.start_steps[i] <= self.step <= self.end_steps[i] or self.end_steps[i] == -1): + if(self.scale_factors[i] > 0): + model_options["model_function_wrapper"] = create_adaptive_low_res_model_wrapper( + min_scale_factor=self.scale_factors[i], + max_scale_factor=self.scale_factors[i] + ) + ensure_model(self.models[i],"normal") + pred.append(comfy.samplers.calc_cond_batch(self.prepared_models[i][0], cond, x, timestep, model_options)) + del model_options["model_function_wrapper"] + else: + ensure_model(self.models[i],"normal") + + pred.append(comfy.samplers.calc_cond_batch(self.prepared_models[i][0], cond, x, timestep, model_options)) + + self.active_masks.append(self.masks[i]) + self.active_weights.append(self.weights[i]) + self.active_cfgs.append(self.cfgs[i]) + self.active_uncond_decays.append(self.uncond_decays[i]) + + if not pred: + logging.warning(f"No models were active at step {self.step}, using default model 0") + cond = [self.conds.get(f"positive_1", None), self.conds.get(f"negative_1", None)] + pred.append(comfy.samplers.calc_cond_batch(self.prepared_models[0][0], cond, x, timestep, model_options)) + self.active_masks.append(self.masks[0]) + self.active_weights.append(self.weights[0]) + self.active_cfgs.append(self.cfgs[0]) + self.active_uncond_decays.append(self.uncond_decays[0]) + return self.mix_function(self.prepared_models[0][0], pred, x, timestep, model_options, seed) + + def mix_function(self, model, pred, x, timestep, model_options, seed): + haspostfunctions = False + + postuncond = torch.zeros_like(pred[0][1]) + postcond = torch.zeros_like(pred[0][0]) + + self.active_weights = sumweights(self.active_weights) + if(self.mode == "team"): + if any(mask is not None for mask in self.active_masks): + mix, postuncond = mix_masked(pred, self.active_weights, self.active_cfgs, self.active_masks,) + else: + mix, postuncond = mix_standard(pred, self.active_weights, self.active_cfgs, self.active_uncond_decays) + + elif self.mode == "styleinjection": + mix, postuncond = mix_standard(pred, self.active_weights, self.active_cfgs) + + elif self.mode == "dynamic_mask": + mix, postuncond, new_masks = mix_dynamic_mask(pred, self.active_weights, self.active_cfgs, self.step, self.mask_history, self.decay_factor, self.threshold, self.kwargs.get("blur_sigma", 0.5)) + self.mask_history = new_masks # Store the updated masks for next step + + elif self.mode == "dynamic_mask_alternative": + mix, postuncond = mix_dynamic_mask_alternative(pred, self.active_weights, self.active_cfgs, self.depth_model, self.threshold, self.kwargs.get("invert", False), self.kwargs.get("blur_sigma", 1.0)) + + elif self.mode == "foreground_background": + if(self.step == 0): + self.foreground_background_masks = [] + + masks, mix, postuncond = mix_foreground_background(pred, self.active_weights, self.last_latent, self.active_cfgs, vae=self.vae, depth_model=self.depth_model, sharpness=self.kwargs.get("sharpness", 0.0), invert=self.kwargs.get("invert", False), step=self.step, masks=self.foreground_background_masks, startmask=self.kwargs.get("startmask", 5)) + + if any(mask is not None for mask in masks): + self.foreground_background_masks = masks + + elif self.mode == "bandfft": + mix, postuncond = mix_bandfft(pred, self.active_weights, self.active_cfgs, self.bands) + + elif self.mode == "2m2f": + mix, postuncond = mix_2m2f(pred, self.active_weights, self.active_cfgs, self.ratio) + + elif self.mode == "teamfft": + mix, postuncond = mix_teamfft(pred, self.active_weights, self.active_cfgs) + + elif self.mode == "teamfft4freq": + mix, postuncond = mix_teamfft4freq(pred, self.active_weights, self.active_cfgs) + + elif self.mode == "2model_fft": + mix, postuncond = mix_2model_fft(pred, self.active_weights, self.active_cfgs, self.ratio) + + elif self.mode == "fft_overlap": + mix, postuncond = mix_fft_overlap(pred, self.active_weights, self.active_cfgs, self.ratio) + + elif self.mode == "fft_full": + mix, postuncond = mix_fft_full(pred, self.active_weights, self.active_cfgs) + else: + # Fallback to standard mixing if mode is not recognized + mix, postuncond = mix_standard(pred, self.active_weights, self.active_cfgs) + + for fn in model_options.get("sampler_post_cfg_function", []): + args = {"denoised": mix, "cond": self.conds.get(f"positive_1", None), "uncond": self.conds.get(f"negative_1", None), "cond_scale": self.active_cfgs[0], "model": model, "uncond_denoised": postuncond, "cond_denoised": postcond, + "sigma": timestep, "model_options": model_options, "input": x} + mix = fn(args) + self.step += 1 + self.seed = seed + return mix + + def outer_sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None): + # Get the device from the first model + device = self.models[0].load_device + + # Prepare the models, conditions, and loaded models for each model + + all_loaded_models = [] + + for i, model in enumerate(self.models): + # Get relevant conditions for this model + model_conds = {k: v for k, v in self.conds.items() if k.endswith(f"_{i+1}")} + + + # Prepare sampling for this model + inner_model, prepared_conds, loaded_models = comfy.sampler_helpers.prepare_sampling(model, noise.shape, model_conds, model.model_options) + + # Save the prepared data + self.prepared_models.append((inner_model, prepared_conds)) + + all_loaded_models.extend(loaded_models) + + # Update the conditions for this model + for k, v in prepared_conds.items(): + self.conds[k] = v + + self.inner_model = self.prepared_models[0][0] + # Prepare the device-related tensors + if denoise_mask is not None: + denoise_mask = comfy.sampler_helpers.prepare_mask(denoise_mask, noise.shape, device) + + noise = noise.to(device) + latent_image = latent_image.to(device) + sigmas = sigmas.to(device) + + + # Set model options for the main model patcher + comfy.samplers.cast_to_load_options(self.model_options, device=device, dtype=self.model_patcher.model_dtype()) + + try: + # Pre-run all models + for model in self.models: + model.pre_run() + + # Set the inner_model to be a list of prepared models + + # Run the inner sample with prepared models + output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed) + finally: + # Clean up all models + for model in self.models: + model.cleanup() + # Clean up models and loaded models + comfy.sampler_helpers.cleanup_models(self.conds, all_loaded_models) + self.prepared_models = [] + del all_loaded_models + + + return output + + def inner_sample(self, noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed): + # Create a wrapper callback that captures intermediate denoised results (x0) + original_callback = callback + + def wrapped_callback(step, x0, x, total_steps): + # Store the current x0 (denoised latent) + self.step = step + self.total_steps = total_steps + self.last_latent = x0.detach().clone() + # Call the original callback if provided + if original_callback is not None: + original_callback(step, x0, x, total_steps) + + if latent_image is not None and torch.count_nonzero(latent_image) > 0: #Don't shift the empty latent image. + latent_image = self.inner_model.process_latent_in(latent_image) + + # Process each cond with each prepared inner model separately + processed_conds = {} + for i, (inner_model, _) in enumerate(self.prepared_models): + # Get the conditions for this model (positive_i, negative_i) + model_conds = {k: v for k, v in self.conds.items() if k.endswith(f"_{i+1}")} + + # Process conditions with this model + processed_model_conds = comfy.samplers.process_conds(inner_model, noise, model_conds, device, latent_image, denoise_mask, seed) + + # Add to processed conditions + processed_conds.update(processed_model_conds) + + # Merge back into self.conds + self.conds = processed_conds + + extra_model_options = comfy.model_patcher.create_model_options_clone(self.model_options) + extra_model_options.setdefault("transformer_options", {})["sample_sigmas"] = sigmas + extra_args = {"model_options": extra_model_options, "seed": seed} + + executor = comfy.patcher_extension.WrapperExecutor.new_class_executor( + sampler.sample, + sampler, + comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.SAMPLER_SAMPLE, extra_args["model_options"], is_model_options=True) + ) + samples = executor.execute(self, sigmas, extra_args, wrapped_callback, noise, latent_image, denoise_mask, disable_pbar) + return samples + + def sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None): + if sigmas.shape[-1] == 0: + return latent_image + + self.conds = {} + for k in self.original_conds: + self.conds[k] = list(map(lambda a: a.copy(), self.original_conds[k])) + comfy.samplers.preprocess_conds_hooks(self.conds) + + try: + orig_model_options = self.model_options + self.model_options = comfy.model_patcher.create_model_options_clone(self.model_options) + # if one hook type (or just None), then don't bother caching weights for hooks (will never change after first step) + orig_hook_mode = self.model_patcher.hook_mode + if comfy.samplers.get_total_hook_groups_in_conds(self.conds) <= 1: + self.model_patcher.hook_mode = comfy.hooks.EnumHookMode.MinVram + + comfy.sampler_helpers.prepare_model_patcher(self.model_patcher, self.conds, self.model_options) + + comfy.samplers.filter_registered_hooks_on_conds(self.conds, self.model_options) + executor = comfy.patcher_extension.WrapperExecutor.new_class_executor( + self.outer_sample, + self, + comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.OUTER_SAMPLE, self.model_options, is_model_options=True) + ) + output = executor.execute(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed) + finally: + comfy.samplers.cast_to_load_options(self.model_options, device=self.model_patcher.offload_device) + self.model_options = orig_model_options + self.model_patcher.hook_mode = orig_hook_mode + self.model_patcher.restore_hook_patches() + + del self.conds + return output \ No newline at end of file diff --git a/nodes.py b/nodes.py index 1aee070..55a49f3 100644 --- a/nodes.py +++ b/nodes.py @@ -65,7 +65,7 @@ class MixModGuiderComponentPipelineNode: "weight": base_weight, "cfg": cfg, "options": options, - "prev": prev_component + "prev": prev_component, } return (component,) @@ -76,7 +76,7 @@ class MixModPipelineNode: return {"required": {"model": ("MODEL",), "positive": ("CONDITIONING",), - "negative": ("CONDITIONING",), + "negative": ("CONDITIONING",) } } @@ -387,6 +387,7 @@ NODE_CLASS_MAPPINGS = { "MixModGuiderNode": MixModGuiderNode, "MixModFFTGuiderNode": MixModFFTGuiderNode, "MixModBandFFTGuiderNode": MixModBandFFTGuiderNode, + "MixModHighResGuiderNode": MixModHighResGuiderNode, "MixModOptionsMaskNode": MixModOptionsMaskNode, "MixModOptionsSchedulerNode": MixModOptionsSchedulerNode, "MixModOptionsScaleNode": MixModOptionsScaleNode, @@ -402,6 +403,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "MixModGuiderNode": "MixMod Guider", "MixModFFTGuiderNode": "MixMod FFT Guider", "MixModBandFFTGuiderNode": "MixMod Band FFT Guider", + "MixModHighResGuiderNode": "MixMod High Res Guider", "MixModOptionsMaskNode": "MixMod Options Mask", "MixModOptionsSchedulerNode": "MixMod Options Scheduler", "MixModOptionsScaleNode": "MixMod Options Scale", @@ -409,4 +411,4 @@ NODE_DISPLAY_NAME_MAPPINGS = { "MixModDynamicMaskGuiderNode": "MixMod Dynamic Mask Guider", "MixModDynamicMaskAlternativeGuiderNode": "MixMod Dynamic Mask Alternative Guider", "MixModDepthGuiderNode": "MixMod Depth Guider", -} +} \ No newline at end of file diff --git a/utils.py b/utils.py index 5f3dc7c..727c896 100644 --- a/utils.py +++ b/utils.py @@ -1,392 +1,528 @@ -from __future__ import annotations -import folder_paths -import comfy.utils -import torch -import numpy as np -import os -from PIL import Image -import torchvision.transforms.functional as F -import math -from comfy import model_management -from comfy import latent_formats - - -def ensure_model(model_patcher, mode = None): - if not mode or mode == "disable": - # Ensure model mode disabled. - return - mp = model_patcher - lm = model_management.LoadedModel(mp) - found_lm = None - for list_lm in model_management.current_loaded_models: - if lm == list_lm: - found_lm = list_lm - break - if getattr(found_lm, "currently_used", False): - # Model already exists and appears to be loaded. - return - if found_lm is not None: - lm = found_lm - mp = lm.model - if mode.startswith("normal"): - if mode == "normal_unload": - model_management.unload_all_models() - model_management.load_models_gpu((mp,)) - return - model_management.unload_all_models() - if mode == "lowvram": - # Logic from comfy.model_management.load_models_gpu - min_inference_memory = model_management.minimum_inference_memory() - minimum_memory_required = max( - min_inference_memory, - model_management.extra_reserved_memory(), - ) - loaded_memory = lm.model_loaded_memory() - current_free_mem = model_management.get_free_memory(lm.device) + loaded_memory - - lowvram_model_memory = max( - 64 * 1024 * 1024, - (current_free_mem - minimum_memory_required), - min( - current_free_mem * model_management.MIN_WEIGHT_MEMORY_RATIO, - current_free_mem - model_management.minimum_inference_memory(), - ), - ) - lowvram_model_memory = max(0.1, lowvram_model_memory - loaded_memory) - elif mode == "novram": - lowvram_model_memory = 0.1 - else: - raise ValueError("Bad ensure_model_mode") - lm.model_load(lowvram_model_memory) - model_management.current_loaded_models.insert(0, lm) - - -def fallback(val, default, *, exclude=None, default_is_fun=False): - return val if val is not exclude else (default() if default_is_fun else default) - - -def scale_dim(n, factor=1.0, *, increment=64) -> int: - return math.ceil((n * factor) / increment) * increment - - -def sigma_to_float(sigma): - return sigma.detach().cpu().max().item() - -def load_taesd(name): - sd = {} - approx_vaes = folder_paths.get_filename_list("vae_approx") - - encoder = next(filter(lambda a: a.startswith("{}_encoder.".format(name)), approx_vaes)) - decoder = next(filter(lambda a: a.startswith("{}_decoder.".format(name)), approx_vaes)) - - enc = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", encoder)) - for k in enc: - sd["taesd_encoder.{}".format(k)] = enc[k] - - dec = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", decoder),device=model_management.get_torch_device()) - for k in dec: - sd["taesd_decoder.{}".format(k)] = dec[k] - - if name == "taesd": - sd["vae_scale"] = torch.tensor(0.18215) - sd["vae_shift"] = torch.tensor(0.0) - elif name == "taesdxl": - sd["vae_scale"] = torch.tensor(0.13025) - sd["vae_shift"] = torch.tensor(0.0) - elif name == "taesd3": - sd["vae_scale"] = torch.tensor(1.5305) - sd["vae_shift"] = torch.tensor(0.0609) - elif name == "taef1": - sd["vae_scale"] = torch.tensor(0.3611) - sd["vae_shift"] = torch.tensor(0.1159) - return sd - -def load_vae(name): - vae = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae", name)) - return vae - -def create_depth_map(latent_image, depth_model, resolution, **kwargs): - batch_size = latent_image.shape[0] - out_tensor = None - for i, image in enumerate(latent_image): - np_image = np.asarray(image.cpu() * 255., dtype=np.uint8) - np_result = depth_model(np_image, output_type="np", detect_resolution=resolution, **kwargs) - out = torch.from_numpy(np_result.astype(np.float32) / 255.0) - if out_tensor is None: - out_tensor = torch.zeros(batch_size, *out.shape, dtype=torch.float32) - out_tensor[i] = out - - return out_tensor - -def save_debug_image(tensor, filename, scale=True): - try: - # Ensure tensor is detached and on CPU - tensor = tensor.detach().cpu().float() - print(tensor.shape) - # Remove batch dimension if present - if tensor.ndim == 4 and tensor.shape[0] == 1: - tensor = tensor.squeeze(0) - - # Handle 2D tensors (grayscale masks) by converting to RGB - if tensor.ndim == 2: - # Convert 2D tensor to 3D by repeating across 3 channels - tensor = tensor.unsqueeze(-1).repeat(1, 1, 3) - - # If scale is True, scale to [0, 255] - if scale: - array = (tensor.numpy() * 255).clip(0, 255).astype('uint8') - else: - array = tensor.numpy().clip(0, 255).astype('uint8') - # Ensure shape is (H, W, 3) - if array.shape[-1] != 3: - raise ValueError(f"Expected last dimension to be 3 (RGB), got {array.shape[-1]}") - img = Image.fromarray(array) - - # Ensure debug directory exists - os.makedirs("debug", exist_ok=True) - - img.save("debug/" + filename) - return True - except Exception as e: - print(f"Error saving debug image {filename}: {e}") - return False - -def process_mask(mask, H, W, device): - # Convert to tensor if not already - if not isinstance(mask, torch.Tensor): - mask = torch.tensor(mask, device=device, dtype=torch.float32) - mask = mask.to(device).float() - - # Squeeze singleton dimensions - mask = mask.squeeze() - - # If mask is RGB, convert to grayscale - if mask.ndim == 3 and mask.shape[-1] == 3: - # [H, W, 3] -> [3, H, W] - mask = mask.permute(2, 0, 1) - mask = mask.mean(dim=0) # [H, W] - elif mask.ndim == 4 and mask.shape[-1] == 3: - # [B, H, W, 3] -> [B, 3, H, W] - mask = mask.permute(0, 3, 1, 2) - mask = mask.mean(dim=1) # [B, H, W] - mask = mask.squeeze(0) # Remove batch if present - - # Now mask should be [H, W] - if mask.shape != (H, W): - mask = torch.nn.functional.interpolate(mask.unsqueeze(0).unsqueeze(0), size=(H, W), mode='bilinear')[0, 0] - - return mask - - -def create_adaptive_low_res_model_wrapper( min_scale_factor=0.5, max_scale_factor=1.0, start_step=0, end_step=None): - """ - Creates a model function wrapper that runs the model at progressively higher resolution as sampling proceeds. - - Args: - min_scale_factor: Minimum scale factor to use (for early steps) - max_scale_factor: Maximum scale factor to use (for final steps) - start_step: Step number to start increasing resolution (0 means immediately) - end_step: Step number when scaling should reach max_scale_factor (None means the last step) - - Returns: - A function wrapper that can be used as model_options['model_function_wrapper'] - """ - def wrapper(apply_model_func, args_dict): - input_x = args_dict["input"] - timestep = args_dict["timestep"] - c = args_dict["c"] - - # Get sampling progress from sigmas in transformer_options - sample_sigmas = c['transformer_options'].get('sample_sigmas', None) - current_sigma = c['transformer_options'].get('sigmas', None) - - # Determine scale factor based on current step - if sample_sigmas is not None and current_sigma is not None: - # Get total steps from the length of sample_sigmas - total_steps = len(sample_sigmas) - - # Find current step by matching current_sigma to sample_sigmas - current_sigma_val = float(current_sigma[0]) - - # Find the index of the closest sigma value in sample_sigmas - # This gives us the current step number - sigma_diffs = [abs(float(s) - current_sigma_val) for s in sample_sigmas] - current_step = sigma_diffs.index(min(sigma_diffs)) - - # If end_step is not provided, use the last step - actual_end_step = end_step if end_step is not None else (total_steps - 1) - - # Clamp start and end steps to valid range - start_step_clamped = max(0, min(start_step, total_steps - 1)) - end_step_clamped = max(start_step_clamped, min(actual_end_step, total_steps - 1)) - - # Calculate scale factor based on current step - if current_step < start_step_clamped: - # Before start step - use minimum scale - scale_factor = min_scale_factor - elif current_step > end_step_clamped: - # After end step - use maximum scale - scale_factor = max_scale_factor - else: - # In between - interpolate between min and max - step_range = end_step_clamped - start_step_clamped - if step_range > 0: - progress = (current_step - start_step_clamped) / step_range - scale_factor = min_scale_factor + progress * (max_scale_factor - min_scale_factor) - else: - # If start_step equals end_step - scale_factor = min_scale_factor - else: - # Fallback if we can't determine steps - scale_factor = max_scale_factor - - # Original shape - original_shape = input_x.shape - - # Skip resizing if we're at full resolution - if scale_factor >= 0.99: - return apply_model_func(input_x, timestep, **c) - print(f"Scaling down to {scale_factor}x for step {current_step} using bislerp") - # Calculate new dimensions - new_h = int(original_shape[2] * scale_factor) - new_w = int(original_shape[3] * scale_factor) - - # Ensure dimensions are at least 8 (minimum size for stable diffusion) - new_h = max(8, new_h) - new_w = max(8, new_w) - - # Downscale using comfy.utils.common_upscale with bislerp mode - downscaled_x = comfy.utils.common_upscale( - input_x, - new_w, - new_h, - 'bislerp', - 'disabled' - ) - - # Run model at lower resolution - output = apply_model_func(downscaled_x, timestep, **c) - - # Upscale result back to original resolution using bislerp - upscaled_output = comfy.utils.common_upscale( - output, - original_shape[3], # width - original_shape[2], # height - 'bislerp', - 'disabled' - ) - - """ # #blurring the output - # Apply a small amount of Gaussian blur to smooth the upscaled output - # This helps reduce artifacts that can appear from the upscaling process - kernel_size = 3 # Small kernel for subtle blurring - sigma = 0.5 # Low sigma for gentle blur - - # Create Gaussian kernel - kernel = torch.tensor([ - [0.0625, 0.125, 0.0625], - [0.125, 0.25, 0.125], - [0.0625, 0.125, 0.0625] - ], device=upscaled_output.device).view(1, 1, kernel_size, kernel_size) - - Expand kernel for all channels - channels = upscaled_output.shape[1] - kernel = kernel.expand(channels, 1, kernel_size, kernel_size) - - # Group kernel for batch processing - kernel = kernel.repeat(1, 1, 1, 1) - - # Apply convolution for blurring - # Use padding='same' to maintain dimensions - padding = kernel_size // 2 - blurred_output = torch.nn.functional.conv2d( - upscaled_output, - kernel.to(upscaled_output.dtype), - padding=padding, - groups=channels - ) - - # Use the blurred output instead of the direct upscaled output - upscaled_output = blurred_output """ - return upscaled_output - - return wrapper - -def create_noise_injection_wrapper(latent, noise_weight=1.0, callback=None): - """ - Creates a wrapper that takes an input image, converts it to latent space, - noises it to match the current timestep/sigma, and injects its noise prediction. - - Args: - image: Input image to inject (tensor in pixel space) - vae: VAE model for encoding to latent space - noise_weight: Weight of the injected noise prediction (0-1) - - Returns: - A function wrapper to use as model_options['model_function_wrapper'] - - - """ - def wrapper(apply_model_func, args_dict): - # Extract args - input_x = args_dict["input"] - timestep = args_dict["timestep"] - c = args_dict["c"] - - - # Get current sampling parameters - sigmas = c['transformer_options'].get('sigmas', None) - sigma_idx = c['transformer_options'].get('sigma_idx', 0) - current_sigma = sigmas[sigma_idx] if sigmas is not None else None - - # Noise the image latent to match current noise level - if current_sigma is not None: - # Move tensors to the same device as input_x - device = input_x.device - image_latent = latent.to(device) - current_sigma = current_sigma.to(device) - - image_latent = latent_formats.SDXL().process_in(image_latent) - - # Resize image_latent to match input_x dimensions - if image_latent.shape[2:] != input_x.shape[2:]: - image_latent = torch.nn.functional.interpolate( - image_latent, - size=(input_x.shape[2], input_x.shape[3]), - mode='bilinear', - align_corners=False - ) - - noise = torch.randn_like(image_latent) - noised_image_latent = image_latent + noise * current_sigma - - noised_image_latent = (1 - noise_weight) * input_x + noise_weight * noised_image_latent - - # Get noise prediction on the image latent - image_args = args_dict.copy() - image_args["input"] = noised_image_latent - combined_pred = apply_model_func(noised_image_latent, timestep, **c) - - # Run original model to get base noise prediction - #original_noise_pred = apply_model_func(input_x, timestep, **c) - - - - # Combine predictions - #combined_pred = (1 - noise_weight) * original_noise_pred + noise_weight * image_noise_pred - #if callback is not None: - # callback(combined_pred) - return combined_pred - - # Fall back to original model if we can't determine sigma - return apply_model_func(input_x, timestep, **c) - - return wrapper - -def sumweights(weights): - weight_sum = sum(weights) - - # Normalize weights if sum is not zero - if weight_sum > 0: - weights = [w / weight_sum for w in weights] - return weights +from __future__ import annotations +import folder_paths +import comfy.utils +import torch +import numpy as np +import os +from PIL import Image +import torchvision.transforms.functional as F +import math +from comfy import model_management +from comfy import latent_formats + + +def ensure_model(model_patcher, mode = None): + if not mode or mode == "disable": + # Ensure model mode disabled. + return + mp = model_patcher + lm = model_management.LoadedModel(mp) + found_lm = None + for list_lm in model_management.current_loaded_models: + if lm == list_lm: + found_lm = list_lm + break + if getattr(found_lm, "currently_used", False): + # Model already exists and appears to be loaded. + return + if found_lm is not None: + lm = found_lm + mp = lm.model + if mode.startswith("normal"): + if mode == "normal_unload": + model_management.unload_all_models() + model_management.load_models_gpu((mp,)) + return + model_management.unload_all_models() + if mode == "lowvram": + # Logic from comfy.model_management.load_models_gpu + min_inference_memory = model_management.minimum_inference_memory() + minimum_memory_required = max( + min_inference_memory, + model_management.extra_reserved_memory(), + ) + loaded_memory = lm.model_loaded_memory() + current_free_mem = model_management.get_free_memory(lm.device) + loaded_memory + + lowvram_model_memory = max( + 64 * 1024 * 1024, + (current_free_mem - minimum_memory_required), + min( + current_free_mem * model_management.MIN_WEIGHT_MEMORY_RATIO, + current_free_mem - model_management.minimum_inference_memory(), + ), + ) + lowvram_model_memory = max(0.1, lowvram_model_memory - loaded_memory) + elif mode == "novram": + lowvram_model_memory = 0.1 + else: + raise ValueError("Bad ensure_model_mode") + lm.model_load(lowvram_model_memory) + model_management.current_loaded_models.insert(0, lm) + + +def fallback(val, default, *, exclude=None, default_is_fun=False): + return val if val is not exclude else (default() if default_is_fun else default) + + +def scale_dim(n, factor=1.0, *, increment=64) -> int: + return math.ceil((n * factor) / increment) * increment + + +def sigma_to_float(sigma): + return sigma.detach().cpu().max().item() + +def load_taesd(name): + sd = {} + approx_vaes = folder_paths.get_filename_list("vae_approx") + + encoder = next(filter(lambda a: a.startswith("{}_encoder.".format(name)), approx_vaes)) + decoder = next(filter(lambda a: a.startswith("{}_decoder.".format(name)), approx_vaes)) + + enc = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", encoder)) + for k in enc: + sd["taesd_encoder.{}".format(k)] = enc[k] + + dec = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", decoder),device=model_management.get_torch_device()) + for k in dec: + sd["taesd_decoder.{}".format(k)] = dec[k] + + if name == "taesd": + sd["vae_scale"] = torch.tensor(0.18215) + sd["vae_shift"] = torch.tensor(0.0) + elif name == "taesdxl": + sd["vae_scale"] = torch.tensor(0.13025) + sd["vae_shift"] = torch.tensor(0.0) + elif name == "taesd3": + sd["vae_scale"] = torch.tensor(1.5305) + sd["vae_shift"] = torch.tensor(0.0609) + elif name == "taef1": + sd["vae_scale"] = torch.tensor(0.3611) + sd["vae_shift"] = torch.tensor(0.1159) + + vae = comfy.sd.VAE(sd=sd) + return vae + +def load_vae(name): + vae = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae", name),device=model_management.get_torch_device()) + vae = comfy.sd.VAE(sd=vae) + return vae + +def create_depth_map(latent_image, depth_model, resolution, **kwargs): + batch_size = latent_image.shape[0] + out_tensor = None + for i, image in enumerate(latent_image): + np_image = np.asarray(image.cpu() * 255., dtype=np.uint8) + np_result = depth_model(np_image, output_type="np", detect_resolution=resolution, **kwargs) + out = torch.from_numpy(np_result.astype(np.float32) / 255.0) + if out_tensor is None: + out_tensor = torch.zeros(batch_size, *out.shape, dtype=torch.float32) + out_tensor[i] = out + + return out_tensor + +def save_debug_image(tensor, filename, scale=True): + try: + # Ensure tensor is detached and on CPU + tensor = tensor.detach().cpu().float() + print(tensor.shape) + # Remove batch dimension if present + if tensor.ndim == 4 and tensor.shape[0] == 1: + tensor = tensor.squeeze(0) + + # Handle 2D tensors (grayscale masks) by converting to RGB + if tensor.ndim == 2: + # Convert 2D tensor to 3D by repeating across 3 channels + tensor = tensor.unsqueeze(-1).repeat(1, 1, 3) + + # If scale is True, scale to [0, 255] + if scale: + array = (tensor.numpy() * 255).clip(0, 255).astype('uint8') + else: + array = tensor.numpy().clip(0, 255).astype('uint8') + # Ensure shape is (H, W, 3) + if array.shape[-1] != 3: + raise ValueError(f"Expected last dimension to be 3 (RGB), got {array.shape[-1]}") + img = Image.fromarray(array) + + # Ensure debug directory exists + os.makedirs("debug", exist_ok=True) + + img.save("debug/" + filename) + return True + except Exception as e: + print(f"Error saving debug image {filename}: {e}") + return False + +def process_mask(mask, H, W, device): + # Convert to tensor if not already + if not isinstance(mask, torch.Tensor): + mask = torch.tensor(mask, device=device, dtype=torch.float32) + mask = mask.to(device).float() + + # Squeeze singleton dimensions + mask = mask.squeeze() + + # If mask is RGB, convert to grayscale + if mask.ndim == 3 and mask.shape[-1] == 3: + # [H, W, 3] -> [3, H, W] + mask = mask.permute(2, 0, 1) + mask = mask.mean(dim=0) # [H, W] + elif mask.ndim == 4 and mask.shape[-1] == 3: + # [B, H, W, 3] -> [B, 3, H, W] + mask = mask.permute(0, 3, 1, 2) + mask = mask.mean(dim=1) # [B, H, W] + mask = mask.squeeze(0) # Remove batch if present + + # Now mask should be [H, W] + if mask.shape != (H, W): + mask = torch.nn.functional.interpolate(mask.unsqueeze(0).unsqueeze(0), size=(H, W), mode='bilinear')[0, 0] + + return mask + + +def create_adaptive_low_res_model_wrapper( min_scale_factor=0.5, max_scale_factor=1.0, start_step=0, end_step=None): + """ + Creates a model function wrapper that runs the model at progressively higher resolution as sampling proceeds. + + Args: + min_scale_factor: Minimum scale factor to use (for early steps) + max_scale_factor: Maximum scale factor to use (for final steps) + start_step: Step number to start increasing resolution (0 means immediately) + end_step: Step number when scaling should reach max_scale_factor (None means the last step) + + Returns: + A function wrapper that can be used as model_options['model_function_wrapper'] + """ + def wrapper(apply_model_func, args_dict): + input_x = args_dict["input"] + timestep = args_dict["timestep"] + c = args_dict["c"] + + # Get sampling progress from sigmas in transformer_options + sample_sigmas = c['transformer_options'].get('sample_sigmas', None) + current_sigma = c['transformer_options'].get('sigmas', None) + + # Determine scale factor based on current step + if sample_sigmas is not None and current_sigma is not None: + # Get total steps from the length of sample_sigmas + total_steps = len(sample_sigmas) + + # Find current step by matching current_sigma to sample_sigmas + current_sigma_val = float(current_sigma[0]) + + # Find the index of the closest sigma value in sample_sigmas + # This gives us the current step number + sigma_diffs = [abs(float(s) - current_sigma_val) for s in sample_sigmas] + current_step = sigma_diffs.index(min(sigma_diffs)) + + # If end_step is not provided, use the last step + actual_end_step = end_step if end_step is not None else (total_steps - 1) + + # Clamp start and end steps to valid range + start_step_clamped = max(0, min(start_step, total_steps - 1)) + end_step_clamped = max(start_step_clamped, min(actual_end_step, total_steps - 1)) + + # Calculate scale factor based on current step + if current_step < start_step_clamped: + # Before start step - use minimum scale + scale_factor = min_scale_factor + elif current_step > end_step_clamped: + # After end step - use maximum scale + scale_factor = max_scale_factor + else: + # In between - interpolate between min and max + step_range = end_step_clamped - start_step_clamped + if step_range > 0: + progress = (current_step - start_step_clamped) / step_range + scale_factor = min_scale_factor + progress * (max_scale_factor - min_scale_factor) + else: + # If start_step equals end_step + scale_factor = min_scale_factor + else: + # Fallback if we can't determine steps + scale_factor = max_scale_factor + + # Original shape + original_shape = input_x.shape + + # Skip resizing if we're at full resolution + if scale_factor >= 0.99: + return apply_model_func(input_x, timestep, **c) + print(f"Scaling down to {scale_factor}x for step {current_step} using bislerp") + # Calculate new dimensions + new_h = int(original_shape[2] * scale_factor) + new_w = int(original_shape[3] * scale_factor) + + # Ensure dimensions are at least 8 (minimum size for stable diffusion) + new_h = max(8, new_h) + new_w = max(8, new_w) + + # Downscale using comfy.utils.common_upscale with bislerp mode + downscaled_x = comfy.utils.common_upscale( + input_x, + new_w, + new_h, + 'bislerp', + 'disabled' + ) + + # Run model at lower resolution + output = apply_model_func(downscaled_x, timestep, **c) + + # Upscale result back to original resolution using bislerp + upscaled_output = comfy.utils.common_upscale( + output, + original_shape[3], # width + original_shape[2], # height + 'bislerp', + 'disabled' + ) + + """ # #blurring the output + # Apply a small amount of Gaussian blur to smooth the upscaled output + # This helps reduce artifacts that can appear from the upscaling process + kernel_size = 3 # Small kernel for subtle blurring + sigma = 0.5 # Low sigma for gentle blur + + # Create Gaussian kernel + kernel = torch.tensor([ + [0.0625, 0.125, 0.0625], + [0.125, 0.25, 0.125], + [0.0625, 0.125, 0.0625] + ], device=upscaled_output.device).view(1, 1, kernel_size, kernel_size) + + Expand kernel for all channels + channels = upscaled_output.shape[1] + kernel = kernel.expand(channels, 1, kernel_size, kernel_size) + + # Group kernel for batch processing + kernel = kernel.repeat(1, 1, 1, 1) + + # Apply convolution for blurring + # Use padding='same' to maintain dimensions + padding = kernel_size // 2 + blurred_output = torch.nn.functional.conv2d( + upscaled_output, + kernel.to(upscaled_output.dtype), + padding=padding, + groups=channels + ) + + # Use the blurred output instead of the direct upscaled output + upscaled_output = blurred_output """ + return upscaled_output + + return wrapper + +def create_noise_injection_wrapper(latent, noise_weight=1.0, callback=None): + """ + Creates a wrapper that takes an input image, converts it to latent space, + noises it to match the current timestep/sigma, and injects its noise prediction. + + Args: + image: Input image to inject (tensor in pixel space) + vae: VAE model for encoding to latent space + noise_weight: Weight of the injected noise prediction (0-1) + + Returns: + A function wrapper to use as model_options['model_function_wrapper'] + + + """ + def wrapper(apply_model_func, args_dict): + # Extract args + input_x = args_dict["input"] + timestep = args_dict["timestep"] + c = args_dict["c"] + + + # Get current sampling parameters + sigmas = c['transformer_options'].get('sigmas', None) + sigma_idx = c['transformer_options'].get('sigma_idx', 0) + current_sigma = sigmas[sigma_idx] if sigmas is not None else None + + # Noise the image latent to match current noise level + if current_sigma is not None: + # Move tensors to the same device as input_x + device = input_x.device + image_latent = latent.to(device) + current_sigma = current_sigma.to(device) + + image_latent = latent_formats.SDXL().process_in(image_latent) + + # Resize image_latent to match input_x dimensions + if image_latent.shape[2:] != input_x.shape[2:]: + image_latent = torch.nn.functional.interpolate( + image_latent, + size=(input_x.shape[2], input_x.shape[3]), + mode='bilinear', + align_corners=False + ) + + noise = torch.randn_like(image_latent) + noised_image_latent = image_latent + noise * current_sigma + + noised_image_latent = (1 - noise_weight) * input_x + noise_weight * noised_image_latent + + # Get noise prediction on the image latent + image_args = args_dict.copy() + image_args["input"] = noised_image_latent + combined_pred = apply_model_func(noised_image_latent, timestep, **c) + + # Run original model to get base noise prediction + #original_noise_pred = apply_model_func(input_x, timestep, **c) + + + + # Combine predictions + #combined_pred = (1 - noise_weight) * original_noise_pred + noise_weight * image_noise_pred + #if callback is not None: + # callback(combined_pred) + return combined_pred + + # Fall back to original model if we can't determine sigma + return apply_model_func(input_x, timestep, **c) + + return wrapper + +def create_model_wrapper(type, vae15, vaexl): + """ + Creates a wrapper that takes an input image, converts it to latent space, + noises it to match the current timestep/sigma, and injects its noise prediction. + + Args: + image: Input image to inject (tensor in pixel space) + vae: VAE model for encoding to latent space + noise_weight: Weight of the injected noise prediction (0-1) + + Returns: + A function wrapper to use as model_options['model_function_wrapper'] + + + """ + def wrapper(apply_model_func, args_dict): + # Extract args + input_x = args_dict["input"] + timestep = args_dict["timestep"] + c = args_dict["c"] + + # Noise the image latent to match current noise level + if type == "sd15": + # Get device from input_x + device = input_x.device + + image_latent = convert_latent_format(input_x, "sdxl", "sd15", device) + + # + model_output = apply_model_func(image_latent, timestep, **c) + model_output = model_output.to(device) + + final_output = convert_latent_format(model_output, "sd15", "sdxl", device) + + return final_output + + # Fall back to original model if type is not SD15 + return apply_model_func(input_x, timestep, **c) + + return wrapper + +def sumweights(weights): + weight_sum = sum(weights) + + # Normalize weights if sum is not zero + if weight_sum > 0: + weights = [w / weight_sum for w in weights] + return weights + +def convert_latent_format(latent, source_format="sd15", target_format="sdxl", device=None): + """ + Converts latents between SD1.5 and SDXL formats using their RGB factors and scales. + + Args: + latent: Input latent tensor (B, 4, H, W) + source_format: Source format ("sd15" or "sdxl") + target_format: Target format ("sd15" or "sdxl") + device: Device to use (if None, uses the latent's device) + + Returns: + Converted latent tensor (B, 4, H, W) + """ + if source_format == target_format: + return latent + + if device is None: + device = latent.device + + # Define format parameters + sd15_scale = 0.18215 + sdxl_scale = 0.13025 + + sd15_factors = torch.tensor([ + [ 0.3512, 0.2297, 0.3227], + [ 0.3250, 0.4974, 0.2350], + [-0.2829, 0.1762, 0.2721], + [-0.2120, -0.2616, -0.7177] + ], device=device) + + sdxl_factors = torch.tensor([ + [ 0.3651, 0.4232, 0.4341], + [-0.2533, -0.0042, 0.1068], + [ 0.1076, 0.1111, -0.0362], + [-0.3165, -0.2492, -0.2188] + ], device=device) + + sdxl_bias = torch.tensor([0.1084, -0.0175, -0.0011], device=device) + + # Unscale the latent based on source format + if source_format == "sd15": + unscaled_latent = latent * sd15_scale + else: # sdxl + unscaled_latent = latent * sdxl_scale + + # Get dimensions + B, C, H, W = unscaled_latent.shape + + # For simplicity, we'll process each sample in the batch separately + target_latent = torch.zeros_like(unscaled_latent) + + for b in range(B): + # Get the current sample + sample = unscaled_latent[b] # [4, H, W] + + # Reshape to [4, H*W] + sample_flat = sample.reshape(4, -1) + + # Convert to RGB representation + if source_format == "sd15": + # Calculate pseudo-inverse (4x3 matrix becomes 3x4) + sd15_pinv = torch.pinverse(sd15_factors) + # Transform from latent to RGB: [3x4] x [4, H*W] -> [3, H*W] + rgb = torch.matmul(sd15_pinv, sample_flat) + else: # sdxl + sdxl_pinv = torch.pinverse(sdxl_factors) + # Transform from latent to RGB: [3x4] x [4, H*W] -> [3, H*W] + rgb = torch.matmul(sdxl_pinv, sample_flat) + # Remove bias (for each channel) + for i in range(3): + rgb[i] = rgb[i] - sdxl_bias[i] + + # Convert RGB to target latent space + if target_format == "sd15": + # RGB to SD1.5 latent: [4x3] x [3, H*W] -> [4, H*W] + out = torch.matmul(sd15_factors, rgb) + # Apply scale + out = out / sd15_scale + else: # sdxl + # Add bias + for i in range(3): + rgb[i] = rgb[i] + sdxl_bias[i] + # RGB to SDXL latent: [4x3] x [3, H*W] -> [4, H*W] + out = torch.matmul(sdxl_factors, rgb) + # Apply scale + out = out / sdxl_scale + + # Reshape back to [4, H, W] and store + target_latent[b] = out.reshape(4, H, W) + + return target_latent