fix: flex latent blend
deprecate audio latent blend nodes
This commit is contained in:
RyanOnTheInside
+42
-94
@@ -24,31 +24,26 @@ class AudioLatentBlend:
|
||||
CATEGORY = "latent/audio"
|
||||
|
||||
def blend(self, latent1, latent2, blend_factor: float, blend_mode: str):
|
||||
# Ensure we're working with latent dictionaries
|
||||
if not isinstance(latent1, dict) or not isinstance(latent2, dict):
|
||||
raise ValueError("Inputs must be latent dictionaries")
|
||||
|
||||
|
||||
if "samples" not in latent1 or "samples" not in latent2:
|
||||
raise ValueError("Inputs must contain 'samples' key")
|
||||
|
||||
# Check if the latents have the correct shape for audio (8 channels)
|
||||
s1 = latent1["samples"]
|
||||
s2 = latent2["samples"]
|
||||
|
||||
if s1.shape[1] != 8 or s2.shape[1] != 8:
|
||||
raise ValueError("Both inputs must be audio latents (8 channels)")
|
||||
|
||||
samples_out = latent1.copy()
|
||||
|
||||
|
||||
# Ensure both latents have the same shape
|
||||
if s1.shape != s2.shape:
|
||||
# Reshape s2 to match s1's dimensions
|
||||
s2 = torch.nn.functional.interpolate(
|
||||
s2,
|
||||
size=s1.shape[2:],
|
||||
mode='linear',
|
||||
align_corners=False
|
||||
)
|
||||
# Match time dimension (last dim) via interpolation
|
||||
if s1.dim() == 3:
|
||||
# ACE-Step 1.5: [B, C, T]
|
||||
s2 = torch.nn.functional.interpolate(s2, size=s1.shape[2], mode='linear', align_corners=False)
|
||||
else:
|
||||
# ACE-Step 1.0: [B, C, H, T]
|
||||
s2 = torch.nn.functional.interpolate(s2, size=s1.shape[2:], mode='bilinear', align_corners=False)
|
||||
|
||||
# Apply blend mode
|
||||
if blend_mode == "normal":
|
||||
@@ -101,107 +96,66 @@ class FlexlatentAudioBlend:
|
||||
CATEGORY = "latent/audio"
|
||||
|
||||
def blend(self, latent1, latent2, blend_factors, blend_mode: str, interpolation: str):
|
||||
# Ensure we're working with latent dictionaries
|
||||
if not isinstance(latent1, dict) or not isinstance(latent2, dict):
|
||||
raise ValueError("Inputs must be latent dictionaries")
|
||||
|
||||
|
||||
if "samples" not in latent1 or "samples" not in latent2:
|
||||
raise ValueError("Inputs must contain 'samples' key")
|
||||
|
||||
# Check if the latents have the correct shape for audio (8 channels)
|
||||
s1 = latent1["samples"]
|
||||
s2 = latent2["samples"]
|
||||
|
||||
print(f"Input latent shapes - s1: {s1.shape}, s2: {s2.shape}")
|
||||
|
||||
if s1.shape[1] != 8 or s2.shape[1] != 8:
|
||||
raise ValueError("Both inputs must be audio latents (8 channels)")
|
||||
is_3d = s1.dim() == 3 # ACE-Step 1.5: [B, C, T] vs 1.0: [B, C, H, T]
|
||||
|
||||
samples_out = latent1.copy()
|
||||
|
||||
|
||||
# Ensure both latents have the same shape
|
||||
if s1.shape != s2.shape:
|
||||
# Reshape s2 to match s1's dimensions
|
||||
s2 = torch.nn.functional.interpolate(
|
||||
s2,
|
||||
size=s1.shape[2:],
|
||||
mode='linear',
|
||||
align_corners=False
|
||||
)
|
||||
print(f"Resized s2 shape: {s2.shape}")
|
||||
|
||||
if is_3d:
|
||||
s2 = torch.nn.functional.interpolate(s2, size=s1.shape[2], mode='linear', align_corners=False)
|
||||
else:
|
||||
s2 = torch.nn.functional.interpolate(s2, size=s1.shape[2:], mode='bilinear', align_corners=False)
|
||||
|
||||
# Convert blend_factors to a list if it's a single float
|
||||
if isinstance(blend_factors, float):
|
||||
blend_factors_list = [blend_factors]
|
||||
else:
|
||||
# Already a list of floats
|
||||
blend_factors_list = blend_factors
|
||||
|
||||
|
||||
# Clamp values to 0.0-1.0 range
|
||||
blend_factors_list = [max(0.0, min(1.0, f)) for f in blend_factors_list]
|
||||
|
||||
print(f"Blend factors: {blend_factors_list}")
|
||||
|
||||
|
||||
if len(blend_factors_list) == 0:
|
||||
raise ValueError("No blend factors provided")
|
||||
|
||||
# Get the temporal dimension size
|
||||
temporal_length = s1.shape[3]
|
||||
print(f"Temporal length: {temporal_length}")
|
||||
|
||||
|
||||
# Temporal dimension is last: shape[2] for 3D, shape[3] for 4D
|
||||
temporal_length = s1.shape[-1]
|
||||
|
||||
# Create blend factor tensor across temporal dimension
|
||||
if len(blend_factors_list) == 1:
|
||||
# If only one factor, use it uniformly
|
||||
blend_factor_tensor = torch.ones(temporal_length, device=s1.device) * blend_factors_list[0]
|
||||
print(f"Using single blend factor: {blend_factors_list[0]}")
|
||||
else:
|
||||
# Create interpolation points
|
||||
points = torch.linspace(0, temporal_length-1, len(blend_factors_list))
|
||||
values = torch.tensor(blend_factors_list, device=s1.device)
|
||||
|
||||
print(f"Interpolation points: {points.tolist()}")
|
||||
print(f"Values at points: {values.tolist()}")
|
||||
|
||||
# Create a tensor for all temporal positions
|
||||
positions = torch.arange(temporal_length, device=s1.device)
|
||||
|
||||
|
||||
if interpolation == "step":
|
||||
# For step interpolation, find nearest point
|
||||
indices = torch.searchsorted(points, positions) - 1
|
||||
indices = torch.clamp(indices, 0, len(blend_factors_list) - 1)
|
||||
blend_factor_tensor = values[indices]
|
||||
print(f"Using step interpolation")
|
||||
else: # linear interpolation
|
||||
# For positions before first point
|
||||
blend_factor_tensor = torch.zeros(temporal_length, device=s1.device)
|
||||
|
||||
# For each segment between points
|
||||
|
||||
for i in range(len(points) - 1):
|
||||
# Find positions in this segment
|
||||
start, end = points[i], points[i+1]
|
||||
mask = (positions >= start) & (positions <= end)
|
||||
|
||||
# Calculate interpolation weight
|
||||
weight = (positions[mask] - start) / (end - start)
|
||||
|
||||
# Interpolate values
|
||||
blend_factor_tensor[mask] = values[i] * (1 - weight) + values[i+1] * weight
|
||||
print(f"Using linear interpolation")
|
||||
|
||||
# For positions before first point, use first value
|
||||
mask = positions < points[0]
|
||||
blend_factor_tensor[mask] = values[0]
|
||||
|
||||
# For positions after last point, use last value
|
||||
mask = positions > points[-1]
|
||||
blend_factor_tensor[mask] = values[-1]
|
||||
|
||||
# Print blend factor tensor stats
|
||||
print(f"Blend factor tensor stats - min: {blend_factor_tensor.min().item()}, max: {blend_factor_tensor.max().item()}, mean: {blend_factor_tensor.mean().item()}")
|
||||
print(f"First few blend factors: {blend_factor_tensor[:5].tolist()}")
|
||||
print(f"Last few blend factors: {blend_factor_tensor[-5:].tolist()}")
|
||||
|
||||
# Ensure blend_factor_tensor has the shape [temporal_length]
|
||||
|
||||
# Fill edges
|
||||
blend_factor_tensor[positions < points[0]] = values[0]
|
||||
blend_factor_tensor[positions > points[-1]] = values[-1]
|
||||
|
||||
# Ensure blend_factor_tensor has the correct temporal length
|
||||
if blend_factor_tensor.shape[0] != temporal_length:
|
||||
blend_factor_tensor = torch.nn.functional.interpolate(
|
||||
blend_factor_tensor.unsqueeze(0).unsqueeze(0),
|
||||
@@ -209,26 +163,20 @@ class FlexlatentAudioBlend:
|
||||
mode='linear',
|
||||
align_corners=False
|
||||
).squeeze(0).squeeze(0)
|
||||
print(f"Resized blend factor tensor shape: {blend_factor_tensor.shape}")
|
||||
|
||||
# Expand blend_factor_tensor to match latent dimensions [batch, channels, 16, time]
|
||||
# First create the expanded shape manually to ensure compatibility
|
||||
batch_size, num_channels = s1.shape[0], s1.shape[1]
|
||||
blend_factor_tensor = blend_factor_tensor.unsqueeze(0).unsqueeze(0).unsqueeze(0) # [1, 1, 1, temporal_length]
|
||||
|
||||
print(f"Blend factor tensor shape after unsqueeze: {blend_factor_tensor.shape}")
|
||||
|
||||
# Create explicit shape for expansion
|
||||
|
||||
# Expand blend_factor_tensor to match latent dimensions
|
||||
# 3D [B, C, T]: unsqueeze to [1, 1, T]
|
||||
# 4D [B, C, H, T]: unsqueeze to [1, 1, 1, T]
|
||||
if is_3d:
|
||||
blend_factor_tensor = blend_factor_tensor.unsqueeze(0).unsqueeze(0)
|
||||
else:
|
||||
blend_factor_tensor = blend_factor_tensor.unsqueeze(0).unsqueeze(0).unsqueeze(0)
|
||||
|
||||
target_shape = list(s1.shape)
|
||||
print(f"Target shape for expansion: {target_shape}")
|
||||
print(f"Blend factor tensor shape before expansion: {blend_factor_tensor.shape}")
|
||||
|
||||
# Apply blend mode with time-varying factors
|
||||
if blend_mode == "normal":
|
||||
# Expand blend_factor_tensor to match all dimensions
|
||||
blend_factor_expanded = blend_factor_tensor.expand(target_shape)
|
||||
print(f"Blend factor expanded shape: {blend_factor_expanded.shape}")
|
||||
print(f"Blend factor expanded stats - min: {blend_factor_expanded.min().item()}, max: {blend_factor_expanded.max().item()}, mean: {blend_factor_expanded.mean().item()}")
|
||||
blended = s1 * blend_factor_expanded + s2 * (1 - blend_factor_expanded)
|
||||
elif blend_mode == "add":
|
||||
blend_factor_expanded = blend_factor_tensor.expand(target_shape)
|
||||
@@ -263,6 +211,6 @@ NODE_CLASS_MAPPINGS = {
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AudioLatentBlend": "Audio Latent Blend",
|
||||
"FlexlatentAudioBlend": "FlexlatentAudioBlend (BETA)",
|
||||
"AudioLatentBlend": "Audio Latent Blend (DEPRECATED USE FLEX LATENT BLEND)",
|
||||
"FlexlatentAudioBlend": "FlexlatentAudioBlend (DEPRECATED USE FLEX LATENT BLEND)",
|
||||
}
|
||||
Reference in New Issue
Block a user