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

refact project

Browse Source
This commit is contained in:
chuxij
2025-04-25 15:45:24 +00:00
parent f91ce4867f
commit db32300e75
20 changed files with 2031 additions and 4027 deletions
Download Patch File Download Diff File Expand all files Collapse all files
music_dcae/__init__.py
-2
View File
@@ -1,2 +0,0 @@
from .music_dcae import MusicDCAE
from .music_log_mel import LogMelSpectrogram, get_mel_transform
music_dcae/music_dcae_pipeline.py
+15 -20
View File
@@ -1,37 +1,32 @@
import os
import torch
import torch.nn as nn
from diffusers import AutoencoderDC
import torchaudio
import torchvision.transforms as transforms
import torchaudio
from diffusers.models.modeling_utils import ModelMixin
from diffusers.loaders import FromOriginalModelMixin
from diffusers.configuration_utils import ConfigMixin, register_to_config
try:
from .music_log_mel import get_mel_transform
from .music_vocoder import ADaMoSHiFiGANV1
except ImportError:
from music_log_mel import get_mel_transform
from music_vocoder import ADaMoSHiFiGANV1
import os
root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_PRETRAINED_PATH = os.path.join(root_dir, "checkpoints", "music_dcae_f8c8")
VOCODER_PRETRAINED_PATH = os.path.join(root_dir, "checkpoints", "music_vocoder.pt")
VOCODER_PRETRAINED_PATH = os.path.join(root_dir, "checkpoints", "music_vocoder")
class MusicDCAE(nn.Module):
def __init__(self, pretrained_path=DEFAULT_PRETRAINED_PATH, encoder_only=False, source_sample_rate=None):
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):
super(MusicDCAE, self).__init__()
dcae = AutoencoderDC.from_pretrained(pretrained_path)
self.encoder_only = encoder_only
self.mel_transform = get_mel_transform()
if encoder_only:
self.encoder = dcae.encoder
else:
self.encoder = dcae.encoder
self.decoder = dcae.decoder
self.vocoder = ADaMoSHiFiGANV1(VOCODER_PRETRAINED_PATH).eval()
self.dcae = AutoencoderDC.from_pretrained(dcae_checkpoint_path)
self.vocoder = ADaMoSHiFiGANV1.from_pretrained(vocoder_checkpoint_path)
if source_sample_rate is None:
source_sample_rate = 48000
@@ -57,7 +52,7 @@ class MusicDCAE(nn.Module):
def forward_mel(self, audios):
mels = []
for i in range(len(audios)):
image = self.mel_transform(audios[i])
image = self.vocoder.mel_transform(audios[i])
mels.append(image)
mels = torch.stack(mels)
return mels
@@ -89,7 +84,7 @@ class MusicDCAE(nn.Module):
mels = self.transform(mels)
latents = []
for mel in mels:
latent = self.encoder(mel.unsqueeze(0))
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()
@@ -103,7 +98,7 @@ class MusicDCAE(nn.Module):
mels = []
for latent in latents:
mel = self.decoder(latent.unsqueeze(0))
mel = self.dcae.decoder(latent.unsqueeze(0))
mels.append(mel)
mels = torch.cat(mels, dim=0)
@@ -152,4 +147,4 @@ if __name__ == "__main__":
print("latent_lengths: ", latent_lengths)
print("sr: ", sr)
torchaudio.save("test_reconstructed.flac", pred_wavs[0], sr)
print("/test_reconstructed.flac")
print("test_reconstructed.flac")
music_dcae/music_log_mel.py
-12
View File
@@ -105,15 +105,3 @@ class LogMelSpectrogram(nn.Module):
return x, self.compress(linear)
return x
def get_mel_transform():
return LogMelSpectrogram(
sample_rate=44100,
n_fft=2048,
win_length=2048,
hop_length=512,
f_min=40,
f_max=16000,
n_mels=128,
)
music_dcae/music_vocoder.py
+56 -45
View File
@@ -11,6 +11,10 @@ import torch.nn.functional as F
from torch.nn import Conv1d
from torch.nn.utils import weight_norm
from torch.nn.utils.parametrize import remove_parametrizations as remove_weight_norm
from diffusers.models.modeling_utils import ModelMixin
from diffusers.loaders import FromOriginalModelMixin
from diffusers.configuration_utils import ConfigMixin, register_to_config
try:
from music_log_mel import LogMelSpectrogram
@@ -480,60 +484,67 @@ class HiFiGANGenerator(nn.Module):
remove_weight_norm(self.conv_post)
class ADaMoSHiFiGANV1(nn.Module):
class ADaMoSHiFiGANV1(ModelMixin, ConfigMixin, FromOriginalModelMixin):
@register_to_config
def __init__(
self,
checkpoint_path: str = "checkpoints/adamos-generator-1640000.pth",
input_channels: int = 128,
depths: List[int] = [3, 3, 9, 3],
dims: List[int] = [128, 256, 384, 512],
drop_path_rate: float = 0.0,
kernel_sizes: Tuple[int] = (7,),
upsample_rates: Tuple[int] = (4, 4, 2, 2, 2, 2, 2),
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)),
num_mels: int = 512,
upsample_initial_channel: int = 1024,
use_template: bool = False,
pre_conv_kernel_size: int = 13,
post_conv_kernel_size: int = 13,
sampling_rate: int = 44100,
n_fft: int = 2048,
win_length: int = 2048,
hop_length: int = 512,
f_min: int = 40,
f_max: int = 16000,
n_mels: int = 128,
):
super().__init__()
self.backbone = ConvNeXtEncoder(
input_channels=128,
depths=[3, 3, 9, 3],
dims=[128, 256, 384, 512],
drop_path_rate=0,
kernel_sizes=(7,),
input_channels=input_channels,
depths=depths,
dims=dims,
drop_path_rate=drop_path_rate,
kernel_sizes=kernel_sizes,
)
self.head = HiFiGANGenerator(
hop_length=512,
upsample_rates=(4, 4, 2, 2, 2, 2, 2),
upsample_kernel_sizes=(8, 8, 4, 4, 4, 4, 4),
resblock_kernel_sizes=(3, 7, 11, 13),
resblock_dilation_sizes=(
(1, 3, 5), (1, 3, 5), (1, 3, 5), (1, 3, 5)),
num_mels=512,
upsample_initial_channel=1024,
use_template=False,
pre_conv_kernel_size=13,
post_conv_kernel_size=13,
hop_length=hop_length,
upsample_rates=upsample_rates,
upsample_kernel_sizes=upsample_kernel_sizes,
resblock_kernel_sizes=resblock_kernel_sizes,
resblock_dilation_sizes=resblock_dilation_sizes,
num_mels=num_mels,
upsample_initial_channel=upsample_initial_channel,
use_template=use_template,
pre_conv_kernel_size=pre_conv_kernel_size,
post_conv_kernel_size=post_conv_kernel_size,
)
self.sampling_rate = 44100
ckpt_state = torch.load(checkpoint_path, map_location="cpu")
if "state_dict" in ckpt_state:
ckpt_state = ckpt_state["state_dict"]
if any(k.startswith("generator.") for k in ckpt_state):
ckpt_state = {
k.replace("generator.", ""): v
for k, v in ckpt_state.items()
if k.startswith("generator.")
}
self.load_state_dict(ckpt_state)
self.eval()
self.sampling_rate = sampling_rate
self.mel_transform = LogMelSpectrogram(
sample_rate=44100,
n_fft=2048,
win_length=2048,
hop_length=512,
f_min=40,
f_max=16000,
n_mels=128,
sample_rate=sampling_rate,
n_fft=n_fft,
win_length=win_length,
hop_length=hop_length,
f_min=f_min,
f_max=f_max,
n_mels=n_mels,
)
self.eval()
@torch.no_grad()
def decode(self, mel):
@@ -554,12 +565,12 @@ class ADaMoSHiFiGANV1(nn.Module):
if __name__ == "__main__":
import soundfile as sf
x = "./test.wav"
model = ADaMoSHiFiGANV1(checkpoint_path='./step_001640000.pth')
x = "test_audio.flac"
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]
mel = model.encode(wav)
wav = model.decode(mel)[0].mT
sf.write("test_out.wav", wav.cpu().numpy(), 44100)
sf.write("test_audio_vocoder_rec.flac", wav.cpu().numpy(), 44100)