118 lines
4.0 KiB
Python
118 lines
4.0 KiB
Python
import os
|
|
import torch
|
|
import torchaudio
|
|
from torchaudio.functional import resample
|
|
import json
|
|
import av
|
|
import io
|
|
|
|
class LoopAudioUtils:
|
|
"""Utility class for managing audio files"""
|
|
|
|
DEFAULT_SAMPLE_RATE = 44100
|
|
|
|
@staticmethod
|
|
def load_or_create_audio(audio: dict | None, path: str, load: bool = True, target_sample_rate: int = DEFAULT_SAMPLE_RATE) -> dict:
|
|
"""
|
|
Load an existing audio file or create the file.
|
|
"""
|
|
if os.path.exists(path) and load:
|
|
return LoopAudioUtils.load_audio(path, target_sample_rate)
|
|
else:
|
|
if audio is None:
|
|
waveform = torch.zeros((1, 1, target_sample_rate), dtype=torch.float32)
|
|
audio = {"waveform": waveform, "sample_rate": target_sample_rate}
|
|
LoopAudioUtils.save_audio(audio, path)
|
|
return audio
|
|
|
|
@staticmethod
|
|
def load_audio(path: str, target_sample_rate: int = DEFAULT_SAMPLE_RATE) -> dict:
|
|
"""
|
|
Load an existing audio file and return a dict with 'waveform' and 'sample_rate'.
|
|
eventually resample to target_sample_rate.
|
|
"""
|
|
waveform, sample_rate = torchaudio.load(path)
|
|
if waveform.ndim == 2:
|
|
waveform = waveform.unsqueeze(0)
|
|
if sample_rate != target_sample_rate:
|
|
waveform = resample(waveform, sample_rate, target_sample_rate)
|
|
sample_rate = target_sample_rate
|
|
return {"waveform": waveform, "sample_rate": sample_rate}
|
|
|
|
# @staticmethod
|
|
# def save_audio(audio: dict, path: str, metadata: dict | None = None) -> str:
|
|
# """
|
|
# Save an audio file as wav and return input path.
|
|
# """
|
|
# waveform = audio["waveform"]
|
|
# sample_rate = audio["sample_rate"]
|
|
# if waveform.ndim == 3:
|
|
# waveform = waveform[0]
|
|
# os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
# torchaudio.save(path, waveform.cpu(), sample_rate)
|
|
# return path
|
|
|
|
@staticmethod
|
|
def save_audio(audio: dict, path: str, metadata: dict | None = None) -> str:
|
|
"""
|
|
Save an audio file as flac (mono to up to 6 channels). return input path.
|
|
"""
|
|
waveform = audio["waveform"]
|
|
sample_rate = audio["sample_rate"]
|
|
|
|
if waveform.ndim == 3:
|
|
waveform = waveform[0] # Shape: (channels, samples)
|
|
|
|
num_channels = waveform.shape[0]
|
|
|
|
# layout
|
|
if num_channels == 1:
|
|
layout = 'mono'
|
|
elif num_channels == 2:
|
|
layout = 'stereo'
|
|
else:
|
|
layout = f'{num_channels}.0' # up to 6 channels
|
|
|
|
# Transform for PyAV
|
|
if num_channels == 1:
|
|
frame_tensor = waveform # (1, samples)
|
|
else:
|
|
frame_tensor = waveform.movedim(0, 1).reshape(1, -1) # (channels, samples) → (1, samples * channels)
|
|
|
|
frame_data = frame_tensor.cpu().numpy()
|
|
|
|
output_buffer = io.BytesIO()
|
|
output_container = av.open(output_buffer, mode='w', format="flac")
|
|
|
|
if metadata:
|
|
for key, value in metadata.items():
|
|
output_container.metadata[key] = str(value)
|
|
|
|
out_stream = output_container.add_stream("flac", rate=sample_rate)
|
|
|
|
frame = av.AudioFrame.from_ndarray(frame_data, format='flt', layout=layout)
|
|
frame.sample_rate = sample_rate
|
|
frame.pts = 0
|
|
|
|
output_container.mux(out_stream.encode(frame))
|
|
output_container.mux(out_stream.encode(None))
|
|
output_container.close()
|
|
|
|
with open(path, 'wb') as f:
|
|
f.write(output_buffer.getbuffer())
|
|
|
|
return path
|
|
|
|
@staticmethod
|
|
def prepare_metadata(prompt: str, extra_pnginfo: str) -> dict[str, str]:
|
|
"""
|
|
Feed a dict with metadata.
|
|
"""
|
|
metadata = {}
|
|
if prompt is not None:
|
|
metadata["prompt"] = json.dumps(prompt)
|
|
if extra_pnginfo is not None:
|
|
for x in extra_pnginfo:
|
|
metadata[x] = json.dumps(extra_pnginfo[x])
|
|
return metadata
|