nodes/ryanontheinside-ComfyUI_RyanOnTheInside
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

acep wip

Browse Source
This commit is contained in:
RyanOnTheInside
2026-02-06 05:32:13 -05:00
parent 418dd26946
commit c9776ac9de
7 changed files with 1899 additions and 331 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/acestep/ace_step_guiders.py
+779 -91
View File
File diff suppressed because it is too large Load Diff
nodes/acestep/ace_step_utils.py
+156 -108
View File
@@ -4,113 +4,156 @@ import math
class ACEStepLatentUtils:
"""Utility functions for ACEStep audio latent manipulation"""
# Version constants
V1_0 = "v1.0"
V1_5 = "v1.5"
# Frame rates: frames per second
V1_0_FPS = 44100 / 512 / 8 # ≈ 10.77 fps
V1_5_FPS = 48000 / 1920 # = 25 fps
@staticmethod
def time_to_frame_index(time_seconds):
"""Convert time in seconds to ACE latent frame index
ACE audio: 1 second = 44100 samples / 512 / 8 ≈ 10.8 frames
def detect_version(latent):
"""Detect ACE-Step version from latent shape
v1.0: (batch, 8, 16, length) - 4D
v1.5: (batch, 64, length) - 3D
"""
return int(time_seconds * 44100 / 512 / 8)
if len(latent.shape) == 4 and latent.shape[1] == 8 and latent.shape[2] == 16:
return ACEStepLatentUtils.V1_0
elif len(latent.shape) == 3 and latent.shape[1] == 64:
return ACEStepLatentUtils.V1_5
else:
return None
@staticmethod
def frame_index_to_time(frame_index):
def get_fps(version):
"""Get frames per second for a given version"""
if version == ACEStepLatentUtils.V1_5:
return ACEStepLatentUtils.V1_5_FPS
return ACEStepLatentUtils.V1_0_FPS
@staticmethod
def time_to_frame_index(time_seconds, version=None):
"""Convert time in seconds to ACE latent frame index
v1.0: 1 second = 44100 / 512 / 8 ≈ 10.77 frames
v1.5: 1 second = 48000 / 1920 = 25 frames
"""
fps = ACEStepLatentUtils.get_fps(version)
return int(time_seconds * fps)
@staticmethod
def frame_index_to_time(frame_index, version=None):
"""Convert ACE latent frame index to time in seconds"""
return frame_index * 512 * 8 / 44100
fps = ACEStepLatentUtils.get_fps(version)
return frame_index / fps
@staticmethod
def get_latent_length(latent):
"""Get the time dimension length from latent (last dim for both versions)"""
return latent.shape[-1]
@staticmethod
def validate_ace_latent_shape(latent):
"""Validate that latent has correct ACE audio shape"""
if len(latent.shape) != 4:
raise ValueError(f"ACE latent must be 4D, got {len(latent.shape)}D")
batch_size, channels, height, length = latent.shape
if channels != 8:
raise ValueError(f"ACE latent must have 8 channels, got {channels}")
if height != 16:
raise ValueError(f"ACE latent must have height 16, got {height}")
return True
"""Validate that latent has correct ACE audio shape (v1.0 or v1.5)"""
version = ACEStepLatentUtils.detect_version(latent)
if version is None:
if len(latent.shape) == 4:
raise ValueError(f"ACE v1.0 latent must have shape (batch, 8, 16, length), got {latent.shape}")
elif len(latent.shape) == 3:
raise ValueError(f"ACE v1.5 latent must have shape (batch, 64, length), got {latent.shape}")
else:
raise ValueError(f"ACE latent must be 3D (v1.5) or 4D (v1.0), got {len(latent.shape)}D")
return version
@staticmethod
def create_repaint_mask(latent_shape, start_frame, end_frame):
def create_repaint_mask(latent_shape, start_frame, end_frame, version=None):
"""Create a mask for repainting a specific region
Args:
latent_shape: (batch, 8, 16, length)
latent_shape: (batch, 8, 16, length) for v1.0 or (batch, 64, length) for v1.5
start_frame: Frame index to start repainting
end_frame: Frame index to end repainting
version: ACE-Step version (auto-detected if None)
Returns:
torch.Tensor: Binary mask with 1.0 in repaint region
"""
batch_size, channels, height, length = latent_shape
length = latent_shape[-1]
# Clamp frame indices to valid range
start_frame = max(0, start_frame)
end_frame = min(length, end_frame)
if start_frame >= end_frame:
raise ValueError(f"start_frame ({start_frame}) must be less than end_frame ({end_frame})")
mask = torch.zeros(latent_shape)
mask[:, :, :, start_frame:end_frame] = 1.0
if len(latent_shape) == 4: # v1.0
mask[:, :, :, start_frame:end_frame] = 1.0
else: # v1.5
mask[:, :, start_frame:end_frame] = 1.0
return mask
@staticmethod
def create_extend_mask(source_shape, left_frames, right_frames):
"""Create a mask for extension regions
Args:
source_shape: (batch, 8, 16, length) shape of source latent
source_shape: (batch, 8, 16, length) for v1.0 or (batch, 64, length) for v1.5
left_frames: Number of frames to extend on the left
right_frames: Number of frames to extend on the right
Returns:
torch.Tensor: Binary mask with 1.0 in extension regions
"""
batch_size, channels, height, source_length = source_shape
source_length = source_shape[-1]
total_length = source_length + left_frames + right_frames
mask = torch.zeros((batch_size, channels, height, total_length))
# Mark left extension region
if left_frames > 0:
mask[:, :, :, :left_frames] = 1.0
# Mark right extension region
if right_frames > 0:
mask[:, :, :, -right_frames:] = 1.0
if len(source_shape) == 4: # v1.0
batch_size, channels, height, _ = source_shape
mask = torch.zeros((batch_size, channels, height, total_length))
if left_frames > 0:
mask[:, :, :, :left_frames] = 1.0
if right_frames > 0:
mask[:, :, :, -right_frames:] = 1.0
else: # v1.5
batch_size, channels, _ = source_shape
mask = torch.zeros((batch_size, channels, total_length))
if left_frames > 0:
mask[:, :, :left_frames] = 1.0
if right_frames > 0:
mask[:, :, -right_frames:] = 1.0
return mask
@staticmethod
def create_extend_region_mask(extended_shape, left_frames, right_frames, source_length):
"""Create a mask indicating which regions to generate vs preserve in extended latent
Args:
extended_shape: Shape of the extended latent (batch, 8, 16, total_length)
extended_shape: Shape of the extended latent
left_frames: Number of frames added on the left
right_frames: Number of frames added on the right
source_length: Length of the original source latent
Returns:
torch.Tensor: Binary mask with 1.0 = generate (extension regions), 0.0 = preserve (source region)
"""
batch_size, channels, height, total_length = extended_shape
mask = torch.zeros(extended_shape)
# Mark left extension region for generation
if left_frames > 0:
mask[:, :, :, :left_frames] = 1.0
# Mark right extension region for generation
if right_frames > 0:
mask[:, :, :, -right_frames:] = 1.0
# The middle region (source_length) remains 0.0 for preservation
if len(extended_shape) == 4: # v1.0
if left_frames > 0:
mask[:, :, :, :left_frames] = 1.0
if right_frames > 0:
mask[:, :, :, -right_frames:] = 1.0
else: # v1.5
if left_frames > 0:
mask[:, :, :left_frames] = 1.0
if right_frames > 0:
mask[:, :, -right_frames:] = 1.0
return mask
@staticmethod
@@ -161,47 +204,52 @@ class ACEStepLatentUtils:
@staticmethod
def create_feather_mask(mask, feather_frames=2):
"""Create a feathered (soft) version of a binary mask
Args:
mask: Binary mask tensor
mask: Binary mask tensor (3D for v1.5 or 4D for v1.0)
feather_frames: Number of frames to feather on each side
Returns:
torch.Tensor: Feathered mask with smooth transitions
"""
if feather_frames <= 0:
return mask
# Create gaussian kernel for smoothing
kernel_size = feather_frames * 2 + 1
sigma = feather_frames / 3.0
x = torch.arange(kernel_size, dtype=torch.float32) - feather_frames
gaussian_kernel = torch.exp(-0.5 * (x / sigma) ** 2)
gaussian_kernel = gaussian_kernel / gaussian_kernel.sum()
# Reshape for 1D convolution along time dimension
kernel = gaussian_kernel.view(1, 1, 1, -1)
# Apply convolution to each channel separately
batch_size, channels, height, length = mask.shape
result = torch.zeros_like(mask)
for b in range(batch_size):
for c in range(channels):
for h in range(height):
# Extract 1D signal
signal = mask[b:b+1, c:c+1, h:h+1, :] # Shape: (1, 1, 1, length)
# Pad with constant values (edge values)
is_v1_5 = len(mask.shape) == 3
if is_v1_5:
# v1.5: (batch, channels, length)
batch_size, channels, length = mask.shape
kernel = gaussian_kernel.view(1, 1, -1)
for b in range(batch_size):
for c in range(channels):
signal = mask[b:b+1, c:c+1, :]
padded_signal = F.pad(signal, (feather_frames, feather_frames), 'constant', 0.0)
# Apply 1D convolution
smoothed = F.conv2d(padded_signal, kernel, padding=0)
# Store result
result[b, c, h, :] = smoothed[0, 0, 0, :]
smoothed = F.conv1d(padded_signal, kernel, padding=0)
result[b, c, :] = smoothed[0, 0, :]
else:
# v1.0: (batch, channels, height, length)
batch_size, channels, height, length = mask.shape
kernel = gaussian_kernel.view(1, 1, 1, -1)
for b in range(batch_size):
for c in range(channels):
for h in range(height):
signal = mask[b:b+1, c:c+1, h:h+1, :]
padded_signal = F.pad(signal, (feather_frames, feather_frames), 'constant', 0.0)
smoothed = F.conv2d(padded_signal, kernel, padding=0)
result[b, c, h, :] = smoothed[0, 0, 0, :]
return torch.clamp(result, 0.0, 1.0)
@staticmethod
@@ -282,29 +330,29 @@ class ACEStepLatentUtils:
class ACEStepRepaintHelper:
"""Helper class for repaint operations"""
def __init__(self, source_latents, start_time, end_time,
def __init__(self, source_latents, start_time, end_time,
repaint_strength=0.7, feather_time=0.1):
"""
Args:
source_latents: Original audio latents (batch, 8, 16, length)
source_latents: Original audio latents (v1.0 or v1.5 shape)
start_time: Start time in seconds for repaint region
end_time: End time in seconds for repaint region
repaint_strength: Strength of repainting (0.0 = no repaint, 1.0 = full repaint)
feather_time: Time in seconds for feathering edges
"""
ACEStepLatentUtils.validate_ace_latent_shape(source_latents)
self.version = ACEStepLatentUtils.validate_ace_latent_shape(source_latents)
self.source_latents = source_latents
self.start_frame = ACEStepLatentUtils.time_to_frame_index(start_time)
self.end_frame = ACEStepLatentUtils.time_to_frame_index(end_time)
self.start_frame = ACEStepLatentUtils.time_to_frame_index(start_time, self.version)
self.end_frame = ACEStepLatentUtils.time_to_frame_index(end_time, self.version)
self.repaint_strength = repaint_strength
self.feather_frames = ACEStepLatentUtils.time_to_frame_index(feather_time)
self.feather_frames = ACEStepLatentUtils.time_to_frame_index(feather_time, self.version)
# Create masks
self.repaint_mask = ACEStepLatentUtils.create_repaint_mask(
source_latents.shape, self.start_frame, self.end_frame)
source_latents.shape, self.start_frame, self.end_frame, self.version)
# Apply feathering for smooth transitions
if self.feather_frames > 0:
self.repaint_mask = ACEStepLatentUtils.create_feather_mask(
@@ -325,24 +373,24 @@ class ACEStepRepaintHelper:
class ACEStepExtendHelper:
"""Helper class for extend operations"""
def __init__(self, source_latents, extend_left_time=0, extend_right_time=0):
"""
Args:
source_latents: Original audio latents (batch, 8, 16, length)
source_latents: Original audio latents (v1.0 or v1.5 shape)
extend_left_time: Time in seconds to extend before the audio
extend_right_time: Time in seconds to extend after the audio
"""
ACEStepLatentUtils.validate_ace_latent_shape(source_latents)
self.version = ACEStepLatentUtils.validate_ace_latent_shape(source_latents)
self.source_latents = source_latents
self.left_frames = ACEStepLatentUtils.time_to_frame_index(extend_left_time)
self.right_frames = ACEStepLatentUtils.time_to_frame_index(extend_right_time)
self.left_frames = ACEStepLatentUtils.time_to_frame_index(extend_left_time, self.version)
self.right_frames = ACEStepLatentUtils.time_to_frame_index(extend_right_time, self.version)
# Create extended latent with zero padding
self.extended_latents = ACEStepLatentUtils.pad_latent_for_extend(
source_latents, self.left_frames, self.right_frames)
# Create extension mask
self.extension_mask = ACEStepLatentUtils.create_extend_mask(
source_latents.shape, self.left_frames, self.right_frames)
nodes/acestep/audio_mask_nodes.py
+125 -74
View File
@@ -4,29 +4,40 @@ from .ace_step_utils import ACEStepLatentUtils
class AudioMaskBase:
"""Base class for audio mask nodes with common functionality"""
def _extract_and_validate_latent(self, audio_latents):
"""Extract and validate ACE audio latent tensor"""
"""Extract and validate ACE audio latent tensor, returns (tensor, version)"""
if "samples" in audio_latents:
latent_tensor = audio_latents["samples"]
else:
raise ValueError("Invalid latent format")
ACEStepLatentUtils.validate_ace_latent_shape(latent_tensor)
return latent_tensor
def _create_base_mask(self, latent_tensor):
"""Create base mask with correct shape for CondPairSetProps"""
batch_size, channels, height, length = latent_tensor.shape
return torch.zeros((batch_size, height, length)), (batch_size, height, length)
version = ACEStepLatentUtils.validate_ace_latent_shape(latent_tensor)
return latent_tensor, version
def _create_base_mask(self, latent_tensor, version=None):
"""Create base mask with correct shape for conditioning
v1.0: (batch, height, length) where height=16 - matches noise.shape[2:] = (16, length)
Note: Temporal mask conditioning only works with ACE-Step 1.0
"""
if version == ACEStepLatentUtils.V1_5:
# v1.5 uses discrete audio_codes tokens that cannot be interpolated
# Temporal mask conditioning is not supported for v1.5
raise ValueError("Temporal mask conditioning is only supported for ACE-Step 1.0. "
"ACE-Step 1.5 uses discrete audio tokens that cannot be blended temporally.")
else:
batch_size, channels, height, length = latent_tensor.shape
mask = torch.zeros((batch_size, height, length))
dims = (batch_size, height, length)
return mask, dims
def _finalize_mask(self, mask):
"""Ensure mask values are in valid range"""
return torch.clamp(mask, 0.0, 1.0)
class AudioTemporalMask(AudioMaskBase):
"""Create temporal crossfade mask from a list of float values"""
@classmethod
def INPUT_TYPES(s):
return {
@@ -36,16 +47,18 @@ class AudioTemporalMask(AudioMaskBase):
"length_mismatch": (["repeat", "loop"], {"default": "repeat"}),
}
}
RETURN_TYPES = ("MASK",)
FUNCTION = "create_temporal_mask"
CATEGORY = "conditioning/audio"
def create_temporal_mask(self, audio_latents, value, length_mismatch):
"""Create temporal mask from custom float values"""
latent_tensor = self._extract_and_validate_latent(audio_latents)
mask, (batch_size, height, length) = self._create_base_mask(latent_tensor)
latent_tensor, version = self._extract_and_validate_latent(audio_latents)
mask, dims = self._create_base_mask(latent_tensor, version)
batch_size = dims[0]
length = dims[-1] # Always last dimension
# Handle both single float and list of floats
if isinstance(value, (int, float)):
@@ -75,11 +88,14 @@ class AudioTemporalMask(AudioMaskBase):
while len(values) < length:
remaining = length - len(values)
values.extend(values[:min(remaining, len(values))])
# Apply values to mask
# Apply values to mask - v1.5 is 2D (batch, length), v1.0 is 3D (batch, height, length)
for i in range(length):
mask[:, :, i] = values[i]
if len(mask.shape) == 2:
mask[:, i] = values[i]
else:
mask[:, :, i] = values[i]
return (self._finalize_mask(mask),)
class AudioRegionMask(AudioMaskBase):
@@ -103,33 +119,44 @@ class AudioRegionMask(AudioMaskBase):
def create_region_mask(self, audio_latents, start_time, end_time, mask_value, feather_seconds):
"""Create mask for specific time region with optional feathering"""
latent_tensor, version = self._extract_and_validate_latent(audio_latents)
mask, dims = self._create_base_mask(latent_tensor, version)
length = dims[-1] # Always last dimension
# Convert to frames using version-specific FPS
start_frame = ACEStepLatentUtils.time_to_frame_index(start_time, version)
end_frame = ACEStepLatentUtils.time_to_frame_index(end_time, version)
feather_frames = ACEStepLatentUtils.time_to_frame_index(feather_seconds, version)
latent_tensor = self._extract_and_validate_latent(audio_latents)
mask, (batch_size, height, length) = self._create_base_mask(latent_tensor)
# Convert to frames
start_frame = ACEStepLatentUtils.time_to_frame_index(start_time)
end_frame = ACEStepLatentUtils.time_to_frame_index(end_time)
feather_frames = ACEStepLatentUtils.time_to_frame_index(feather_seconds)
# Set region value
# Set region value - v1.5 is 2D (batch, length), v1.0 is 3D (batch, height, length)
is_2d = len(mask.shape) == 2
if start_frame < end_frame:
mask[:, :, start_frame:end_frame] = mask_value
if is_2d:
mask[:, start_frame:end_frame] = mask_value
else:
mask[:, :, start_frame:end_frame] = mask_value
# Apply feathering if requested
if feather_frames > 0:
# Feather start
for i in range(max(0, start_frame - feather_frames), start_frame):
if i >= 0 and i < length:
progress = (i - (start_frame - feather_frames)) / feather_frames
mask[:, :, i] = mask_value * progress
if is_2d:
mask[:, i] = mask_value * progress
else:
mask[:, :, i] = mask_value * progress
# Feather end
for i in range(end_frame, min(length, end_frame + feather_frames)):
if i >= 0 and i < length:
progress = 1.0 - (i - end_frame) / feather_frames
mask[:, :, i] = mask_value * progress
if is_2d:
mask[:, i] = mask_value * progress
else:
mask[:, :, i] = mask_value * progress
return (self._finalize_mask(mask),)
class AudioMaskAnalyzer:
@@ -149,31 +176,44 @@ class AudioMaskAnalyzer:
def analyze_mask(self, mask):
"""Show mask info and visualization"""
info = f"Audio Mask Analysis:\n"
info += f"Shape: {mask.shape}\n"
if len(mask.shape) >= 3:
batch_size, height, length = mask.shape[0], mask.shape[1], mask.shape[-1]
duration = ACEStepLatentUtils.frame_index_to_time(length)
if len(mask.shape) >= 2:
batch_size = mask.shape[0]
length = mask.shape[-1]
# Detect version based on shape - v1.5 is 2D, v1.0 is 3D
if len(mask.shape) == 2:
version = ACEStepLatentUtils.V1_5
info += f"Version: v1.5 (2D mask)\n"
else:
version = ACEStepLatentUtils.V1_0
info += f"Version: v1.0 (3D mask)\n"
duration = ACEStepLatentUtils.frame_index_to_time(length, version)
info += f"Duration: {duration:.2f} seconds\n"
info += f"Frames: {length}\n"
# Get mask values along time dimension (first batch, first height)
mask_values = mask[0, 0, :].cpu().numpy()
# Get mask values along time dimension
if len(mask.shape) == 2:
mask_values = mask[0, :].cpu().numpy()
else:
mask_values = mask[0, 0, :].cpu().numpy()
info += f"Value range: {mask_values.min():.3f} to {mask_values.max():.3f}\n"
# Create simple visualization
viz_length = min(60, length)
step = length / viz_length if viz_length > 0 else 1
viz_line = ""
for i in range(viz_length):
idx = int(i * step)
val = mask_values[idx] if idx < len(mask_values) else 0
if val > 0.8:
viz_line += "█"
elif val > 0.6:
@@ -184,11 +224,11 @@ class AudioMaskAnalyzer:
viz_line += "░"
else:
viz_line += "·"
info += f"\nTemporal pattern:\n{viz_line}\n"
info += "█=0.8+ ▓=0.6+ ▒=0.4+ ░=0.2+ ·=<0.2\n"
info += f"Time: 0s{'':>{viz_length-10}}{duration:.1f}s"
return (info,)
class MaskToAudioMask(AudioMaskBase):
@@ -210,9 +250,11 @@ class MaskToAudioMask(AudioMaskBase):
def convert_spatial_to_audio_mask(self, audio_latents, spatial_mask, frame_summary):
"""Convert spatial mask to audio temporal mask - batch=time, spatial→single value per frame"""
latent_tensor = self._extract_and_validate_latent(audio_latents)
audio_mask, (batch_size, height, target_length) = self._create_base_mask(latent_tensor)
latent_tensor, version = self._extract_and_validate_latent(audio_latents)
audio_mask, dims = self._create_base_mask(latent_tensor, version)
batch_size = dims[0]
target_length = dims[-1]
# Handle different spatial mask shapes
if len(spatial_mask.shape) == 2:
@@ -261,11 +303,15 @@ class MaskToAudioMask(AudioMaskBase):
else:
final_pattern = torch.tensor(temporal_values, dtype=torch.float32)
# Apply temporal pattern to all batches and height channels
# Apply temporal pattern to all batches - v1.5 is 2D, v1.0 is 3D
is_2d = len(audio_mask.shape) == 2
for b in range(batch_size):
for h in range(height):
audio_mask[b, h, :] = final_pattern
if is_2d:
audio_mask[b, :] = final_pattern
else:
for h in range(audio_mask.shape[1]):
audio_mask[b, h, :] = final_pattern
return (self._finalize_mask(audio_mask),)
class AudioLatentInfo(AudioMaskBase):
@@ -287,27 +333,32 @@ class AudioLatentInfo(AudioMaskBase):
def get_audio_info(self, audio_latents, frame_rate):
"""Extract batch size, length, and duration from audio latents"""
latent_tensor = self._extract_and_validate_latent(audio_latents)
batch_size, channels, height, actual_length = latent_tensor.shape
# Calculate duration using the default frame rate (10.77)
default_frame_rate = 10.77
duration = actual_length / default_frame_rate
latent_tensor, version = self._extract_and_validate_latent(audio_latents)
actual_length = latent_tensor.shape[-1]
batch_size = latent_tensor.shape[0]
channels = latent_tensor.shape[1]
# Calculate duration using version-specific frame rate
native_fps = ACEStepLatentUtils.get_fps(version)
duration = actual_length / native_fps
# Calculate length_frames based on the specified frame rate
length_frames = int(duration * frame_rate)
# Create info string
info = f"Audio Latent Info:\n"
info += f"Version: {version or 'unknown'}\n"
info += f"Batch size: {batch_size}\n"
info += f"Channels: {channels}\n"
info += f"Height: {height}\n"
if version != ACEStepLatentUtils.V1_5:
info += f"Height: {latent_tensor.shape[2]}\n"
info += f"Actual length: {actual_length} frames\n"
info += f"Frame rate: {frame_rate:.2f} fps\n"
info += f"Length: {length_frames} frames\n"
info += f"Native frame rate: {native_fps:.2f} fps\n"
info += f"Target frame rate: {frame_rate:.2f} fps\n"
info += f"Length at target rate: {length_frames} frames\n"
info += f"Duration: {duration:.2f} seconds"
return (batch_size, length_frames, duration, info)
# Node mappings
@@ -320,9 +371,9 @@ AUDIO_MASK_NODE_CLASS_MAPPINGS = {
}
AUDIO_MASK_NODE_DISPLAY_NAME_MAPPINGS = {
"AudioTemporalMask": "Audio Temporal Mask",
"AudioRegionMask": "Audio Region Mask",
"AudioTemporalMask": "Audio Temporal Mask (ACE 1.0)",
"AudioRegionMask": "Audio Region Mask (ACE 1.0)",
"AudioMaskAnalyzer": "Audio Mask Analyzer",
"MaskToAudioMask": "Mask to Audio Mask",
"MaskToAudioMask": "Mask to Audio Mask (ACE 1.0)",
"AudioLatentInfo": "Audio Latent Info",
}
nodes/acestep/nodes.py
+542 -46
View File
@@ -1,7 +1,17 @@
import torch
import os
import comfy.model_management
import comfy.samplers
from .ace_step_guiders import ACEStepRepaintGuider, ACEStepExtendGuider, ACEStepHybridGuider
import folder_paths
# Import patches module to apply ACE-Step 1.5 patches
from . import patches
from .ace_step_guiders import (
ACEStepRepaintGuider, ACEStepExtendGuider, ACEStepHybridGuider,
ACEStep15NativeEditGuider,
ACEStep15NativeCoverGuider, ACEStep15NativeExtractGuider, ACEStep15NativeLegoGuider
)
from .ace_step_utils import ACEStepLatentUtils
from .audio_mask_nodes import (
AudioTemporalMask, AudioRegionMask, AudioMaskAnalyzer,
@@ -9,19 +19,109 @@ from .audio_mask_nodes import (
)
def validate_audio_latent(latents):
"""Validate that latents are audio latents, handling both typed and untyped cases"""
"""Validate that latents are audio latents, handling both typed and untyped cases
Supports both v1.0 (batch, 8, 16, length) and v1.5 (batch, 64, length) shapes
"""
# Check if it's explicitly marked as audio type
if latents.get("type") == "audio":
return True
# If no type field, check if it has the ACE audio latent shape (batch, 8, 16, length)
# If no type field, check shape
if "samples" in latents:
tensor = latents["samples"]
# v1.0: (batch, 8, 16, length)
if len(tensor.shape) == 4 and tensor.shape[1] == 8 and tensor.shape[2] == 16:
return True
# v1.5: (batch, 64, length)
if len(tensor.shape) == 3 and tensor.shape[1] == 64:
return True
return False
class ACEStep15SilenceLatentLoader:
"""
Loads the ACE-Step 1.5 silence_latent tensor required for repaint and lego tasks.
The silence_latent is a LEARNED tensor that represents "silence" in the latent space.
It is NOT zeros - it's a specific tensor that the model was trained with.
Downloads automatically from HuggingFace on first use.
"""
# HuggingFace URL for silence_latent.pt
SILENCE_LATENT_URL = "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/main/acestep-v15-turbo/silence_latent.pt"
@classmethod
def INPUT_TYPES(s):
return {
"required": {}
}
RETURN_TYPES = ("SILENCE_LATENT",)
RETURN_NAMES = ("silence_latent",)
FUNCTION = "load"
CATEGORY = "audio/acestep"
@classmethod
def get_silence_latent_path(cls):
"""Get the path where silence_latent.pt should be stored."""
# Use ComfyUI's models directory structure
models_dir = folder_paths.models_dir
ace_step_dir = os.path.join(models_dir, "ace_step")
return os.path.join(ace_step_dir, "silence_latent.pt")
@classmethod
def download_silence_latent(cls):
"""Download silence_latent.pt from HuggingFace."""
import urllib.request
silence_latent_path = cls.get_silence_latent_path()
ace_step_dir = os.path.dirname(silence_latent_path)
# Create directory if it doesn't exist
os.makedirs(ace_step_dir, exist_ok=True)
print(f"[ACE-Step] Downloading silence_latent.pt from HuggingFace...")
print(f"[ACE-Step] URL: {cls.SILENCE_LATENT_URL}")
print(f"[ACE-Step] Destination: {silence_latent_path}")
try:
urllib.request.urlretrieve(cls.SILENCE_LATENT_URL, silence_latent_path)
print(f"[ACE-Step] Download complete!")
return True
except Exception as e:
print(f"[ACE-Step] ERROR: Failed to download silence_latent.pt: {e}")
return False
def load(self):
"""Load the silence_latent tensor, downloading if necessary."""
silence_latent_path = self.get_silence_latent_path()
# Download if not exists
if not os.path.exists(silence_latent_path):
success = self.download_silence_latent()
if not success:
raise RuntimeError(
f"Failed to download silence_latent.pt. Please manually download from:\n"
f" {self.SILENCE_LATENT_URL}\n"
f"and place it at:\n"
f" {silence_latent_path}"
)
# Load the tensor
print(f"[ACE-Step] Loading silence_latent from {silence_latent_path}")
silence_latent = torch.load(silence_latent_path, map_location="cpu", weights_only=True)
# Transpose from [1, D, T] to [1, T, D] as per reference implementation
# handler.py line 465: self.silence_latent = torch.load(silence_latent_path).transpose(1, 2)
silence_latent = silence_latent.transpose(1, 2)
print(f"[ACE-Step] silence_latent loaded, shape: {silence_latent.shape}")
return (silence_latent,)
class ACEStepRepaintGuiderNode:
"""Node that creates a repaint guider for use with SamplerCustomAdvanced"""
@@ -32,7 +132,7 @@ class ACEStepRepaintGuiderNode:
"model": ("MODEL",),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
"cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
"source_latents": ("LATENT",),
"start_time": ("FLOAT", {"default": 10.0, "min": 0.0, "max": 1000.0, "step": 0.1}),
"end_time": ("FLOAT", {"default": 20.0, "min": 0.0, "max": 1000.0, "step": 0.1}),
@@ -77,7 +177,7 @@ class ACEStepExtendGuiderNode:
"model": ("MODEL",),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
"cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
"source_latents": ("LATENT",),
"extend_left_time": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.1}),
"extend_right_time": ("FLOAT", {"default": 10.0, "min": 0.0, "max": 100.0, "step": 0.1}),
@@ -90,22 +190,31 @@ class ACEStepExtendGuiderNode:
def get_guider(self, model, positive, negative, cfg, source_latents,
extend_left_time, extend_right_time):
print(f"[EXTEND_GUIDER_NODE] get_guider called")
print(f"[EXTEND_GUIDER_NODE] extend_left_time: {extend_left_time}")
print(f"[EXTEND_GUIDER_NODE] extend_right_time: {extend_right_time}")
if extend_left_time == 0 and extend_right_time == 0:
raise ValueError("At least one of extend_left_time or extend_right_time must be > 0")
if not validate_audio_latent(source_latents):
raise ValueError("source_latents must be audio latents (from VAEEncodeAudio or EmptyAceStepLatentAudio)")
# Extract latent tensor
latent_tensor = source_latents["samples"]
print(f"[EXTEND_GUIDER_NODE] latent_tensor.shape: {latent_tensor.shape}")
version = ACEStepLatentUtils.detect_version(latent_tensor)
print(f"[EXTEND_GUIDER_NODE] detected version: {version}")
# Create and return the extend guider
guider = ACEStepExtendGuider(
model, positive, negative, cfg,
latent_tensor, extend_left_time, extend_right_time
)
print(f"[EXTEND_GUIDER_NODE] guider.is_v1_5: {guider.is_v1_5}")
print(f"[EXTEND_GUIDER_NODE] guider.extended_latent.shape: {guider.extended_latent.shape}")
return (guider,)
class ACEStepHybridGuiderNode:
@@ -118,7 +227,7 @@ class ACEStepHybridGuiderNode:
"model": ("MODEL",),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
"cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
"source_latents": ("LATENT",),
"extend_left_time": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.1}),
"extend_right_time": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.1}),
@@ -186,24 +295,32 @@ class ACEStepAnalyzeLatentNode:
def analyze_latent(self, latents):
if not validate_audio_latent(latents):
return ("Not audio latents", 0.0, 0, 0)
latent_tensor = latents["samples"]
batch_size, channels, height, length = latent_tensor.shape
# Calculate duration
duration = ACEStepLatentUtils.frame_index_to_time(length)
version = ACEStepLatentUtils.detect_version(latent_tensor)
length = latent_tensor.shape[-1]
# Calculate duration using version-specific FPS
duration = ACEStepLatentUtils.frame_index_to_time(length, version)
# Create info string
info = f"ACE Audio Latent Analysis:\n"
info += f"Version: {version or 'unknown'}\n"
info += f"Shape: {latent_tensor.shape}\n"
info += f"Batch size: {batch_size}\n"
info += f"Channels: {channels}\n"
info += f"Height: {height}\n"
info += f"Batch size: {latent_tensor.shape[0]}\n"
if version == ACEStepLatentUtils.V1_5:
info += f"Channels: {latent_tensor.shape[1]}\n"
else:
info += f"Channels: {latent_tensor.shape[1]}\n"
info += f"Height: {latent_tensor.shape[2]}\n"
info += f"Length (frames): {length}\n"
info += f"Duration: {duration:.2f} seconds\n"
info += f"Frame rate: ~{length/duration:.1f} frames/second"
return (info, duration, length, channels)
fps = ACEStepLatentUtils.get_fps(version)
info += f"Frame rate: {fps:.2f} frames/second"
return (info, duration, length, latent_tensor.shape[1])
class ACEStepTimeRangeNode:
"""Node for converting time ranges to frame indices for debugging"""
@@ -258,24 +375,29 @@ class ACEStepMaskVisualizerNode:
def visualize_mask(self, latents, start_time, end_time, feather_time):
if not validate_audio_latent(latents):
return ("Not audio latents",)
latent_tensor = latents["samples"]
version = ACEStepLatentUtils.detect_version(latent_tensor)
# Create mask
start_frame = ACEStepLatentUtils.time_to_frame_index(start_time)
end_frame = ACEStepLatentUtils.time_to_frame_index(end_time)
feather_frames = ACEStepLatentUtils.time_to_frame_index(feather_time)
mask = ACEStepLatentUtils.create_repaint_mask(latent_tensor.shape, start_frame, end_frame)
start_frame = ACEStepLatentUtils.time_to_frame_index(start_time, version)
end_frame = ACEStepLatentUtils.time_to_frame_index(end_time, version)
feather_frames = ACEStepLatentUtils.time_to_frame_index(feather_time, version)
mask = ACEStepLatentUtils.create_repaint_mask(latent_tensor.shape, start_frame, end_frame, version)
if feather_frames > 0:
feathered_mask = ACEStepLatentUtils.create_feather_mask(mask, feather_frames)
else:
feathered_mask = mask
# Create visualization
total_frames = latent_tensor.shape[-1]
mask_summary = feathered_mask[0, 0, 0, :].cpu().numpy() # Take first channel of first batch
# Handle both 3D (v1.5) and 4D (v1.0) shapes
if len(feathered_mask.shape) == 3:
mask_summary = feathered_mask[0, 0, :].cpu().numpy()
else:
mask_summary = feathered_mask[0, 0, 0, :].cpu().numpy()
info = f"Mask Visualization (Length: {total_frames} frames):\n"
info += f"Repaint region: Frame {start_frame} to {end_frame}\n"
@@ -361,7 +483,11 @@ class ACEStepAudioPostProcessor:
# Check if this is an extend guider to determine padding strategy
if hasattr(guider, 'left_frames') and hasattr(guider, 'right_frames'):
# We know the extend amounts, pad accordingly
left_samples = int(guider.left_frames * 512 * 8 / 44100 * orig_sr)
# Use version-aware conversion if available
if hasattr(guider, 'version') and guider.version == ACEStepLatentUtils.V1_5:
left_samples = int(guider.left_frames / 25.0 * orig_sr) # v1.5: 25 fps
else:
left_samples = int(guider.left_frames * 512 * 8 / 44100 * orig_sr) # v1.0
right_samples = proc_length - orig_length - left_samples
@@ -431,10 +557,14 @@ class ACEStepAudioPostProcessor:
# Latent frames correspond to audio segments
latent_frames = latent_mask.shape[-1]
audio_frames_per_latent = audio_length // latent_frames
# Upsample mask to audio resolution
# Take the first channel/batch element for the mask pattern
mask_pattern = latent_mask[0, 0, 0, :].cpu() # Shape: [latent_frames]
# Handle both v1.5 (3D) and v1.0 (4D) shapes
if len(latent_mask.shape) == 3:
mask_pattern = latent_mask[0, 0, :].cpu() # v1.5: Shape: [latent_frames]
else:
mask_pattern = latent_mask[0, 0, 0, :].cpu() # v1.0: Shape: [latent_frames]
# Repeat each latent frame mask value for corresponding audio samples
audio_mask = mask_pattern.repeat_interleave(audio_frames_per_latent)
@@ -497,26 +627,392 @@ class ACEStepAudioPostProcessor:
return smoothed_mask
# =============================================================================
# ACE-Step 1.5 Native Guider Nodes
# These use the model's native mask input via model wrapping
# =============================================================================
class ACEStep15NativeEditGuiderNode:
"""
Unified ACE-Step 1.5 edit guider for extend and/or repaint operations.
This guider injects chunk_masks and src_latents directly into the model's forward()
method using model wrapping, rather than manipulating latents during sampling.
Operations:
- Extend: Set extend_left_seconds and/or extend_right_seconds > 0
- Repaint: Set repaint_start_seconds and repaint_end_seconds (repaint_end > repaint_start)
- Both: Combine extend and repaint in a single operation
IMPORTANT: Requires silence_latent from ACEStep15SilenceLatentLoader.
The silence_latent is a learned tensor that tells the model to generate
new content in the specified regions.
"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL",),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"source_latents": ("LATENT",),
"silence_latent": ("SILENCE_LATENT",),
"cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
},
"optional": {
"extend_left_seconds": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 240.0, "step": 0.1}),
"extend_right_seconds": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 240.0, "step": 0.1}),
"repaint_start_seconds": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 1000.0, "step": 0.1}),
"repaint_end_seconds": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 1000.0, "step": 0.1}),
}
}
RETURN_TYPES = ("GUIDER", "LATENT")
RETURN_NAMES = ("guider", "output_latent")
FUNCTION = "get_guider"
CATEGORY = "sampling/custom_sampling/guiders"
def get_guider(self, model, positive, negative, source_latents, silence_latent, cfg,
extend_left_seconds=0.0, extend_right_seconds=0.0,
repaint_start_seconds=-1.0, repaint_end_seconds=-1.0):
source_tensor = source_latents["samples"]
# Validate v1.5 shape: (batch, 64, length)
if len(source_tensor.shape) != 3 or source_tensor.shape[1] != 64:
raise ValueError(f"ACE-Step 1.5 requires latent shape (batch, 64, length), got {source_tensor.shape}. Use the v1.0 guiders for v1.0 latents.")
# Determine which operations are enabled
has_extend = extend_left_seconds > 0 or extend_right_seconds > 0
has_repaint = repaint_start_seconds >= 0 and repaint_end_seconds > repaint_start_seconds
if not has_extend and not has_repaint:
raise ValueError("At least one operation must be enabled: set extend_left/right_seconds > 0 or set valid repaint_start/end_seconds")
# Convert -1 values to None for the guider
repaint_start = repaint_start_seconds if repaint_start_seconds >= 0 else None
repaint_end = repaint_end_seconds if repaint_end_seconds >= 0 else None
print(f"[ACE15_EDIT_NODE] Creating guider")
print(f"[ACE15_EDIT_NODE] source_tensor.shape: {source_tensor.shape}")
print(f"[ACE15_EDIT_NODE] silence_latent.shape: {silence_latent.shape}")
if has_extend:
print(f"[ACE15_EDIT_NODE] extend: left={extend_left_seconds}s, right={extend_right_seconds}s")
if has_repaint:
print(f"[ACE15_EDIT_NODE] repaint: {repaint_start_seconds}s - {repaint_end_seconds}s")
# Create the unified guider
guider = ACEStep15NativeEditGuider(
model, positive, negative, cfg,
source_tensor,
silence_latent,
extend_left_seconds=extend_left_seconds,
extend_right_seconds=extend_right_seconds,
repaint_start_seconds=repaint_start,
repaint_end_seconds=repaint_end
)
# Return the working latent (may be extended or same as source)
output_latent = {"samples": guider.working_latent, "type": "audio"}
return (guider, output_latent)
class ACEStep15NativeCoverGuiderNode:
"""
Creates an ACE-Step 1.5 cover guider for style transfer/regeneration.
Uses the source audio as context while generating new content across
the entire duration. Useful for cover songs or regenerating with a new prompt.
"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL",),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
"source_latents": ("LATENT",),
}
}
RETURN_TYPES = ("GUIDER",)
FUNCTION = "get_guider"
CATEGORY = "sampling/custom_sampling/guiders"
def get_guider(self, model, positive, negative, cfg, source_latents):
source_tensor = source_latents["samples"]
# Validate v1.5 shape: (batch, 64, length)
if len(source_tensor.shape) != 3 or source_tensor.shape[1] != 64:
raise ValueError(f"ACE-Step 1.5 requires latent shape (batch, 64, length), got {source_tensor.shape}")
print(f"[ACE15_COVER_NODE] Creating guider")
print(f"[ACE15_COVER_NODE] source_tensor.shape: {source_tensor.shape}")
guider = ACEStep15NativeCoverGuider(
model, positive, negative, cfg, source_tensor
)
return (guider,)
class ACEStep15NativeExtractGuiderNode:
"""
Creates an ACE-Step 1.5 extract guider for extracting specific tracks.
Extracts a specific track (vocals, drums, bass, etc.) from the source audio.
Use with the ACEStep15TaskTextEncode node with task_type="extract".
"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL",),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
"source_latents": ("LATENT",),
"track_name": (["vocals", "drums", "bass", "guitar", "keyboard", "strings",
"percussion", "synth", "fx", "brass", "woodwinds", "backing_vocals"],),
}
}
RETURN_TYPES = ("GUIDER",)
FUNCTION = "get_guider"
CATEGORY = "sampling/custom_sampling/guiders"
def get_guider(self, model, positive, negative, cfg, source_latents, track_name):
source_tensor = source_latents["samples"]
# Validate v1.5 shape: (batch, 64, length)
if len(source_tensor.shape) != 3 or source_tensor.shape[1] != 64:
raise ValueError(f"ACE-Step 1.5 requires latent shape (batch, 64, length), got {source_tensor.shape}")
print(f"[ACE15_EXTRACT_NODE] Creating guider")
print(f"[ACE15_EXTRACT_NODE] source_tensor.shape: {source_tensor.shape}")
print(f"[ACE15_EXTRACT_NODE] track_name: {track_name}")
guider = ACEStep15NativeExtractGuider(
model, positive, negative, cfg, source_tensor, track_name
)
return (guider,)
class ACEStep15NativeLegoGuiderNode:
"""
Creates an ACE-Step 1.5 lego guider for generating specific tracks in a region.
Generates a specific track (vocals, drums, etc.) within a time region while
preserving the rest. Useful for adding instruments to existing audio.
Use with the ACEStep15TaskTextEncode node with task_type="lego".
IMPORTANT: Requires silence_latent from ACEStep15SilenceLatentLoader.
The silence_latent is a learned tensor that tells the model to generate
new content in the lego region.
"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL",),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
"source_latents": ("LATENT",),
"silence_latent": ("SILENCE_LATENT",),
"track_name": (["vocals", "drums", "bass", "guitar", "keyboard", "strings",
"percussion", "synth", "fx", "brass", "woodwinds", "backing_vocals"],),
"start_seconds": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.1}),
"end_seconds": ("FLOAT", {"default": 30.0, "min": 0.0, "max": 1000.0, "step": 0.1}),
}
}
RETURN_TYPES = ("GUIDER",)
FUNCTION = "get_guider"
CATEGORY = "sampling/custom_sampling/guiders"
def get_guider(self, model, positive, negative, cfg, source_latents,
silence_latent, track_name, start_seconds, end_seconds):
if start_seconds >= end_seconds:
raise ValueError(f"start_seconds ({start_seconds}) must be less than end_seconds ({end_seconds})")
source_tensor = source_latents["samples"]
# Validate v1.5 shape: (batch, 64, length)
if len(source_tensor.shape) != 3 or source_tensor.shape[1] != 64:
raise ValueError(f"ACE-Step 1.5 requires latent shape (batch, 64, length), got {source_tensor.shape}")
print(f"[ACE15_LEGO_NODE] Creating guider")
print(f"[ACE15_LEGO_NODE] source_tensor.shape: {source_tensor.shape}")
print(f"[ACE15_LEGO_NODE] silence_latent.shape: {silence_latent.shape}")
print(f"[ACE15_LEGO_NODE] track_name: {track_name}")
print(f"[ACE15_LEGO_NODE] region: {start_seconds}s - {end_seconds}s")
guider = ACEStep15NativeLegoGuider(
model, positive, negative, cfg, source_tensor,
silence_latent, track_name, start_seconds, end_seconds
)
return (guider,)
class ACEStep15TaskTextEncodeNode:
"""
Task-aware text encoder for ACE-Step 1.5.
Encodes text with task-specific instructions:
- text2music: Generate audio from text description
- repaint: Regenerate a region while preserving the rest
- cover: Style transfer/regeneration with source audio as context
- extract: Extract a specific track (requires track_name)
- lego: Generate a specific track in a region (requires track_name)
The task_type determines the instruction prefix used by the model.
For extract/lego tasks, track_name specifies which track to extract/generate.
Returns single conditioning output. Use ConditioningZeroOut node for negative
conditioning (standard ACE-Step workflow pattern).
"""
# Valid keyscales from reference: 7 notes × 3 accidentals (plain, #, b) × 2 modes = 42 combinations
# Using ASCII accidentals for compatibility
KEYSCALE_NOTES = ['C', 'D', 'E', 'F', 'G', 'A', 'B']
KEYSCALE_ACCIDENTALS = ['', '#', 'b']
KEYSCALE_MODES = ['major', 'minor']
# Generate all valid keyscales
VALID_KEYSCALES = []
for note in KEYSCALE_NOTES:
for acc in KEYSCALE_ACCIDENTALS:
for mode in KEYSCALE_MODES:
VALID_KEYSCALES.append(f"{note}{acc} {mode}")
# Valid time signatures from reference
VALID_TIME_SIGNATURES = ["2", "3", "4", "6"]
# Valid languages from reference
VALID_LANGUAGES = [
'en', 'zh', 'ja', 'ko', 'es', 'fr', 'de', 'it', 'pt', 'ru',
'ar', 'hi', 'vi', 'th', 'id', 'ms', 'tl', 'nl', 'pl', 'tr',
'sv', 'da', 'no', 'fi', 'cs', 'sk', 'hu', 'ro', 'bg', 'hr',
'sr', 'uk', 'el', 'he', 'fa', 'bn', 'ta', 'te', 'pa', 'ur',
'ne', 'sw', 'ht', 'is', 'lt', 'la', 'az', 'ca', 'sa', 'yue',
'unknown'
]
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"clip": ("CLIP",),
"text": ("STRING", {"multiline": True, "dynamicPrompts": True, "default": "A melodic electronic track with soft synths"}),
"task_type": (["text2music", "repaint", "cover", "extract", "lego"],),
},
"optional": {
"track_name": (["", "vocals", "drums", "bass", "guitar", "keyboard", "strings",
"percussion", "synth", "fx", "brass", "woodwinds", "backing_vocals"],
{"default": ""}),
"lyrics": ("STRING", {"multiline": True, "default": ""}),
"bpm": ("INT", {"default": 120, "min": 30, "max": 300}),
"duration": ("INT", {"default": 60, "min": 10, "max": 600}),
# Use COMBO type to accept connections from AudioInfo detected_key output
"keyscale": ("COMBO", {"default": "C major", "options": s.VALID_KEYSCALES}),
"timesignature": ("COMBO", {"default": "4", "options": s.VALID_TIME_SIGNATURES}),
"language": ("COMBO", {"default": "en", "options": s.VALID_LANGUAGES}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff}),
}
}
RETURN_TYPES = ("CONDITIONING",)
RETURN_NAMES = ("conditioning",)
FUNCTION = "encode"
CATEGORY = "conditioning"
def encode(self, clip, text, task_type, track_name="", lyrics="", bpm=120,
duration=60, keyscale="C major", timesignature="4", language="en", seed=0):
# Validate track_name for extract/lego tasks
if task_type in ["extract", "lego"] and not track_name:
print(f"[ACE15_TEXT_ENCODE] Warning: track_name not specified for {task_type} task, using default instruction")
# Convert timesignature from string to int
timesig_int = int(timesignature)
# Build kwargs for tokenizer
kwargs = {
"lyrics": lyrics,
"bpm": bpm,
"duration": duration,
"keyscale": keyscale,
"timesignature": timesig_int,
"language": language,
"seed": seed,
"task_type": task_type,
"track_name": track_name if track_name else None,
}
print(f"[ACE15_TEXT_ENCODE] Encoding with task_type={task_type}, track_name={track_name or 'N/A'}")
# Use the patched tokenize_with_weights which accepts task_type
# encode_from_tokens_scheduled returns conditioning directly (list format)
tokens = clip.tokenize(text, **kwargs)
conditioning = clip.encode_from_tokens_scheduled(tokens)
# For negative conditioning, use ComfyUI's ConditioningZeroOut node
# This matches the standard ACE-Step workflow pattern
return (conditioning,)
# Node class mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
# ACE-Step 1.0 guiders (latent-level ODE blending)
"ACEStepRepaintGuider": ACEStepRepaintGuiderNode,
"ACEStepExtendGuider": ACEStepExtendGuiderNode,
"ACEStepHybridGuider": ACEStepHybridGuiderNode,
# Utility nodes
"ACEStepAnalyzeLatent": ACEStepAnalyzeLatentNode,
"ACEStepTimeRange": ACEStepTimeRangeNode,
"ACEStepMaskVisualizer": ACEStepMaskVisualizerNode,
"ACEStepAudioPostProcessor": ACEStepAudioPostProcessor,
# ACE-Step 1.5 loaders
"ACEStep15SilenceLatentLoader": ACEStep15SilenceLatentLoader,
# ACE-Step 1.5 guiders (model-level mask input)
"ACEStep15NativeEditGuider": ACEStep15NativeEditGuiderNode,
"ACEStep15NativeCoverGuider": ACEStep15NativeCoverGuiderNode,
"ACEStep15NativeExtractGuider": ACEStep15NativeExtractGuiderNode,
"ACEStep15NativeLegoGuider": ACEStep15NativeLegoGuiderNode,
# ACE-Step 1.5 text encoder
"ACEStep15TaskTextEncode": ACEStep15TaskTextEncodeNode,
**AUDIO_MASK_NODE_CLASS_MAPPINGS, # Add audio mask nodes
}
# Display name mappings for ComfyUI
NODE_DISPLAY_NAME_MAPPINGS = {
"ACEStepRepaintGuider": "ACEStep Repaint Guider BETA",
"ACEStepExtendGuider": "ACEStep Extend Guider BETA",
"ACEStepHybridGuider": "ACEStep Hybrid Guider BETA",
"ACEStepAnalyzeLatent": "ACEStep Analyze Latent BETA",
"ACEStepTimeRange": "ACEStep Time Range BETA",
"ACEStepMaskVisualizer": "ACEStep Mask Visualizer BETA",
"ACEStepAudioPostProcessor": "ACEStep Audio Post Processor BETA",
**AUDIO_MASK_NODE_DISPLAY_NAME_MAPPINGS, # Add audio mask node display names
# ACE-Step 1.0 guiders
"ACEStepRepaintGuider": "ACEStep 1.0 Repaint Guider BETA",
"ACEStepExtendGuider": "ACEStep 1.0 Extend Guider BETA",
"ACEStepHybridGuider": "ACEStep 1.0 Hybrid Guider BETA",
# Utility nodes
"ACEStepAnalyzeLatent": "ACEStep Analyze Latent",
"ACEStepTimeRange": "ACEStep Time Range",
"ACEStepMaskVisualizer": "ACEStep Mask Visualizer",
"ACEStepAudioPostProcessor": "ACEStep Audio Post Processor",
# ACE-Step 1.5 loaders
"ACEStep15SilenceLatentLoader": "ACE-Step 1.5 Load Silence Latent",
# ACE-Step 1.5 guiders
"ACEStep15NativeEditGuider": "ACE-Step 1.5 Edit Guider (Extend/Repaint)",
"ACEStep15NativeCoverGuider": "ACE-Step 1.5 Cover Guider",
"ACEStep15NativeExtractGuider": "ACE-Step 1.5 Extract Guider",
"ACEStep15NativeLegoGuider": "ACE-Step 1.5 Lego Guider",
# ACE-Step 1.5 text encoder
"ACEStep15TaskTextEncode": "ACE-Step 1.5 Task Text Encode",
**AUDIO_MASK_NODE_DISPLAY_NAME_MAPPINGS,
}
nodes/acestep/patches.py
+232
View File
@@ -0,0 +1,232 @@
"""
Monkey patches for ACE-Step 1.5 support in ComfyUI.
These patches extend ComfyUI's ACE-Step implementation to support all task types:
- text2music: Pure generation from text
- repaint: Regenerate a region while preserving the rest
- extend: Add new audio before/after existing audio
- cover: Style transfer/regeneration with target audio as context
- extract: Extract specific track (e.g., vocals, drums) from audio
- lego: Generate specific track within a region
Patches are applied at module import time.
"""
import torch
import math
import functools
# Flag to track if patches have been applied
_patches_applied = False
# Task instructions from official ACE-Step 1.5 constants
TASK_INSTRUCTIONS = {
"text2music": "Fill the audio semantic mask based on the given conditions:",
"repaint": "Repaint the mask area based on the given conditions:",
"cover": "Generate audio semantic tokens based on the given conditions:",
"extract": "Extract the {TRACK_NAME} track from the audio:",
"extract_default": "Extract the track from the audio:",
"lego": "Generate the {TRACK_NAME} track based on the audio context:",
"lego_default": "Generate the track based on the audio context:",
"complete": "Complete the input track with {TRACK_CLASSES}:",
"complete_default": "Complete the input track:",
}
# Valid track names for extract/lego tasks
TRACK_NAMES = [
"woodwinds", "brass", "fx", "synth", "strings", "percussion",
"keyboard", "guitar", "bass", "drums", "backing_vocals", "vocals"
]
def get_task_instruction(task_type: str, track_name: str = None) -> str:
"""
Get the appropriate instruction for a task type.
Args:
task_type: One of text2music, repaint, cover, extract, lego, complete
track_name: Required for extract/lego tasks (e.g., "vocals", "drums")
Returns:
The instruction string for the task
"""
if task_type not in TASK_INSTRUCTIONS and f"{task_type}_default" not in TASK_INSTRUCTIONS:
raise ValueError(f"Unknown task type: {task_type}. Valid types: {list(TASK_INSTRUCTIONS.keys())}")
instruction = TASK_INSTRUCTIONS.get(task_type, TASK_INSTRUCTIONS.get(f"{task_type}_default", ""))
# Handle track name placeholder for extract/lego
if "{TRACK_NAME}" in instruction:
if track_name:
instruction = instruction.format(TRACK_NAME=track_name.upper())
else:
# Fall back to default instruction without track name
instruction = TASK_INSTRUCTIONS.get(f"{task_type}_default", instruction.replace("{TRACK_NAME}", ""))
return instruction
def _create_patched_ace_step15_forward(original_forward):
"""
Create a patched forward method for AceStepConditionGenerationModel.
The patched version:
1. Accepts chunk_masks and src_latents as kwargs
2. Sets is_covers=0 when src_latents is provided (to use provided latents instead of lm_hints)
3. Properly transposes user-provided tensors from (batch, channels, length) to (batch, length, channels)
This enables extend, repaint, cover, extract, and lego task types.
"""
@functools.wraps(original_forward)
def patched_forward(self, x, timestep, context, lyric_embed=None, refer_audio=None,
audio_codes=None, chunk_masks=None, src_latents=None, **kwargs):
text_attention_mask = None
lyric_attention_mask = None
refer_audio_order_mask = None
attention_mask = None
# chunk_masks and src_latents now come from kwargs for extend/repaint support
# When src_latents is explicitly provided (for extend/repaint), we need is_covers=0
# so prepare_condition() uses our src_latents instead of replacing with lm_hints
is_covers = None
if src_latents is not None:
# Set is_covers to 0 to indicate we want to use the provided src_latents
is_covers = torch.zeros((x.shape[0],), device=x.device, dtype=torch.long)
precomputed_lm_hints_25Hz = None
lyric_hidden_states = lyric_embed
text_hidden_states = context
refer_audio_acoustic_hidden_states_packed = refer_audio.movedim(-1, -2)
x = x.movedim(-1, -2)
if refer_audio_order_mask is None:
refer_audio_order_mask = torch.zeros((x.shape[0],), device=x.device, dtype=torch.long)
# Transpose user-provided src_latents/chunk_masks from (batch, channels, length) to (batch, length, channels)
if src_latents is not None:
src_latents = src_latents.movedim(-1, -2)
if chunk_masks is not None:
chunk_masks = chunk_masks.movedim(-1, -2)
if src_latents is None and is_covers is None:
src_latents = x
if chunk_masks is None:
chunk_masks = torch.ones_like(x)
enc_hidden, enc_mask, context_latents = self.prepare_condition(
text_hidden_states, text_attention_mask,
lyric_hidden_states, lyric_attention_mask,
refer_audio_acoustic_hidden_states_packed, refer_audio_order_mask,
src_latents, chunk_masks, is_covers,
precomputed_lm_hints_25Hz=precomputed_lm_hints_25Hz,
audio_codes=audio_codes
)
out = self.decoder(
hidden_states=x,
timestep=timestep,
timestep_r=timestep,
attention_mask=attention_mask,
encoder_hidden_states=enc_hidden,
encoder_attention_mask=enc_mask,
context_latents=context_latents
)
return out.movedim(-1, -2)
return patched_forward
def _create_patched_tokenizer(original_tokenize_with_weights):
"""
Create a patched tokenize_with_weights method for ACE15Tokenizer.
The patched version:
1. Accepts task_type kwarg to determine the instruction
2. Accepts track_name kwarg for extract/lego tasks
3. Uses appropriate instruction template based on task type
"""
@functools.wraps(original_tokenize_with_weights)
def patched_tokenize_with_weights(self, text, return_word_ids=False, **kwargs):
out = {}
lyrics = kwargs.get("lyrics", "")
bpm = kwargs.get("bpm", 120)
duration = kwargs.get("duration", 120)
keyscale = kwargs.get("keyscale", "C major")
timesignature = kwargs.get("timesignature", 2)
language = kwargs.get("language", "en")
seed = kwargs.get("seed", 0)
# New kwargs for task type support
task_type = kwargs.get("task_type", "text2music")
track_name = kwargs.get("track_name", None)
# Get appropriate instruction for task type
instruction = get_task_instruction(task_type, track_name)
duration = math.ceil(duration)
meta_lm = 'bpm: {}\nduration: {}\nkeyscale: {}\ntimesignature: {}'.format(bpm, duration, keyscale, timesignature)
# Use task-specific instruction in template
lm_template = "<|im_start|>system\n# Instruction\n{}\n\n<|im_end|>\n<|im_start|>user\n# Caption\n{}\n{}\n<|im_end|>\n<|im_start|>assistant\n<think>\n{}\n</think>\n\n<|im_end|>\n"
meta_cap = '- bpm: {}\n- timesignature: {}\n- keyscale: {}\n- duration: {}\n'.format(bpm, timesignature, keyscale, duration)
out["lm_prompt"] = self.qwen3_06b.tokenize_with_weights(lm_template.format(instruction, text, lyrics, meta_lm), disable_weights=True)
out["lm_prompt_negative"] = self.qwen3_06b.tokenize_with_weights(lm_template.format(instruction, text, lyrics, ""), disable_weights=True)
out["lyrics"] = self.qwen3_06b.tokenize_with_weights("# Languages\n{}\n\n# Lyric{}<|endoftext|><|endoftext|>".format(language, lyrics), return_word_ids, disable_weights=True, **kwargs)
out["qwen3_06b"] = self.qwen3_06b.tokenize_with_weights("# Instruction\n{}\n\n# Caption\n{}# Metas\n{}<|endoftext|>\n<|endoftext|>".format(instruction, text, meta_cap), return_word_ids, **kwargs)
out["lm_metadata"] = {"min_tokens": duration * 5, "seed": seed}
# Store task info for downstream use
out["task_type"] = task_type
out["track_name"] = track_name
return out
return patched_tokenize_with_weights
def apply_patches():
"""
Apply all ACE-Step 1.5 patches.
This should be called once at module load time.
Safe to call multiple times (will only apply once).
"""
global _patches_applied
if _patches_applied:
return
try:
# Patch AceStepConditionGenerationModel.forward
from comfy.ldm.ace.ace_step15 import AceStepConditionGenerationModel
original_forward = AceStepConditionGenerationModel.forward
AceStepConditionGenerationModel.forward = _create_patched_ace_step15_forward(original_forward)
print("[ACE-Step Patches] Patched AceStepConditionGenerationModel.forward")
except ImportError as e:
print(f"[ACE-Step Patches] Warning: Could not patch ace_step15: {e}")
try:
# Patch ACE15Tokenizer.tokenize_with_weights
from comfy.text_encoders.ace15 import ACE15Tokenizer
original_tokenize = ACE15Tokenizer.tokenize_with_weights
ACE15Tokenizer.tokenize_with_weights = _create_patched_tokenizer(original_tokenize)
print("[ACE-Step Patches] Patched ACE15Tokenizer.tokenize_with_weights")
except ImportError as e:
print(f"[ACE-Step Patches] Warning: Could not patch ace15 tokenizer: {e}")
_patches_applied = True
print("[ACE-Step Patches] All patches applied successfully")
def is_patched():
"""Check if patches have been applied."""
return _patches_applied
# Apply patches when this module is imported
apply_patches()
nodes/audio/audio_nodes_utility.py
+65 -11
View File
@@ -203,6 +203,14 @@ class AudioSubtract(AudioUtility):
#TODO: TOO SLOW
@apply_tooltips
class AudioInfo(AudioUtility):
# Krumhansl-Kessler key profiles for major and minor keys
MAJOR_PROFILE = np.array([6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88])
MINOR_PROFILE = np.array([6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17])
# Note names matching the keyscale format in TextEncodeAceStepAudio1.5
# Using sharps for detection, which covers all keys
NOTE_NAMES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]
@classmethod
def INPUT_TYPES(cls):
return {
@@ -212,47 +220,91 @@ class AudioInfo(AudioUtility):
}
}
RETURN_TYPES = ("INT", "INT", "INT", "INT", "INT", "FLOAT", "FLOAT", "FLOAT", "INT", "INT", "INT", "FLOAT", "FLOAT", "FLOAT", "STRING")
RETURN_TYPES = ("INT", "INT", "INT", "INT", "INT", "FLOAT", "FLOAT", "FLOAT", "INT", "INT", "INT", "INT", "FLOAT", "FLOAT", "FLOAT", "STRING", "COMBO")
RETURN_NAMES = (
"total_frames", "frames_per_beat", "frames_per_bar", "frames_per_quarter", "frames_per_eighth",
"audio_duration", "beats_per_second", "detected_bpm",
"audio_duration", "beats_per_second", "detected_bpm", "detected_bpm_rounded",
"sample_rate", "num_channels", "num_samples",
"max_amplitude", "mean_amplitude", "rms_amplitude", "bit_depth"
"max_amplitude", "mean_amplitude", "rms_amplitude", "bit_depth",
"detected_key"
)
FUNCTION = "get_audio_info"
def _detect_key(self, audio_mono, sample_rate):
"""Detect musical key using chromagram and Krumhansl-Kessler key profiles."""
# Extract chromagram using Constant-Q Transform (better for pitch detection)
chromagram = librosa.feature.chroma_cqt(y=audio_mono, sr=sample_rate)
# Average chroma values across time to get a single 12-element profile
chroma_profile = np.mean(chromagram, axis=1)
# Normalize the profile
chroma_profile = chroma_profile / (np.linalg.norm(chroma_profile) + 1e-8)
best_correlation = -1
best_key = "C major"
# Test all 24 keys (12 major + 12 minor)
for i in range(12):
# Rotate the key profiles to match each root note
major_rotated = np.roll(self.MAJOR_PROFILE, i)
minor_rotated = np.roll(self.MINOR_PROFILE, i)
# Normalize profiles
major_rotated = major_rotated / (np.linalg.norm(major_rotated) + 1e-8)
minor_rotated = minor_rotated / (np.linalg.norm(minor_rotated) + 1e-8)
# Calculate correlation with major key
major_corr = np.corrcoef(chroma_profile, major_rotated)[0, 1]
if major_corr > best_correlation:
best_correlation = major_corr
best_key = f"{self.NOTE_NAMES[i]} major"
# Calculate correlation with minor key
minor_corr = np.corrcoef(chroma_profile, minor_rotated)[0, 1]
if minor_corr > best_correlation:
best_correlation = minor_corr
best_key = f"{self.NOTE_NAMES[i]} minor"
return best_key
def get_audio_info(self, audio, frame_rate):
# Get basic audio info
waveform = audio['waveform']
sample_rate = audio['sample_rate']
# Calculate original audio info first
num_channels = waveform.shape[1] if waveform.dim() > 2 else 1
num_samples = waveform.shape[-1]
audio_duration = num_samples / sample_rate
# Calculate total frames
total_frames = int(audio_duration * frame_rate)
# Detect BPM using librosa
audio_mono = waveform.squeeze(0).mean(axis=0).cpu().numpy()
tempo, _ = librosa.beat.beat_track(y=audio_mono, sr=sample_rate)
tempo_scalar = float(tempo.item() if hasattr(tempo, 'item') else tempo)
detected_bpm_rounded = int(round(tempo_scalar))
beats_per_second = tempo / 60.0
# Calculate frames per beat and musical divisions
frames_per_beat = int(frame_rate / beats_per_second)
frames_per_bar = frames_per_beat * 4 # Assuming 4/4 time signature
frames_per_quarter = frames_per_beat
frames_per_eighth = frames_per_beat // 2
# Calculate amplitude statistics
max_amplitude = float(torch.max(torch.abs(waveform)))
mean_amplitude = float(torch.mean(torch.abs(waveform)))
rms_amplitude = float(torch.sqrt(torch.mean(waveform ** 2)))
# Get bit depth from dtype
bit_depth = str(waveform.dtype)
# Detect musical key
detected_key = self._detect_key(audio_mono, sample_rate)
return (
total_frames,
frames_per_beat,
@@ -262,13 +314,15 @@ class AudioInfo(AudioUtility):
audio_duration,
beats_per_second,
tempo, # detected_bpm
detected_bpm_rounded,
sample_rate,
num_channels,
num_samples,
max_amplitude,
mean_amplitude,
rms_amplitude,
bit_depth
bit_depth,
detected_key
)
@apply_tooltips
requirements.txt
-1
View File
@@ -6,6 +6,5 @@ pymunk>=6.8.0,<7.0.0
matplotlib
openunmix
mido
torchaudio
pygame
scikit-image