ckinpdx
68f16d6826
Updated the LoadAudioandSplit class to handle a maximum of 8 audio frames instead of 50 and added overlap functionality in the split_audio method.
107 lines
3.5 KiB
Python
107 lines
3.5 KiB
Python
import torch
|
|
|
|
|
|
class LoadAudioandSplit:
|
|
"""
|
|
Splits audio into chunks where each chunk's duration is calculated from
|
|
connected frame counts and FPS. Duration = frames / fps
|
|
Auto-detects how many frames_N inputs are connected.
|
|
"""
|
|
RETURN_TYPES = ("INT", "FLOAT") + tuple(["AUDIO"] * 8)
|
|
RETURN_NAMES = ("scene_count", "final_timestamp") + tuple([f"audio_{i}" for i in range(1, 9)])
|
|
FUNCTION = "split_audio"
|
|
CATEGORY = "wankeyframebuilder"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
optional_inputs = {}
|
|
for i in range(1, 9):
|
|
optional_inputs[f"frames_{i}"] = ("INT", {"default": None, "min": 1, "forceInput": True})
|
|
|
|
return {
|
|
"required": {
|
|
"audio": ("AUDIO",),
|
|
"fps": ("FLOAT", {"default": 25.0, "min": 0.001, "max": 120.0}),
|
|
"overlap_frames": ("INT", {"default": 0, "min": 0, "max": 500}),
|
|
},
|
|
"optional": optional_inputs
|
|
}
|
|
|
|
def split_audio(self, audio, fps, overlap_frames=0, **kwargs):
|
|
make_stereo = True
|
|
|
|
waveform = audio["waveform"]
|
|
src_sample_rate = int(audio.get("sample_rate", 44100))
|
|
|
|
if waveform.ndim == 2:
|
|
waveform = waveform.unsqueeze(0)
|
|
|
|
total_samples = waveform.shape[-1]
|
|
audio_total_duration = float(total_samples) / float(src_sample_rate)
|
|
|
|
# auto-detect connected frame inputs (contiguous from frames_1)
|
|
frame_counts = []
|
|
for i in range(1, 9):
|
|
frame_key = f"frames_{i}"
|
|
val = kwargs.get(frame_key, None)
|
|
if val is None:
|
|
break
|
|
frame_counts.append(int(val))
|
|
|
|
scene_count = len(frame_counts)
|
|
|
|
# if nothing connected, default to 1 scene with 97 frames
|
|
if scene_count == 0:
|
|
frame_counts = [97]
|
|
scene_count = 1
|
|
|
|
# calculate durations from frame counts and fps
|
|
durations = [float(fc) / float(fps) for fc in frame_counts]
|
|
overlap_seconds = float(overlap_frames) / float(fps)
|
|
|
|
# compute start times with overlap
|
|
starts = []
|
|
current_time = 0.0
|
|
for idx, d in enumerate(durations):
|
|
starts.append(current_time)
|
|
# next chunk starts at (duration - overlap) after this one
|
|
if idx < len(durations) - 1:
|
|
current_time += d - overlap_seconds
|
|
|
|
segments = []
|
|
|
|
for idx, start_time in enumerate(starts):
|
|
requested_duration = durations[idx]
|
|
start_samp = max(0, int(start_time * src_sample_rate))
|
|
end_samp = min(total_samples, int(start_samp + requested_duration * src_sample_rate))
|
|
|
|
seg = waveform[..., start_samp:end_samp]
|
|
|
|
# if segment is empty, create silence
|
|
if seg.numel() == 0:
|
|
silence_samples = int(requested_duration * src_sample_rate)
|
|
seg = torch.zeros((1, 2, silence_samples), dtype=waveform.dtype)
|
|
|
|
if make_stereo and seg.shape[1] == 1:
|
|
seg = seg.repeat(1, 2, 1)
|
|
|
|
segments.append({"waveform": seg, "sample_rate": src_sample_rate})
|
|
|
|
# pad with None for unused outputs
|
|
while len(segments) < 8:
|
|
segments.append(None)
|
|
|
|
# calculate end timestamp of last chunk
|
|
last_chunk_end = starts[-1] + durations[-1]
|
|
|
|
return (scene_count, float(last_chunk_end), *tuple(segments))
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LoadAudioandSplit": LoadAudioandSplit
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LoadAudioandSplit": "Load Audio and Split"
|
|
}
|