495 lines
20 KiB
Python
495 lines
20 KiB
Python
import torch
|
|
|
|
from comfy import latent_formats
|
|
from ..utils import create_depth_map, save_debug_image, process_mask
|
|
import comfy.sd
|
|
import numpy as np
|
|
import comfy.samplers
|
|
from ..utils import sumweights
|
|
|
|
|
|
def mix_standard(pred, weights, cfgs, uncond_decays=[1.0]):
|
|
"""
|
|
Standard mixing function that applies CFG for each model based on weights.
|
|
|
|
Args:
|
|
pred: List of model predictions, each containing [conditional, unconditional]
|
|
weights: List of weights for each model
|
|
cfgs: List of cfg values for each model
|
|
|
|
Returns:
|
|
mix: The mixed prediction
|
|
postuncond: The unconditioned prediction for post-processing
|
|
"""
|
|
print("mix_standard")
|
|
# Initialize with the shape of the first model's prediction
|
|
mix = torch.zeros_like(pred[0][1])
|
|
|
|
|
|
# First, combine all unconditional predictions based on weights
|
|
for i in range(len(pred)):
|
|
mix = mix + pred[i][1] * weights[i]
|
|
|
|
# Save the combined unconditional prediction for post-processing
|
|
postuncond = mix.clone()
|
|
|
|
# For each model, apply the guidance according to weights and cfg values
|
|
for i in range(len(pred)):
|
|
# Extract conditional and unconditional predictions for this model
|
|
c_pred = pred[i][0]
|
|
u_pred = pred[i][1]
|
|
|
|
# Apply the CFG formula with the corresponding weight and cfg
|
|
# Add the weighted difference between conditional and unconditional
|
|
mix = mix + (c_pred - u_pred) * cfgs[i]
|
|
|
|
return mix, postuncond
|
|
|
|
def mix_masked(pred, weights, cfgs, masks=None):
|
|
"""
|
|
Mask-based mixing function that applies different models to different regions.
|
|
|
|
Args:
|
|
pred: List of model predictions, each containing [conditional, unconditional]
|
|
weights: List of weights for each model
|
|
cfgs: List of cfg values for each model
|
|
masks: List of masks (one per model) - should already be properly formatted/feathered
|
|
|
|
Returns:
|
|
mix: The mixed prediction
|
|
postuncond: The unconditioned prediction for post-processing
|
|
"""
|
|
print("mix_masked")
|
|
# Check if we have valid mask in the list
|
|
#if not any(mask is not None for mask in masks):
|
|
# # Fall back to standard mixing if no valid masks
|
|
# return mix_standard(pred, weights, cfgs)
|
|
|
|
# Initialize with zeros
|
|
B, C, H, W = pred[0][0].shape
|
|
mix = torch.zeros_like(pred[0][0])
|
|
postuncond = torch.zeros_like(pred[0][1])
|
|
# Prepare masks
|
|
processed_masks = []
|
|
for i, mask in enumerate(masks):
|
|
try:
|
|
if(mask is not None):
|
|
processed_mask = process_mask(mask, H, W, device=pred[0][0].device)
|
|
processed_masks.append(processed_mask)
|
|
else:
|
|
processed_masks.append(torch.ones((H, W), device=pred[0][0].device))
|
|
except Exception as e:
|
|
print(f"Error processing mask {i}: {e}")
|
|
processed_masks.append(torch.zeros((H, W), device=pred[0][0].device))
|
|
|
|
# If we couldn't process any masks, fall back to standard mixing
|
|
if len(processed_masks) == 0:
|
|
return mix_standard(pred, weights, cfgs)
|
|
|
|
# If we have fewer masks than models, pad with zeros
|
|
while len(processed_masks) < len(pred):
|
|
processed_masks.append(torch.ones_like(processed_masks[0], device=pred[0][0].device))
|
|
|
|
print(len(processed_masks))
|
|
# Stack and normalize masks
|
|
stacked_masks = torch.stack(processed_masks[:len(pred)])
|
|
mask_sum = torch.sum(stacked_masks, dim=0, keepdim=True)
|
|
mask_sum = torch.clamp(mask_sum, min=1e-6) # Avoid division by zero
|
|
normalized_masks = stacked_masks / mask_sum
|
|
|
|
# Apply each model according to its normalized mask
|
|
for i in range(len(pred)):
|
|
if i >= len(normalized_masks):
|
|
continue
|
|
|
|
# Get mask for this model and add batch and channel dimensions
|
|
model_mask = normalized_masks[i].unsqueeze(0).unsqueeze(0).repeat(B, C, 1, 1)
|
|
|
|
# Extract predictions
|
|
c_pred = pred[i][0]
|
|
u_pred = pred[i][1]
|
|
|
|
# Apply CFG
|
|
model_result = u_pred + (c_pred - u_pred) * cfgs[i]
|
|
|
|
# Add to result based on mask
|
|
mix = mix + model_result * model_mask
|
|
postuncond = postuncond + u_pred * model_mask
|
|
|
|
return mix, postuncond
|
|
|
|
def mix_dynamic_mask(pred, weights, cfgs, step=0, history=None, decay_factor=0.9, threshold=0.1, blur_sigma=0.5, max_update_step=15):
|
|
"""
|
|
Dynamic mask-based mixing that evolves masks based on prediction magnitudes over time.
|
|
|
|
Args:
|
|
pred: List of model predictions, each containing [conditional, unconditional]
|
|
weights: List of weights for each model
|
|
cfgs: List of cfg values for each model
|
|
step: Current generation step (used for mask evolution)
|
|
history: Previous mask history to evolve from (None for first step)
|
|
decay_factor: How much to retain from previous mask (0-1)
|
|
threshold: Threshold for magnitude difference to affect mask
|
|
blur_sigma: Standard deviation for Gaussian blur applied to masks (0 for no blur)
|
|
max_update_step: If set, stops updating masks after this step and uses existing masks
|
|
|
|
Returns:
|
|
mix: The mixed prediction
|
|
postuncond: The unconditioned prediction for post-processing
|
|
masks: The updated masks to pass to the next step
|
|
"""
|
|
# Get dimensions from the first prediction
|
|
B, C, H, W = pred[0][0].shape
|
|
device = pred[0][0].device
|
|
# Initialize output tensors
|
|
mix = torch.zeros_like(pred[0][0])
|
|
postuncond = torch.zeros_like(pred[0][1])
|
|
|
|
# Initialize or update masks
|
|
if history is None or step == 0:
|
|
# First step: initialize with equal masks or based on weights
|
|
masks = []
|
|
for i in range(len(pred)):
|
|
masks.append(torch.ones((H, W), device=device) * weights[i])
|
|
else:
|
|
# Use previous masks as starting point
|
|
masks = history
|
|
# Save masks to debug folder
|
|
for i, mask in enumerate(masks):
|
|
save_debug_image(mask, f"mask_{i}.png")
|
|
|
|
# Check if we should update masks based on max_update_step
|
|
should_update_masks = max_update_step is None or step <= max_update_step
|
|
|
|
if should_update_masks:
|
|
# Calculate magnitudes for each model's prediction
|
|
magnitudes = []
|
|
for i in range(len(pred)):
|
|
# Use the difference between conditional and unconditional as a measure of "importance"
|
|
model_contrib = torch.abs(pred[i][0] - pred[i][1])
|
|
# Average across channels to get a 2D magnitude map
|
|
magnitude = torch.mean(model_contrib, dim=1)[0] # Take first batch
|
|
magnitudes.append(magnitude)
|
|
|
|
# Stack magnitudes for comparison
|
|
stacked_magnitudes = torch.stack(magnitudes)
|
|
|
|
# Find which model has highest magnitude at each pixel
|
|
max_indices = torch.argmax(stacked_magnitudes, dim=0)
|
|
|
|
# Update masks based on magnitudes
|
|
updated_masks = []
|
|
for i in range(len(pred)):
|
|
# Create binary mask where this model has highest magnitude
|
|
model_dominant = (max_indices == i).float()
|
|
|
|
# Apply threshold to only update where difference is significant
|
|
magnitude_threshold = torch.max(stacked_magnitudes) * threshold
|
|
significant_diff = (stacked_magnitudes[i] > magnitude_threshold).float()
|
|
|
|
# Combine with threshold
|
|
update_mask = model_dominant * significant_diff
|
|
|
|
# Evolve mask gradually (blend old and new)
|
|
if i < len(masks):
|
|
# Decay old mask and add new information
|
|
new_mask = masks[i] * decay_factor + update_mask * (1 - decay_factor)
|
|
updated_masks.append(new_mask)
|
|
else:
|
|
updated_masks.append(update_mask)
|
|
|
|
# Use the updated masks
|
|
mask_list = updated_masks
|
|
else:
|
|
# Just use the existing masks without updating
|
|
if isinstance(masks, list):
|
|
mask_list = masks
|
|
else:
|
|
# Convert tensor masks back to list if needed
|
|
mask_list = [masks[i] for i in range(masks.shape[0])]
|
|
|
|
# Normalize masks to sum to 1 at each pixel
|
|
stacked_masks = torch.stack(mask_list)
|
|
mask_sum = torch.sum(stacked_masks, dim=0, keepdim=True)
|
|
mask_sum = torch.clamp(mask_sum, min=1e-6) # Avoid division by zero
|
|
normalized_masks = stacked_masks / mask_sum
|
|
|
|
# Apply Gaussian blur to the masks if blur_sigma > 0
|
|
if blur_sigma > 0:
|
|
import torch.nn.functional as F
|
|
import math
|
|
blurred_masks = []
|
|
|
|
# Create a custom Gaussian blur function since F.gaussian_blur isn't available
|
|
def custom_gaussian_blur(input_tensor, kernel_size, sigma):
|
|
# Create Gaussian kernel
|
|
channels = input_tensor.shape[1]
|
|
kernel_size = [kernel_size, kernel_size]
|
|
sigma = [sigma, sigma]
|
|
|
|
# Create a grid of coordinates
|
|
x_coord = torch.arange(kernel_size[0], device=input_tensor.device)
|
|
x_grid = x_coord.repeat(kernel_size[1]).view(kernel_size[1], kernel_size[0])
|
|
y_grid = x_grid.t()
|
|
|
|
xy_grid = torch.stack([x_grid, y_grid], dim=-1).float()
|
|
|
|
# Calculate the center of the filter
|
|
mean = (torch.tensor(kernel_size, device=input_tensor.device) - 1) / 2.
|
|
variance = torch.tensor(sigma, device=input_tensor.device) ** 2
|
|
|
|
# Calculate the 2d gaussian kernel
|
|
gaussian_kernel = torch.exp(
|
|
-torch.sum((xy_grid - mean) ** 2., dim=-1) / (2 * variance[0])
|
|
)
|
|
|
|
# Make sure sum of values in gaussian kernel equals 1.
|
|
gaussian_kernel = gaussian_kernel / torch.sum(gaussian_kernel)
|
|
|
|
# Reshape to 2d depth-wise convolutional weight
|
|
gaussian_kernel = gaussian_kernel.view(1, 1, kernel_size[0], kernel_size[1])
|
|
gaussian_kernel = gaussian_kernel.repeat(channels, 1, 1, 1)
|
|
|
|
# Apply padding
|
|
pad_size = (kernel_size[0] // 2, kernel_size[1] // 2)
|
|
padded_input = F.pad(input_tensor, (pad_size[0], pad_size[0], pad_size[1], pad_size[1]), mode='reflect')
|
|
|
|
# Apply Gaussian blur using depthwise convolution
|
|
return F.conv2d(padded_input, gaussian_kernel, groups=channels)
|
|
|
|
for i in range(len(normalized_masks)):
|
|
mask = normalized_masks[i].unsqueeze(0).unsqueeze(0) # Add batch and channel dimensions
|
|
# Apply Gaussian blur (kernel size should be odd and >= 3)
|
|
kernel_size = max(3, int(blur_sigma * 4 + 0.5) * 2 + 1)
|
|
blurred_mask = custom_gaussian_blur(mask, kernel_size, blur_sigma)
|
|
blurred_masks.append(blurred_mask.squeeze(0).squeeze(0))
|
|
|
|
# Re-normalize after blurring to ensure they still sum to 1
|
|
stacked_blurred_masks = torch.stack(blurred_masks)
|
|
blurred_sum = torch.sum(stacked_blurred_masks, dim=0, keepdim=True)
|
|
blurred_sum = torch.clamp(blurred_sum, min=1e-6)
|
|
normalized_masks = stacked_blurred_masks / blurred_sum
|
|
|
|
# Apply each model according to its normalized mask
|
|
for i in range(len(pred)):
|
|
# Get mask for this model and add batch and channel dimensions
|
|
model_mask = normalized_masks[i].unsqueeze(0).unsqueeze(0).repeat(B, C, 1, 1)
|
|
|
|
# Extract predictions
|
|
c_pred = pred[i][0]
|
|
u_pred = pred[i][1]
|
|
|
|
# Apply CFG
|
|
model_result = u_pred + (c_pred - u_pred) * cfgs[i]
|
|
|
|
# Add to result based on mask
|
|
mix = mix + model_result * model_mask
|
|
postuncond = postuncond + u_pred * model_mask
|
|
|
|
# Return the mixed result, unconditioned prediction, and updated masks for next step
|
|
return mix, postuncond, normalized_masks.detach()
|
|
|
|
def mix_dynamic_mask_alternative(pred, weights, cfgs, depth_model=None, threshold=0.5, invert=False, blur_sigma=1.0):
|
|
"""
|
|
Mixing function that automatically detects foreground/background and applies different models to each.
|
|
|
|
Args:
|
|
pred: List of model predictions, each containing [conditional, unconditional]
|
|
weights: List of weights for each model
|
|
cfgs: List of cfg values for each model
|
|
depth_model: Optional depth estimation model. If None, uses simple heuristics
|
|
threshold: Threshold for foreground/background separation (0-1)
|
|
invert: If True, inverts the mask (foreground becomes background)
|
|
blur_sigma: Blur amount to apply to mask edges
|
|
|
|
Returns:
|
|
mix: The mixed prediction
|
|
postuncond: The unconditioned prediction for post-processing
|
|
"""
|
|
# Need at least 2 models to do foreground/background mixing
|
|
if len(pred) < 2:
|
|
return mix_standard(pred, weights, cfgs)
|
|
|
|
# Get dimensions from the first prediction
|
|
B, C, H, W = pred[0][0].shape
|
|
device = pred[0][0].device
|
|
|
|
# Generate foreground/background mask using prediction features
|
|
if depth_model is None:
|
|
# Simple heuristic: use the magnitude of the conditional prediction
|
|
# as an approximation of feature importance
|
|
c_pred = pred[0][0] # Use first model's conditional prediction
|
|
|
|
# Calculate magnitude across channels (approximate depth/saliency)
|
|
magnitude = torch.mean(torch.abs(c_pred), dim=1, keepdim=True)
|
|
|
|
# Normalize to 0-1 range
|
|
min_val = torch.min(magnitude)
|
|
max_val = torch.max(magnitude)
|
|
normalized = (magnitude - min_val) / (max_val - min_val + 1e-8)
|
|
|
|
# Apply threshold
|
|
mask = (normalized > threshold).float()
|
|
|
|
# Invert if requested
|
|
if invert:
|
|
mask = 1.0 - mask
|
|
else:
|
|
# Use provided depth model (not implemented here, would need integration)
|
|
# This would replace the simple heuristic above with actual depth estimation
|
|
mask = torch.ones((B, 1, H, W), device=device) * 0.5
|
|
# mask = depth_model(c_pred)
|
|
|
|
# Apply Gaussian blur to smooth mask edges
|
|
if blur_sigma > 0:
|
|
import torch.nn.functional as F
|
|
|
|
# Create a Gaussian kernel
|
|
kernel_size = max(3, int(blur_sigma * 4 + 0.5) * 2 + 1)
|
|
|
|
# Apply Gaussian blur
|
|
mask = F.pad(mask, (kernel_size//2, kernel_size//2, kernel_size//2, kernel_size//2), mode='reflect')
|
|
mask = F.avg_pool2d(mask, kernel_size, stride=1, padding=0)
|
|
|
|
# Create two masks: foreground and background
|
|
mask_fg = mask
|
|
mask_bg = 1.0 - mask
|
|
|
|
# Ensure masks have the correct shape (H, W) without extra dimensions
|
|
if mask_fg.dim() > 2:
|
|
# Extract a 2D mask, squeezing out extra dimensions
|
|
if mask_fg.dim() == 3: # [C, H, W] or [1, H, W]
|
|
mask_fg = mask_fg.squeeze(0)
|
|
elif mask_fg.dim() == 4: # [B, C, H, W]
|
|
mask_fg = mask_fg.squeeze(0).squeeze(0)
|
|
|
|
# If mask still has more than 2 dimensions, use only the first 2D slice
|
|
if mask_fg.dim() > 2:
|
|
mask_fg = mask_fg[0]
|
|
|
|
# Do the same for background mask
|
|
if mask_bg.dim() > 2:
|
|
if mask_bg.dim() == 3:
|
|
mask_bg = mask_bg.squeeze(0)
|
|
elif mask_bg.dim() == 4:
|
|
mask_bg = mask_bg.squeeze(0).squeeze(0)
|
|
|
|
if mask_bg.dim() > 2:
|
|
mask_bg = mask_bg[0]
|
|
|
|
# Create list of masks for each model
|
|
masks = [mask_fg, mask_bg]
|
|
|
|
# Create directory if it doesn't exist
|
|
for i, mask in enumerate(masks):
|
|
save_debug_image(mask, f"mask_{i}.png")
|
|
|
|
# If more than 2 models, create zero masks for the rest
|
|
for i in range(2, len(pred)):
|
|
masks.append(torch.zeros_like(mask_fg))
|
|
|
|
# Use the existing masked mixing function with our generated masks
|
|
return mix_masked(pred, weights, cfgs, masks)
|
|
|
|
def mix_foreground_background(pred, weights, last_latent, cfgs, step=0, invert=False, masks=None, vae=None, sharpness=0.0, depth_model=None, startmask=5):
|
|
"""
|
|
Mixing function that uses Depth Anything V2 to separate foreground and background.
|
|
|
|
Args:
|
|
pred: List of model predictions, each containing [conditional, unconditional]
|
|
weights: List of weights for each model
|
|
cfgs: List of cfg values for each model
|
|
invert: If True, inverts the mask (foreground becomes background)
|
|
ckpt_name: The Depth Anything V2 checkpoint to use
|
|
threshold: Threshold for foreground/background separation (0-1)
|
|
blur_sigma: Blur amount to apply to mask edges
|
|
|
|
Returns:
|
|
mix: The mixed prediction
|
|
postuncond: The unconditioned prediction for post-processing
|
|
"""
|
|
|
|
# Get dimensions from the first prediction
|
|
B, C, H, W = pred[0][0].shape
|
|
#get the higes resolution of the latent image make sure its a power of 2
|
|
res = 2**(max(H, W).bit_length() - 1)
|
|
device = pred[0][0].device
|
|
|
|
# Get the unconditioned prediction from the first model to generate a depth map
|
|
u_pred = pred[0][1]
|
|
|
|
# Convert latent to image for depth estimation
|
|
# This is a simplified approach - assumes we're working with latents
|
|
# In a real implementation, you'd need proper latent-to-image conversion
|
|
|
|
if not any(mask is not None for mask in masks):
|
|
if step == startmask:
|
|
# Use Depth Anything V2 to generate a depth map
|
|
latent_image = last_latent
|
|
try:
|
|
import comfy.model_management as model_management
|
|
from custom_nodes.comfyui_controlnet_aux.utils import common_annotator_call
|
|
|
|
|
|
|
|
# Convert latent to proper image using TaeSDXL Mini VAE
|
|
try:
|
|
latent_image = latent_formats.SDXL().process_out(latent_image)
|
|
latent_image = vae.decode(latent_image)
|
|
|
|
except Exception as e:
|
|
print(f"Error using TAESD: {e}")
|
|
|
|
|
|
out_tensor = create_depth_map(latent_image, depth_model, res)
|
|
|
|
depth_norm = out_tensor
|
|
|
|
|
|
except Exception as e:
|
|
print(f"Error using Depth Anything V2: {e}")
|
|
print("Falling back to simple heuristic for foreground/background separation")
|
|
|
|
# Fallback: use the magnitude of the conditional prediction as approximation
|
|
c_pred = pred[0][0]
|
|
depth_norm = torch.mean(torch.abs(c_pred), dim=1, keepdim=True)
|
|
depth_norm = (depth_norm - depth_norm.min()) / (depth_norm.max() - depth_norm.min() + 1e-8)
|
|
mask = depth_norm * (1.0 + sharpness) - (depth_norm * sharpness).mean()
|
|
mask = torch.clamp(mask, 0.0, 1.0)
|
|
|
|
mask = mask
|
|
# Invert if requested
|
|
if invert:
|
|
mask = 1.0 - mask
|
|
|
|
# Create two masks: foreground and background
|
|
mask_fg = mask
|
|
mask_bg = 1.0 - mask
|
|
|
|
|
|
|
|
# Create list of masks for each model
|
|
masks = [mask_fg, mask_bg]
|
|
|
|
# If more than 2 models, create zero masks for the rest
|
|
for i in range(2, len(pred)):
|
|
masks.append(torch.zeros_like(mask_fg))
|
|
|
|
mix, postuncond = mix_masked(pred, weights, cfgs, masks)
|
|
# Debug: Save masks
|
|
try:
|
|
# Save various debug images
|
|
save_debug_image(latent_image, "last_latent.png")
|
|
|
|
save_debug_image(mask, "mask.png")
|
|
except Exception as e:
|
|
print(f"Failed to save debug masks: {e}")
|
|
else:
|
|
mix, postuncond = mix_standard(pred, weights, cfgs)
|
|
else:
|
|
mix, postuncond = mix_masked(pred, weights, cfgs, masks)
|
|
torch.cuda.empty_cache()
|
|
# Use the existing masked mixing function with our generated masks
|
|
return masks, mix, postuncond
|
|
|