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 normalize_model_predictions(pred, model_types, sigma=None): """ Normalize predictions between different model types (flow, eps, v-prediction). Args: pred: List of model predictions, each containing [conditional, unconditional] model_types: List of model types ("FLOW", "EPS", "V_PREDICTION", etc.) sigma: Current noise level (required for proper scaling) Returns: normalized_pred: Normalized predictions that can be mixed together """ normalized_pred = [] if sigma is None: print("Warning: sigma not provided for model prediction normalization. Using default scaling.") sigma = torch.tensor(1.0, device=pred[0][0].device) for i, (cond, uncond) in enumerate(pred): model_type = model_types[i].lower() if i < len(model_types) else "eps" if model_type == "flow": # Flow models predict score function (gradient of log-likelihood) # For proper mixing with noise-prediction models, we use: # noise = -score * sigma sigma_factor = sigma.view(sigma.shape[:1] + (1,) * (cond.ndim - 1)) # Convert score to noise prediction with correct scaling relationship normalized_cond = -cond * sigma_factor normalized_uncond = -uncond * sigma_factor elif model_type == "v_prediction": # V-prediction models predict velocity # v = (x_t - α_t·x_0)/σ_t # To convert v-prediction to noise prediction (ε): # ε = v * σ sigma_factor = sigma.view(sigma.shape[:1] + (1,) * (cond.ndim - 1)) normalized_cond = cond * sigma_factor normalized_uncond = uncond * sigma_factor else: # EPS is the default # EPS models predict noise directly - use as is normalized_cond = cond normalized_uncond = uncond normalized_pred.append([normalized_cond, normalized_uncond]) return normalized_pred 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