Merge pull request #73 from ryanontheinside/fix/init-reqs
ACE Step 1.5 Repaint, Extend, Cover support
This commit is contained in:
RyanOnTheInside
File diff suppressed because one or more lines are too long
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1 @@
|
||||
# Empty init file for nodes package
|
||||
@@ -0,0 +1,26 @@
|
||||
"""
|
||||
ACE-Step nodes for ComfyUI.
|
||||
|
||||
This package provides nodes for ACE-Step 1.0 and 1.5 audio generation,
|
||||
including extend, repaint, cover, and extract operations.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
# Create logger for the acestep module
|
||||
logger = logging.getLogger("comfyui_ryanontheinside.acestep")
|
||||
|
||||
# Set default level from environment variable or default to INFO
|
||||
# Use ACESTEP_LOG_LEVEL=DEBUG for verbose output
|
||||
log_level = os.environ.get("ACESTEP_LOG_LEVEL", "INFO").upper()
|
||||
logger.setLevel(getattr(logging, log_level, logging.INFO))
|
||||
|
||||
# Only add handler if none exist (avoid duplicate handlers on reload)
|
||||
if not logger.handlers:
|
||||
handler = logging.StreamHandler()
|
||||
handler.setFormatter(logging.Formatter("[%(name)s] %(message)s"))
|
||||
logger.addHandler(handler)
|
||||
|
||||
# Prevent propagation to root logger to avoid duplicate messages
|
||||
logger.propagate = False
|
||||
File diff suppressed because it is too large
Load Diff
+323
-108
@@ -1,116 +1,326 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import math
|
||||
import os
|
||||
import comfy.model_management
|
||||
import folder_paths
|
||||
from . import logger
|
||||
|
||||
# Module-level cache for silence_latent (loaded once per session)
|
||||
_silence_latent_cache = None
|
||||
|
||||
# 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"
|
||||
|
||||
|
||||
def get_silence_latent_path():
|
||||
"""Get the path where silence_latent.pt should be stored."""
|
||||
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")
|
||||
|
||||
|
||||
def download_silence_latent():
|
||||
"""Download silence_latent.pt from HuggingFace."""
|
||||
import urllib.request
|
||||
|
||||
silence_latent_path = 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)
|
||||
|
||||
logger.info(f"[ACE-Step] Downloading silence_latent.pt from HuggingFace...")
|
||||
logger.info(f"[ACE-Step] URL: {SILENCE_LATENT_URL}")
|
||||
logger.info(f"[ACE-Step] Destination: {silence_latent_path}")
|
||||
|
||||
try:
|
||||
urllib.request.urlretrieve(SILENCE_LATENT_URL, silence_latent_path)
|
||||
logger.info(f"[ACE-Step] Download complete!")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.info(f"[ACE-Step] ERROR: Failed to download silence_latent.pt: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def load_silence_latent(verbose=True):
|
||||
"""
|
||||
Load the silence_latent tensor, downloading if necessary.
|
||||
|
||||
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.
|
||||
This tells the model to "generate new content here".
|
||||
|
||||
Returns:
|
||||
torch.Tensor: silence_latent in shape [1, T, D] (transposed from file)
|
||||
"""
|
||||
global _silence_latent_cache
|
||||
|
||||
# Return cached version if available
|
||||
if _silence_latent_cache is not None:
|
||||
if verbose:
|
||||
logger.info(f"[ACE-Step] Using cached silence_latent, shape: {_silence_latent_cache.shape}")
|
||||
return _silence_latent_cache
|
||||
|
||||
silence_latent_path = get_silence_latent_path()
|
||||
|
||||
# Download if not exists
|
||||
if not os.path.exists(silence_latent_path):
|
||||
success = download_silence_latent()
|
||||
if not success:
|
||||
raise RuntimeError(
|
||||
f"Failed to download silence_latent.pt. Please manually download from:\n"
|
||||
f" {SILENCE_LATENT_URL}\n"
|
||||
f"and place it at:\n"
|
||||
f" {silence_latent_path}"
|
||||
)
|
||||
|
||||
# Load the tensor
|
||||
if verbose:
|
||||
logger.info(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)
|
||||
|
||||
if verbose:
|
||||
logger.info(f"[ACE-Step] silence_latent loaded, shape: {silence_latent.shape}")
|
||||
|
||||
# Cache for future use
|
||||
_silence_latent_cache = silence_latent
|
||||
|
||||
return silence_latent
|
||||
|
||||
|
||||
def extract_semantic_hints(model, source_latent, verbose=True):
|
||||
"""
|
||||
Extract semantic hints from source audio latents for Cover/Extract tasks.
|
||||
|
||||
This extracts the semantic structure of the audio (rhythm, melody, harmony)
|
||||
while abstracting away the timbre. Used by Cover and Extract guiders.
|
||||
|
||||
Args:
|
||||
model: ComfyUI ModelPatcher wrapping the ACE-Step 1.5 model
|
||||
source_latent: Source audio latent in ComfyUI format [B, 64, T]
|
||||
verbose: Whether to print diagnostic information
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Semantic hints in ComfyUI format [B, 64, T]
|
||||
"""
|
||||
prefix = "[SEMANTIC_EXTRACT]"
|
||||
|
||||
# Validate v1.5 shape: (batch, 64, length)
|
||||
if len(source_latent.shape) != 3 or source_latent.shape[1] != 64:
|
||||
raise ValueError(f"ACE-Step 1.5 requires latent shape (batch, 64, length), got {source_latent.shape}")
|
||||
|
||||
# Get the diffusion model (AceStepConditionGenerationModel)
|
||||
diffusion_model = model.model.diffusion_model
|
||||
|
||||
if not hasattr(diffusion_model, 'tokenizer') or not hasattr(diffusion_model, 'detokenizer'):
|
||||
raise RuntimeError(
|
||||
"Model does not have tokenizer/detokenizer. "
|
||||
"Make sure you're using an ACE-Step 1.5 model."
|
||||
)
|
||||
|
||||
# Load model to GPU
|
||||
comfy.model_management.load_model_gpu(model)
|
||||
|
||||
device = comfy.model_management.get_torch_device()
|
||||
dtype = model.model.get_dtype()
|
||||
|
||||
# Move source tensor to model device/dtype
|
||||
# Source is in ComfyUI format: [B, D, T] = [B, 64, length]
|
||||
# Tokenizer expects: [B, T, D] = [B, length, 64]
|
||||
source_on_device = source_latent.to(device=device, dtype=dtype)
|
||||
source_transposed = source_on_device.movedim(-1, -2) # [B, T, D]
|
||||
|
||||
if verbose:
|
||||
logger.debug(f"{prefix} Extracting from source shape: {source_latent.shape}")
|
||||
logger.debug(f"{prefix} source stats: mean={source_transposed.mean():.4f}, std={source_transposed.std():.4f}")
|
||||
|
||||
# Verify weights are on GPU
|
||||
tokenizer_device = next(diffusion_model.tokenizer.parameters()).device
|
||||
if verbose:
|
||||
logger.debug(f"{prefix} tokenizer on: {tokenizer_device}, target: {device}")
|
||||
|
||||
# Step 1: Tokenize - get quantized embeddings at 5Hz
|
||||
with torch.no_grad():
|
||||
quantized, indices = diffusion_model.tokenizer.tokenize(source_transposed)
|
||||
|
||||
if verbose:
|
||||
logger.debug(f"{prefix} quantized shape: {quantized.shape}")
|
||||
logger.debug(f"{prefix} quantized stats: mean={quantized.mean():.4f}, std={quantized.std():.4f}")
|
||||
|
||||
# Step 2: Detokenize - upsample from 5Hz to 25Hz
|
||||
# Use quantized directly (not get_output_from_indices) to avoid redundant computation
|
||||
with torch.no_grad():
|
||||
lm_hints = diffusion_model.detokenizer(quantized)
|
||||
|
||||
if verbose:
|
||||
logger.debug(f"{prefix} lm_hints shape: {lm_hints.shape}")
|
||||
logger.debug(f"{prefix} lm_hints stats: mean={lm_hints.mean():.4f}, std={lm_hints.std():.4f}")
|
||||
if lm_hints.std() < 0.01:
|
||||
logger.warning(f"{prefix} Very low variance in semantic hints!")
|
||||
|
||||
# Transpose back to ComfyUI format: [B, T, D] → [B, D, T]
|
||||
semantic_hints = lm_hints.movedim(-1, -2)
|
||||
|
||||
return semantic_hints
|
||||
|
||||
|
||||
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 +371,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 +497,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 +540,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)
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
+591
-46
@@ -1,7 +1,18 @@
|
||||
import torch
|
||||
import comfy.model_management
|
||||
import comfy.samplers
|
||||
from .ace_step_guiders import ACEStepRepaintGuider, ACEStepExtendGuider, ACEStepHybridGuider
|
||||
|
||||
# Import patches module to apply ACE-Step 1.5 patches
|
||||
from . import patches
|
||||
from . import logger
|
||||
|
||||
from .ace_step_guiders import (
|
||||
ACEStepRepaintGuider, ACEStepExtendGuider, ACEStepHybridGuider,
|
||||
ACEStep15NativeEditGuider,
|
||||
ACEStep15NativeCoverGuider,
|
||||
# TODO: Re-enable when ready
|
||||
# ACEStep15NativeExtractGuider, ACEStep15NativeLegoGuider
|
||||
)
|
||||
from .ace_step_utils import ACEStepLatentUtils
|
||||
from .audio_mask_nodes import (
|
||||
AudioTemporalMask, AudioRegionMask, AudioMaskAnalyzer,
|
||||
@@ -9,19 +20,123 @@ 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 ACEStep15SemanticExtractor:
|
||||
"""
|
||||
Extracts semantic tokens from source audio latents for use with Cover/Extract tasks.
|
||||
|
||||
The ACE-Step 1.5 model uses semantic tokens (lm_hints) as structural guidance
|
||||
for tasks like Cover and Extract. This node extracts those tokens from VAE latents.
|
||||
|
||||
Flow: VAE latents → tokenizer → quantized → detokenizer → semantic_hints
|
||||
|
||||
Use the output with ACE-Step 1.5 Cover Guider or Extract Guider for creative
|
||||
workflows like blending semantic hints from multiple sources.
|
||||
|
||||
Note: Cover and Extract guiders will auto-extract semantic hints if not provided.
|
||||
This node is useful for advanced workflows like semantic hint blending.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"source_latents": ("LATENT",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SEMANTIC_HINTS",)
|
||||
RETURN_NAMES = ("semantic_hints",)
|
||||
FUNCTION = "extract"
|
||||
CATEGORY = "audio/acestep"
|
||||
|
||||
def extract(self, model, source_latents):
|
||||
from .ace_step_utils import extract_semantic_hints
|
||||
|
||||
source_tensor = source_latents["samples"]
|
||||
logger.debug(f"[ACE15_SEMANTIC_EXTRACTOR] Extracting semantic hints from shape {source_tensor.shape}")
|
||||
|
||||
# Use shared utility function
|
||||
semantic_hints = extract_semantic_hints(model, source_tensor, verbose=True)
|
||||
|
||||
# Move to CPU for storage
|
||||
semantic_hints = semantic_hints.cpu()
|
||||
|
||||
return (semantic_hints,)
|
||||
|
||||
|
||||
class ACEStep15SemanticHintsBlend:
|
||||
"""
|
||||
Blends semantic hints from two sources for creative mashups.
|
||||
|
||||
Use this to:
|
||||
- Create song mashups by blending structure from two sources
|
||||
- Interpolate between musical structures
|
||||
- Create hybrid compositions
|
||||
|
||||
blend_factor: 0.0 = 100% hints_a, 1.0 = 100% hints_b, 0.5 = 50/50 mix
|
||||
|
||||
TODO: Enhance with feature system for temporal blending, crossfades,
|
||||
and region-based masking.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"hints_a": ("SEMANTIC_HINTS",),
|
||||
"hints_b": ("SEMANTIC_HINTS",),
|
||||
"blend_factor": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SEMANTIC_HINTS",)
|
||||
RETURN_NAMES = ("blended_hints",)
|
||||
FUNCTION = "blend"
|
||||
CATEGORY = "audio/acestep"
|
||||
|
||||
def blend(self, hints_a, hints_b, blend_factor):
|
||||
# Validate shapes match
|
||||
if hints_a.shape != hints_b.shape:
|
||||
# Try to match lengths by truncating to shorter
|
||||
min_length = min(hints_a.shape[-1], hints_b.shape[-1])
|
||||
hints_a = hints_a[..., :min_length]
|
||||
hints_b = hints_b[..., :min_length]
|
||||
logger.debug(f"[SEMANTIC_BLEND] Truncated to common length: {min_length}")
|
||||
|
||||
logger.debug(f"[SEMANTIC_BLEND] Blending hints")
|
||||
logger.debug(f"[SEMANTIC_BLEND] hints_a shape: {hints_a.shape}, stats: mean={hints_a.mean():.4f}, std={hints_a.std():.4f}")
|
||||
logger.debug(f"[SEMANTIC_BLEND] hints_b shape: {hints_b.shape}, stats: mean={hints_b.mean():.4f}, std={hints_b.std():.4f}")
|
||||
logger.debug(f"[SEMANTIC_BLEND] blend_factor: {blend_factor} (0=A, 1=B)")
|
||||
|
||||
# Simple linear interpolation
|
||||
blended = (1.0 - blend_factor) * hints_a + blend_factor * hints_b
|
||||
|
||||
logger.debug(f"[SEMANTIC_BLEND] result stats: mean={blended.mean():.4f}, std={blended.std():.4f}")
|
||||
|
||||
return (blended,)
|
||||
|
||||
|
||||
class ACEStepRepaintGuiderNode:
|
||||
"""Node that creates a repaint guider for use with SamplerCustomAdvanced"""
|
||||
|
||||
@@ -32,7 +147,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 +192,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 +205,31 @@ class ACEStepExtendGuiderNode:
|
||||
|
||||
def get_guider(self, model, positive, negative, cfg, source_latents,
|
||||
extend_left_time, extend_right_time):
|
||||
|
||||
|
||||
logger.debug(f"[EXTEND_GUIDER_NODE] get_guider called")
|
||||
logger.debug(f"[EXTEND_GUIDER_NODE] extend_left_time: {extend_left_time}")
|
||||
logger.debug(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"]
|
||||
|
||||
logger.debug(f"[EXTEND_GUIDER_NODE] latent_tensor.shape: {latent_tensor.shape}")
|
||||
version = ACEStepLatentUtils.detect_version(latent_tensor)
|
||||
logger.debug(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
|
||||
)
|
||||
|
||||
logger.debug(f"[EXTEND_GUIDER_NODE] guider.is_v1_5: {guider.is_v1_5}")
|
||||
logger.debug(f"[EXTEND_GUIDER_NODE] guider.extended_latent.shape: {guider.extended_latent.shape}")
|
||||
|
||||
return (guider,)
|
||||
|
||||
class ACEStepHybridGuiderNode:
|
||||
@@ -118,7 +242,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 +310,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 +390,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 +498,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 +572,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 +642,426 @@ 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
|
||||
|
||||
The silence_latent is automatically loaded internally (downloaded from HuggingFace
|
||||
on first use). It's a learned tensor that tells the model to generate new content.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"source_latents": ("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, 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
|
||||
|
||||
logger.debug(f"[ACE15_EDIT_NODE] Creating guider")
|
||||
logger.debug(f"[ACE15_EDIT_NODE] source_tensor.shape: {source_tensor.shape}")
|
||||
if has_extend:
|
||||
logger.debug(f"[ACE15_EDIT_NODE] extend: left={extend_left_seconds}s, right={extend_right_seconds}s")
|
||||
if has_repaint:
|
||||
logger.debug(f"[ACE15_EDIT_NODE] repaint: {repaint_start_seconds}s - {repaint_end_seconds}s")
|
||||
|
||||
# Create the unified guider (silence_latent is loaded internally)
|
||||
guider = ACEStep15NativeEditGuider(
|
||||
model, positive, negative, cfg,
|
||||
source_tensor,
|
||||
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 semantic tokens from the source audio as structural guidance while
|
||||
generating new content with a different style/timbre.
|
||||
|
||||
Semantic hints are automatically extracted if not provided. For advanced
|
||||
workflows like blending hints from multiple sources, use the Semantic
|
||||
Extractor and Semantic Hints Blend nodes.
|
||||
"""
|
||||
|
||||
@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",),
|
||||
},
|
||||
"optional": {
|
||||
"semantic_hints": ("SEMANTIC_HINTS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GUIDER",)
|
||||
FUNCTION = "get_guider"
|
||||
CATEGORY = "sampling/custom_sampling/guiders"
|
||||
|
||||
def get_guider(self, model, positive, negative, cfg, source_latents, semantic_hints=None):
|
||||
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}")
|
||||
|
||||
logger.debug(f"[ACE15_COVER_NODE] Creating guider")
|
||||
logger.debug(f"[ACE15_COVER_NODE] source_tensor.shape: {source_tensor.shape}")
|
||||
if semantic_hints is not None:
|
||||
logger.debug(f"[ACE15_COVER_NODE] semantic_hints provided externally: {semantic_hints.shape}")
|
||||
else:
|
||||
logger.debug(f"[ACE15_COVER_NODE] semantic_hints will be auto-extracted")
|
||||
|
||||
guider = ACEStep15NativeCoverGuider(
|
||||
model, positive, negative, cfg, source_tensor, semantic_hints
|
||||
)
|
||||
|
||||
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".
|
||||
|
||||
Semantic hints are automatically extracted if not provided. For advanced
|
||||
workflows, use the Semantic Extractor node.
|
||||
"""
|
||||
|
||||
@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"],),
|
||||
},
|
||||
"optional": {
|
||||
"semantic_hints": ("SEMANTIC_HINTS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GUIDER",)
|
||||
FUNCTION = "get_guider"
|
||||
CATEGORY = "sampling/custom_sampling/guiders"
|
||||
|
||||
def get_guider(self, model, positive, negative, cfg, source_latents, track_name, semantic_hints=None):
|
||||
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}")
|
||||
|
||||
logger.debug(f"[ACE15_EXTRACT_NODE] Creating guider")
|
||||
logger.debug(f"[ACE15_EXTRACT_NODE] source_tensor.shape: {source_tensor.shape}")
|
||||
logger.debug(f"[ACE15_EXTRACT_NODE] track_name: {track_name}")
|
||||
if semantic_hints is not None:
|
||||
logger.debug(f"[ACE15_EXTRACT_NODE] semantic_hints provided externally: {semantic_hints.shape}")
|
||||
else:
|
||||
logger.debug(f"[ACE15_EXTRACT_NODE] semantic_hints will be auto-extracted")
|
||||
|
||||
guider = ACEStep15NativeExtractGuider(
|
||||
model, positive, negative, cfg, source_tensor, track_name, semantic_hints
|
||||
)
|
||||
|
||||
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".
|
||||
|
||||
The silence_latent is automatically loaded internally (downloaded from HuggingFace
|
||||
on first use). It's a learned tensor that tells the model to generate new content.
|
||||
"""
|
||||
|
||||
@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"],),
|
||||
"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,
|
||||
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}")
|
||||
|
||||
logger.debug(f"[ACE15_LEGO_NODE] Creating guider")
|
||||
logger.debug(f"[ACE15_LEGO_NODE] source_tensor.shape: {source_tensor.shape}")
|
||||
logger.debug(f"[ACE15_LEGO_NODE] track_name: {track_name}")
|
||||
logger.debug(f"[ACE15_LEGO_NODE] region: {start_seconds}s - {end_seconds}s")
|
||||
|
||||
# Create guider (silence_latent is loaded internally)
|
||||
guider = ACEStep15NativeLegoGuider(
|
||||
model, positive, negative, cfg, source_tensor,
|
||||
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"],), # TODO: Re-enable "extract", "lego" when ready
|
||||
},
|
||||
"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:
|
||||
logger.debug(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,
|
||||
}
|
||||
|
||||
logger.debug(f"[ACE15_TEXT_ENCODE] Encoding with task_type={task_type}, track_name={track_name or 'N/A'}")
|
||||
|
||||
# Import to get task instruction for verification
|
||||
from .patches import get_task_instruction, is_patched
|
||||
instruction = get_task_instruction(task_type, track_name if track_name else None)
|
||||
logger.debug(f"[ACE15_TEXT_ENCODE] patches applied: {is_patched()}")
|
||||
logger.debug(f"[ACE15_TEXT_ENCODE] task instruction: {instruction}")
|
||||
|
||||
# Use the patched tokenize_with_weights which accepts task_type
|
||||
# encode_from_tokens_scheduled returns conditioning directly (list format)
|
||||
tokens = clip.tokenize(text, **kwargs)
|
||||
|
||||
# Verify the tokens contain task_type info (patched tokenizer should include this)
|
||||
if isinstance(tokens, dict):
|
||||
if "task_type" in tokens:
|
||||
logger.debug(f"[ACE15_TEXT_ENCODE] tokens contain task_type: {tokens.get('task_type')}")
|
||||
else:
|
||||
logger.debug(f"[ACE15_TEXT_ENCODE] WARNING: tokens dict does not contain task_type - patch may not be applied!")
|
||||
logger.debug(f"[ACE15_TEXT_ENCODE] tokens keys: {list(tokens.keys())}")
|
||||
|
||||
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 semantic extraction and blending
|
||||
"ACEStep15SemanticExtractor": ACEStep15SemanticExtractor,
|
||||
"ACEStep15SemanticHintsBlend": ACEStep15SemanticHintsBlend,
|
||||
# ACE-Step 1.5 guiders (model-level mask input)
|
||||
"ACEStep15NativeEditGuider": ACEStep15NativeEditGuiderNode,
|
||||
"ACEStep15NativeCoverGuider": ACEStep15NativeCoverGuiderNode,
|
||||
# TODO: Re-enable when ready
|
||||
# "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 semantic extraction and blending
|
||||
"ACEStep15SemanticExtractor": "ACE-Step 1.5 Semantic Extractor",
|
||||
"ACEStep15SemanticHintsBlend": "ACE-Step 1.5 Semantic Hints Blend",
|
||||
# ACE-Step 1.5 guiders
|
||||
"ACEStep15NativeEditGuider": "ACE-Step 1.5 Edit Guider (Extend/Repaint)",
|
||||
"ACEStep15NativeCoverGuider": "ACE-Step 1.5 Cover Guider",
|
||||
# TODO: Re-enable when ready
|
||||
# "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,
|
||||
}
|
||||
@@ -0,0 +1,260 @@
|
||||
"""
|
||||
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
|
||||
from . import logger
|
||||
|
||||
# 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, src_latents, is_covers, and precomputed_lm_hints_25Hz as kwargs
|
||||
2. Allows guiders to explicitly control is_covers flag for different task types
|
||||
3. Properly transposes user-provided tensors from (batch, channels, length) to (batch, length, channels)
|
||||
|
||||
Task type behavior via is_covers:
|
||||
- is_covers=0: Use provided src_latents directly (extend/repaint)
|
||||
- is_covers=1: Use lm_hints from precomputed_lm_hints_25Hz or audio_codes (cover/extract)
|
||||
- is_covers=None: Default behavior (src_latents = x)
|
||||
|
||||
This enables extend, repaint, cover, extract, and lego task types.
|
||||
"""
|
||||
# One-time logging flag
|
||||
_forward_logged = [False]
|
||||
|
||||
@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,
|
||||
is_covers=None, precomputed_lm_hints_25Hz=None, **kwargs):
|
||||
text_attention_mask = None
|
||||
lyric_attention_mask = None
|
||||
refer_audio_order_mask = None
|
||||
attention_mask = None
|
||||
|
||||
# One-time diagnostic log for cover task
|
||||
if not _forward_logged[0] and precomputed_lm_hints_25Hz is not None:
|
||||
logger.debug(f"[ACE15_PATCHED_FORWARD] Cover mode detected")
|
||||
logger.debug(f"[ACE15_PATCHED_FORWARD] x.shape (before transpose): {x.shape}")
|
||||
logger.debug(f"[ACE15_PATCHED_FORWARD] is_covers: {is_covers}")
|
||||
logger.debug(f"[ACE15_PATCHED_FORWARD] precomputed_lm_hints_25Hz.shape (before transpose): {precomputed_lm_hints_25Hz.shape}")
|
||||
logger.debug(f"[ACE15_PATCHED_FORWARD] precomputed_lm_hints_25Hz stats: mean={precomputed_lm_hints_25Hz.mean():.4f}, std={precomputed_lm_hints_25Hz.std():.4f}")
|
||||
_forward_logged[0] = True
|
||||
|
||||
# is_covers can now be passed explicitly by guiders:
|
||||
# - Cover/Extract tasks: is_covers=1 (use lm_hints from precomputed_lm_hints_25Hz)
|
||||
# - Extend/Repaint tasks: is_covers=0 (use provided src_latents)
|
||||
# - Default (not provided): fallback to old behavior
|
||||
if is_covers is None and src_latents is not None:
|
||||
# Backwards compatibility: default to 0 when src_latents provided but is_covers not specified
|
||||
is_covers = torch.zeros((x.shape[0],), device=x.device, dtype=torch.long)
|
||||
|
||||
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 tensors from ComfyUI format (batch, channels, length)
|
||||
# to model format (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 precomputed_lm_hints_25Hz is not None:
|
||||
# precomputed_lm_hints_25Hz comes in ComfyUI format [B, D, T], transpose to [B, T, D]
|
||||
precomputed_lm_hints_25Hz = precomputed_lm_hints_25Hz.movedim(-1, -2)
|
||||
# Log post-transpose shape (one-time)
|
||||
if _forward_logged[0] and not hasattr(patched_forward, '_post_transpose_logged'):
|
||||
logger.debug(f"[ACE15_PATCHED_FORWARD] precomputed_lm_hints_25Hz.shape (after transpose): {precomputed_lm_hints_25Hz.shape}")
|
||||
logger.debug(f"[ACE15_PATCHED_FORWARD] x.shape (after transpose): {x.shape}")
|
||||
logger.debug(f"[ACE15_PATCHED_FORWARD] src_latents.shape (after transpose): {src_latents.shape if src_latents is not None else 'None'}")
|
||||
patched_forward._post_transpose_logged = True
|
||||
|
||||
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)
|
||||
logger.info("[ACE-Step Patches] Patched AceStepConditionGenerationModel.forward")
|
||||
except ImportError as e:
|
||||
logger.info(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)
|
||||
logger.info("[ACE-Step Patches] Patched ACE15Tokenizer.tokenize_with_weights")
|
||||
except ImportError as e:
|
||||
logger.info(f"[ACE-Step Patches] Warning: Could not patch ace15 tokenizer: {e}")
|
||||
|
||||
_patches_applied = True
|
||||
logger.info("[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()
|
||||
@@ -0,0 +1 @@
|
||||
# Empty init file for audio package
|
||||
@@ -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
|
||||
|
||||
+2
-2
@@ -1,10 +1,10 @@
|
||||
[project]
|
||||
name = "comfyui_ryanonyheinside"
|
||||
description = "Custom nodes introducing everything-reactivity and timelines to ComfyUI"
|
||||
version = "2.0.12"
|
||||
version = "2.1.0"
|
||||
license = {file = "LICENSE"}
|
||||
|
||||
dependencies = ["pygame","opencv-python>=4.10.0","scipy","torchaudio", "pillow","librosa>=0.10.2","pymunk==6.8.1","matplotlib","openunmix","mido","scikit-image"]
|
||||
dependencies = ["pygame", "opencv-python>=4.5.0", "scipy", "torchaudio", "pillow", "librosa>=0.10.1", "pymunk>=6.8.0,<7.0.0", "matplotlib", "openunmix", "mido", "scikit-image"]
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/ryanontheinside/ComfyUI_RyanOnTheInside"
|
||||
|
||||
@@ -12,7 +12,7 @@
|
||||
- 🎵 **Audio & MIDI Processing**: Separate instruments, manipulate audio, work with ACEStep and other audio models
|
||||
- 🎆 **Particle Systems**: Create mesmerizing, fluid-like effects
|
||||
- 🌊 **Optical Flow**: Generate masks based on motion in videos
|
||||
- 🎨 **ACEStep Native**: Native ComfyUI implementation for audio repaint and extend functionality
|
||||
- 🎨 **ACE-Step 1.5 Native**: Native ComfyUI guiders for ACE-Step 1.5 — extend/repaint and cover
|
||||
- 🌊 **DEPTH Flow** 🌊: Use flex features to control depthflow nodes, adding parallax animations to your workflows!
|
||||
- 👤 **AdvancedLivePortrait** 👤: Use flex features to control facial animation expressions!
|
||||
- 🎨 **Advanced Controlnet** 🎨: Direct integration with ComfyUI-AdvancedControlnet!
|
||||
@@ -42,19 +42,6 @@ If you are not getting updates, or are having trouble updating from ComfyUI-Mana
|
||||
- **And much more!**
|
||||
|
||||
|
||||
### ⚠️ Breaking Changes Notice
|
||||
Due to ComfyUI's workflow loading mechanism, **existing workflows using these nodes ~~may~~ *will* break after updating**. I did consider this carefully, as I have yet to introduce breaking changes to this node system, but this extensive update neccesitated a complete overhaul.
|
||||
|
||||
If you need to run an older workflow, you can revert to the previous version of these nodes by using the Manager, or by running this command in your ComfyUI_RyanOnTheInside directory:
|
||||
|
||||
```bash
|
||||
git checkout dab96492ac7d906368ac9c7a17cb0dbd670923d9
|
||||
```
|
||||
|
||||
To return to the latest version later, use:
|
||||
```bash
|
||||
git checkout main
|
||||
```
|
||||
|
||||
<table style="border-collapse: collapse; border: none;">
|
||||
<tr>
|
||||
@@ -135,10 +122,11 @@ Dynamic control over various aspects of your workflow:
|
||||
- Extract features from audio and MIDI for visual effects
|
||||
- Create audio-reactive animations and transformations
|
||||
|
||||
### 🎨 ACEStep Native
|
||||
- **Repaint**: Selectively regenerate specific time ranges in existing audio
|
||||
- **Extend**: Add new content before or after existing audio
|
||||
- **Hybrid**: Combine repaint and extend operations simultaneously
|
||||
### 🎨 ACE-Step 1.5 Native
|
||||
Native ComfyUI guiders for the new [ACE-Step 1.5](https://github.com/ace-step/ACE-Step) music generation model.
|
||||
- **Extend/Repaint**: Selectively regenerate specific regions or extend audio with new content
|
||||
- **Cover**: Style transfer/regeneration using semantic guidance from source audio
|
||||
- *Extract and Lego guiders coming soon*
|
||||
|
||||
### 🌊 Optical Flow
|
||||
- Generate masks based on movement in video sequences
|
||||
@@ -301,42 +289,41 @@ These audio and MIDI processing capabilities enable the creation of music-driven
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><h3>ACEStep Native</h3></summary>
|
||||
<summary><h3>ACE-Step 1.5 Native</h3></summary>
|
||||
|
||||
Native ComfyUI implementation of ACEStep audio repaint and extend functionality. This provides seamless integration with ComfyUI's existing sampling infrastructure to enable selective audio regeneration and extension without requiring custom types.
|
||||
Native ComfyUI guiders for the [ACE-Step 1.5](https://github.com/ace-step/ACE-Step) music generation model. These integrate directly with ComfyUI's sampling infrastructure (SamplerCustomAdvanced) using model-level mask injection.
|
||||
|
||||
#### Core Functionality Nodes
|
||||
#### Available Guiders
|
||||
|
||||
**ACEStep Repaint Guider**
|
||||
Creates a guider for repainting specific audio regions.
|
||||
- **Inputs**: Model, conditioning, source latents, time range, strength settings
|
||||
- **Output**: GUIDER (for use with SamplerCustomAdvanced)
|
||||
**ACE-Step 1.5 Edit Guider (Extend/Repaint)**
|
||||
Unified guider for extend and/or repaint operations.
|
||||
- Extend audio by adding new content before or after existing audio
|
||||
- Repaint specific time regions while preserving the rest
|
||||
- Combine both operations in a single pass
|
||||
- Silence latent is automatically loaded from HuggingFace on first use
|
||||
|
||||
**ACEStep Extend Guider**
|
||||
Creates a guider for extending audio before/after existing content.
|
||||
- **Inputs**: Model, conditioning, source latents, extend times
|
||||
- **Output**: GUIDER (for use with SamplerCustomAdvanced)
|
||||
|
||||
**ACEStep Hybrid Guider**
|
||||
Combines repaint and extend functionality in a single node.
|
||||
- **Inputs**: All repaint/extend parameters with optional controls
|
||||
- **Output**: GUIDER (for use with SamplerCustomAdvanced)
|
||||
**ACE-Step 1.5 Cover Guider**
|
||||
Style transfer/regeneration using semantic guidance from source audio.
|
||||
- Uses semantic tokens extracted from source audio as structural guidance
|
||||
- Generates new content with different style/timbre while preserving musical structure
|
||||
- Semantic hints are auto-extracted or can be provided manually for advanced workflows
|
||||
|
||||
#### Utility Nodes
|
||||
|
||||
- **ACEStep Analyze Latent**: Analyzes ACE latent properties for debugging
|
||||
- **ACEStep Time Range**: Converts time ranges to frame indices
|
||||
- **ACEStep Mask Visualizer**: Visualizes repaint masks for debugging
|
||||
- **ACEStep Audio Post Processor**: Preserves original audio fidelity in unchanged regions
|
||||
- **ACE-Step 1.5 Task Text Encode**: Task-aware text encoder with support for lyrics, BPM, key, time signature, and language
|
||||
- **ACE-Step 1.5 Semantic Extractor**: Extract semantic tokens for use with Cover guider
|
||||
- **ACE-Step 1.5 Semantic Hints Blend**: Blend semantic hints from multiple sources for mashups
|
||||
- **ACEStep Analyze Latent**: Analyze audio latent properties
|
||||
- **ACEStep Audio Post Processor**: Preserve original audio fidelity in unchanged regions
|
||||
|
||||
#### Basic Workflows
|
||||
#### Coming Soon
|
||||
|
||||
**Repaint Workflow**: Replace audio content from 10-20 seconds with new generation.
|
||||
See: [examples/acestep_repaint.json](examples/acestep_repaint.json)
|
||||
- **Extract Guider**: Extract specific tracks (vocals, drums, bass, etc.) from source audio
|
||||
- **Lego Guider**: Generate specific tracks within a time region
|
||||
|
||||
**Extend Workflow**: Add 10 seconds of new content after existing audio.
|
||||
See: [examples/acestep_extend.json](examples/acestep_extend.json)
|
||||
#### Legacy (ACE-Step 1.0)
|
||||
|
||||
The v1.0 Repaint, Extend, and Hybrid guiders remain available for use with ACE-Step 1.0 latents.
|
||||
|
||||
</details>
|
||||
|
||||
|
||||
+3
-4
@@ -1,11 +1,10 @@
|
||||
opencv-python>=4.10.0
|
||||
opencv-python>=4.5.0
|
||||
scipy
|
||||
pillow
|
||||
librosa>=0.10.2
|
||||
pymunk==6.8.1
|
||||
librosa>=0.10.1
|
||||
pymunk>=6.8.0,<7.0.0
|
||||
matplotlib
|
||||
openunmix
|
||||
mido
|
||||
torchaudio
|
||||
pygame
|
||||
scikit-image
|
||||
@@ -1,5 +1,4 @@
|
||||
import { app } from "../../scripts/app.js";
|
||||
import { ComfyWidgets } from "../../scripts/widgets.js";
|
||||
import { FeatureTemplates } from "./feature_templates.js";
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user