facebookresearch--audiocraft
Alexandre Défossez (autodoc)
619 行
34 KiB
HTML
619 行
34 KiB
HTML
<!doctype html>
|
||
<html lang="en">
|
||
<head>
|
||
<meta charset="utf-8">
|
||
<meta name="viewport" content="width=device-width, initial-scale=1, minimum-scale=1">
|
||
<meta name="generator" content="pdoc3 0.11.5">
|
||
<title>audiocraft.models.multibanddiffusion API documentation</title>
|
||
<meta name="description" content="Multi Band Diffusion models as described in
|
||
"From Discrete Tokens to High-Fidelity Audio Using Multi-Band Diffusion"
|
||
(paper link).">
|
||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/13.0.0/sanitize.min.css" integrity="sha512-y1dtMcuvtTMJc1yPgEqF0ZjQbhnc/bFhyvIyVNb9Zk5mIGtqVaAB1Ttl28su8AvFMOY0EwRbAe+HCLqj6W7/KA==" crossorigin>
|
||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/13.0.0/typography.min.css" integrity="sha512-Y1DYSb995BAfxobCkKepB1BqJJTPrOp3zPL74AWFugHHmmdcvO+C48WLrUOlhGMc0QG7AE3f7gmvvcrmX2fDoA==" crossorigin>
|
||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/styles/default.min.css" crossorigin>
|
||
<style>:root{--highlight-color:#fe9}.flex{display:flex !important}body{line-height:1.5em}#content{padding:20px}#sidebar{padding:1.5em;overflow:hidden}#sidebar > *:last-child{margin-bottom:2cm}.http-server-breadcrumbs{font-size:130%;margin:0 0 15px 0}#footer{font-size:.75em;padding:5px 30px;border-top:1px solid #ddd;text-align:right}#footer p{margin:0 0 0 1em;display:inline-block}#footer p:last-child{margin-right:30px}h1,h2,h3,h4,h5{font-weight:300}h1{font-size:2.5em;line-height:1.1em}h2{font-size:1.75em;margin:2em 0 .50em 0}h3{font-size:1.4em;margin:1.6em 0 .7em 0}h4{margin:0;font-size:105%}h1:target,h2:target,h3:target,h4:target,h5:target,h6:target{background:var(--highlight-color);padding:.2em 0}a{color:#058;text-decoration:none;transition:color .2s ease-in-out}a:visited{color:#503}a:hover{color:#b62}.title code{font-weight:bold}h2[id^="header-"]{margin-top:2em}.ident{color:#900;font-weight:bold}pre code{font-size:.8em;line-height:1.4em;padding:1em;display:block}code{background:#f3f3f3;font-family:"DejaVu Sans Mono",monospace;padding:1px 4px;overflow-wrap:break-word}h1 code{background:transparent}pre{border-top:1px solid #ccc;border-bottom:1px solid #ccc;margin:1em 0}#http-server-module-list{display:flex;flex-flow:column}#http-server-module-list div{display:flex}#http-server-module-list dt{min-width:10%}#http-server-module-list p{margin-top:0}.toc ul,#index{list-style-type:none;margin:0;padding:0}#index code{background:transparent}#index h3{border-bottom:1px solid #ddd}#index ul{padding:0}#index h4{margin-top:.6em;font-weight:bold}@media (min-width:200ex){#index .two-column{column-count:2}}@media (min-width:300ex){#index .two-column{column-count:3}}dl{margin-bottom:2em}dl dl:last-child{margin-bottom:4em}dd{margin:0 0 1em 3em}#header-classes + dl > dd{margin-bottom:3em}dd dd{margin-left:2em}dd p{margin:10px 0}.name{background:#eee;font-size:.85em;padding:5px 10px;display:inline-block;min-width:40%}.name:hover{background:#e0e0e0}dt:target .name{background:var(--highlight-color)}.name > span:first-child{white-space:nowrap}.name.class > span:nth-child(2){margin-left:.4em}.inherited{color:#999;border-left:5px solid #eee;padding-left:1em}.inheritance em{font-style:normal;font-weight:bold}.desc h2{font-weight:400;font-size:1.25em}.desc h3{font-size:1em}.desc dt code{background:inherit}.source > summary,.git-link-div{color:#666;text-align:right;font-weight:400;font-size:.8em;text-transform:uppercase}.source summary > *{white-space:nowrap;cursor:pointer}.git-link{color:inherit;margin-left:1em}.source pre{max-height:500px;overflow:auto;margin:0}.source pre code{font-size:12px;overflow:visible;min-width:max-content}.hlist{list-style:none}.hlist li{display:inline}.hlist li:after{content:',\2002'}.hlist li:last-child:after{content:none}.hlist .hlist{display:inline;padding-left:1em}img{max-width:100%}td{padding:0 .5em}.admonition{padding:.1em 1em;margin:1em 0}.admonition-title{font-weight:bold}.admonition.note,.admonition.info,.admonition.important{background:#aef}.admonition.todo,.admonition.versionadded,.admonition.tip,.admonition.hint{background:#dfd}.admonition.warning,.admonition.versionchanged,.admonition.deprecated{background:#fd4}.admonition.error,.admonition.danger,.admonition.caution{background:lightpink}</style>
|
||
<style media="screen and (min-width: 700px)">@media screen and (min-width:700px){#sidebar{width:30%;height:100vh;overflow:auto;position:sticky;top:0}#content{width:70%;max-width:100ch;padding:3em 4em;border-left:1px solid #ddd}pre code{font-size:1em}.name{font-size:1em}main{display:flex;flex-direction:row-reverse;justify-content:flex-end}.toc ul ul,#index ul ul{padding-left:1em}.toc > ul > li{margin-top:.5em}}</style>
|
||
<style media="print">@media print{#sidebar h1{page-break-before:always}.source{display:none}}@media print{*{background:transparent !important;color:#000 !important;box-shadow:none !important;text-shadow:none !important}a[href]:after{content:" (" attr(href) ")";font-size:90%}a[href][title]:after{content:none}abbr[title]:after{content:" (" attr(title) ")"}.ir a:after,a[href^="javascript:"]:after,a[href^="#"]:after{content:""}pre,blockquote{border:1px solid #999;page-break-inside:avoid}thead{display:table-header-group}tr,img{page-break-inside:avoid}img{max-width:100% !important}@page{margin:0.5cm}p,h2,h3{orphans:3;widows:3}h1,h2,h3,h4,h5,h6{page-break-after:avoid}}</style>
|
||
<script defer src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/highlight.min.js" integrity="sha512-D9gUyxqja7hBtkWpPWGt9wfbfaMGVt9gnyCvYa+jojwwPHLCzUm5i8rpk7vD7wNee9bA35eYIjobYPaQuKS1MQ==" crossorigin></script>
|
||
<script>window.addEventListener('DOMContentLoaded', () => {
|
||
hljs.configure({languages: ['bash', 'css', 'diff', 'graphql', 'ini', 'javascript', 'json', 'plaintext', 'python', 'python-repl', 'rust', 'shell', 'sql', 'typescript', 'xml', 'yaml']});
|
||
hljs.highlightAll();
|
||
/* Collapse source docstrings */
|
||
setTimeout(() => {
|
||
[...document.querySelectorAll('.hljs.language-python > .hljs-string')]
|
||
.filter(el => el.innerHTML.length > 200 && ['"""', "'''"].includes(el.innerHTML.substring(0, 3)))
|
||
.forEach(el => {
|
||
let d = document.createElement('details');
|
||
d.classList.add('hljs-string');
|
||
d.innerHTML = '<summary>"""</summary>' + el.innerHTML.substring(3);
|
||
el.replaceWith(d);
|
||
});
|
||
}, 100);
|
||
})</script>
|
||
</head>
|
||
<body>
|
||
<main>
|
||
<article id="content">
|
||
<header>
|
||
<h1 class="title">Module <code>audiocraft.models.multibanddiffusion</code></h1>
|
||
</header>
|
||
<section id="section-intro">
|
||
<p>Multi Band Diffusion models as described in
|
||
"From Discrete Tokens to High-Fidelity Audio Using Multi-Band Diffusion"
|
||
(paper link).</p>
|
||
</section>
|
||
<section>
|
||
</section>
|
||
<section>
|
||
</section>
|
||
<section>
|
||
</section>
|
||
<section>
|
||
<h2 class="section-title" id="header-classes">Classes</h2>
|
||
<dl>
|
||
<dt id="audiocraft.models.multibanddiffusion.DiffusionProcess"><code class="flex name class">
|
||
<span>class <span class="ident">DiffusionProcess</span></span>
|
||
<span>(</span><span>model: <a title="audiocraft.models.unet.DiffusionUnet" href="unet.html#audiocraft.models.unet.DiffusionUnet">DiffusionUnet</a>,<br>noise_schedule: <a title="audiocraft.modules.diffusion_schedule.NoiseSchedule" href="../modules/diffusion_schedule.html#audiocraft.modules.diffusion_schedule.NoiseSchedule">NoiseSchedule</a>)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">class DiffusionProcess:
|
||
"""Sampling for a diffusion Model.
|
||
|
||
Args:
|
||
model (DiffusionUnet): Diffusion U-Net model.
|
||
noise_schedule (NoiseSchedule): Noise schedule for diffusion process.
|
||
"""
|
||
def __init__(self, model: DiffusionUnet, noise_schedule: NoiseSchedule) -> None:
|
||
self.model = model
|
||
self.schedule = noise_schedule
|
||
|
||
def generate(self, condition: torch.Tensor, initial_noise: torch.Tensor,
|
||
step_list: tp.Optional[tp.List[int]] = None):
|
||
"""Perform one diffusion process to generate one of the bands.
|
||
|
||
Args:
|
||
condition (torch.Tensor): The embeddings from the compression model.
|
||
initial_noise (torch.Tensor): The initial noise to start the process.
|
||
"""
|
||
return self.schedule.generate_subsampled(model=self.model, initial=initial_noise, step_list=step_list,
|
||
condition=condition)</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Sampling for a diffusion Model.</p>
|
||
<h2 id="args">Args</h2>
|
||
<dl>
|
||
<dt><strong><code>model</code></strong> : <code>DiffusionUnet</code></dt>
|
||
<dd>Diffusion U-Net model.</dd>
|
||
<dt><strong><code>noise_schedule</code></strong> : <code>NoiseSchedule</code></dt>
|
||
<dd>Noise schedule for diffusion process.</dd>
|
||
</dl></div>
|
||
<h3>Methods</h3>
|
||
<dl>
|
||
<dt id="audiocraft.models.multibanddiffusion.DiffusionProcess.generate"><code class="name flex">
|
||
<span>def <span class="ident">generate</span></span>(<span>self,<br>condition: torch.Tensor,<br>initial_noise: torch.Tensor,<br>step_list: List[int] | None = None)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">def generate(self, condition: torch.Tensor, initial_noise: torch.Tensor,
|
||
step_list: tp.Optional[tp.List[int]] = None):
|
||
"""Perform one diffusion process to generate one of the bands.
|
||
|
||
Args:
|
||
condition (torch.Tensor): The embeddings from the compression model.
|
||
initial_noise (torch.Tensor): The initial noise to start the process.
|
||
"""
|
||
return self.schedule.generate_subsampled(model=self.model, initial=initial_noise, step_list=step_list,
|
||
condition=condition)</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Perform one diffusion process to generate one of the bands.</p>
|
||
<h2 id="args">Args</h2>
|
||
<dl>
|
||
<dt><strong><code>condition</code></strong> : <code>torch.Tensor</code></dt>
|
||
<dd>The embeddings from the compression model.</dd>
|
||
<dt><strong><code>initial_noise</code></strong> : <code>torch.Tensor</code></dt>
|
||
<dd>The initial noise to start the process.</dd>
|
||
</dl></div>
|
||
</dd>
|
||
</dl>
|
||
</dd>
|
||
<dt id="audiocraft.models.multibanddiffusion.MultiBandDiffusion"><code class="flex name class">
|
||
<span>class <span class="ident">MultiBandDiffusion</span></span>
|
||
<span>(</span><span>DPs: List[<a title="audiocraft.models.multibanddiffusion.DiffusionProcess" href="#audiocraft.models.multibanddiffusion.DiffusionProcess">DiffusionProcess</a>],<br>codec_model: <a title="audiocraft.models.encodec.CompressionModel" href="encodec.html#audiocraft.models.encodec.CompressionModel">CompressionModel</a>)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">class MultiBandDiffusion:
|
||
"""Sample from multiple diffusion models.
|
||
|
||
Args:
|
||
DPs (list of DiffusionProcess): Diffusion processes.
|
||
codec_model (CompressionModel): Underlying compression model used to obtain discrete tokens.
|
||
"""
|
||
def __init__(self, DPs: tp.List[DiffusionProcess], codec_model: CompressionModel) -> None:
|
||
self.DPs = DPs
|
||
self.codec_model = codec_model
|
||
self.device = next(self.codec_model.parameters()).device
|
||
|
||
@property
|
||
def sample_rate(self) -> int:
|
||
return self.codec_model.sample_rate
|
||
|
||
@staticmethod
|
||
def get_mbd_musicgen(device=None):
|
||
"""Load our diffusion models trained for MusicGen."""
|
||
if device is None:
|
||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||
path = 'facebook/multiband-diffusion'
|
||
filename = 'mbd_musicgen_32khz.th'
|
||
name = 'facebook/musicgen-small'
|
||
codec_model = load_compression_model(name, device=device)
|
||
models, processors, cfgs = load_diffusion_models(path, filename=filename, device=device)
|
||
DPs = []
|
||
for i in range(len(models)):
|
||
schedule = NoiseSchedule(**cfgs[i].schedule, sample_processor=processors[i], device=device)
|
||
DPs.append(DiffusionProcess(model=models[i], noise_schedule=schedule))
|
||
return MultiBandDiffusion(DPs=DPs, codec_model=codec_model)
|
||
|
||
@staticmethod
|
||
def get_mbd_24khz(bw: float = 3.0,
|
||
device: tp.Optional[tp.Union[torch.device, str]] = None,
|
||
n_q: tp.Optional[int] = None):
|
||
"""Get the pretrained Models for MultibandDiffusion.
|
||
|
||
Args:
|
||
bw (float): Bandwidth of the compression model.
|
||
device (torch.device or str, optional): Device on which the models are loaded.
|
||
n_q (int, optional): Number of quantizers to use within the compression model.
|
||
"""
|
||
if device is None:
|
||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||
assert bw in [1.5, 3.0, 6.0], f"bandwidth {bw} not available"
|
||
if n_q is not None:
|
||
assert n_q in [2, 4, 8]
|
||
assert {1.5: 2, 3.0: 4, 6.0: 8}[bw] == n_q, \
|
||
f"bandwidth and number of codebooks missmatch to use n_q = {n_q} bw should be {n_q * (1.5 / 2)}"
|
||
n_q = {1.5: 2, 3.0: 4, 6.0: 8}[bw]
|
||
codec_model = CompressionSolver.model_from_checkpoint(
|
||
'//pretrained/facebook/encodec_24khz', device=device)
|
||
codec_model.set_num_codebooks(n_q)
|
||
codec_model = codec_model.to(device)
|
||
path = 'facebook/multiband-diffusion'
|
||
filename = f'mbd_comp_{n_q}.pt'
|
||
models, processors, cfgs = load_diffusion_models(path, filename=filename, device=device)
|
||
DPs = []
|
||
for i in range(len(models)):
|
||
schedule = NoiseSchedule(**cfgs[i].schedule, sample_processor=processors[i], device=device)
|
||
DPs.append(DiffusionProcess(model=models[i], noise_schedule=schedule))
|
||
return MultiBandDiffusion(DPs=DPs, codec_model=codec_model)
|
||
|
||
@torch.no_grad()
|
||
def get_condition(self, wav: torch.Tensor, sample_rate: int) -> torch.Tensor:
|
||
"""Get the conditioning (i.e. latent representations of the compression model) from a waveform.
|
||
Args:
|
||
wav (torch.Tensor): The audio that we want to extract the conditioning from.
|
||
sample_rate (int): Sample rate of the audio."""
|
||
if sample_rate != self.sample_rate:
|
||
wav = julius.resample_frac(wav, sample_rate, self.sample_rate)
|
||
codes, scale = self.codec_model.encode(wav)
|
||
assert scale is None, "Scaled compression models not supported."
|
||
emb = self.get_emb(codes)
|
||
return emb
|
||
|
||
@torch.no_grad()
|
||
def get_emb(self, codes: torch.Tensor):
|
||
"""Get latent representation from the discrete codes.
|
||
Args:
|
||
codes (torch.Tensor): Discrete tokens."""
|
||
emb = self.codec_model.decode_latent(codes)
|
||
return emb
|
||
|
||
def generate(self, emb: torch.Tensor, size: tp.Optional[torch.Size] = None,
|
||
step_list: tp.Optional[tp.List[int]] = None):
|
||
"""Generate waveform audio from the latent embeddings of the compression model.
|
||
Args:
|
||
emb (torch.Tensor): Conditioning embeddings
|
||
size (None, torch.Size): Size of the output
|
||
if None this is computed from the typical upsampling of the model.
|
||
step_list (list[int], optional): list of Markov chain steps, defaults to 50 linearly spaced step.
|
||
"""
|
||
if size is None:
|
||
upsampling = int(self.codec_model.sample_rate / self.codec_model.frame_rate)
|
||
size = torch.Size([emb.size(0), self.codec_model.channels, emb.size(-1) * upsampling])
|
||
assert size[0] == emb.size(0)
|
||
out = torch.zeros(size).to(self.device)
|
||
for DP in self.DPs:
|
||
out += DP.generate(condition=emb, step_list=step_list, initial_noise=torch.randn_like(out))
|
||
return out
|
||
|
||
def re_eq(self, wav: torch.Tensor, ref: torch.Tensor, n_bands: int = 32, strictness: float = 1):
|
||
"""Match the eq to the encodec output by matching the standard deviation of some frequency bands.
|
||
Args:
|
||
wav (torch.Tensor): Audio to equalize.
|
||
ref (torch.Tensor): Reference audio from which we match the spectrogram.
|
||
n_bands (int): Number of bands of the eq.
|
||
strictness (float): How strict the matching. 0 is no matching, 1 is exact matching.
|
||
"""
|
||
split = julius.SplitBands(n_bands=n_bands, sample_rate=self.codec_model.sample_rate).to(wav.device)
|
||
bands = split(wav)
|
||
bands_ref = split(ref)
|
||
out = torch.zeros_like(ref)
|
||
for i in range(n_bands):
|
||
out += bands[i] * (bands_ref[i].std() / bands[i].std()) ** strictness
|
||
return out
|
||
|
||
def regenerate(self, wav: torch.Tensor, sample_rate: int):
|
||
"""Regenerate a waveform through compression and diffusion regeneration.
|
||
Args:
|
||
wav (torch.Tensor): Original 'ground truth' audio.
|
||
sample_rate (int): Sample rate of the input (and output) wav.
|
||
"""
|
||
if sample_rate != self.codec_model.sample_rate:
|
||
wav = julius.resample_frac(wav, sample_rate, self.codec_model.sample_rate)
|
||
emb = self.get_condition(wav, sample_rate=self.codec_model.sample_rate)
|
||
size = wav.size()
|
||
out = self.generate(emb, size=size)
|
||
if sample_rate != self.codec_model.sample_rate:
|
||
out = julius.resample_frac(out, self.codec_model.sample_rate, sample_rate)
|
||
return out
|
||
|
||
def tokens_to_wav(self, tokens: torch.Tensor, n_bands: int = 32):
|
||
"""Generate Waveform audio with diffusion from the discrete codes.
|
||
Args:
|
||
tokens (torch.Tensor): Discrete codes.
|
||
n_bands (int): Bands for the eq matching.
|
||
"""
|
||
wav_encodec = self.codec_model.decode(tokens)
|
||
condition = self.get_emb(tokens)
|
||
wav_diffusion = self.generate(emb=condition, size=wav_encodec.size())
|
||
return self.re_eq(wav=wav_diffusion, ref=wav_encodec, n_bands=n_bands)</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Sample from multiple diffusion models.</p>
|
||
<h2 id="args">Args</h2>
|
||
<dl>
|
||
<dt><strong><code>DPs</code></strong> : <code>list</code> of <code><a title="audiocraft.models.multibanddiffusion.DiffusionProcess" href="#audiocraft.models.multibanddiffusion.DiffusionProcess">DiffusionProcess</a></code></dt>
|
||
<dd>Diffusion processes.</dd>
|
||
<dt><strong><code>codec_model</code></strong> : <code>CompressionModel</code></dt>
|
||
<dd>Underlying compression model used to obtain discrete tokens.</dd>
|
||
</dl></div>
|
||
<h3>Static methods</h3>
|
||
<dl>
|
||
<dt id="audiocraft.models.multibanddiffusion.MultiBandDiffusion.get_mbd_24khz"><code class="name flex">
|
||
<span>def <span class="ident">get_mbd_24khz</span></span>(<span>bw: float = 3.0,<br>device: str | torch.device | None = None,<br>n_q: int | None = None)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@staticmethod
|
||
def get_mbd_24khz(bw: float = 3.0,
|
||
device: tp.Optional[tp.Union[torch.device, str]] = None,
|
||
n_q: tp.Optional[int] = None):
|
||
"""Get the pretrained Models for MultibandDiffusion.
|
||
|
||
Args:
|
||
bw (float): Bandwidth of the compression model.
|
||
device (torch.device or str, optional): Device on which the models are loaded.
|
||
n_q (int, optional): Number of quantizers to use within the compression model.
|
||
"""
|
||
if device is None:
|
||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||
assert bw in [1.5, 3.0, 6.0], f"bandwidth {bw} not available"
|
||
if n_q is not None:
|
||
assert n_q in [2, 4, 8]
|
||
assert {1.5: 2, 3.0: 4, 6.0: 8}[bw] == n_q, \
|
||
f"bandwidth and number of codebooks missmatch to use n_q = {n_q} bw should be {n_q * (1.5 / 2)}"
|
||
n_q = {1.5: 2, 3.0: 4, 6.0: 8}[bw]
|
||
codec_model = CompressionSolver.model_from_checkpoint(
|
||
'//pretrained/facebook/encodec_24khz', device=device)
|
||
codec_model.set_num_codebooks(n_q)
|
||
codec_model = codec_model.to(device)
|
||
path = 'facebook/multiband-diffusion'
|
||
filename = f'mbd_comp_{n_q}.pt'
|
||
models, processors, cfgs = load_diffusion_models(path, filename=filename, device=device)
|
||
DPs = []
|
||
for i in range(len(models)):
|
||
schedule = NoiseSchedule(**cfgs[i].schedule, sample_processor=processors[i], device=device)
|
||
DPs.append(DiffusionProcess(model=models[i], noise_schedule=schedule))
|
||
return MultiBandDiffusion(DPs=DPs, codec_model=codec_model)</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Get the pretrained Models for MultibandDiffusion.</p>
|
||
<h2 id="args">Args</h2>
|
||
<dl>
|
||
<dt><strong><code>bw</code></strong> : <code>float</code></dt>
|
||
<dd>Bandwidth of the compression model.</dd>
|
||
<dt><strong><code>device</code></strong> : <code>torch.device</code> or <code>str</code>, optional</dt>
|
||
<dd>Device on which the models are loaded.</dd>
|
||
<dt><strong><code>n_q</code></strong> : <code>int</code>, optional</dt>
|
||
<dd>Number of quantizers to use within the compression model.</dd>
|
||
</dl></div>
|
||
</dd>
|
||
<dt id="audiocraft.models.multibanddiffusion.MultiBandDiffusion.get_mbd_musicgen"><code class="name flex">
|
||
<span>def <span class="ident">get_mbd_musicgen</span></span>(<span>device=None)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@staticmethod
|
||
def get_mbd_musicgen(device=None):
|
||
"""Load our diffusion models trained for MusicGen."""
|
||
if device is None:
|
||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||
path = 'facebook/multiband-diffusion'
|
||
filename = 'mbd_musicgen_32khz.th'
|
||
name = 'facebook/musicgen-small'
|
||
codec_model = load_compression_model(name, device=device)
|
||
models, processors, cfgs = load_diffusion_models(path, filename=filename, device=device)
|
||
DPs = []
|
||
for i in range(len(models)):
|
||
schedule = NoiseSchedule(**cfgs[i].schedule, sample_processor=processors[i], device=device)
|
||
DPs.append(DiffusionProcess(model=models[i], noise_schedule=schedule))
|
||
return MultiBandDiffusion(DPs=DPs, codec_model=codec_model)</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Load our diffusion models trained for MusicGen.</p></div>
|
||
</dd>
|
||
</dl>
|
||
<h3>Instance variables</h3>
|
||
<dl>
|
||
<dt id="audiocraft.models.multibanddiffusion.MultiBandDiffusion.sample_rate"><code class="name">prop <span class="ident">sample_rate</span> : int</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@property
|
||
def sample_rate(self) -> int:
|
||
return self.codec_model.sample_rate</code></pre>
|
||
</details>
|
||
<div class="desc"></div>
|
||
</dd>
|
||
</dl>
|
||
<h3>Methods</h3>
|
||
<dl>
|
||
<dt id="audiocraft.models.multibanddiffusion.MultiBandDiffusion.generate"><code class="name flex">
|
||
<span>def <span class="ident">generate</span></span>(<span>self,<br>emb: torch.Tensor,<br>size: torch.Size | None = None,<br>step_list: List[int] | None = None)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">def generate(self, emb: torch.Tensor, size: tp.Optional[torch.Size] = None,
|
||
step_list: tp.Optional[tp.List[int]] = None):
|
||
"""Generate waveform audio from the latent embeddings of the compression model.
|
||
Args:
|
||
emb (torch.Tensor): Conditioning embeddings
|
||
size (None, torch.Size): Size of the output
|
||
if None this is computed from the typical upsampling of the model.
|
||
step_list (list[int], optional): list of Markov chain steps, defaults to 50 linearly spaced step.
|
||
"""
|
||
if size is None:
|
||
upsampling = int(self.codec_model.sample_rate / self.codec_model.frame_rate)
|
||
size = torch.Size([emb.size(0), self.codec_model.channels, emb.size(-1) * upsampling])
|
||
assert size[0] == emb.size(0)
|
||
out = torch.zeros(size).to(self.device)
|
||
for DP in self.DPs:
|
||
out += DP.generate(condition=emb, step_list=step_list, initial_noise=torch.randn_like(out))
|
||
return out</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Generate waveform audio from the latent embeddings of the compression model.</p>
|
||
<h2 id="args">Args</h2>
|
||
<dl>
|
||
<dt><strong><code>emb</code></strong> : <code>torch.Tensor</code></dt>
|
||
<dd>Conditioning embeddings</dd>
|
||
<dt><strong><code>size</code></strong> : <code>None, torch.Size</code></dt>
|
||
<dd>Size of the output
|
||
if None this is computed from the typical upsampling of the model.</dd>
|
||
<dt><strong><code>step_list</code></strong> : <code>list[int]</code>, optional</dt>
|
||
<dd>list of Markov chain steps, defaults to 50 linearly spaced step.</dd>
|
||
</dl></div>
|
||
</dd>
|
||
<dt id="audiocraft.models.multibanddiffusion.MultiBandDiffusion.get_condition"><code class="name flex">
|
||
<span>def <span class="ident">get_condition</span></span>(<span>self, wav: torch.Tensor, sample_rate: int) ‑> torch.Tensor</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@torch.no_grad()
|
||
def get_condition(self, wav: torch.Tensor, sample_rate: int) -> torch.Tensor:
|
||
"""Get the conditioning (i.e. latent representations of the compression model) from a waveform.
|
||
Args:
|
||
wav (torch.Tensor): The audio that we want to extract the conditioning from.
|
||
sample_rate (int): Sample rate of the audio."""
|
||
if sample_rate != self.sample_rate:
|
||
wav = julius.resample_frac(wav, sample_rate, self.sample_rate)
|
||
codes, scale = self.codec_model.encode(wav)
|
||
assert scale is None, "Scaled compression models not supported."
|
||
emb = self.get_emb(codes)
|
||
return emb</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Get the conditioning (i.e. latent representations of the compression model) from a waveform.</p>
|
||
<h2 id="args">Args</h2>
|
||
<dl>
|
||
<dt><strong><code>wav</code></strong> : <code>torch.Tensor</code></dt>
|
||
<dd>The audio that we want to extract the conditioning from.</dd>
|
||
<dt><strong><code>sample_rate</code></strong> : <code>int</code></dt>
|
||
<dd>Sample rate of the audio.</dd>
|
||
</dl></div>
|
||
</dd>
|
||
<dt id="audiocraft.models.multibanddiffusion.MultiBandDiffusion.get_emb"><code class="name flex">
|
||
<span>def <span class="ident">get_emb</span></span>(<span>self, codes: torch.Tensor)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@torch.no_grad()
|
||
def get_emb(self, codes: torch.Tensor):
|
||
"""Get latent representation from the discrete codes.
|
||
Args:
|
||
codes (torch.Tensor): Discrete tokens."""
|
||
emb = self.codec_model.decode_latent(codes)
|
||
return emb</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Get latent representation from the discrete codes.</p>
|
||
<h2 id="args">Args</h2>
|
||
<dl>
|
||
<dt><strong><code>codes</code></strong> : <code>torch.Tensor</code></dt>
|
||
<dd>Discrete tokens.</dd>
|
||
</dl></div>
|
||
</dd>
|
||
<dt id="audiocraft.models.multibanddiffusion.MultiBandDiffusion.re_eq"><code class="name flex">
|
||
<span>def <span class="ident">re_eq</span></span>(<span>self,<br>wav: torch.Tensor,<br>ref: torch.Tensor,<br>n_bands: int = 32,<br>strictness: float = 1)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">def re_eq(self, wav: torch.Tensor, ref: torch.Tensor, n_bands: int = 32, strictness: float = 1):
|
||
"""Match the eq to the encodec output by matching the standard deviation of some frequency bands.
|
||
Args:
|
||
wav (torch.Tensor): Audio to equalize.
|
||
ref (torch.Tensor): Reference audio from which we match the spectrogram.
|
||
n_bands (int): Number of bands of the eq.
|
||
strictness (float): How strict the matching. 0 is no matching, 1 is exact matching.
|
||
"""
|
||
split = julius.SplitBands(n_bands=n_bands, sample_rate=self.codec_model.sample_rate).to(wav.device)
|
||
bands = split(wav)
|
||
bands_ref = split(ref)
|
||
out = torch.zeros_like(ref)
|
||
for i in range(n_bands):
|
||
out += bands[i] * (bands_ref[i].std() / bands[i].std()) ** strictness
|
||
return out</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Match the eq to the encodec output by matching the standard deviation of some frequency bands.</p>
|
||
<h2 id="args">Args</h2>
|
||
<dl>
|
||
<dt><strong><code>wav</code></strong> : <code>torch.Tensor</code></dt>
|
||
<dd>Audio to equalize.</dd>
|
||
<dt><strong><code>ref</code></strong> : <code>torch.Tensor</code></dt>
|
||
<dd>Reference audio from which we match the spectrogram.</dd>
|
||
<dt><strong><code>n_bands</code></strong> : <code>int</code></dt>
|
||
<dd>Number of bands of the eq.</dd>
|
||
<dt><strong><code>strictness</code></strong> : <code>float</code></dt>
|
||
<dd>How strict the matching. 0 is no matching, 1 is exact matching.</dd>
|
||
</dl></div>
|
||
</dd>
|
||
<dt id="audiocraft.models.multibanddiffusion.MultiBandDiffusion.regenerate"><code class="name flex">
|
||
<span>def <span class="ident">regenerate</span></span>(<span>self, wav: torch.Tensor, sample_rate: int)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">def regenerate(self, wav: torch.Tensor, sample_rate: int):
|
||
"""Regenerate a waveform through compression and diffusion regeneration.
|
||
Args:
|
||
wav (torch.Tensor): Original 'ground truth' audio.
|
||
sample_rate (int): Sample rate of the input (and output) wav.
|
||
"""
|
||
if sample_rate != self.codec_model.sample_rate:
|
||
wav = julius.resample_frac(wav, sample_rate, self.codec_model.sample_rate)
|
||
emb = self.get_condition(wav, sample_rate=self.codec_model.sample_rate)
|
||
size = wav.size()
|
||
out = self.generate(emb, size=size)
|
||
if sample_rate != self.codec_model.sample_rate:
|
||
out = julius.resample_frac(out, self.codec_model.sample_rate, sample_rate)
|
||
return out</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Regenerate a waveform through compression and diffusion regeneration.</p>
|
||
<h2 id="args">Args</h2>
|
||
<dl>
|
||
<dt><strong><code>wav</code></strong> : <code>torch.Tensor</code></dt>
|
||
<dd>Original 'ground truth' audio.</dd>
|
||
<dt><strong><code>sample_rate</code></strong> : <code>int</code></dt>
|
||
<dd>Sample rate of the input (and output) wav.</dd>
|
||
</dl></div>
|
||
</dd>
|
||
<dt id="audiocraft.models.multibanddiffusion.MultiBandDiffusion.tokens_to_wav"><code class="name flex">
|
||
<span>def <span class="ident">tokens_to_wav</span></span>(<span>self, tokens: torch.Tensor, n_bands: int = 32)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">def tokens_to_wav(self, tokens: torch.Tensor, n_bands: int = 32):
|
||
"""Generate Waveform audio with diffusion from the discrete codes.
|
||
Args:
|
||
tokens (torch.Tensor): Discrete codes.
|
||
n_bands (int): Bands for the eq matching.
|
||
"""
|
||
wav_encodec = self.codec_model.decode(tokens)
|
||
condition = self.get_emb(tokens)
|
||
wav_diffusion = self.generate(emb=condition, size=wav_encodec.size())
|
||
return self.re_eq(wav=wav_diffusion, ref=wav_encodec, n_bands=n_bands)</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Generate Waveform audio with diffusion from the discrete codes.</p>
|
||
<h2 id="args">Args</h2>
|
||
<dl>
|
||
<dt><strong><code>tokens</code></strong> : <code>torch.Tensor</code></dt>
|
||
<dd>Discrete codes.</dd>
|
||
<dt><strong><code>n_bands</code></strong> : <code>int</code></dt>
|
||
<dd>Bands for the eq matching.</dd>
|
||
</dl></div>
|
||
</dd>
|
||
</dl>
|
||
</dd>
|
||
</dl>
|
||
</section>
|
||
</article>
|
||
<nav id="sidebar">
|
||
<div class="toc">
|
||
<ul></ul>
|
||
</div>
|
||
<ul id="index">
|
||
<li><h3>Super-module</h3>
|
||
<ul>
|
||
<li><code><a title="audiocraft.models" href="index.html">audiocraft.models</a></code></li>
|
||
</ul>
|
||
</li>
|
||
<li><h3><a href="#header-classes">Classes</a></h3>
|
||
<ul>
|
||
<li>
|
||
<h4><code><a title="audiocraft.models.multibanddiffusion.DiffusionProcess" href="#audiocraft.models.multibanddiffusion.DiffusionProcess">DiffusionProcess</a></code></h4>
|
||
<ul class="">
|
||
<li><code><a title="audiocraft.models.multibanddiffusion.DiffusionProcess.generate" href="#audiocraft.models.multibanddiffusion.DiffusionProcess.generate">generate</a></code></li>
|
||
</ul>
|
||
</li>
|
||
<li>
|
||
<h4><code><a title="audiocraft.models.multibanddiffusion.MultiBandDiffusion" href="#audiocraft.models.multibanddiffusion.MultiBandDiffusion">MultiBandDiffusion</a></code></h4>
|
||
<ul class="two-column">
|
||
<li><code><a title="audiocraft.models.multibanddiffusion.MultiBandDiffusion.generate" href="#audiocraft.models.multibanddiffusion.MultiBandDiffusion.generate">generate</a></code></li>
|
||
<li><code><a title="audiocraft.models.multibanddiffusion.MultiBandDiffusion.get_condition" href="#audiocraft.models.multibanddiffusion.MultiBandDiffusion.get_condition">get_condition</a></code></li>
|
||
<li><code><a title="audiocraft.models.multibanddiffusion.MultiBandDiffusion.get_emb" href="#audiocraft.models.multibanddiffusion.MultiBandDiffusion.get_emb">get_emb</a></code></li>
|
||
<li><code><a title="audiocraft.models.multibanddiffusion.MultiBandDiffusion.get_mbd_24khz" href="#audiocraft.models.multibanddiffusion.MultiBandDiffusion.get_mbd_24khz">get_mbd_24khz</a></code></li>
|
||
<li><code><a title="audiocraft.models.multibanddiffusion.MultiBandDiffusion.get_mbd_musicgen" href="#audiocraft.models.multibanddiffusion.MultiBandDiffusion.get_mbd_musicgen">get_mbd_musicgen</a></code></li>
|
||
<li><code><a title="audiocraft.models.multibanddiffusion.MultiBandDiffusion.re_eq" href="#audiocraft.models.multibanddiffusion.MultiBandDiffusion.re_eq">re_eq</a></code></li>
|
||
<li><code><a title="audiocraft.models.multibanddiffusion.MultiBandDiffusion.regenerate" href="#audiocraft.models.multibanddiffusion.MultiBandDiffusion.regenerate">regenerate</a></code></li>
|
||
<li><code><a title="audiocraft.models.multibanddiffusion.MultiBandDiffusion.sample_rate" href="#audiocraft.models.multibanddiffusion.MultiBandDiffusion.sample_rate">sample_rate</a></code></li>
|
||
<li><code><a title="audiocraft.models.multibanddiffusion.MultiBandDiffusion.tokens_to_wav" href="#audiocraft.models.multibanddiffusion.MultiBandDiffusion.tokens_to_wav">tokens_to_wav</a></code></li>
|
||
</ul>
|
||
</li>
|
||
</ul>
|
||
</li>
|
||
</ul>
|
||
</nav>
|
||
</main>
|
||
<footer id="footer">
|
||
<p>Generated by <a href="https://pdoc3.github.io/pdoc" title="pdoc: Python API documentation generator"><cite>pdoc</cite> 0.11.5</a>.</p>
|
||
</footer>
|
||
</body>
|
||
</html>
|