merge main

Author: julian-parker

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name='stable-audio-tools-dev',
-    version='0.0.9',
+    version='0.0.10',
     url='https://github.com/Stability-AI/stable-audio-tools-dev.git',
     author='Stability AI',
     description='Training and inference tools for generative audio models from Stability AI',
@@ -17,7 +17,8 @@
         'einops-exts==0.0.4',
         'ema-pytorch==0.2.3',
         'encodec==0.1.1',
-        'gradio==3.42.0',
+        'flash-attn>=2.5.0',
+        'gradio>=3.42.0',
         'huggingface_hub',
         'importlib-resources==5.12.0',
         'k-diffusion==0.1.1',

stable_audio_tools/data/dataset.py

@@ -589,7 +589,7 @@ def create_dataloader_from_config(dataset_config, batch_size, sample_size, sampl
             sample_rate=sample_rate,
             sample_size=sample_size,
             batch_size=batch_size,
-            random_crop=True,
+            random_crop=dataset_config.get("random_crop", True),
             num_workers=num_workers,
             persistent_workers=True,
             force_channels=force_channels,

stable_audio_tools/inference/generation.py

@@ -211,6 +211,8 @@ def generate_diffusion_cond(
     # Denoising process done. 
     # If this is latent diffusion, decode latents back into audio
     if model.pretransform is not None and not return_latents:
+        #cast sampled latents to pretransform dtype
+        sampled = sampled.to(next(model.pretransform.parameters()).dtype)
         sampled = model.pretransform.decode(sampled)
 
     # Return audio

stable_audio_tools/interface/gradio.py

@@ -64,10 +64,10 @@ def generate_cond(
         steps=250,
         preview_every=None,
         seed=-1,
-        sampler_type="dpmpp-2m-sde",
+        sampler_type="dpmpp-3m-sde",
         sigma_min=0.03,
-        sigma_max=50,
-        cfg_rescale=0.4,
+        sigma_max=1000,
+        cfg_rescale=0.0,
         use_init=False,
         init_audio=None,
         init_noise_level=1.0,
@@ -86,6 +86,8 @@ def generate_cond(
         torch.cuda.empty_cache()
     gc.collect()
 
+    print(f"Prompt: {prompt}")
+
     global preview_images
     preview_images = []
     if preview_every == 0:
@@ -185,7 +187,7 @@ def progress_callback(callback_info):
 
     # Convert to WAV file
     audio = rearrange(audio, "b d n -> d (b n)")
-    audio = audio.clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+    audio = audio.to(torch.float32).div(torch.max(torch.abs(audio))).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
     torchaudio.save("output.wav", audio, sample_rate)
 
     # Let's look at a nice spectrogram too
@@ -196,9 +198,9 @@ def progress_callback(callback_info):
 def generate_uncond(
         steps=250,
         seed=-1,
-        sampler_type="dpmpp-2m-sde",
+        sampler_type="dpmpp-3m-sde",
         sigma_min=0.03,
-        sigma_max=50,
+        sigma_max=1000,
         use_init=False,
         init_audio=None,
         init_noise_level=1.0,
@@ -282,7 +284,7 @@ def progress_callback(callback_info):
 
     audio = rearrange(audio, "b d n -> d (b n)")
 
-    audio = audio.clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+    audio = audio.to(torch.float32).div(torch.max(torch.abs(audio))).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
 
     torchaudio.save("output.wav", audio, sample_rate)
 
@@ -296,7 +298,6 @@ def generate_lm(
         top_k=0,    
         batch_size=1,
         ):
-
 
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
@@ -317,7 +318,7 @@ def generate_lm(
 
     audio = rearrange(audio, "b d n -> d (b n)")
 
-    audio = audio.clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+    audio = audio.to(torch.float32).div(torch.max(torch.abs(audio))).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
 
     torchaudio.save("output.wav", audio, sample_rate)
 
@@ -342,9 +343,9 @@ def create_uncond_sampling_ui(model_config):
 
             # Sampler params
                 with gr.Row():
-                    sampler_type_dropdown = gr.Dropdown(["dpmpp-2m-sde", "dpmpp-3m-sde", "k-heun", "k-lms", "k-dpmpp-2s-ancestral", "k-dpm-2", "k-dpm-fast"], label="Sampler type", value="dpmpp-2m-sde")
+                    sampler_type_dropdown = gr.Dropdown(["dpmpp-2m-sde", "dpmpp-3m-sde", "k-heun", "k-lms", "k-dpmpp-2s-ancestral", "k-dpm-2", "k-dpm-fast"], label="Sampler type", value="dpmpp-3m-sde")
                     sigma_min_slider = gr.Slider(minimum=0.0, maximum=2.0, step=0.01, value=0.03, label="Sigma min")
-                    sigma_max_slider = gr.Slider(minimum=0.0, maximum=200.0, step=0.1, value=80, label="Sigma max")
+                    sigma_max_slider = gr.Slider(minimum=0.0, maximum=1000.0, step=0.1, value=1000, label="Sigma max")
 
             with gr.Accordion("Init audio", open=False):
                 init_audio_checkbox = gr.Checkbox(label="Use init audio")
@@ -417,10 +418,10 @@ def create_sampling_ui(model_config, inpainting=False):
 
                 # Sampler params
                 with gr.Row():
-                    sampler_type_dropdown = gr.Dropdown(["dpmpp-2m-sde", "dpmpp-3m-sde", "k-heun", "k-lms", "k-dpmpp-2s-ancestral", "k-dpm-2", "k-dpm-fast"], label="Sampler type", value="dpmpp-2m-sde")
+                    sampler_type_dropdown = gr.Dropdown(["dpmpp-2m-sde", "dpmpp-3m-sde", "k-heun", "k-lms", "k-dpmpp-2s-ancestral", "k-dpm-2", "k-dpm-fast"], label="Sampler type", value="dpmpp-3m-sde")
                     sigma_min_slider = gr.Slider(minimum=0.0, maximum=2.0, step=0.01, value=0.03, label="Sigma min")
-                    sigma_max_slider = gr.Slider(minimum=0.0, maximum=200.0, step=0.1, value=80, label="Sigma max")
-                    cfg_rescale_slider = gr.Slider(minimum=0.0, maximum=1, step=0.01, value=0.2, label="CFG rescale amount")
+                    sigma_max_slider = gr.Slider(minimum=0.0, maximum=1000.0, step=0.1, value=1000, label="Sigma max")
+                    cfg_rescale_slider = gr.Slider(minimum=0.0, maximum=1, step=0.01, value=0.0, label="CFG rescale amount")
 
             if inpainting: 
                 # Inpainting Tab
@@ -547,18 +548,18 @@ def autoencoder_process(audio, latent_noise, n_quantizers):
     audio = audio.to(dtype)
 
     if n_quantizers > 0:
-        latents = model.encode_audio(audio, chunked=True, n_quantizers=n_quantizers)
+        latents = model.encode_audio(audio, chunked=False, n_quantizers=n_quantizers)
     else:
-        latents = model.encode_audio(audio, chunked=True)
+        latents = model.encode_audio(audio, chunked=False)
 
     if latent_noise > 0:
         latents = latents + torch.randn_like(latents) * latent_noise
 
-    audio = model.decode_audio(latents, chunked=True)
+    audio = model.decode_audio(latents, chunked=False)
 
     audio = rearrange(audio, "b d n -> d (b n)")
 
-    audio = audio.clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+    audio = audio.to(torch.float32).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
 
     torchaudio.save("output.wav", audio, sample_rate)
 
@@ -607,7 +608,7 @@ def diffusion_prior_process(audio, steps, sampler_type, sigma_min, sigma_max):
 
     audio = rearrange(audio, "b d n -> d (b n)")
 
-    audio = audio.clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+    audio = audio.to(torch.float32).div(torch.max(torch.abs(audio))).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
 
     torchaudio.save("output.wav", audio, sample_rate)
 
@@ -620,9 +621,9 @@ def create_diffusion_prior_ui(model_config):
         # Sampler params
         with gr.Row():
             steps_slider = gr.Slider(minimum=1, maximum=500, step=1, value=100, label="Steps")
-            sampler_type_dropdown = gr.Dropdown(["dpmpp-2m-sde", "dpmpp-3m-sde", "k-heun", "k-lms", "k-dpmpp-2s-ancestral", "k-dpm-2", "k-dpm-fast"], label="Sampler type", value="dpmpp-2m-sde")
+            sampler_type_dropdown = gr.Dropdown(["dpmpp-2m-sde", "dpmpp-3m-sde", "k-heun", "k-lms", "k-dpmpp-2s-ancestral", "k-dpm-2", "k-dpm-fast"], label="Sampler type", value="dpmpp-3m-sde")
             sigma_min_slider = gr.Slider(minimum=0.0, maximum=2.0, step=0.01, value=0.03, label="Sigma min")
-            sigma_max_slider = gr.Slider(minimum=0.0, maximum=200.0, step=0.1, value=80, label="Sigma max")
+            sigma_max_slider = gr.Slider(minimum=0.0, maximum=1000.0, step=0.1, value=1000, label="Sigma max")
         process_button = gr.Button("Process", variant='primary', scale=1)
         process_button.click(fn=diffusion_prior_process, inputs=[input_audio, steps_slider, sampler_type_dropdown, sigma_min_slider, sigma_max_slider], outputs=output_audio, api_name="process")    
 

stable_audio_tools/models/adp.py

@@ -1,6 +1,7 @@
 # Copied and modified from https://github.com/archinetai/audio-diffusion-pytorch/blob/v0.0.94/audio_diffusion_pytorch/modules.py under MIT License
 # License can be found in LICENSES/LICENSE_ADP.txt
 
+import math
 from inspect import isfunction
 from math import ceil, floor, log, pi, log2
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, TypeVar, Union