AEmotionStudio
37bf766c1b
Fixes two issues where shape matching/correction handles channel dimensions before spatial dimensions, causing errors when both differ: - blending.py: _match_tensor_shape and _match_video_tensor_shape now handle spatial interpolation (Step 3) before channel copying (Step 4) - sampler.py: _correct_frame_shape and _correct_image_shape now handle spatial interpolation before channel copying The issue was that when creating torch.zeros_like(target) for channel adjustment, the target has target's spatial dimensions, but the source slice being copied has source's spatial dimensions. If both differ, the assignment fails with incompatible shapes. By handling spatial interpolation first, the source tensor's spatial dimensions match the target before any channel slicing occurs.
596 lines
21 KiB
Python
596 lines
21 KiB
Python
"""
|
|
Core sampling functionality for shader noise KSampler.
|
|
|
|
This module contains the main sampling logic, stage calculations,
|
|
and shader noise generation coordination.
|
|
"""
|
|
import torch
|
|
import math
|
|
import contextlib
|
|
from typing import Optional, Dict, Any, List, Tuple, Callable
|
|
|
|
import torch.nn.functional as F
|
|
|
|
from .constants import SUPPORTED_DISTRIBUTIONS
|
|
from .transforms import normalize_noise
|
|
from .params import ShaderParams
|
|
|
|
|
|
def calculate_stage_strengths(
|
|
base_strength: float,
|
|
num_stages: int,
|
|
distribution: str
|
|
) -> List[float]:
|
|
"""
|
|
Calculate strength for each stage based on distribution type.
|
|
|
|
Args:
|
|
base_strength: Base shader strength
|
|
num_stages: Number of stages
|
|
distribution: Distribution type (one of SUPPORTED_DISTRIBUTIONS)
|
|
|
|
Returns:
|
|
List of strength values for each stage
|
|
"""
|
|
if num_stages <= 0:
|
|
return []
|
|
|
|
strengths = []
|
|
|
|
if distribution == "uniform":
|
|
strengths = [base_strength] * num_stages
|
|
|
|
elif distribution == "linear_decrease":
|
|
min_factor = 0.25 # Ensure a minimum shader strength in each stage
|
|
for i in range(num_stages):
|
|
factor = 1.0 - (i / (num_stages - 1 if num_stages > 1 else 1))
|
|
factor = max(factor, min_factor)
|
|
strengths.append(base_strength * factor)
|
|
|
|
elif distribution == "linear_increase":
|
|
min_factor = 0.25 # Ensure a minimum shader strength (same as linear_decrease)
|
|
for i in range(num_stages):
|
|
factor = i / (num_stages - 1 if num_stages > 1 else 1)
|
|
# For single stage, use full strength instead of zero
|
|
if num_stages == 1:
|
|
factor = 1.0
|
|
factor = max(factor, min_factor)
|
|
strengths.append(base_strength * factor)
|
|
|
|
elif distribution == "gaussian":
|
|
# Create a bell curve with peak in the middle
|
|
mid_point = (num_stages - 1) / 2
|
|
for i in range(num_stages):
|
|
# Standard deviation as 1/3 of the range
|
|
std_dev = num_stages / 3
|
|
factor = math.exp(-((i - mid_point) ** 2) / (2 * std_dev ** 2))
|
|
strengths.append(base_strength * factor)
|
|
|
|
elif distribution == "first_stronger":
|
|
# First stage strongest, rapid falloff
|
|
for i in range(num_stages):
|
|
factor = math.exp(-i) # Exponential decay
|
|
strengths.append(base_strength * factor)
|
|
|
|
elif distribution == "last_stronger":
|
|
# Last stage strongest, exponential increase
|
|
for i in range(num_stages):
|
|
factor = math.exp(i - num_stages + 1) # Exponential increase
|
|
strengths.append(base_strength * factor)
|
|
|
|
else:
|
|
# Default to uniform if unknown distribution
|
|
strengths = [base_strength] * num_stages
|
|
|
|
return strengths
|
|
|
|
|
|
def calculate_step_ranges(total_steps: int, num_stages: int) -> List[Tuple[int, int]]:
|
|
"""
|
|
Calculate step ranges for sequential stages.
|
|
|
|
Args:
|
|
total_steps: Total number of sampling steps
|
|
num_stages: Number of stages
|
|
|
|
Returns:
|
|
List of (start_step, end_step) tuples
|
|
"""
|
|
if num_stages <= 0:
|
|
return []
|
|
|
|
step_ranges = []
|
|
step_size = total_steps / num_stages
|
|
|
|
for i in range(num_stages):
|
|
start_step = int(i * step_size)
|
|
end_step = int((i + 1) * step_size) if i < num_stages - 1 else total_steps
|
|
step_ranges.append((start_step, end_step))
|
|
|
|
return step_ranges
|
|
|
|
|
|
def calculate_step_points(total_steps: int, num_stages: int) -> List[int]:
|
|
"""
|
|
Calculate at which steps to apply different shader noises for injection stages.
|
|
|
|
Args:
|
|
total_steps: Total number of sampling steps
|
|
num_stages: Number of injection stages
|
|
|
|
Returns:
|
|
List of step points
|
|
"""
|
|
# Match behavior of calculate_stage_strengths and calculate_step_ranges
|
|
if num_stages <= 0:
|
|
return []
|
|
if num_stages == 1:
|
|
return [0]
|
|
|
|
step_points = []
|
|
for i in range(num_stages):
|
|
step_point = int(i * total_steps / (num_stages - 1)) if num_stages > 1 else 0
|
|
step_points.append(min(step_point, total_steps - 1))
|
|
|
|
return step_points
|
|
|
|
|
|
def generate_shader_noise(
|
|
latent_samples: torch.Tensor,
|
|
target_noise_shape: Tuple[int, ...],
|
|
shader_params: Dict[str, Any],
|
|
shader_type: str,
|
|
seed: int,
|
|
device: str = "cuda",
|
|
model: Optional[Any] = None,
|
|
model_name: Optional[str] = None,
|
|
frame_count: int = 1,
|
|
frame_dim_idx: int = -1,
|
|
generator_func: Optional[Callable] = None,
|
|
debugger: Optional[object] = None
|
|
) -> torch.Tensor:
|
|
"""
|
|
Generate noise using the specified shader.
|
|
|
|
Args:
|
|
latent_samples: The input latent samples (for context/shape)
|
|
target_noise_shape: Expected shape of the final noise tensor
|
|
shader_params: Parameters for shader generation
|
|
shader_type: Type of shader generator to use
|
|
seed: Random seed
|
|
device: Device to create tensor on
|
|
model: The model object (for channel detection)
|
|
model_name: Optional model name for customized generation
|
|
frame_count: Number of frames to generate
|
|
frame_dim_idx: Index of the frame dimension (-1 if not video)
|
|
generator_func: Shader generator function to use
|
|
debugger: Optional debugger for logging
|
|
|
|
Returns:
|
|
Generated noise tensor with shape matching target_noise_shape
|
|
"""
|
|
debug_enabled = debugger is not None and getattr(debugger, 'enabled', False)
|
|
debug_level = getattr(debugger, 'debug_level', 0) if debug_enabled else 0
|
|
|
|
# Set random seed for reproducibility
|
|
torch.manual_seed(seed)
|
|
|
|
# Extract shape information from the target_noise_shape
|
|
is_video = len(target_noise_shape) == 5
|
|
batch_size = target_noise_shape[0]
|
|
height = target_noise_shape[-2]
|
|
width = target_noise_shape[-1]
|
|
|
|
# Get channels from target_noise_shape
|
|
if is_video:
|
|
# Handle default frame_dim_idx of -1 by detecting format from shape
|
|
if frame_dim_idx == -1:
|
|
# Common latent channel counts vs frame counts
|
|
common_channels = {4, 8, 12, 16, 32, 64, 128}
|
|
dim1, dim2 = target_noise_shape[1], target_noise_shape[2]
|
|
# Heuristic: smaller dimension that looks like channels is the channel dim
|
|
if dim1 in common_channels and dim2 not in common_channels:
|
|
channel_dim_idx = 1 # B,C,F,H,W
|
|
frame_dim_idx = 2
|
|
elif dim2 in common_channels and dim1 not in common_channels:
|
|
channel_dim_idx = 2 # B,F,C,H,W
|
|
frame_dim_idx = 1
|
|
else:
|
|
# Default assumption: B,F,C,H,W format (most common in video models)
|
|
# This must match the default in blending.py _detect_video_format()
|
|
channel_dim_idx = 2
|
|
frame_dim_idx = 1
|
|
else:
|
|
# Explicit frame_dim_idx provided
|
|
channel_dim_idx = 2 if frame_dim_idx == 1 else 1
|
|
channels = target_noise_shape[channel_dim_idx]
|
|
else:
|
|
channels = target_noise_shape[1]
|
|
channel_dim_idx = 1
|
|
|
|
# Add model information to shader params
|
|
shader_params = shader_params.copy()
|
|
shader_params["target_channels"] = channels
|
|
|
|
if model is not None:
|
|
_add_model_info_to_params(shader_params, model, model_name, debugger)
|
|
|
|
if debug_enabled and debug_level >= 2:
|
|
print(f"🎨 Generating {shader_type} shader noise for {'video' if is_video else 'image'}")
|
|
print(f" Target Shape: {target_noise_shape}, Device: {device}, Seed: {seed}")
|
|
if model_name:
|
|
print(f" Model: {model_name}")
|
|
print(f" Target channels: {channels}")
|
|
|
|
# Generate noise based on tensor type
|
|
if is_video:
|
|
noise = _generate_video_noise(
|
|
target_noise_shape, shader_params, seed, device,
|
|
frame_count, frame_dim_idx, channels, height, width,
|
|
batch_size, generator_func, latent_samples.dtype, debugger
|
|
)
|
|
else:
|
|
noise = _generate_image_noise(
|
|
target_noise_shape, shader_params, seed, device,
|
|
channels, height, width, batch_size,
|
|
generator_func, latent_samples.dtype, debugger
|
|
)
|
|
|
|
# Normalize the noise tensor
|
|
noise = normalize_noise(noise)
|
|
|
|
# Analyze and log the generated noise for debugging
|
|
if debug_enabled and hasattr(debugger, 'analyze_tensor'):
|
|
debugger.analyze_tensor(noise, "shader_noise")
|
|
|
|
# Final shape check and correction
|
|
if noise.shape != target_noise_shape:
|
|
noise = _correct_noise_shape(noise, target_noise_shape, channel_dim_idx, device, latent_samples.dtype)
|
|
|
|
return noise
|
|
|
|
|
|
def _add_model_info_to_params(
|
|
shader_params: Dict[str, Any],
|
|
model: Any,
|
|
model_name: Optional[str],
|
|
debugger: Optional[object]
|
|
) -> None:
|
|
"""Add model information to shader parameters."""
|
|
debug_enabled = debugger is not None and getattr(debugger, 'enabled', False)
|
|
debug_level = getattr(debugger, 'debug_level', 0) if debug_enabled else 0
|
|
|
|
# Get and store model name
|
|
if model_name is None:
|
|
model_name = getattr(model, 'model_name', None)
|
|
if model_name is None and hasattr(model, 'model') and hasattr(model.model, 'model_name'):
|
|
model_name = model.model.model_name
|
|
|
|
if model_name:
|
|
shader_params["model_name"] = model_name
|
|
if debug_enabled and debug_level >= 2:
|
|
print(f"📝 Added model_name='{model_name}' to shader_params")
|
|
|
|
# Add model class name
|
|
shader_params["model_class"] = model.__class__.__name__
|
|
|
|
# Add inner model class name if it exists
|
|
if hasattr(model, 'model'):
|
|
shader_params["inner_model_class"] = model.model.__class__.__name__
|
|
inner_model_name = getattr(model.model, 'model_name', None)
|
|
if inner_model_name:
|
|
shader_params["inner_model_name"] = inner_model_name
|
|
|
|
|
|
def _generate_video_noise(
|
|
target_noise_shape: Tuple[int, ...],
|
|
shader_params: Dict[str, Any],
|
|
seed: int,
|
|
device: str,
|
|
frame_count: int,
|
|
frame_dim_idx: int,
|
|
channels: int,
|
|
height: int,
|
|
width: int,
|
|
batch_size: int,
|
|
generator_func: Optional[Callable],
|
|
dtype: torch.dtype,
|
|
debugger: Optional[object]
|
|
) -> torch.Tensor:
|
|
"""Generate noise for video latents."""
|
|
debug_enabled = debugger is not None and getattr(debugger, 'enabled', False)
|
|
debug_level = getattr(debugger, 'debug_level', 0) if debug_enabled else 0
|
|
|
|
if debug_enabled and debug_level >= 2:
|
|
print(f" Video dimensions: Target={target_noise_shape}, Frames={frame_count}")
|
|
|
|
if generator_func is None:
|
|
# Return random noise if no generator
|
|
return torch.randn(target_noise_shape, device=device, dtype=dtype)
|
|
|
|
# Generate noise frame by frame
|
|
noise_frames = []
|
|
|
|
for frame_idx in range(frame_count):
|
|
# Update time parameter for temporal variation
|
|
frame_params = shader_params.copy()
|
|
frame_time_increment = 0.1 if frame_count <= 1 else (1.0 / (frame_count - 1)) * frame_idx
|
|
frame_params["time"] = shader_params.get("time", 0.0) + frame_time_increment
|
|
|
|
if debug_enabled and debug_level >= 3:
|
|
print(f" Generating frame {frame_idx} with time={frame_params['time']:.2f}")
|
|
|
|
# Prepare arguments - use "params" to match BaseNoiseGenerator.generate() signature
|
|
# Wrap dict in ShaderParams for attribute-style access and validate
|
|
generator_args = {
|
|
"params": ShaderParams(frame_params).validate(),
|
|
"height": height,
|
|
"width": width,
|
|
"batch_size": batch_size,
|
|
"device": device,
|
|
"seed": seed + frame_idx if not shader_params.get("useTemporalCoherence", False) else seed,
|
|
"target_channels": channels
|
|
}
|
|
|
|
# Generate the noise for this frame
|
|
frame_noise = generator_func(**generator_args)
|
|
|
|
# Validate and correct frame shape
|
|
expected_frame_shape = (batch_size, channels, height, width)
|
|
if frame_noise.shape != expected_frame_shape:
|
|
frame_noise = _correct_frame_shape(frame_noise, expected_frame_shape, device, dtype)
|
|
|
|
noise_frames.append(frame_noise)
|
|
|
|
# Stack frames along the correct dimension
|
|
if noise_frames:
|
|
return torch.stack(noise_frames, dim=frame_dim_idx)
|
|
else:
|
|
return torch.zeros(target_noise_shape, device=device, dtype=dtype)
|
|
|
|
|
|
def _generate_image_noise(
|
|
target_noise_shape: Tuple[int, ...],
|
|
shader_params: Dict[str, Any],
|
|
seed: int,
|
|
device: str,
|
|
channels: int,
|
|
height: int,
|
|
width: int,
|
|
batch_size: int,
|
|
generator_func: Optional[Callable],
|
|
dtype: torch.dtype,
|
|
debugger: Optional[object]
|
|
) -> torch.Tensor:
|
|
"""Generate noise for image latents."""
|
|
debug_enabled = debugger is not None and getattr(debugger, 'enabled', False)
|
|
debug_level = getattr(debugger, 'debug_level', 0) if debug_enabled else 0
|
|
|
|
if debug_enabled and debug_level >= 2:
|
|
print(f" Image dimensions: Target={target_noise_shape}")
|
|
|
|
if generator_func is None:
|
|
# Return random noise if no generator
|
|
return torch.randn(target_noise_shape, device=device, dtype=dtype)
|
|
|
|
# Prepare arguments - use "params" to match BaseNoiseGenerator.generate() signature
|
|
# Wrap dict in ShaderParams for attribute-style access and validate
|
|
generator_args = {
|
|
"params": ShaderParams(shader_params).validate(),
|
|
"height": height,
|
|
"width": width,
|
|
"batch_size": batch_size,
|
|
"device": device,
|
|
"seed": seed,
|
|
"target_channels": channels
|
|
}
|
|
|
|
# Generate the noise
|
|
noise = generator_func(**generator_args)
|
|
|
|
# Ensure image has correct final shape
|
|
if noise.shape != target_noise_shape:
|
|
noise = _correct_image_shape(noise, target_noise_shape, channels, height, width, device, dtype)
|
|
|
|
return noise
|
|
|
|
|
|
def _correct_frame_shape(
|
|
frame_noise: torch.Tensor,
|
|
expected_shape: Tuple[int, ...],
|
|
device: str,
|
|
dtype: torch.dtype
|
|
) -> torch.Tensor:
|
|
"""Correct frame noise shape to match expected shape."""
|
|
try:
|
|
# FIRST fix spatial dimensions with interpolation
|
|
# This must happen before channel handling to prevent shape mismatches
|
|
if frame_noise.shape[2:] != expected_shape[2:]:
|
|
frame_noise = F.interpolate(
|
|
frame_noise,
|
|
size=(expected_shape[2], expected_shape[3]),
|
|
mode='bilinear',
|
|
align_corners=False
|
|
)
|
|
|
|
# THEN fix channel count - now spatial dims match
|
|
if frame_noise.shape[1] != expected_shape[1]:
|
|
corrected_frame = torch.zeros(expected_shape, device=device, dtype=dtype)
|
|
min_c = min(frame_noise.shape[1], expected_shape[1])
|
|
corrected_frame[:, :min_c] = frame_noise[:, :min_c]
|
|
frame_noise = corrected_frame
|
|
|
|
# Final check
|
|
if frame_noise.shape != expected_shape:
|
|
raise ValueError("Shape correction failed")
|
|
|
|
except Exception:
|
|
# Create a zero tensor as fallback
|
|
frame_noise = torch.zeros(expected_shape, device=device, dtype=dtype)
|
|
|
|
return frame_noise
|
|
|
|
|
|
def _correct_image_shape(
|
|
noise: torch.Tensor,
|
|
target_shape: Tuple[int, ...],
|
|
channels: int,
|
|
height: int,
|
|
width: int,
|
|
device: str,
|
|
dtype: torch.dtype
|
|
) -> torch.Tensor:
|
|
"""Correct image noise shape to match target shape."""
|
|
try:
|
|
# FIRST fix spatial dimensions with interpolation
|
|
# This must happen before channel handling to prevent shape mismatches
|
|
if noise.shape[2:] != target_shape[2:]:
|
|
noise = F.interpolate(
|
|
noise,
|
|
size=(height, width),
|
|
mode='bilinear',
|
|
align_corners=False
|
|
)
|
|
|
|
# THEN fix channel count - now spatial dims match
|
|
if noise.shape[1] != channels:
|
|
corrected_noise = torch.zeros(target_shape, device=device, dtype=dtype)
|
|
min_c = min(noise.shape[1], channels)
|
|
corrected_noise[:, :min_c] = noise[:, :min_c]
|
|
noise = corrected_noise
|
|
|
|
if noise.shape != target_shape:
|
|
raise ValueError("Shape correction failed")
|
|
|
|
except Exception:
|
|
noise = torch.zeros(target_shape, device=device, dtype=dtype)
|
|
|
|
return noise
|
|
|
|
|
|
def _correct_noise_shape(
|
|
noise: torch.Tensor,
|
|
target_shape: Tuple[int, ...],
|
|
channel_dim_idx: int,
|
|
device: str,
|
|
dtype: torch.dtype
|
|
) -> torch.Tensor:
|
|
"""Final correction of noise shape to match target."""
|
|
|
|
try:
|
|
if len(noise.shape) == 5:
|
|
# 5D video tensor - handle both B,C,F,H,W and B,F,C,H,W formats
|
|
# Determine frame dimension from channel_dim_idx
|
|
frame_dim_idx = 2 if channel_dim_idx == 1 else 1
|
|
|
|
# For trilinear interpolation, PyTorch expects (N, C, D, H, W) format
|
|
# If we have B,F,C,H,W format, we need to convert to B,C,F,H,W first
|
|
if channel_dim_idx == 2: # B,F,C,H,W format
|
|
# Transpose to B,C,F,H,W for interpolation
|
|
noise = noise.transpose(1, 2).contiguous()
|
|
|
|
# Now noise is in B,C,F,H,W format, extract target dimensions
|
|
# target_shape is in the original format, so extract appropriately
|
|
if channel_dim_idx == 2: # Original was B,F,C,H,W
|
|
tgt_frames = target_shape[1]
|
|
tgt_channels = target_shape[2]
|
|
else: # B,C,F,H,W
|
|
tgt_channels = target_shape[1]
|
|
tgt_frames = target_shape[2]
|
|
tgt_h, tgt_w = target_shape[3], target_shape[4]
|
|
|
|
# Use trilinear interpolation for (frames, height, width)
|
|
noise = F.interpolate(
|
|
noise,
|
|
size=(tgt_frames, tgt_h, tgt_w),
|
|
mode='trilinear',
|
|
align_corners=False
|
|
)
|
|
|
|
# Handle channel mismatch
|
|
if noise.shape[1] != tgt_channels:
|
|
b, c, f, h, w = noise.shape
|
|
new_noise = torch.zeros((b, tgt_channels, f, h, w), device=device, dtype=dtype)
|
|
min_c = min(c, tgt_channels)
|
|
new_noise[:, :min_c] = noise[:, :min_c]
|
|
noise = new_noise
|
|
|
|
# Convert back to original format if needed
|
|
if channel_dim_idx == 2: # Original was B,F,C,H,W
|
|
noise = noise.transpose(1, 2).contiguous()
|
|
else:
|
|
# 4D image tensor - use bilinear interpolation
|
|
noise = F.interpolate(noise, size=target_shape[2:], mode='bilinear', align_corners=False)
|
|
|
|
# Handle channel mismatch
|
|
if noise.shape[1] != target_shape[1]:
|
|
final_noise = torch.zeros(target_shape, device=device, dtype=dtype)
|
|
min_c = min(noise.shape[1], target_shape[1])
|
|
final_noise[:, :min_c, ...] = noise[:, :min_c, ...]
|
|
noise = final_noise
|
|
|
|
except Exception:
|
|
# Fallback: for 5D tensors, try frame-by-frame bilinear interpolation
|
|
if len(noise.shape) == 5 and len(target_shape) == 5:
|
|
try:
|
|
noise = _interpolate_video_framewise(noise, target_shape, channel_dim_idx, device, dtype)
|
|
except Exception:
|
|
noise = torch.zeros(target_shape, device=device, dtype=dtype)
|
|
else:
|
|
noise = torch.zeros(target_shape, device=device, dtype=dtype)
|
|
|
|
return noise
|
|
|
|
|
|
def _interpolate_video_framewise(
|
|
noise: torch.Tensor,
|
|
target_shape: Tuple[int, ...],
|
|
channel_dim_idx: int,
|
|
device: str,
|
|
dtype: torch.dtype
|
|
) -> torch.Tensor:
|
|
"""Fallback frame-by-frame interpolation for video tensors."""
|
|
# Determine frame dimension
|
|
frame_dim = 2 if channel_dim_idx == 1 else 1
|
|
target_h, target_w = target_shape[3], target_shape[4]
|
|
|
|
frames_list = []
|
|
num_frames = noise.shape[frame_dim]
|
|
|
|
for i in range(num_frames):
|
|
if frame_dim == 1:
|
|
frame = noise[:, i] # B,C,H,W
|
|
else:
|
|
frame = noise[:, :, i] # B,C,H,W
|
|
|
|
frame_resized = F.interpolate(
|
|
frame,
|
|
size=(target_h, target_w),
|
|
mode='bilinear',
|
|
align_corners=False
|
|
)
|
|
frames_list.append(frame_resized)
|
|
|
|
if frame_dim == 1:
|
|
result = torch.stack(frames_list, dim=1)
|
|
else:
|
|
result = torch.stack(frames_list, dim=2)
|
|
|
|
# Handle frame count mismatch
|
|
target_frames = target_shape[frame_dim]
|
|
if result.shape[frame_dim] != target_frames:
|
|
# Simple repeat/truncate for frame adjustment
|
|
if result.shape[frame_dim] < target_frames:
|
|
repeats = (target_frames // result.shape[frame_dim]) + 1
|
|
if frame_dim == 1:
|
|
result = result.repeat(1, repeats, 1, 1, 1)[:, :target_frames]
|
|
else:
|
|
result = result.repeat(1, 1, repeats, 1, 1)[:, :, :target_frames]
|
|
else:
|
|
if frame_dim == 1:
|
|
result = result[:, :target_frames]
|
|
else:
|
|
result = result[:, :, :target_frames]
|
|
|
|
return result
|