Hugging Face's logo

扩散课程文档

音频扩散

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

Open In Colab

音频扩散

在本笔记本中，我们将简要了解如何使用扩散模型生成音频。

你将学到什么：

  • 音频在计算机中的表示方式
  • 在原始音频数据和频谱图之间转换的方法
  • 如何使用自定义的 collate 函数准备数据加载器，以将音频片段转换为频谱图
  • 在特定音乐类型上微调现有的音频扩散模型
  • 将你的自定义 pipeline 上传到 Hugging Face Hub

注意：这主要用于教育目的 - 不保证我们的模型听起来会很好 😉。

让我们开始吧！

设置和导入

%pip install -q datasets diffusers torchaudio accelerate
import torch, random
import numpy as np
import torch.nn.functional as F
from tqdm.auto import tqdm
from IPython.display import Audio
from matplotlib import pyplot as plt
from diffusers import DiffusionPipeline
from torchaudio import transforms as AT
from torchvision import transforms as IT

从预训练音频 pipeline 中采样

让我们首先按照音频扩散文档加载一个预先存在的音频扩散模型 pipeline

# Load a pre-trained audio diffusion pipeline
device = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
pipe = DiffusionPipeline.from_pretrained("teticio/audio-diffusion-instrumental-hiphop-256").to(device)

与我们在之前的单元中使用的 pipelines 一样，我们可以通过调用 pipeline 来创建样本，如下所示

>>> # Sample from the pipeline and display the outputs
>>> output = pipe()
>>> display(output.images[0])
>>> display(Audio(output.audios[0], rate=pipe.mel.get_sample_rate()))

在这里，rate 参数指定音频的采样率；我们稍后会深入探讨这一点。你还会注意到 pipeline 返回了多个内容。这是怎么回事？让我们仔细看看这两个输出。

第一个是数据数组，表示生成的音频

# The audio array
output.audios[0].shape

第二个看起来像灰度图像

# The output image (spectrogram)
output.images[0].size

这给了我们一个关于此 pipeline 如何工作的提示。音频不是直接通过扩散生成的 - 相反，pipeline 具有与我们在单元 1中看到的无条件图像生成 pipelines 相同的 2D UNet，用于生成频谱图，然后将其后处理为最终音频。

该 pipe 有一个额外的组件来处理这些转换，我们可以通过 pipe.mel 访问它

pipe.mel

从音频到图像再到音频

音频“波形”编码了随时间推移的原始音频样本 - 例如，这可以是来自麦克风的电信号。处理这种“时域”表示可能很棘手，因此通常的做法是将其转换为其他形式，通常称为频谱图。频谱图显示了不同频率（y 轴）与时间（x 轴）的强度

>>> # Calculate and show a spectrogram for our generated audio sample using torchaudio
>>> spec_transform = AT.Spectrogram(power=2)
>>> spectrogram = spec_transform(torch.tensor(output.audios[0]))
>>> print(spectrogram.min(), spectrogram.max())
>>> log_spectrogram = spectrogram.log()
>>> plt.imshow(log_spectrogram[0], cmap="gray")
tensor(0.) tensor(6.0842)

我们刚刚制作的频谱图的值介于 0.0000000000001 和 1 之间，其中大多数值接近该范围的低端。这对于可视化或建模来说并不理想 - 事实上，我们必须对这些值取对数才能获得显示任何细节的灰度图。因此，我们通常使用一种特殊的频谱图，称为 Mel 频谱图，它旨在通过对信号的不同频率分量应用一些变换来捕获对人类听觉重要的信息。

torchaudio docs diagram 来自 torchaudio 文档的一些音频变换

幸运的是，我们甚至不必太担心这些变换 - pipeline 的 mel 功能为我们处理了这些细节。使用它，我们可以像这样将频谱图图像转换为音频

a = pipe.mel.image_to_audio(output.images[0])
a.shape

我们可以通过首先加载原始音频数据，然后调用 audio_slice_to_image() 函数，将音频数据数组转换为频谱图图像。更长的剪辑会自动切片成正确长度的块，以生成 256x256 频谱图图像

>>> pipe.mel.load_audio(raw_audio=a)
>>> im = pipe.mel.audio_slice_to_image(0)
>>> im

音频表示为一个长数字数组。要大声播放它，我们需要另一个关键信息：采样率。我们使用多少个样本（单个值）来表示一秒钟的音频？

我们可以使用以下命令查看此 pipeline 训练期间使用的采样率

sample_rate_pipeline = pipe.mel.get_sample_rate()
sample_rate_pipeline

如果我们错误地指定采样率，我们会得到加速或减速的音频

display(Audio(output.audios[0], rate=44100))  # 2x speed

微调 pipeline

现在我们对 pipeline 的工作原理有了大致的了解，让我们在新音频数据上对其进行微调！

数据集是不同类型的音频剪辑集合，我们可以像这样从 hub 加载它

from datasets import load_dataset

dataset = load_dataset("lewtun/music_genres", split="train")
dataset

你可以使用下面的代码来查看数据集中的不同类型以及每种类型包含多少个样本

>>> for g in list(set(dataset["genre"])):
...     print(g, sum(x == g for x in dataset["genre"]))
Pop 945
Blues 58
Punk 2582
Old-Time / Historic 408
Experimental 1800
Folk 1214
Electronic 3071
Spoken 94
Classical 495
Country 142
Instrumental 1044
Chiptune / Glitch 1181
International 814
Ambient Electronic 796
Jazz 306
Soul-RnB 94
Hip-Hop 1757
Easy Listening 13
Rock 3095

数据集的音频为数组形式

>>> audio_array = dataset[0]["audio"]["array"]
>>> sample_rate_dataset = dataset[0]["audio"]["sampling_rate"]
>>> print("Audio array shape:", audio_array.shape)
>>> print("Sample rate:", sample_rate_dataset)
>>> display(Audio(audio_array, rate=sample_rate_dataset))
Audio array shape: (1323119,)
Sample rate: 44100

请注意，此音频的采样率更高 - 如果我们要使用现有的 pipeline，我们需要“重采样”它以匹配。剪辑也比 pipeline 设置的剪辑长。幸运的是，当我们使用 pipe.mel 加载音频时，它会自动将剪辑切片成更小的部分

>>> a = dataset[0]["audio"]["array"]  # Get the audio array
>>> pipe.mel.load_audio(raw_audio=a)  # Load it with pipe.mel
>>> pipe.mel.audio_slice_to_image(0)  # View the first 'slice' as a spectrogram

我们需要记住调整采样率，因为此数据集中的数据每秒的样本数是原来的两倍

sample_rate_dataset = dataset[0]["audio"]["sampling_rate"]
sample_rate_dataset

这里我们使用 torchaudio 的变换（作为 AT 导入）进行重采样，使用 pipe 的 mel 将音频转换为图像，并使用 torchvision 的变换（作为 IT 导入）将图像转换为张量。这为我们提供了一个函数，可以将音频剪辑转换为频谱图张量，我们可以将其用于训练

resampler = AT.Resample(sample_rate_dataset, sample_rate_pipeline, dtype=torch.float32)
to_t = IT.ToTensor()


def to_image(audio_array):
    audio_tensor = torch.tensor(audio_array).to(torch.float32)
    audio_tensor = resampler(audio_tensor)
    pipe.mel.load_audio(raw_audio=np.array(audio_tensor))
    num_slices = pipe.mel.get_number_of_slices()
    slice_idx = random.randint(0, num_slices - 1)  # Pic a random slice each time (excluding the last short slice)
    im = pipe.mel.audio_slice_to_image(slice_idx)
    return im

我们将使用我们的 to_image() 函数作为自定义 collate 函数的一部分，以将我们的数据集转换为可用于训练的数据加载器。collate 函数定义了如何将数据集中的一批示例转换为准备好用于训练的最终数据批次。在这种情况下，我们将每个音频样本转换为频谱图图像，并将生成的张量堆叠在一起

>>> def collate_fn(examples):
...     # to image -> to tensor -> rescale to (-1, 1) -> stack into batch
...     audio_ims = [to_t(to_image(x["audio"]["array"])) * 2 - 1 for x in examples]
...     return torch.stack(audio_ims)


>>> # Create a dataset with only the 'Chiptune / Glitch' genre of songs
>>> batch_size = 4  # 4 on colab, 12 on A100
>>> chosen_genre = "Electronic"  # <<< Try training on different genres <<<
>>> indexes = [i for i, g in enumerate(dataset["genre"]) if g == chosen_genre]
>>> filtered_dataset = dataset.select(indexes)
>>> dl = torch.utils.data.DataLoader(
...     filtered_dataset.shuffle(), batch_size=batch_size, collate_fn=collate_fn, shuffle=True
... )
>>> batch = next(iter(dl))
>>> print(batch.shape)
torch.Size([4, 1, 256, 256])

注意：除非你有足够的 GPU vRAM 可用，否则你需要使用较低的批次大小（例如，4）。

训练循环

这是一个简单的训练循环，它在数据加载器上运行几个 epoch，以微调 pipeline 的 UNet。你也可以跳过此单元格，并使用以下单元格中的代码加载 pipeline。

epochs = 3
lr = 1e-4

pipe.unet.train()
pipe.scheduler.set_timesteps(1000)
optimizer = torch.optim.AdamW(pipe.unet.parameters(), lr=lr)

for epoch in range(epochs):
    for step, batch in tqdm(enumerate(dl), total=len(dl)):

        # Prepare the input images
        clean_images = batch.to(device)
        bs = clean_images.shape[0]

        # Sample a random timestep for each image
        timesteps = torch.randint(0, pipe.scheduler.num_train_timesteps, (bs,), device=clean_images.device).long()

        # Add noise to the clean images according to the noise magnitude at each timestep
        noise = torch.randn(clean_images.shape).to(clean_images.device)
        noisy_images = pipe.scheduler.add_noise(clean_images, noise, timesteps)

        # Get the model prediction
        noise_pred = pipe.unet(noisy_images, timesteps, return_dict=False)[0]

        # Calculate the loss
        loss = F.mse_loss(noise_pred, noise)
        loss.backward(loss)

        # Update the model parameters with the optimizer
        optimizer.step()
        optimizer.zero_grad()
# OR: Load the version I trained earlier
pipe = DiffusionPipeline.from_pretrained("johnowhitaker/Electronic_test").to(device)
>>> output = pipe()
>>> display(output.images[0])
>>> display(Audio(output.audios[0], rate=22050))
>>> # Make a longer sample by passing in a starting noise tensor with a different shape
>>> noise = torch.randn(1, 1, pipe.unet.sample_size[0], pipe.unet.sample_size[1] * 4).to(device)
>>> output = pipe(noise=noise)
>>> display(output.images[0])
>>> display(Audio(output.audios[0], rate=22050))

不是最动听的输出，但这是一个开始 :) 探索调整学习率和 epoch 数量，并在 Discord 上分享你的最佳结果，以便我们一起改进！

一些需要考虑的事项

  • 我们正在处理 256px 的方形频谱图图像，这限制了我们的批次大小。你能从 128x128 的频谱图中恢复足够质量的音频吗？
  • 我们每次都选择音频剪辑的不同片段来代替随机图像增强，但是当训练多个 epoch 时，可以通过一些不同类型的增强来改进吗？
  • 我们还可以如何使用它来生成更长的剪辑？也许你可以生成一个 5 秒的起始剪辑，然后使用受 inpainting 启发的想法继续生成从初始剪辑开始的其他音频片段……
  • 在这种频谱图扩散上下文中，图像到图像的等效物是什么？

推送到 Hub

一旦你对你的模型感到满意，你可以保存它并将其推送到 hub，供其他人享用

from huggingface_hub import get_full_repo_name, HfApi, create_repo, ModelCard
# Pick a name for the model
model_name = "audio-diffusion-electronic"
hub_model_id = get_full_repo_name(model_name)
# Save the pipeline locally
pipe.save_pretrained(model_name)
>>> # Inspect the folder contents
>>> !ls {model_name}
mel  model_index.json  scheduler  unet

# Create a repository
create_repo(hub_model_id)
# Upload the files
api = HfApi()
api.upload_folder(folder_path=f"{model_name}/scheduler", path_in_repo="scheduler", repo_id=hub_model_id)
api.upload_folder(folder_path=f"{model_name}/mel", path_in_repo="mel", repo_id=hub_model_id)
api.upload_folder(folder_path=f"{model_name}/unet", path_in_repo="unet", repo_id=hub_model_id)
api.upload_file(
    path_or_fileobj=f"{model_name}/model_index.json",
    path_in_repo="model_index.json",
    repo_id=hub_model_id,
)
# Push a model card
content = f"""
---
license: mit
tags:
- pytorch
- diffusers
- unconditional-audio-generation
- diffusion-models-class
---

# Model Card for Unit 4 of the [Diffusion Models Class 🧨](https://github.com/huggingface/diffusion-models-class)

This model is a diffusion model for unconditional audio generation of music in the genre {chosen_genre}

## Usage

<pre>
from IPython.display import Audio
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("{hub_model_id}")
output = pipe()
display(output.images[0])
display(Audio(output.audios[0], rate=pipe.mel.get_sample_rate()))
</pre>
"""

card = ModelCard(content)
card.push_to_hub(hub_model_id)

结论

希望本笔记本能让你对音频生成的潜力有所了解。查看本单元介绍中链接的一些参考资料，了解一些更高级的方法以及它们可以创建的惊人样本！

< > 在 GitHub 上更新

←DDIM 反演 概述→