mrleo1nid/ACE-Step
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

work on pip package

Browse Source
This commit is contained in:
mrfakename
2025-05-06 18:59:32 -07:00
parent 54da683d36
commit a5746eaab6
36 changed files with 2928 additions and 1251 deletions
Download Patch File Download Diff File Expand all files Collapse all files
acestep/music_dcae/music_dcae_pipeline.py
+37 -10
View File
@@ -1,3 +1,11 @@
"""
ACE-Step: A Step Towards Music Generation Foundation Model
https://github.com/ace-step/ACE-Step
Apache 2.0 License
"""
import os
import torch
from diffusers import AutoencoderDC
@@ -21,7 +29,12 @@ VOCODER_PRETRAINED_PATH = os.path.join(root_dir, "checkpoints", "music_vocoder")
class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
@register_to_config
def __init__(self, source_sample_rate=None, dcae_checkpoint_path=DEFAULT_PRETRAINED_PATH, vocoder_checkpoint_path=VOCODER_PRETRAINED_PATH):
def __init__(
self,
source_sample_rate=None,
dcae_checkpoint_path=DEFAULT_PRETRAINED_PATH,
vocoder_checkpoint_path=VOCODER_PRETRAINED_PATH,
):
super(MusicDCAE, self).__init__()
self.dcae = AutoencoderDC.from_pretrained(dcae_checkpoint_path)
@@ -32,9 +45,11 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
self.resampler = torchaudio.transforms.Resample(source_sample_rate, 44100)
self.transform = transforms.Compose([
transforms.Normalize(0.5, 0.5),
])
self.transform = transforms.Compose(
[
transforms.Normalize(0.5, 0.5),
]
)
self.min_mel_value = -11.0
self.max_mel_value = 3.0
self.audio_chunk_size = int(round((1024 * 512 / 44100 * 48000)))
@@ -76,7 +91,9 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
max_audio_len = audio.shape[-1]
if max_audio_len % (8 * 512) != 0:
audio = torch.nn.functional.pad(audio, (0, 8 * 512 - max_audio_len % (8 * 512)))
audio = torch.nn.functional.pad(
audio, (0, 8 * 512 - max_audio_len % (8 * 512))
)
mels = self.forward_mel(audio)
mels = (mels - self.min_mel_value) / (self.max_mel_value - self.min_mel_value)
@@ -86,7 +103,9 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
latent = self.dcae.encoder(mel.unsqueeze(0))
latents.append(latent)
latents = torch.cat(latents, dim=0)
latent_lengths = (audio_lengths / sr * 44100 / 512 / self.time_dimention_multiple).long()
latent_lengths = (
audio_lengths / sr * 44100 / 512 / self.time_dimention_multiple
).long()
latents = (latents - self.shift_factor) * self.scale_factor
return latents, latent_lengths
@@ -103,18 +122,26 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
wav = self.vocoder.decode(mels[0]).squeeze(1)
if sr is not None:
resampler = torchaudio.transforms.Resample(44100, sr).to(latents.device).to(latents.dtype)
resampler = (
torchaudio.transforms.Resample(44100, sr)
.to(latents.device)
.to(latents.dtype)
)
wav = resampler(wav)
else:
sr = 44100
pred_wavs.append(wav)
if audio_lengths is not None:
pred_wavs = [wav[:, :length].cpu() for wav, length in zip(pred_wavs, audio_lengths)]
pred_wavs = [
wav[:, :length].cpu() for wav, length in zip(pred_wavs, audio_lengths)
]
return sr, pred_wavs
def forward(self, audios, audio_lengths=None, sr=None):
latents, latent_lengths = self.encode(audios=audios, audio_lengths=audio_lengths, sr=sr)
latents, latent_lengths = self.encode(
audios=audios, audio_lengths=audio_lengths, sr=sr
)
sr, pred_wavs = self.decode(latents=latents, audio_lengths=audio_lengths, sr=sr)
return sr, pred_wavs, latents, latent_lengths
@@ -124,7 +151,7 @@ if __name__ == "__main__":
audio, sr = torchaudio.load("test.wav")
audio_lengths = torch.tensor([audio.shape[1]])
audios = audio.unsqueeze(0)
# test encode only
model = MusicDCAE()
# latents, latent_lengths = model.encode(audios, audio_lengths)
acestep/music_dcae/music_log_mel.py
+8
View File
@@ -1,3 +1,11 @@
"""
ACE-Step: A Step Towards Music Generation Foundation Model
https://github.com/ace-step/ACE-Step
Apache 2.0 License
"""
import torch
import torch.nn as nn
from torch import Tensor
acestep/music_dcae/music_vocoder.py
+22 -11
View File
@@ -1,3 +1,11 @@
"""
ACE-Step: A Step Towards Music Generation Foundation Model
https://github.com/ace-step/ACE-Step
Apache 2.0 License
"""
import librosa
import torch
from torch import nn
@@ -132,13 +140,11 @@ class ConvNeXtBlock(nn.Module):
self.act = nn.GELU()
self.pwconv2 = nn.Linear(int(mlp_ratio * dim), dim)
self.gamma = (
nn.Parameter(layer_scale_init_value *
torch.ones((dim)), requires_grad=True)
nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
if layer_scale_init_value > 0
else None
)
self.drop_path = DropPath(
drop_path) if drop_path > 0.0 else nn.Identity()
self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
def forward(self, x, apply_residual: bool = True):
input = x
@@ -367,8 +373,7 @@ class HiFiGANGenerator(nn.Module):
upsample_rates: Tuple[int] = (8, 8, 2, 2, 2),
upsample_kernel_sizes: Tuple[int] = (16, 16, 8, 2, 2),
resblock_kernel_sizes: Tuple[int] = (3, 7, 11),
resblock_dilation_sizes: Tuple[Tuple[int]] = (
(1, 3, 5), (1, 3, 5), (1, 3, 5)),
resblock_dilation_sizes: Tuple[Tuple[int]] = ((1, 3, 5), (1, 3, 5), (1, 3, 5)),
num_mels: int = 128,
upsample_initial_channel: int = 512,
use_template: bool = True,
@@ -417,7 +422,7 @@ class HiFiGANGenerator(nn.Module):
continue
if i + 1 < len(upsample_rates):
stride_f0 = np.prod(upsample_rates[i + 1:])
stride_f0 = np.prod(upsample_rates[i + 1 :])
self.noise_convs.append(
Conv1d(
1,
@@ -485,7 +490,7 @@ class HiFiGANGenerator(nn.Module):
class ADaMoSHiFiGANV1(ModelMixin, ConfigMixin, FromOriginalModelMixin):
@register_to_config
def __init__(
self,
@@ -498,7 +503,11 @@ class ADaMoSHiFiGANV1(ModelMixin, ConfigMixin, FromOriginalModelMixin):
upsample_kernel_sizes: Tuple[int] = (8, 8, 4, 4, 4, 4, 4),
resblock_kernel_sizes: Tuple[int] = (3, 7, 11, 13),
resblock_dilation_sizes: Tuple[Tuple[int]] = (
(1, 3, 5), (1, 3, 5), (1, 3, 5), (1, 3, 5)),
(1, 3, 5),
(1, 3, 5),
(1, 3, 5),
(1, 3, 5),
),
num_mels: int = 512,
upsample_initial_channel: int = 1024,
use_template: bool = False,
@@ -555,7 +564,7 @@ class ADaMoSHiFiGANV1(ModelMixin, ConfigMixin, FromOriginalModelMixin):
@torch.no_grad()
def encode(self, x):
return self.mel_transform(x)
def forward(self, mel):
y = self.backbone(mel)
y = self.head(y)
@@ -566,7 +575,9 @@ if __name__ == "__main__":
import soundfile as sf
x = "test_audio.flac"
model = ADaMoSHiFiGANV1.from_pretrained("./checkpoints/music_vocoder", local_files_only=True)
model = ADaMoSHiFiGANV1.from_pretrained(
"./checkpoints/music_vocoder", local_files_only=True
)
wav, sr = librosa.load(x, sr=44100, mono=True)
wav = torch.from_numpy(wav).float()[None]