Stable Audio Open + progbars + mp3 + batched forward + cleanup

.gitattributes

@@ -34,3 +34,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 *.wav filter=lfs diff=lfs merge=lfs -text
+*.mp3 filter=lfs diff=lfs merge=lfs -text

Examples/{Beethoven.wav → Beethoven.mp3}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:30a6a087a9e0eb87422aa3b48ad966eabb1dfe105d73a25d356b71d3aee31493
-size 4828972
+oid sha256:3dcc79fe071d118df3caaeeb85d7944f93a5df40bbdb72a26b67bd57da2af7c5
+size 1097142

Examples/{Cat_dog.wav → Beethoven_arcade.mp3}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:90a97dc229eeccef307dd40db97fb09cc439ce0b45a320fd84b2ea6b03d0deb2
-size 327822
+oid sha256:542bd61d9cc1723ccfd9bfc06b0818e77fc763013827ff1f9289e2ac6a912904
+size 563040

Examples/{Beethoven_arcade.wav → Beethoven_piano.mp3}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ccd929b93c15706f2102a27973d490a84ce0eb97faba6a92ece0c6d81ed2c26e
-size 1794746
+oid sha256:000d82c39d8c41b10188d328e29cb1baa948232bacd693f22e297cc54f4bb707
+size 563040

Examples/{Beethoven_piano.wav → Beethoven_rock.mp3}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5787c31b0b3c78dec33d651d437364785713042e7cfce2290cf4baf01f65ac6f
-size 1794746
+oid sha256:c51d75c9094a50c7892449a013b32ffde266a5abd6dad9f00bf3aeec0ee935ee
+size 1097142

Examples/{Cat.wav → Cat.mp3}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:27b43763a8d9ac90dc78285ed9817b16524f24b4f4d1aa399616f1a04d4a9fd9
-size 1920508
+oid sha256:cff7010e5fb12a57508c7a0941663f1a12bfc8b3b3d01d0973359cd42ae5eb1e
+size 402542

Examples/Cat_dog.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:72ff727243606215c934552e946f7d97b5e2e39c4d6263f7f36659e3f39f3008
+size 207403

Examples/ModalJazz.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:34cf145b84b6b4669050ca42932fb74ac0f28aabbe6c665f12a877c9809fa9c6
+size 4153468

Examples/ModalJazz.wav

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:846a77046d21ebc3996841404eede9d56797c82b3414025e1ccafe586eaf2959
-size 9153322

Examples/ModalJazz_banjo.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:11680068427556981aa6304e6c11bd05debc820ca581c248954c1ffe3cd94569
+size 2128320

Examples/ModalJazz_banjo.wav

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:122e0078c0bf2fc96425071706fe0e8674c93cc1d2787fd02c0e2c0f12de5cc5
-size 6802106

Examples/Shadows.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2e0cab2ebda4507641d6a1b5d9b2d888a7526581b7de48540ebf86ce00579908
+size 1342693

Examples/Shadows_arcade.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:68c84805ea17d0697cd79bc85394754d70fb02f740db4bee4c6ccbb5269a5d84
+size 1342693

README.md

@@ -9,7 +9,10 @@ app_file: app.py
 pinned: false
 license: cc-by-sa-4.0
 short_description: Edit audios with text prompts
+hf_oauth: true
+hf_oauth_scopes:
+- read-repos
 ---
 
 The 30-second limit was introduced to ensure that queue wait times remain reasonable, especially when there are a lot of users.
-For that reason pull-requests that change this limit will not be merged. Please clone or duplicate the space to work locally without limits.
+For that reason pull-requests that change this limit will not be merged. Please clone or duplicate the space to work locally without limits.

app.py

@@ -6,27 +6,26 @@ if os.getenv('SPACES_ZERO_GPU') == "true":
 import gradio as gr
 import random
 import torch
+import os
 from torch import inference_mode
-# from tempfile import NamedTemporaryFile
-from typing import Optional
+from typing import Optional, List
 import numpy as np
 from models import load_model
 import utils
 import spaces
+import huggingface_hub
 from inversion_utils import inversion_forward_process, inversion_reverse_process
 
 
-# current_loaded_model = "cvssp/audioldm2-music"
-# # current_loaded_model = "cvssp/audioldm2-music"
-
-# ldm_stable = load_model(current_loaded_model, device, 200)  # deafult model
 LDM2 = "cvssp/audioldm2"
 MUSIC = "cvssp/audioldm2-music"
 LDM2_LARGE = "cvssp/audioldm2-large"
+STABLEAUD = "stabilityai/stable-audio-open-1.0"
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 ldm2 = load_model(model_id=LDM2, device=device)
 ldm2_large = load_model(model_id=LDM2_LARGE, device=device)
 ldm2_music = load_model(model_id=MUSIC, device=device)
+ldm_stableaud = load_model(model_id=STABLEAUD, device=device, token=os.getenv('PRIV_TOKEN'))
 
 
 def randomize_seed_fn(seed, randomize_seed):
@@ -36,89 +35,136 @@ def randomize_seed_fn(seed, randomize_seed):
     return seed
 
 
-def invert(ldm_stable, x0, prompt_src, num_diffusion_steps, cfg_scale_src):  # , ldm_stable):
+def invert(ldm_stable, x0, prompt_src, num_diffusion_steps, cfg_scale_src, duration, save_compute):
     # ldm_stable.model.scheduler.set_timesteps(num_diffusion_steps, device=device)
 
     with inference_mode():
         w0 = ldm_stable.vae_encode(x0)
 
     # find Zs and wts - forward process
-    _, zs, wts = inversion_forward_process(ldm_stable, w0, etas=1,
-                                           prompts=[prompt_src],
-                                           cfg_scales=[cfg_scale_src],
-                                           prog_bar=True,
-                                           num_inference_steps=num_diffusion_steps,
-                                           numerical_fix=True)
-    return zs, wts
+    _, zs, wts, extra_info = inversion_forward_process(ldm_stable, w0, etas=1,
+                                                       prompts=[prompt_src],
+                                                       cfg_scales=[cfg_scale_src],
+                                                       num_inference_steps=num_diffusion_steps,
+                                                       numerical_fix=True,
+                                                       duration=duration,
+                                                       save_compute=save_compute)
+    return zs, wts, extra_info
 
 
-def sample(ldm_stable, zs, wts, steps, prompt_tar, tstart, cfg_scale_tar):  # , ldm_stable):
+def sample(ldm_stable, zs, wts, extra_info, prompt_tar, tstart, cfg_scale_tar, duration, save_compute):
     # reverse process (via Zs and wT)
     tstart = torch.tensor(tstart, dtype=torch.int)
-    skip = steps - tstart
-    w0, _ = inversion_reverse_process(ldm_stable, xT=wts, skips=steps - skip,
+    w0, _ = inversion_reverse_process(ldm_stable, xT=wts, tstart=tstart,
                                       etas=1., prompts=[prompt_tar],
                                       neg_prompts=[""], cfg_scales=[cfg_scale_tar],
-                                      prog_bar=True,
-                                      zs=zs[:int(steps - skip)])
+                                      zs=zs[:int(tstart)],
+                                      duration=duration,
+                                      extra_info=extra_info,
+                                      save_compute=save_compute)
 
     # vae decode image
     with inference_mode():
         x0_dec = ldm_stable.vae_decode(w0)
-    if x0_dec.dim() < 4:
-        x0_dec = x0_dec[None, :, :, :]
 
-    with torch.no_grad():
-        audio = ldm_stable.decode_to_mel(x0_dec)
+    if 'stable-audio' not in ldm_stable.model_id:
+        if x0_dec.dim() < 4:
+            x0_dec = x0_dec[None, :, :, :]
 
-    return (16000, audio.squeeze().cpu().numpy())
-
-def get_duration(input_audio, model_id: str, do_inversion: bool,
-    wts: Optional[torch.Tensor], zs: Optional[torch.Tensor],
-    saved_inv_model: str, source_prompt="", target_prompt="",
-    steps=200, cfg_scale_src=3.5, cfg_scale_tar=12, t_start=45, randomize_seed=True):
+        with torch.no_grad():
+            audio = ldm_stable.decode_to_mel(x0_dec)
+    else:
+        audio = x0_dec.squeeze(0).T
+
+    return (ldm_stable.get_sr(), audio.squeeze().cpu().numpy())
+
+
+def get_duration(input_audio,
+                 model_id: str,
+                 do_inversion: bool,
+                 wts: Optional[torch.Tensor], zs: Optional[torch.Tensor], extra_info: Optional[List],
+                 saved_inv_model: str,
+                 source_prompt: str = "",
+                 target_prompt: str = "",
+                 steps: int = 200,
+                 cfg_scale_src: float = 3.5,
+                 cfg_scale_tar: float = 12,
+                 t_start: int = 45,
+                 randomize_seed: bool = True,
+                 save_compute: bool = True,
+                 oauth_token: Optional[gr.OAuthToken] = None):
     if model_id == LDM2:
-        factor = 0.8
+        factor = 1
     elif model_id == LDM2_LARGE:
-        factor = 1.5
+        factor = 2.5
+    elif model_id == STABLEAUD:
+        factor = 3.2
     else:  # MUSIC
         factor = 1
 
-    mult = 0
+    forwards = 0
     if do_inversion or randomize_seed:
-        mult = steps
+        forwards = steps if source_prompt == "" else steps * 2  # x2 when there is a prompt text
+    forwards += int(t_start / 100 * steps) * 2
+
+    duration = min(utils.get_duration(input_audio), utils.MAX_DURATION)
+    time_for_maxlength = factor * forwards * 0.15  # 0.25 is the time per forward pass
+    print('expected time:', time_for_maxlength / utils.MAX_DURATION * duration)
+
+    spare_time = 5
+    return max(10, time_for_maxlength / utils.MAX_DURATION * duration + spare_time)
+
 
+def verify_model_params(model_id: str, input_audio, src_prompt: str, tar_prompt: str, cfg_scale_src: float,
+                        oauth_token: gr.OAuthToken | None):
     if input_audio is None:
         raise gr.Error('Input audio missing!')
-    duration = min(utils.get_duration(input_audio), 30)
 
-    time_per_iter_of_full = factor * ((t_start /100 * steps)*2 + mult) * 0.25
-    print('expected time:', time_per_iter_of_full / 30 * duration)
-    return max(15, time_per_iter_of_full / 30 * duration)
-    
+    if tar_prompt == "":
+        raise gr.Error("Please provide a target prompt to edit the audio.")
+
+    if src_prompt != "":
+        if model_id == STABLEAUD and cfg_scale_src != 1:
+            gr.Info("Consider using Source Guidance Scale=1 for Stable Audio Open 1.0.")
+        elif model_id != STABLEAUD and cfg_scale_src != 3:
+            gr.Info(f"Consider using Source Guidance Scale=3 for {model_id}.")
+
+    if model_id == STABLEAUD:
+        if oauth_token is None:
+            raise gr.Error("You must be logged in to use Stable Audio Open 1.0. Please log in and try again.")
+        try:
+            huggingface_hub.get_hf_file_metadata(huggingface_hub.hf_hub_url(STABLEAUD, 'transformer/config.json'),
+                                                 token=oauth_token.token)
+            print('Has Access')
+        # except huggingface_hub.utils._errors.GatedRepoError:
+        except huggingface_hub.errors.GatedRepoError:
+            raise gr.Error("You need to accept the license agreement to use Stable Audio Open 1.0. "
+                           "Visit the <a href='https://huggingface.co/stabilityai/stable-audio-open-1.0'>"
+                           "model page</a> to get access.")
 
-@spaces.GPU(duration=get_duration)
-def edit(
-    # cache_dir,
-    input_audio,
-    model_id: str,
-    do_inversion: bool,
-    #  wtszs_file: str,
-    wts: Optional[torch.Tensor], zs: Optional[torch.Tensor],
-    saved_inv_model: str,
-    source_prompt="",
-    target_prompt="",
-    steps=200,
-    cfg_scale_src=3.5,
-    cfg_scale_tar=12,
-    t_start=45,
-    randomize_seed=True):
 
+@spaces.GPU(duration=get_duration)
+def edit(input_audio,
+         model_id: str,
+         do_inversion: bool,
+         wts: Optional[torch.Tensor], zs: Optional[torch.Tensor], extra_info: Optional[List],
+         saved_inv_model: str,
+         source_prompt: str = "",
+         target_prompt: str = "",
+         steps: int = 200,
+         cfg_scale_src: float = 3.5,
+         cfg_scale_tar: float = 12,
+         t_start: int = 45,
+         randomize_seed: bool = True,
+         save_compute: bool = True,
+         oauth_token: Optional[gr.OAuthToken] = None):
     print(model_id)
     if model_id == LDM2:
         ldm_stable = ldm2
     elif model_id == LDM2_LARGE:
         ldm_stable = ldm2_large
+    elif model_id == STABLEAUD:
+        ldm_stable = ldm_stableaud
     else:  # MUSIC
         ldm_stable = ldm2_music
 
@@ -130,102 +176,126 @@ def edit(
 
     if input_audio is None:
         raise gr.Error('Input audio missing!')
-    x0 = utils.load_audio(input_audio, ldm_stable.get_fn_STFT(), device=device)
-
-    # if not (do_inversion or randomize_seed):
-        # if not os.path.exists(wtszs_file):
-            # do_inversion = True
-            # Too much time has passed
+    x0, _, duration = utils.load_audio(input_audio, ldm_stable.get_fn_STFT(), device=device,
+                                       stft=('stable-audio' not in ldm_stable.model_id), model_sr=ldm_stable.get_sr())
     if wts is None or zs is None:
         do_inversion = True
 
     if do_inversion or randomize_seed:  # always re-run inversion
-        zs_tensor, wts_tensor = invert(ldm_stable=ldm_stable, x0=x0, prompt_src=source_prompt,
-                                       num_diffusion_steps=steps,
-                                       cfg_scale_src=cfg_scale_src)
-        # f = NamedTemporaryFile("wb", dir=cache_dir, suffix=".pth", delete=False)
-        # torch.save({'wts': wts_tensor, 'zs': zs_tensor}, f.name)
-        # wtszs_file = f.name
-        # wtszs_file = gr.State(value=f.name)
-        # wts = gr.State(value=wts_tensor)
+        zs_tensor, wts_tensor, extra_info_list = invert(ldm_stable=ldm_stable, x0=x0, prompt_src=source_prompt,
+                                                        num_diffusion_steps=steps,
+                                                        cfg_scale_src=cfg_scale_src,
+                                                        duration=duration,
+                                                        save_compute=save_compute)
         wts = wts_tensor
         zs = zs_tensor
-        # zs = gr.State(value=zs_tensor)
-        # demo.move_resource_to_block_cache(f.name)
+        extra_info = extra_info_list
         saved_inv_model = model_id
         do_inversion = False
     else:
-        #     wtszs = torch.load(wtszs_file, map_location=device)
-        #     # wtszs = torch.load(wtszs_file.f, map_location=device)
-        #     wts_tensor = wtszs['wts']
-        #     zs_tensor = wtszs['zs']
         wts_tensor = wts.to(device)
         zs_tensor = zs.to(device)
+        extra_info_list = [e.to(device) for e in extra_info if e is not None]
 
-    # make sure t_start is in the right limit
-    # t_start = change_tstart_range(t_start, steps)
-
-    output = sample(ldm_stable, zs_tensor, wts_tensor, steps, prompt_tar=target_prompt,
-                    tstart=int(t_start / 100 * steps), cfg_scale_tar=cfg_scale_tar)
+    output = sample(ldm_stable, zs_tensor, wts_tensor, extra_info_list, prompt_tar=target_prompt,
+                    tstart=int(t_start / 100 * steps), cfg_scale_tar=cfg_scale_tar, duration=duration,
+                    save_compute=save_compute)
 
-    return output, wts.cpu(), zs.cpu(), saved_inv_model, do_inversion
+    return output, wts.cpu(), zs.cpu(), [e.cpu() for e in extra_info if e is not None], saved_inv_model, do_inversion
     # return output, wtszs_file, saved_inv_model, do_inversion
 
 
 def get_example():
     case = [
-        ['Examples/Beethoven.wav',
+        ['Examples/Beethoven.mp3',
          '',
          'A recording of an arcade game soundtrack.',
          45,
          'cvssp/audioldm2-music',
          '27s',
-         'Examples/Beethoven_arcade.wav',
+         'Examples/Beethoven_arcade.mp3',
          ],
-        ['Examples/Beethoven.wav',
+        ['Examples/Beethoven.mp3',
          'A high quality recording of wind instruments and strings playing.',
          'A high quality recording of a piano playing.',
          45,
          'cvssp/audioldm2-music',
          '27s',
-         'Examples/Beethoven_piano.wav',
+         'Examples/Beethoven_piano.mp3',
+         ],
+        ['Examples/Beethoven.mp3',
+         '',
+         'Heavy Rock.',
+         40,
+         'stabilityai/stable-audio-open-1.0',
+         '27s',
+         'Examples/Beethoven_rock.mp3',
          ],
-        ['Examples/ModalJazz.wav',
+        ['Examples/ModalJazz.mp3',
          'Trumpets playing alongside a piano, bass and drums in an upbeat old-timey cool jazz song.',
          'A banjo playing alongside a piano, bass and drums in an upbeat old-timey cool country song.',
          45,
          'cvssp/audioldm2-music',
          '106s',
-         'Examples/ModalJazz_banjo.wav',],
-        ['Examples/Cat.wav',
+         'Examples/ModalJazz_banjo.mp3',],
+        ['Examples/Shadows.mp3',
+         '',
+         '8-bit arcade game soundtrack.',
+         40,
+         'stabilityai/stable-audio-open-1.0',
+         '34s',
+         'Examples/Shadows_arcade.mp3',],
+        ['Examples/Cat.mp3',
          '',
          'A dog barking.',
          75,
          'cvssp/audioldm2-large',
          '10s',
-         'Examples/Cat_dog.wav',]
+         'Examples/Cat_dog.mp3',]
     ]
     return case
 
 
 intro = """
-<h1 style="font-weight: 1400; text-align: center; margin-bottom: 7px;"> ZETA Editing 🎧 </h1>
-<h2 style="font-weight: 1400; text-align: center; margin-bottom: 7px;"> Zero-Shot Text-Based Audio Editing Using DDPM Inversion 🎛️ </h2>
-<h3 style="margin-bottom: 10px; text-align: center;">
+<h1 style="font-weight: 1000; text-align: center; margin: 0px;"> ZETA Editing 🎧 </h1>
+<h2 style="font-weight: 1000; text-align: center; margin: 0px;">
+    Zero-Shot Text-Based Audio Editing Using DDPM Inversion 🎛️ </h2>
+<h3 style="margin-top: 0px; margin-bottom: 10px; text-align: center;">
     <a href="https://arxiv.org/abs/2402.10009">[Paper]</a>&nbsp;|&nbsp;
     <a href="https://hilamanor.github.io/AudioEditing/">[Project page]</a>&nbsp;|&nbsp;
     <a href="https://github.com/HilaManor/AudioEditingCode">[Code]</a>
 </h3>
 
-
-<p style="font-size: 0.9rem; margin: 0rem; line-height: 1.2em; margin-top:1em">
+<p style="font-size: 1rem; line-height: 1.2em;">
 For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings.
 <a href="https://huggingface.co/spaces/hilamanor/audioEditing?duplicate=true">
-<img style="margin-top: 0em; margin-bottom: 0em; display:inline" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" ></a>
+<img style="margin-top: 0em; margin-bottom: 0em; display:inline" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" >
+</a>
+</p>
+<p style="margin: 0px;">
+<b>NEW - 15.10.24:</b> You can now edit using <b>Stable Audio Open 1.0</b>.
+You must be <b>logged in</b> after accepting the
+<b><a href="https://huggingface.co/stabilityai/stable-audio-open-1.0">license agreement</a></b> to use it.</br>
+</p>
+<ul style="padding-left:40px; line-height:normal;">
+<li style="margin: 0px;">Prompts behave differently - e.g.,
+try "8-bit arcade" directly instead of "a recording of...". Check out the new examples below!</li>
+<li style="margin: 0px;">Try to play around <code>T-start=40%</code>.</li>
+<li style="margin: 0px;">Under "More Options": Use <code>Source Guidance Scale=1</code>, 
+and you can try fewer timesteps (even 20!).</li>
+<li style="margin: 0px;">Stable Audio Open is a general-audio model.
+For better music editing, duplicate the space and change to a
+<a href="https://huggingface.co/models?other=base_model:finetune:stabilityai/stable-audio-open-1.0">
+fine-tuned model for music</a>.</li>
+</ul>
+<p>
+<b>NEW - 15.10.24:</b> Parallel editing is enabled by default.
+To disable, uncheck <code>Efficient editing</code> under "More Options".
+Saves a bit of time.
 </p>
-
 """
 
+
 help = """
 <div style="font-size:medium">
 <b>Instructions:</b><br>
@@ -233,22 +303,27 @@ help = """
 <li>You must provide an input audio and a target prompt to edit the audio. </li>
 <li>T<sub>start</sub> is used to control the tradeoff between fidelity to the original signal and text-adhearance.
 Lower value -> favor fidelity. Higher value -> apply a stronger edit.</li>
-<li>Make sure that you use an AudioLDM2 version that is suitable for your input audio.
-For example, use the music version for music and the large version for general audio.
+<li>Make sure that you use a model version that is suitable for your input audio.
+For example, use AudioLDM2-music for music while AudioLDM2-large for general audio.
 </li>
 <li>You can additionally provide a source prompt to guide even further the editing process.</li>
 <li>Longer input will take more time.</li>
 <li><strong>Unlimited length</strong>: This space automatically trims input audio to a maximum length of 30 seconds.
-For unlimited length, duplicated the space, and remove the trimming by changing the code.
-Specifically, in the <code style="display:inline; background-color: lightgrey; ">load_audio</code> function in the <code style="display:inline; background-color: lightgrey; ">utils.py</code> file,
-change <code style="display:inline; background-color: lightgrey; ">duration = min(audioldm.utils.get_duration(audio_path), 30)</code> to 
-<code style="display:inline; background-color: lightgrey; ">duration = audioldm.utils.get_duration(audio_path)</code>.
+For unlimited length, duplicated the space, and change the
+<code style="display:inline; background-color: lightgrey;">MAX_DURATION</code> parameter
+inside <code style="display:inline; background-color: lightgrey;">utils.py</code>
+to <code style="display:inline; background-color: lightgrey;">None</code>.
+</li>
 </ul>
 </div>
 
 """
 
-with gr.Blocks(css='style.css') as demo:  #, delete_cache=(3600, 3600)) as demo:
+css = '.gradio-container {max-width: 1000px !important; padding-top: 1.5rem !important;}' \
+      '.audio-upload .wrap {min-height: 0px;}'
+
+# with gr.Blocks(css='style.css') as demo:
+with gr.Blocks(css=css) as demo:
     def reset_do_inversion(do_inversion_user, do_inversion):
         # do_inversion = gr.State(value=True)
         do_inversion = True
@@ -267,23 +342,22 @@ with gr.Blocks(css='style.css') as demo:  #, delete_cache=(3600, 3600)) as demo:
         return do_inversion_user, do_inversion
 
     gr.HTML(intro)
+
     wts = gr.State()
     zs = gr.State()
-    wtszs = gr.State()
-    # cache_dir = gr.State(demo.GRADIO_CACHE)
+    extra_info = gr.State()
     saved_inv_model = gr.State()
-    # current_loaded_model = gr.State(value="cvssp/audioldm2-music")
-    # ldm_stable = load_model("cvssp/audioldm2-music", device, 200)
-    # ldm_stable = gr.State(value=ldm_stable)
     do_inversion = gr.State(value=True)  # To save some runtime when editing the same thing over and over
     do_inversion_user = gr.State(value=False)
 
     with gr.Group():
-        gr.Markdown("💡 **note**: input longer than **30 sec** is automatically trimmed (for unlimited input, see the Help section below)")
-        with gr.Row():
-            input_audio = gr.Audio(sources=["upload", "microphone"], type="filepath", editable=True, label="Input Audio",
-                                   interactive=True, scale=1)
-            output_audio = gr.Audio(label="Edited Audio", interactive=False, scale=1)
+        gr.Markdown("💡 **note**: input longer than **30 sec** is automatically trimmed "
+                    "(for unlimited input, see the Help section below)")
+        with gr.Row(equal_height=True):
+            input_audio = gr.Audio(sources=["upload", "microphone"], type="filepath",
+                                   editable=True, label="Input Audio", interactive=True, scale=1, format='wav',
+                                   elem_classes=['audio-upload'])
+            output_audio = gr.Audio(label="Edited Audio", interactive=False, scale=1, format='wav')
 
     with gr.Row():
         tar_prompt = gr.Textbox(label="Prompt", info="Describe your desired edited output",
@@ -293,17 +367,16 @@ with gr.Blocks(css='style.css') as demo:  #, delete_cache=(3600, 3600)) as demo:
     with gr.Row():
         t_start = gr.Slider(minimum=15, maximum=85, value=45, step=1, label="T-start (%)", interactive=True, scale=3,
                             info="Lower T-start -> closer to original audio. Higher T-start -> stronger edit.")
-        # model_id = gr.Radio(label="AudioLDM2 Version",
-        model_id = gr.Dropdown(label="AudioLDM2 Version",
-                               choices=["cvssp/audioldm2",
-                                        "cvssp/audioldm2-large",
-                                        "cvssp/audioldm2-music"],
-                               info="Choose a checkpoint suitable for your intended audio and edit",
+        model_id = gr.Dropdown(label="Model Version",
+                               choices=[LDM2,
+                                        LDM2_LARGE,
+                                        MUSIC,
+                                        STABLEAUD],
+                               info="Choose a checkpoint suitable for your audio and edit",
                                value="cvssp/audioldm2-music", interactive=True, type="value", scale=2)
-
     with gr.Row():
-        with gr.Column():
-            submit = gr.Button("Edit")
+        submit = gr.Button("Edit", variant="primary", scale=3)
+        gr.LoginButton(value="Login to HF (For Stable Audio)", scale=1)
 
     with gr.Accordion("More Options", open=False):
         with gr.Row():
@@ -311,58 +384,62 @@ with gr.Blocks(css='style.css') as demo:  #, delete_cache=(3600, 3600)) as demo:
                                     info="Optional: Describe the original audio input",
                                     placeholder="A recording of a happy upbeat classical music piece",)
 
-        with gr.Row():
+        with gr.Row(equal_height=True):
             cfg_scale_src = gr.Number(value=3, minimum=0.5, maximum=25, precision=None,
                                       label="Source Guidance Scale", interactive=True, scale=1)
             cfg_scale_tar = gr.Number(value=12, minimum=0.5, maximum=25, precision=None,
                                       label="Target Guidance Scale", interactive=True, scale=1)
-            steps = gr.Number(value=50, step=1, minimum=20, maximum=300,
+            steps = gr.Number(value=50, step=1, minimum=10, maximum=300,
                               info="Higher values (e.g. 200) yield higher-quality generation.",
-                              label="Num Diffusion Steps", interactive=True, scale=1)
-        with gr.Row():
+                              label="Num Diffusion Steps", interactive=True, scale=2)
+        with gr.Row(equal_height=True):
             seed = gr.Number(value=0, precision=0, label="Seed", interactive=True)
             randomize_seed = gr.Checkbox(label='Randomize seed', value=False)
+            save_compute = gr.Checkbox(label='Efficient editing', value=True)
             length = gr.Number(label="Length", interactive=False, visible=False)
 
     with gr.Accordion("Help💡", open=False):
         gr.HTML(help)
 
     submit.click(
-        fn=randomize_seed_fn,
-        inputs=[seed, randomize_seed],
-        outputs=[seed], queue=False).then(
-            fn=clear_do_inversion_user, inputs=[do_inversion_user], outputs=[do_inversion_user]).then(
-           fn=edit,
-           inputs=[#cache_dir,
-                   input_audio,
-                   model_id,
-                   do_inversion,
-                   #    current_loaded_model, ldm_stable,
-                      wts, zs,
-                #    wtszs,
-                   saved_inv_model,
-                   src_prompt,
-                   tar_prompt,
-                   steps,
-                   cfg_scale_src,
-                   cfg_scale_tar,
-                   t_start,
-                   randomize_seed
-                   ],
-           outputs=[output_audio, wts, zs, # wtszs,
-                    saved_inv_model, do_inversion]  # , current_loaded_model, ldm_stable],
-        ).then(post_match_do_inversion, inputs=[do_inversion_user, do_inversion], outputs=[do_inversion_user, do_inversion]
-               ).then(lambda x: (demo.temp_file_sets.append(set([str(gr.utils.abspath(x))])) if type(x) is str else None),
-                      inputs=wtszs)
-
-    # demo.move_resource_to_block_cache(wtszs.value)
+            fn=verify_model_params,
+            inputs=[model_id, input_audio, src_prompt, tar_prompt, cfg_scale_src],
+            outputs=[]
+        ).success(
+            fn=randomize_seed_fn, inputs=[seed, randomize_seed], outputs=[seed], queue=False
+        ).then(
+            fn=clear_do_inversion_user, inputs=[do_inversion_user], outputs=[do_inversion_user]
+        ).then(
+            fn=edit,
+            inputs=[input_audio,
+                    model_id,
+                    do_inversion,
+                    wts, zs, extra_info,
+                    saved_inv_model,
+                    src_prompt,
+                    tar_prompt,
+                    steps,
+                    cfg_scale_src,
+                    cfg_scale_tar,
+                    t_start,
+                    randomize_seed,
+                    save_compute,
+                    ],
+            outputs=[output_audio, wts, zs, extra_info, saved_inv_model, do_inversion]
+        ).success(
+            fn=post_match_do_inversion,
+            inputs=[do_inversion_user, do_inversion],
+            outputs=[do_inversion_user, do_inversion]
+        )
 
     # If sources changed we have to rerun inversion
-    input_audio.change(fn=reset_do_inversion, inputs=[do_inversion_user, do_inversion], outputs=[do_inversion_user, do_inversion])
-    src_prompt.change(fn=reset_do_inversion, inputs=[do_inversion_user, do_inversion], outputs=[do_inversion_user, do_inversion])
-    model_id.change(fn=reset_do_inversion, inputs=[do_inversion_user, do_inversion], outputs=[do_inversion_user, do_inversion])
-    cfg_scale_src.change(fn=reset_do_inversion, inputs=[do_inversion_user, do_inversion], outputs=[do_inversion_user, do_inversion])
-    steps.change(fn=reset_do_inversion, inputs=[do_inversion_user, do_inversion], outputs=[do_inversion_user, do_inversion])
+    gr.on(
+        triggers=[input_audio.change, src_prompt.change, model_id.change, cfg_scale_src.change,
+                  steps.change, save_compute.change],
+        fn=reset_do_inversion,
+        inputs=[do_inversion_user, do_inversion],
+        outputs=[do_inversion_user, do_inversion]
+    )
 
     gr.Examples(
         label="Examples",

inversion_utils.py

@@ -1,341 +1,135 @@
 import torch
 from tqdm import tqdm
-# from torchvision import transforms as T
-from typing import List, Optional, Dict, Union
+from typing import List, Optional, Tuple
 from models import PipelineWrapper
-
-
-def mu_tilde(model, xt, x0, timestep):
-    "mu_tilde(x_t, x_0) DDPM paper eq. 7"
-    prev_timestep = timestep - model.scheduler.config.num_train_timesteps // model.scheduler.num_inference_steps
-    alpha_prod_t_prev = model.scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 \
-        else model.scheduler.final_alpha_cumprod
-    alpha_t = model.scheduler.alphas[timestep]
-    beta_t = 1 - alpha_t
-    alpha_bar = model.scheduler.alphas_cumprod[timestep]
-    return ((alpha_prod_t_prev ** 0.5 * beta_t) / (1-alpha_bar)) * x0 + \
-        ((alpha_t**0.5 * (1-alpha_prod_t_prev)) / (1 - alpha_bar)) * xt
-
-
-def sample_xts_from_x0(model, x0, num_inference_steps=50, x_prev_mode=False):
-    """
-    Samples from P(x_1:T|x_0)
-    """
-    # torch.manual_seed(43256465436)
-    alpha_bar = model.model.scheduler.alphas_cumprod
-    sqrt_one_minus_alpha_bar = (1-alpha_bar) ** 0.5
-    alphas = model.model.scheduler.alphas
-    # betas = 1 - alphas
-    variance_noise_shape = (
-            num_inference_steps + 1,
-            model.model.unet.config.in_channels,
-            # model.unet.sample_size,
-            # model.unet.sample_size)
-            x0.shape[-2],
-            x0.shape[-1])
-
-    timesteps = model.model.scheduler.timesteps.to(model.device)
-    t_to_idx = {int(v): k for k, v in enumerate(timesteps)}
-    xts = torch.zeros(variance_noise_shape).to(x0.device)
-    xts[0] = x0
-    x_prev = x0
-    for t in reversed(timesteps):
-        # idx = t_to_idx[int(t)]
-        idx = num_inference_steps-t_to_idx[int(t)]
-        if x_prev_mode:
-            xts[idx] = x_prev * (alphas[t] ** 0.5) + torch.randn_like(x0) * ((1-alphas[t]) ** 0.5)
-            x_prev = xts[idx].clone()
-        else:
-            xts[idx] = x0 * (alpha_bar[t] ** 0.5) + torch.randn_like(x0) * sqrt_one_minus_alpha_bar[t]
-    # xts = torch.cat([xts, x0 ],dim = 0)
-
-    return xts
-
-
-def forward_step(model, model_output, timestep, sample):
-    next_timestep = min(model.scheduler.config.num_train_timesteps - 2,
-                        timestep + model.scheduler.config.num_train_timesteps // model.scheduler.num_inference_steps)
-
-    # 2. compute alphas, betas
-    alpha_prod_t = model.scheduler.alphas_cumprod[timestep]
-    # alpha_prod_t_next = self.scheduler.alphas_cumprod[next_timestep] if next_ltimestep >= 0 \
-    #     else self.scheduler.final_alpha_cumprod
-
-    beta_prod_t = 1 - alpha_prod_t
-
-    # 3. compute predicted original sample from predicted noise also called
-    # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
-    pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
-
-    # 5. TODO: simple noising implementatiom
-    next_sample = model.scheduler.add_noise(pred_original_sample, model_output, torch.LongTensor([next_timestep]))
-    return next_sample
+import gradio as gr
 
 
 def inversion_forward_process(model: PipelineWrapper,
                               x0: torch.Tensor,
                               etas: Optional[float] = None,
-                              prog_bar: bool = False,
                               prompts: List[str] = [""],
                               cfg_scales: List[float] = [3.5],
                               num_inference_steps: int = 50,
-                              eps: Optional[float] = None,
-                              cutoff_points: Optional[List[float]] = None,
                               numerical_fix: bool = False,
-                              extract_h_space: bool = False,
-                              extract_skipconns: bool = False,
-                              x_prev_mode: bool = False):
-    if len(prompts) > 1 and extract_h_space:
-        raise NotImplementedError("How do you split cfg_scales for hspace? TODO")
-
+                              duration: Optional[float] = None,
+                              first_order: bool = False,
+                              save_compute: bool = True,
+                              progress=gr.Progress()) -> Tuple:
     if len(prompts) > 1 or prompts[0] != "":
         text_embeddings_hidden_states, text_embeddings_class_labels, \
             text_embeddings_boolean_prompt_mask = model.encode_text(prompts)
-        # text_embeddings = encode_text(model, prompt)
-
-        # # classifier free guidance
-        batch_size = len(prompts)
-        cfg_scales_tensor = torch.ones((batch_size, *x0.shape[1:]), device=model.device, dtype=x0.dtype)
-
-        # if len(prompts) > 1:
-        #     if cutoff_points is None:
-        #         cutoff_points = [i * 1 / batch_size for i in range(1, batch_size)]
-        #     if len(cfg_scales) == 1:
-        #         cfg_scales *= batch_size
-        #     elif len(cfg_scales) < batch_size:
-        #         raise ValueError("Not enough target CFG scales")
-
-        #     cutoff_points = [int(x * cfg_scales_tensor.shape[2]) for x in cutoff_points]
-        #     cutoff_points = [0, *cutoff_points, cfg_scales_tensor.shape[2]]
 
-        #     for i, (start, end) in enumerate(zip(cutoff_points[:-1], cutoff_points[1:])):
-        #         cfg_scales_tensor[i, :, end:] = 0
-        #         cfg_scales_tensor[i, :, :start] = 0
-        #         cfg_scales_tensor[i] *= cfg_scales[i]
-        #         if prompts[i] == "":
-        #             cfg_scales_tensor[i] = 0
-        #     cfg_scales_tensor = T.functional.gaussian_blur(cfg_scales_tensor, kernel_size=15, sigma=1)
-        # else:
-        cfg_scales_tensor *= cfg_scales[0]
+        # In the forward negative prompts are not supported currently (TODO)
+        uncond_embeddings_hidden_states, uncond_embeddings_class_lables, uncond_boolean_prompt_mask = model.encode_text(
+            [""], negative=True, save_compute=save_compute, cond_length=text_embeddings_class_labels.shape[1]
+            if text_embeddings_class_labels is not None else None)
+    else:
+        uncond_embeddings_hidden_states, uncond_embeddings_class_lables, uncond_boolean_prompt_mask = model.encode_text(
+            [""], negative=True, save_compute=False)
 
-    uncond_embedding_hidden_states, uncond_embedding_class_lables, uncond_boolean_prompt_mask = model.encode_text([""])
-    # uncond_embedding = encode_text(model, "")
     timesteps = model.model.scheduler.timesteps.to(model.device)
-    variance_noise_shape = (
-        num_inference_steps,
-        model.model.unet.config.in_channels,
-        # model.unet.sample_size,
-        # model.unet.sample_size)
-        x0.shape[-2],
-        x0.shape[-1])
+    variance_noise_shape = model.get_noise_shape(x0, num_inference_steps)
 
-    if etas is None or (type(etas) in [int, float] and etas == 0):
-        eta_is_zero = True
-        zs = None
-    else:
-        eta_is_zero = False
-        if type(etas) in [int, float]:
-            etas = [etas]*model.model.scheduler.num_inference_steps
-        xts = sample_xts_from_x0(model, x0, num_inference_steps=num_inference_steps, x_prev_mode=x_prev_mode)
-        alpha_bar = model.model.scheduler.alphas_cumprod
-        zs = torch.zeros(size=variance_noise_shape, device=model.device)
-    hspaces = []
-    skipconns = []
-    t_to_idx = {int(v): k for k, v in enumerate(timesteps)}
+    if type(etas) in [int, float]:
+        etas = [etas]*model.model.scheduler.num_inference_steps
+    xts = model.sample_xts_from_x0(x0, num_inference_steps=num_inference_steps)
+    zs = torch.zeros(size=variance_noise_shape, device=model.device)
+    extra_info = [None] * len(zs)
+
+    if timesteps[0].dtype == torch.int64:
+        t_to_idx = {int(v): k for k, v in enumerate(timesteps)}
+    elif timesteps[0].dtype == torch.float32:
+        t_to_idx = {float(v): k for k, v in enumerate(timesteps)}
     xt = x0
-    # op = tqdm(reversed(timesteps)) if prog_bar else reversed(timesteps)
-    op = tqdm(timesteps) if prog_bar else timesteps
+    op = tqdm(timesteps, desc="Inverting")
+    model.setup_extra_inputs(xt, init_timestep=timesteps[0], audio_end_in_s=duration,
+                             save_compute=save_compute and prompts[0] != "")
+    app_op = progress.tqdm(timesteps, desc="Inverting")
+    for t, _ in zip(op, app_op):
+        idx = num_inference_steps - t_to_idx[int(t) if timesteps[0].dtype == torch.int64 else float(t)] - 1
 
-    for t in op:
-        # idx = t_to_idx[int(t)]
-        idx = num_inference_steps - t_to_idx[int(t)] - 1
         # 1. predict noise residual
-        if not eta_is_zero:
-            xt = xts[idx+1][None]
+        xt = xts[idx+1][None]
+        xt_inp = model.model.scheduler.scale_model_input(xt, t)
 
         with torch.no_grad():
-            out, out_hspace, out_skipconns = model.unet_forward(xt, timestep=t,
-                                                                encoder_hidden_states=uncond_embedding_hidden_states,
-                                                                class_labels=uncond_embedding_class_lables,
-                                                                encoder_attention_mask=uncond_boolean_prompt_mask)
-            # out = model.unet.forward(xt, timestep= t, encoder_hidden_states=uncond_embedding)
-            if len(prompts) > 1 or prompts[0] != "":
-                cond_out, cond_out_hspace, cond_out_skipconns = model.unet_forward(
-                    xt.expand(len(prompts), -1, -1, -1), timestep=t,
-                    encoder_hidden_states=text_embeddings_hidden_states,
-                    class_labels=text_embeddings_class_labels,
-                    encoder_attention_mask=text_embeddings_boolean_prompt_mask)
-                # cond_out = model.unet.forward(xt, timestep=t, encoder_hidden_states = text_embeddings)
+            if save_compute and prompts[0] != "":
+                comb_out, _, _ = model.unet_forward(
+                    xt_inp.expand(2, -1, -1, -1) if hasattr(model.model, 'unet') else xt_inp.expand(2, -1, -1),
+                    timestep=t,
+                    encoder_hidden_states=torch.cat([uncond_embeddings_hidden_states, text_embeddings_hidden_states
+                                                     ], dim=0)
+                    if uncond_embeddings_hidden_states is not None else None,
+                    class_labels=torch.cat([uncond_embeddings_class_lables, text_embeddings_class_labels], dim=0)
+                    if uncond_embeddings_class_lables is not None else None,
+                    encoder_attention_mask=torch.cat([uncond_boolean_prompt_mask, text_embeddings_boolean_prompt_mask
+                                                      ], dim=0)
+                    if uncond_boolean_prompt_mask is not None else None,
+                )
+                out, cond_out = comb_out.sample.chunk(2, dim=0)
+            else:
+                out = model.unet_forward(xt_inp, timestep=t,
+                                         encoder_hidden_states=uncond_embeddings_hidden_states,
+                                         class_labels=uncond_embeddings_class_lables,
+                                         encoder_attention_mask=uncond_boolean_prompt_mask)[0].sample
+                if len(prompts) > 1 or prompts[0] != "":
+                    cond_out = model.unet_forward(
+                        xt_inp,
+                        timestep=t,
+                        encoder_hidden_states=text_embeddings_hidden_states,
+                        class_labels=text_embeddings_class_labels,
+                        encoder_attention_mask=text_embeddings_boolean_prompt_mask)[0].sample
 
         if len(prompts) > 1 or prompts[0] != "":
             # # classifier free guidance
-            noise_pred = out.sample + \
-                (cfg_scales_tensor * (cond_out.sample - out.sample.expand(batch_size, -1, -1, -1))
-                 ).sum(axis=0).unsqueeze(0)
-            if extract_h_space or extract_skipconns:
-                noise_h_space = out_hspace + cfg_scales[0] * (cond_out_hspace - out_hspace)
-            if extract_skipconns:
-                noise_skipconns = {k: [out_skipconns[k][j] + cfg_scales[0] *
-                                       (cond_out_skipconns[k][j] - out_skipconns[k][j])
-                                       for j in range(len(out_skipconns[k]))]
-                                   for k in out_skipconns}
-        else:
-            noise_pred = out.sample
-            if extract_h_space or extract_skipconns:
-                noise_h_space = out_hspace
-            if extract_skipconns:
-                noise_skipconns = out_skipconns
-        if extract_h_space or extract_skipconns:
-            hspaces.append(noise_h_space)
-        if extract_skipconns:
-            skipconns.append(noise_skipconns)
-
-        if eta_is_zero:
-            # 2. compute more noisy image and set x_t -> x_t+1
-            xt = forward_step(model.model, noise_pred, t, xt)
+            noise_pred = out + (cfg_scales[0] * (cond_out - out)).sum(axis=0).unsqueeze(0)
         else:
-            # xtm1 =  xts[idx+1][None]
-            xtm1 = xts[idx][None]
-            # pred of x0
-            if model.model.scheduler.config.prediction_type == 'epsilon':
-                pred_original_sample = (xt - (1 - alpha_bar[t]) ** 0.5 * noise_pred) / alpha_bar[t] ** 0.5
-            elif model.model.scheduler.config.prediction_type == 'v_prediction':
-                pred_original_sample = (alpha_bar[t] ** 0.5) * xt - ((1 - alpha_bar[t]) ** 0.5) * noise_pred
-
-            # direction to xt
-            prev_timestep = t - model.model.scheduler.config.num_train_timesteps // \
-                model.model.scheduler.num_inference_steps
-
-            alpha_prod_t_prev = model.get_alpha_prod_t_prev(prev_timestep)
-            variance = model.get_variance(t, prev_timestep)
-
-            if model.model.scheduler.config.prediction_type == 'epsilon':
-                radom_noise_pred = noise_pred
-            elif model.model.scheduler.config.prediction_type == 'v_prediction':
-                radom_noise_pred = (alpha_bar[t] ** 0.5) * noise_pred + ((1 - alpha_bar[t]) ** 0.5) * xt
-
-            pred_sample_direction = (1 - alpha_prod_t_prev - etas[idx] * variance) ** (0.5) * radom_noise_pred
-
-            mu_xt = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
-
-            z = (xtm1 - mu_xt) / (etas[idx] * variance ** 0.5)
-
-            zs[idx] = z
-
-            # correction to avoid error accumulation
-            if numerical_fix:
-                xtm1 = mu_xt + (etas[idx] * variance ** 0.5)*z
-            xts[idx] = xtm1
+            noise_pred = out
+
+        # xtm1 =  xts[idx+1][None]
+        xtm1 = xts[idx][None]
+        z, xtm1, extra = model.get_zs_from_xts(xt, xtm1, noise_pred, t,
+                                               eta=etas[idx], numerical_fix=numerical_fix,
+                                               first_order=first_order)
+        zs[idx] = z
+        # print(f"Fix Xt-1 distance -  NORM:{torch.norm(xts[idx] - xtm1):.4g}, MSE:{((xts[idx] - xtm1)**2).mean():.4g}")
+        xts[idx] = xtm1
+        extra_info[idx] = extra
 
     if zs is not None:
         # zs[-1] = torch.zeros_like(zs[-1])
         zs[0] = torch.zeros_like(zs[0])
         # zs_cycle[0] = torch.zeros_like(zs[0])
 
-    if extract_h_space:
-        hspaces = torch.concat(hspaces, axis=0)
-        return xt, zs, xts, hspaces
-
-    if extract_skipconns:
-        hspaces = torch.concat(hspaces, axis=0)
-        return xt, zs, xts, hspaces, skipconns
-
-    return xt, zs, xts
-
-
-def reverse_step(model, model_output, timestep, sample, eta=0, variance_noise=None):
-    # 1. get previous step value (=t-1)
-    prev_timestep = timestep - model.model.scheduler.config.num_train_timesteps // \
-        model.model.scheduler.num_inference_steps
-    # 2. compute alphas, betas
-    alpha_prod_t = model.model.scheduler.alphas_cumprod[timestep]
-    alpha_prod_t_prev = model.get_alpha_prod_t_prev(prev_timestep)
-    beta_prod_t = 1 - alpha_prod_t
-    # 3. compute predicted original sample from predicted noise also called
-    # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
-    if model.model.scheduler.config.prediction_type == 'epsilon':
-        pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
-    elif model.model.scheduler.config.prediction_type == 'v_prediction':
-        pred_original_sample = (alpha_prod_t ** 0.5) * sample - (beta_prod_t ** 0.5) * model_output
-
-    # 5. compute variance: "sigma_t(η)" -> see formula (16)
-    # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
-    # variance = self.scheduler._get_variance(timestep, prev_timestep)
-    variance = model.get_variance(timestep, prev_timestep)
-    # std_dev_t = eta * variance ** (0.5)
-    # Take care of asymetric reverse process (asyrp)
-    if model.model.scheduler.config.prediction_type == 'epsilon':
-        model_output_direction = model_output
-    elif model.model.scheduler.config.prediction_type == 'v_prediction':
-        model_output_direction = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
-    # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
-    # pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * model_output_direction
-    pred_sample_direction = (1 - alpha_prod_t_prev - eta * variance) ** (0.5) * model_output_direction
-    # 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
-    prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
-    # 8. Add noice if eta > 0
-    if eta > 0:
-        if variance_noise is None:
-            variance_noise = torch.randn(model_output.shape, device=model.device)
-        sigma_z = eta * variance ** (0.5) * variance_noise
-        prev_sample = prev_sample + sigma_z
-
-    return prev_sample
+    del app_op.iterables[0]
+    return xt, zs, xts, extra_info
 
 
 def inversion_reverse_process(model: PipelineWrapper,
                               xT: torch.Tensor,
-                              skips: torch.Tensor,
-                              fix_alpha: float = 0.1,
+                              tstart: torch.Tensor,
                               etas: float = 0,
                               prompts: List[str] = [""],
                               neg_prompts: List[str] = [""],
                               cfg_scales: Optional[List[float]] = None,
-                              prog_bar: bool = False,
                               zs: Optional[List[torch.Tensor]] = None,
-                            #   controller=None,
-                              cutoff_points: Optional[List[float]] = None,
-                              hspace_add: Optional[torch.Tensor] = None,
-                              hspace_replace: Optional[torch.Tensor] = None,
-                              skipconns_replace: Optional[Dict[int, torch.Tensor]] = None,
-                              zero_out_resconns: Optional[Union[int, List]] = None,
-                              asyrp: bool = False,
-                              extract_h_space: bool = False,
-                              extract_skipconns: bool = False):
-
-    batch_size = len(prompts)
+                              duration: Optional[float] = None,
+                              first_order: bool = False,
+                              extra_info: Optional[List] = None,
+                              save_compute: bool = True,
+                              progress=gr.Progress()) -> Tuple[torch.Tensor, torch.Tensor]:
 
     text_embeddings_hidden_states, text_embeddings_class_labels, \
         text_embeddings_boolean_prompt_mask = model.encode_text(prompts)
-    uncond_embedding_hidden_states, uncond_embedding_class_lables, \
-        uncond_boolean_prompt_mask = model.encode_text(neg_prompts)
-    # text_embeddings = encode_text(model, prompts)
-    # uncond_embedding = encode_text(model, [""] * batch_size)
-
-    masks = torch.ones((batch_size, *xT.shape[1:]), device=model.device, dtype=xT.dtype)
-    cfg_scales_tensor = torch.ones((batch_size, *xT.shape[1:]), device=model.device, dtype=xT.dtype)
-
-    # if batch_size > 1:
-    #     if cutoff_points is None:
-    #         cutoff_points = [i * 1 / batch_size for i in range(1, batch_size)]
-    #     if len(cfg_scales) == 1:
-    #         cfg_scales *= batch_size
-    #     elif len(cfg_scales) < batch_size:
-    #         raise ValueError("Not enough target CFG scales")
-
-    #     cutoff_points = [int(x * cfg_scales_tensor.shape[2]) for x in cutoff_points]
-    #     cutoff_points = [0, *cutoff_points, cfg_scales_tensor.shape[2]]
+    uncond_embeddings_hidden_states, uncond_embeddings_class_lables, \
+        uncond_boolean_prompt_mask = model.encode_text(neg_prompts,
+                                                       negative=True,
+                                                       save_compute=save_compute,
+                                                       cond_length=text_embeddings_class_labels.shape[1]
+                                                       if text_embeddings_class_labels is not None else None)
 
-    #     for i, (start, end) in enumerate(zip(cutoff_points[:-1], cutoff_points[1:])):
-    #         cfg_scales_tensor[i, :, end:] = 0
-    #         cfg_scales_tensor[i, :, :start] = 0
-    #         masks[i, :, end:] = 0
-    #         masks[i, :, :start] = 0
-    #         cfg_scales_tensor[i] *= cfg_scales[i]
-    #     cfg_scales_tensor = T.functional.gaussian_blur(cfg_scales_tensor, kernel_size=15, sigma=1)
-    #     masks = T.functional.gaussian_blur(masks, kernel_size=15, sigma=1)
-    # else:
-    cfg_scales_tensor *= cfg_scales[0]
+    xt = xT[tstart.max()].unsqueeze(0)
 
     if etas is None:
         etas = 0
@@ -344,107 +138,71 @@ def inversion_reverse_process(model: PipelineWrapper,
     assert len(etas) == model.model.scheduler.num_inference_steps
     timesteps = model.model.scheduler.timesteps.to(model.device)
 
-    # xt = xT.expand(1, -1, -1, -1)
-    xt = xT[skips.max()].unsqueeze(0)
-    op = tqdm(timesteps[-zs.shape[0]:]) if prog_bar else timesteps[-zs.shape[0]:]
+    op = tqdm(timesteps[-zs.shape[0]:], desc="Editing")
+    if timesteps[0].dtype == torch.int64:
+        t_to_idx = {int(v): k for k, v in enumerate(timesteps[-zs.shape[0]:])}
+    elif timesteps[0].dtype == torch.float32:
+        t_to_idx = {float(v): k for k, v in enumerate(timesteps[-zs.shape[0]:])}
+    model.setup_extra_inputs(xt, extra_info=extra_info, init_timestep=timesteps[-zs.shape[0]],
+                             audio_end_in_s=duration, save_compute=save_compute)
+    app_op = progress.tqdm(timesteps[-zs.shape[0]:], desc="Editing")
+    for it, (t, _) in enumerate(zip(op, app_op)):
+        idx = model.model.scheduler.num_inference_steps - t_to_idx[
+            int(t) if timesteps[0].dtype == torch.int64 else float(t)] - \
+                (model.model.scheduler.num_inference_steps - zs.shape[0] + 1)
+
+        xt_inp = model.model.scheduler.scale_model_input(xt, t)
 
-    t_to_idx = {int(v): k for k, v in enumerate(timesteps[-zs.shape[0]:])}
-    hspaces = []
-    skipconns = []
-
-    for it, t in enumerate(op):
-        # idx = t_to_idx[int(t)]
-        idx = model.model.scheduler.num_inference_steps - t_to_idx[int(t)] - \
-            (model.model.scheduler.num_inference_steps - zs.shape[0] + 1)
         # # Unconditional embedding
         with torch.no_grad():
-            uncond_out, out_hspace, out_skipconns = model.unet_forward(
-                xt, timestep=t,
-                encoder_hidden_states=uncond_embedding_hidden_states,
-                class_labels=uncond_embedding_class_lables,
-                encoder_attention_mask=uncond_boolean_prompt_mask,
-                mid_block_additional_residual=(None if hspace_add is None else
-                                               (1 / (cfg_scales[0] + 1)) *
-                                               (hspace_add[-zs.shape[0]:][it] if hspace_add.shape[0] > 1
-                                                else hspace_add)),
-                replace_h_space=(None if hspace_replace is None else
-                                 (hspace_replace[-zs.shape[0]:][it].unsqueeze(0) if hspace_replace.shape[0] > 1
-                                  else hspace_replace)),
-                zero_out_resconns=zero_out_resconns,
-                replace_skip_conns=(None if skipconns_replace is None else
-                                    (skipconns_replace[-zs.shape[0]:][it] if len(skipconns_replace) > 1
-                                     else skipconns_replace))
-                )  # encoder_hidden_states = uncond_embedding)
-
-        # # Conditional embedding
-        if prompts:
-            with torch.no_grad():
-                cond_out, cond_out_hspace, cond_out_skipconns = model.unet_forward(
-                    xt.expand(batch_size, -1, -1, -1),
+            # print(f'xt_inp.shape: {xt_inp.shape}')
+            # print(f't.shape: {t.shape}')
+            # print(f'uncond_embeddings_hidden_states.shape: {uncond_embeddings_hidden_states.shape}')
+            # print(f'uncond_embeddings_class_lables.shape: {uncond_embeddings_class_lables.shape}')
+            # print(f'uncond_boolean_prompt_mask.shape: {uncond_boolean_prompt_mask.shape}')
+            # print(f'text_embeddings_hidden_states.shape: {text_embeddings_hidden_states.shape}')
+            # print(f'text_embeddings_class_labels.shape: {text_embeddings_class_labels.shape}')
+            # print(f'text_embeddings_boolean_prompt_mask.shape: {text_embeddings_boolean_prompt_mask.shape}')
+
+            if save_compute:
+                comb_out, _, _ = model.unet_forward(
+                    xt_inp.expand(2, -1, -1, -1) if hasattr(model.model, 'unet') else xt_inp.expand(2, -1, -1),
+                    timestep=t,
+                    encoder_hidden_states=torch.cat([uncond_embeddings_hidden_states, text_embeddings_hidden_states
+                                                     ], dim=0)
+                    if uncond_embeddings_hidden_states is not None else None,
+                    class_labels=torch.cat([uncond_embeddings_class_lables, text_embeddings_class_labels], dim=0)
+                    if uncond_embeddings_class_lables is not None else None,
+                    encoder_attention_mask=torch.cat([uncond_boolean_prompt_mask, text_embeddings_boolean_prompt_mask
+                                                      ], dim=0)
+                    if uncond_boolean_prompt_mask is not None else None,
+                )
+                uncond_out, cond_out = comb_out.sample.chunk(2, dim=0)
+            else:
+                uncond_out = model.unet_forward(
+                    xt_inp, timestep=t,
+                    encoder_hidden_states=uncond_embeddings_hidden_states,
+                    class_labels=uncond_embeddings_class_lables,
+                    encoder_attention_mask=uncond_boolean_prompt_mask,
+                    )[0].sample
+
+                # Conditional embedding
+                cond_out = model.unet_forward(
+                    xt_inp,
                     timestep=t,
                     encoder_hidden_states=text_embeddings_hidden_states,
                     class_labels=text_embeddings_class_labels,
                     encoder_attention_mask=text_embeddings_boolean_prompt_mask,
-                    mid_block_additional_residual=(None if hspace_add is None else
-                                                   (cfg_scales[0] / (cfg_scales[0] + 1)) *
-                                                   (hspace_add[-zs.shape[0]:][it] if hspace_add.shape[0] > 1
-                                                    else hspace_add)),
-                    replace_h_space=(None if hspace_replace is None else
-                                     (hspace_replace[-zs.shape[0]:][it].unsqueeze(0) if hspace_replace.shape[0] > 1
-                                      else hspace_replace)),
-                    zero_out_resconns=zero_out_resconns,
-                    replace_skip_conns=(None if skipconns_replace is None else
-                                        (skipconns_replace[-zs.shape[0]:][it] if len(skipconns_replace) > 1
-                                         else skipconns_replace))
-                    )  # encoder_hidden_states = text_embeddings)
+                    )[0].sample
 
         z = zs[idx] if zs is not None else None
-        # print(f'idx: {idx}')
-        # print(f't: {t}')
         z = z.unsqueeze(0)
-        # z = z.expand(batch_size, -1, -1, -1)
-        if prompts:
-            # # classifier free guidance
-            # noise_pred = uncond_out.sample + cfg_scales_tensor * (cond_out.sample - uncond_out.sample)
-            noise_pred = uncond_out.sample + \
-                (cfg_scales_tensor * (cond_out.sample - uncond_out.sample.expand(batch_size, -1, -1, -1))
-                 ).sum(axis=0).unsqueeze(0)
-            if extract_h_space or extract_skipconns:
-                noise_h_space = out_hspace + cfg_scales[0] * (cond_out_hspace - out_hspace)
-            if extract_skipconns:
-                noise_skipconns = {k: [out_skipconns[k][j] + cfg_scales[0] *
-                                       (cond_out_skipconns[k][j] - out_skipconns[k][j])
-                                       for j in range(len(out_skipconns[k]))]
-                                   for k in out_skipconns}
-        else:
-            noise_pred = uncond_out.sample
-            if extract_h_space or extract_skipconns:
-                noise_h_space = out_hspace
-            if extract_skipconns:
-                noise_skipconns = out_skipconns
-
-        if extract_h_space or extract_skipconns:
-            hspaces.append(noise_h_space)
-        if extract_skipconns:
-            skipconns.append(noise_skipconns)
+        # classifier free guidance
+        noise_pred = uncond_out + (cfg_scales[0] * (cond_out - uncond_out)).sum(axis=0).unsqueeze(0)
 
         # 2. compute less noisy image and set x_t -> x_t-1
-        xt = reverse_step(model, noise_pred, t, xt, eta=etas[idx], variance_noise=z)
-        # if controller is not None:
-            # xt = controller.step_callback(xt)
-
-        # "fix" xt
-        apply_fix = ((skips.max() - skips) > it)
-        if apply_fix.any():
-            apply_fix = (apply_fix * fix_alpha).unsqueeze(1).unsqueeze(2).unsqueeze(3).to(xT.device)
-            xt = (masks * (xt.expand(batch_size, -1, -1, -1) * (1 - apply_fix) +
-                           apply_fix * xT[skips.max() - it - 1].expand(batch_size, -1, -1, -1))
-                  ).sum(axis=0).unsqueeze(0)
-
-    if extract_h_space:
-        return xt, zs, torch.concat(hspaces, axis=0)
-
-    if extract_skipconns:
-        return xt, zs, torch.concat(hspaces, axis=0), skipconns
+        xt = model.reverse_step_with_custom_noise(noise_pred, t, xt, variance_noise=z,
+                                                  eta=etas[idx], first_order=first_order)
 
+    del app_op.iterables[0]
     return xt, zs

models.py

@@ -1,46 +1,160 @@
 import torch
-from diffusers import DDIMScheduler
-from diffusers import AudioLDM2Pipeline
-from transformers import RobertaTokenizer, RobertaTokenizerFast
+from diffusers import DDIMScheduler, CosineDPMSolverMultistepScheduler
+from diffusers.schedulers.scheduling_dpmsolver_sde import BrownianTreeNoiseSampler
+from diffusers import AudioLDM2Pipeline, StableAudioPipeline
+from transformers import RobertaTokenizer, RobertaTokenizerFast, VitsTokenizer
 from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput
+from diffusers.models.embeddings import get_1d_rotary_pos_embed
 from typing import Any, Dict, List, Optional, Tuple, Union
+import gradio as gr
 
 
 class PipelineWrapper(torch.nn.Module):
-    def __init__(self, model_id, device, double_precision=False, *args, **kwargs) -> None:
+    def __init__(self, model_id: str,
+                 device: torch.device,
+                 double_precision: bool = False,
+                 token: Optional[str] = None, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self.model_id = model_id
         self.device = device
         self.double_precision = double_precision
+        self.token = token
 
-    def get_sigma(self, timestep) -> float:
+    def get_sigma(self, timestep: int) -> float:
         sqrt_recipm1_alphas_cumprod = torch.sqrt(1.0 / self.model.scheduler.alphas_cumprod - 1)
         return sqrt_recipm1_alphas_cumprod[timestep]
 
-    def load_scheduler(self):
+    def load_scheduler(self) -> None:
         pass
 
-    def get_fn_STFT(self):
+    def get_fn_STFT(self) -> torch.nn.Module:
         pass
 
-    def vae_encode(self, x: torch.Tensor):
+    def get_sr(self) -> int:
+        return 16000
+
+    def vae_encode(self, x: torch.Tensor) -> torch.Tensor:
+        pass
+
+    def vae_decode(self, x: torch.Tensor) -> torch.Tensor:
         pass
 
-    def vae_decode(self, x: torch.Tensor):
+    def decode_to_mel(self, x: torch.Tensor) -> torch.Tensor:
         pass
 
-    def decode_to_mel(self, x: torch.Tensor):
+    def setup_extra_inputs(self, *args, **kwargs) -> None:
         pass
 
-    def encode_text(self, prompts: List[str]) -> Tuple:
+    def encode_text(self, prompts: List[str], **kwargs
+                    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
         pass
 
-    def get_variance(self, timestep, prev_timestep):
+    def get_variance(self, timestep: torch.Tensor, prev_timestep: torch.Tensor) -> torch.Tensor:
         pass
 
-    def get_alpha_prod_t_prev(self, prev_timestep):
+    def get_alpha_prod_t_prev(self, prev_timestep: torch.Tensor) -> torch.Tensor:
         pass
 
+    def get_noise_shape(self, x0: torch.Tensor, num_steps: int) -> Tuple[int, ...]:
+        variance_noise_shape = (num_steps,
+                                self.model.unet.config.in_channels,
+                                x0.shape[-2],
+                                x0.shape[-1])
+        return variance_noise_shape
+
+    def sample_xts_from_x0(self, x0: torch.Tensor, num_inference_steps: int = 50) -> torch.Tensor:
+        """
+        Samples from P(x_1:T|x_0)
+        """
+        alpha_bar = self.model.scheduler.alphas_cumprod
+        sqrt_one_minus_alpha_bar = (1-alpha_bar) ** 0.5
+
+        variance_noise_shape = self.get_noise_shape(x0, num_inference_steps + 1)
+        timesteps = self.model.scheduler.timesteps.to(self.device)
+        t_to_idx = {int(v): k for k, v in enumerate(timesteps)}
+        xts = torch.zeros(variance_noise_shape).to(x0.device)
+        xts[0] = x0
+        for t in reversed(timesteps):
+            idx = num_inference_steps - t_to_idx[int(t)]
+            xts[idx] = x0 * (alpha_bar[t] ** 0.5) + torch.randn_like(x0) * sqrt_one_minus_alpha_bar[t]
+
+        return xts
+
+    def get_zs_from_xts(self, xt: torch.Tensor, xtm1: torch.Tensor, noise_pred: torch.Tensor,
+                        t: torch.Tensor, eta: float = 0, numerical_fix: bool = True, **kwargs
+                        ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        # pred of x0
+        alpha_bar = self.model.scheduler.alphas_cumprod
+        if self.model.scheduler.config.prediction_type == 'epsilon':
+            pred_original_sample = (xt - (1 - alpha_bar[t]) ** 0.5 * noise_pred) / alpha_bar[t] ** 0.5
+        elif self.model.scheduler.config.prediction_type == 'v_prediction':
+            pred_original_sample = (alpha_bar[t] ** 0.5) * xt - ((1 - alpha_bar[t]) ** 0.5) * noise_pred
+
+        # direction to xt
+        prev_timestep = t - self.model.scheduler.config.num_train_timesteps // \
+            self.model.scheduler.num_inference_steps
+
+        alpha_prod_t_prev = self.get_alpha_prod_t_prev(prev_timestep)
+        variance = self.get_variance(t, prev_timestep)
+
+        if self.model.scheduler.config.prediction_type == 'epsilon':
+            radom_noise_pred = noise_pred
+        elif self.model.scheduler.config.prediction_type == 'v_prediction':
+            radom_noise_pred = (alpha_bar[t] ** 0.5) * noise_pred + ((1 - alpha_bar[t]) ** 0.5) * xt
+
+        pred_sample_direction = (1 - alpha_prod_t_prev - eta * variance) ** (0.5) * radom_noise_pred
+
+        mu_xt = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
+
+        z = (xtm1 - mu_xt) / (eta * variance ** 0.5)
+
+        # correction to avoid error accumulation
+        if numerical_fix:
+            xtm1 = mu_xt + (eta * variance ** 0.5)*z
+
+        return z, xtm1, None
+
+    def reverse_step_with_custom_noise(self, model_output: torch.Tensor, timestep: torch.Tensor, sample: torch.Tensor,
+                                       variance_noise: Optional[torch.Tensor] = None, eta: float = 0, **kwargs
+                                       ) -> torch.Tensor:
+        # 1. get previous step value (=t-1)
+        prev_timestep = timestep - self.model.scheduler.config.num_train_timesteps // \
+            self.model.scheduler.num_inference_steps
+        # 2. compute alphas, betas
+        alpha_prod_t = self.model.scheduler.alphas_cumprod[timestep]
+        alpha_prod_t_prev = self.get_alpha_prod_t_prev(prev_timestep)
+        beta_prod_t = 1 - alpha_prod_t
+        # 3. compute predicted original sample from predicted noise also called
+        # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
+        if self.model.scheduler.config.prediction_type == 'epsilon':
+            pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
+        elif self.model.scheduler.config.prediction_type == 'v_prediction':
+            pred_original_sample = (alpha_prod_t ** 0.5) * sample - (beta_prod_t ** 0.5) * model_output
+
+        # 5. compute variance: "sigma_t(η)" -> see formula (16)
+        # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
+        # variance = self.scheduler._get_variance(timestep, prev_timestep)
+        variance = self.get_variance(timestep, prev_timestep)
+        # std_dev_t = eta * variance ** (0.5)
+        # Take care of asymetric reverse process (asyrp)
+        if self.model.scheduler.config.prediction_type == 'epsilon':
+            model_output_direction = model_output
+        elif self.model.scheduler.config.prediction_type == 'v_prediction':
+            model_output_direction = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
+        # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
+        # pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * model_output_direction
+        pred_sample_direction = (1 - alpha_prod_t_prev - eta * variance) ** (0.5) * model_output_direction
+        # 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
+        prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
+        # 8. Add noice if eta > 0
+        if eta > 0:
+            if variance_noise is None:
+                variance_noise = torch.randn(model_output.shape, device=self.device)
+            sigma_z = eta * variance ** (0.5) * variance_noise
+            prev_sample = prev_sample + sigma_z
+
+        return prev_sample
+
     def unet_forward(self,
                      sample: torch.FloatTensor,
                      timestep: Union[torch.Tensor, float, int],
@@ -57,244 +171,27 @@ class PipelineWrapper(torch.nn.Module):
                      replace_skip_conns: Optional[Dict[int, torch.Tensor]] = None,
                      return_dict: bool = True,
                      zero_out_resconns: Optional[Union[int, List]] = None) -> Tuple:
-
-        # By default samples have to be AT least a multiple of the overall upsampling factor.
-        # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).
-        # However, the upsampling interpolation output size can be forced to fit any upsampling size
-        # on the fly if necessary.
-        default_overall_up_factor = 2**self.model.unet.num_upsamplers
-
-        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
-        forward_upsample_size = False
-        upsample_size = None
-
-        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
-            # logger.info("Forward upsample size to force interpolation output size.")
-            forward_upsample_size = True
-
-        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension
-        # expects mask of shape:
-        #   [batch, key_tokens]
-        # adds singleton query_tokens dimension:
-        #   [batch,                    1, key_tokens]
-        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
-        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
-        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
-        if attention_mask is not None:
-            # assume that mask is expressed as:
-            #   (1 = keep,      0 = discard)
-            # convert mask into a bias that can be added to attention scores:
-            #       (keep = +0,     discard = -10000.0)
-            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
-            attention_mask = attention_mask.unsqueeze(1)
-
-        # convert encoder_attention_mask to a bias the same way we do for attention_mask
-        if encoder_attention_mask is not None:
-            encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
-            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
-
-        # 0. center input if necessary
-        if self.model.unet.config.center_input_sample:
-            sample = 2 * sample - 1.0
-
-        # 1. time
-        timesteps = timestep
-        if not torch.is_tensor(timesteps):
-            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
-            # This would be a good case for the `match` statement (Python 3.10+)
-            is_mps = sample.device.type == "mps"
-            if isinstance(timestep, float):
-                dtype = torch.float32 if is_mps else torch.float64
-            else:
-                dtype = torch.int32 if is_mps else torch.int64
-            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
-        elif len(timesteps.shape) == 0:
-            timesteps = timesteps[None].to(sample.device)
-
-        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-        timesteps = timesteps.expand(sample.shape[0])
-
-        t_emb = self.model.unet.time_proj(timesteps)
-
-        # `Timesteps` does not contain any weights and will always return f32 tensors
-        # but time_embedding might actually be running in fp16. so we need to cast here.
-        # there might be better ways to encapsulate this.
-        t_emb = t_emb.to(dtype=sample.dtype)
-
-        emb = self.model.unet.time_embedding(t_emb, timestep_cond)
-
-        if self.model.unet.class_embedding is not None:
-            if class_labels is None:
-                raise ValueError("class_labels should be provided when num_class_embeds > 0")
-
-            if self.model.unet.config.class_embed_type == "timestep":
-                class_labels = self.model.unet.time_proj(class_labels)
-
-                # `Timesteps` does not contain any weights and will always return f32 tensors
-                # there might be better ways to encapsulate this.
-                class_labels = class_labels.to(dtype=sample.dtype)
-
-            class_emb = self.model.unet.class_embedding(class_labels).to(dtype=sample.dtype)
-
-            if self.model.unet.config.class_embeddings_concat:
-                emb = torch.cat([emb, class_emb], dim=-1)
-            else:
-                emb = emb + class_emb
-
-        if self.model.unet.config.addition_embed_type == "text":
-            aug_emb = self.model.unet.add_embedding(encoder_hidden_states)
-            emb = emb + aug_emb
-        elif self.model.unet.config.addition_embed_type == "text_image":
-            # Kadinsky 2.1 - style
-            if "image_embeds" not in added_cond_kwargs:
-                raise ValueError(
-                    f"{self.model.unet.__class__} has the config param `addition_embed_type` set to 'text_image' "
-                    f"which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
-                )
-
-            image_embs = added_cond_kwargs.get("image_embeds")
-            text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states)
-
-            aug_emb = self.model.unet.add_embedding(text_embs, image_embs)
-            emb = emb + aug_emb
-
-        if self.model.unet.time_embed_act is not None:
-            emb = self.model.unet.time_embed_act(emb)
-
-        if self.model.unet.encoder_hid_proj is not None and self.model.unet.config.encoder_hid_dim_type == "text_proj":
-            encoder_hidden_states = self.model.unet.encoder_hid_proj(encoder_hidden_states)
-        elif self.model.unet.encoder_hid_proj is not None and \
-                self.model.unet.config.encoder_hid_dim_type == "text_image_proj":
-            # Kadinsky 2.1 - style
-            if "image_embeds" not in added_cond_kwargs:
-                raise ValueError(
-                    f"{self.model.unet.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' "
-                    f"which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
-                )
-
-            image_embeds = added_cond_kwargs.get("image_embeds")
-            encoder_hidden_states = self.model.unet.encoder_hid_proj(encoder_hidden_states, image_embeds)
-
-        # 2. pre-process
-        sample = self.model.unet.conv_in(sample)
-
-        # 3. down
-        down_block_res_samples = (sample,)
-        for downsample_block in self.model.unet.down_blocks:
-            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
-                sample, res_samples = downsample_block(
-                    hidden_states=sample,
-                    temb=emb,
-                    encoder_hidden_states=encoder_hidden_states,
-                    attention_mask=attention_mask,
-                    cross_attention_kwargs=cross_attention_kwargs,
-                    encoder_attention_mask=encoder_attention_mask,
-                )
-            else:
-                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
-
-            down_block_res_samples += res_samples
-
-        if down_block_additional_residuals is not None:
-            new_down_block_res_samples = ()
-
-            for down_block_res_sample, down_block_additional_residual in zip(
-                down_block_res_samples, down_block_additional_residuals
-            ):
-                down_block_res_sample = down_block_res_sample + down_block_additional_residual
-                new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
-
-            down_block_res_samples = new_down_block_res_samples
-
-        # 4. mid
-        if self.model.unet.mid_block is not None:
-            sample = self.model.unet.mid_block(
-                sample,
-                emb,
-                encoder_hidden_states=encoder_hidden_states,
-                attention_mask=attention_mask,
-                cross_attention_kwargs=cross_attention_kwargs,
-                encoder_attention_mask=encoder_attention_mask,
-            )
-
-        # print(sample.shape)
-
-        if replace_h_space is None:
-            h_space = sample.clone()
-        else:
-            h_space = replace_h_space
-            sample = replace_h_space.clone()
-
-        if mid_block_additional_residual is not None:
-            sample = sample + mid_block_additional_residual
-
-        extracted_res_conns = {}
-        # 5. up
-        for i, upsample_block in enumerate(self.model.unet.up_blocks):
-            is_final_block = i == len(self.model.unet.up_blocks) - 1
-
-            res_samples = down_block_res_samples[-len(upsample_block.resnets):]
-            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
-            if replace_skip_conns is not None and replace_skip_conns.get(i):
-                res_samples = replace_skip_conns.get(i)
-
-            if zero_out_resconns is not None:
-                if (type(zero_out_resconns) is int and i >= (zero_out_resconns - 1)) or \
-                        type(zero_out_resconns) is list and i in zero_out_resconns:
-                    res_samples = [torch.zeros_like(x) for x in res_samples]
-                # down_block_res_samples = [torch.zeros_like(x) for x in down_block_res_samples]
-
-            extracted_res_conns[i] = res_samples
-
-            # if we have not reached the final block and need to forward the
-            # upsample size, we do it here
-            if not is_final_block and forward_upsample_size:
-                upsample_size = down_block_res_samples[-1].shape[2:]
-
-            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
-                sample = upsample_block(
-                    hidden_states=sample,
-                    temb=emb,
-                    res_hidden_states_tuple=res_samples,
-                    encoder_hidden_states=encoder_hidden_states,
-                    cross_attention_kwargs=cross_attention_kwargs,
-                    upsample_size=upsample_size,
-                    attention_mask=attention_mask,
-                    encoder_attention_mask=encoder_attention_mask,
-                )
-            else:
-                sample = upsample_block(
-                    hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size
-                )
-
-        # 6. post-process
-        if self.model.unet.conv_norm_out:
-            sample = self.model.unet.conv_norm_out(sample)
-            sample = self.model.unet.conv_act(sample)
-        sample = self.model.unet.conv_out(sample)
-
-        if not return_dict:
-            return (sample,)
-
-        return UNet2DConditionOutput(sample=sample), h_space, extracted_res_conns
+        pass
 
 
 class AudioLDM2Wrapper(PipelineWrapper):
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         if self.double_precision:
-            self.model = AudioLDM2Pipeline.from_pretrained(self.model_id, torch_dtype=torch.float64).to(self.device)
+            self.model = AudioLDM2Pipeline.from_pretrained(self.model_id, torch_dtype=torch.float64, token=self.token
+                                                           ).to(self.device)
         else:
             try:
-                self.model = AudioLDM2Pipeline.from_pretrained(self.model_id, local_files_only=True).to(self.device)
+                self.model = AudioLDM2Pipeline.from_pretrained(self.model_id, local_files_only=True, token=self.token
+                                                               ).to(self.device)
             except FileNotFoundError:
-                self.model = AudioLDM2Pipeline.from_pretrained(self.model_id, local_files_only=False).to(self.device)
+                self.model = AudioLDM2Pipeline.from_pretrained(self.model_id, local_files_only=False, token=self.token
+                                                               ).to(self.device)
 
-    def load_scheduler(self):
-        # self.model.scheduler = DDIMScheduler.from_config(self.model_id, subfolder="scheduler")
+    def load_scheduler(self) -> None:
         self.model.scheduler = DDIMScheduler.from_pretrained(self.model_id, subfolder="scheduler")
 
-    def get_fn_STFT(self):
+    def get_fn_STFT(self) -> torch.nn.Module:
         from audioldm.audio import TacotronSTFT
         return TacotronSTFT(
             filter_length=1024,
@@ -306,17 +203,17 @@ class AudioLDM2Wrapper(PipelineWrapper):
             mel_fmax=8000,
         )
 
-    def vae_encode(self, x):
+    def vae_encode(self, x: torch.Tensor) -> torch.Tensor:
         # self.model.vae.disable_tiling()
         if x.shape[2] % 4:
             x = torch.nn.functional.pad(x, (0, 0, 4 - (x.shape[2] % 4), 0))
         return (self.model.vae.encode(x).latent_dist.mode() * self.model.vae.config.scaling_factor).float()
         # return (self.encode_no_tiling(x).latent_dist.mode() * self.model.vae.config.scaling_factor).float()
 
-    def vae_decode(self, x):
+    def vae_decode(self, x: torch.Tensor) -> torch.Tensor:
         return self.model.vae.decode(1 / self.model.vae.config.scaling_factor * x).sample
 
-    def decode_to_mel(self, x):
+    def decode_to_mel(self, x: torch.Tensor) -> torch.Tensor:
         if self.double_precision:
             tmp = self.model.mel_spectrogram_to_waveform(x[:, 0].detach().double()).detach()
         tmp = self.model.mel_spectrogram_to_waveform(x[:, 0].detach().float()).detach()
@@ -324,7 +221,9 @@ class AudioLDM2Wrapper(PipelineWrapper):
             tmp = tmp.unsqueeze(0)
         return tmp
 
-    def encode_text(self, prompts: List[str]):
+    def encode_text(self, prompts: List[str], negative: bool = False, 
+                    save_compute: bool = False, cond_length: int = 0, **kwargs
+                    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         tokenizers = [self.model.tokenizer, self.model.tokenizer_2]
         text_encoders = [self.model.text_encoder, self.model.text_encoder_2]
         prompt_embeds_list = []
@@ -333,8 +232,11 @@ class AudioLDM2Wrapper(PipelineWrapper):
         for tokenizer, text_encoder in zip(tokenizers, text_encoders):
             text_inputs = tokenizer(
                 prompts,
-                padding="max_length" if isinstance(tokenizer, (RobertaTokenizer, RobertaTokenizerFast)) else True,
-                max_length=tokenizer.model_max_length,
+                padding="max_length" if (save_compute and negative) or isinstance(tokenizer, (RobertaTokenizer, RobertaTokenizerFast))
+                else True,
+                max_length=tokenizer.model_max_length
+                if (not save_compute) or ((not negative) or isinstance(tokenizer, (RobertaTokenizer, RobertaTokenizerFast, VitsTokenizer)))
+                else cond_length,
                 truncation=True,
                 return_tensors="pt",
             )
@@ -404,7 +306,7 @@ class AudioLDM2Wrapper(PipelineWrapper):
 
         return generated_prompt_embeds, prompt_embeds, attention_mask
 
-    def get_variance(self, timestep, prev_timestep):
+    def get_variance(self, timestep: torch.Tensor, prev_timestep: torch.Tensor) -> torch.Tensor:
         alpha_prod_t = self.model.scheduler.alphas_cumprod[timestep]
         alpha_prod_t_prev = self.get_alpha_prod_t_prev(prev_timestep)
         beta_prod_t = 1 - alpha_prod_t
@@ -412,7 +314,7 @@ class AudioLDM2Wrapper(PipelineWrapper):
         variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
         return variance
 
-    def get_alpha_prod_t_prev(self, prev_timestep):
+    def get_alpha_prod_t_prev(self, prev_timestep: torch.Tensor) -> torch.Tensor:
         return self.model.scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 \
             else self.model.scheduler.final_alpha_cumprod
 
@@ -485,8 +387,6 @@ class AudioLDM2Wrapper(PipelineWrapper):
         # 1. time
         timesteps = timestep
         if not torch.is_tensor(timesteps):
-            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
-            # This would be a good case for the `match` statement (Python 3.10+)
             is_mps = sample.device.type == "mps"
             if isinstance(timestep, float):
                 dtype = torch.float32 if is_mps else torch.float64
@@ -628,12 +528,328 @@ class AudioLDM2Wrapper(PipelineWrapper):
 
         return UNet2DConditionOutput(sample=sample), h_space, extracted_res_conns
 
-    def forward(self, *args, **kwargs):
-        return self
+
+class StableAudWrapper(PipelineWrapper):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        try:
+            self.model = StableAudioPipeline.from_pretrained(self.model_id, token=self.token, local_files_only=True
+                                                             ).to(self.device)
+        except FileNotFoundError:
+            self.model = StableAudioPipeline.from_pretrained(self.model_id, token=self.token, local_files_only=False
+                                                             ).to(self.device)
+        self.model.transformer.eval()
+        self.model.vae.eval()
+
+        if self.double_precision:
+            self.model = self.model.to(torch.float64)
+
+    def load_scheduler(self) -> None:
+        self.model.scheduler = CosineDPMSolverMultistepScheduler.from_pretrained(
+            self.model_id, subfolder="scheduler", token=self.token)
+
+    def encode_text(self, prompts: List[str], negative: bool = False, **kwargs) -> Tuple[torch.Tensor, None, torch.Tensor]:
+        text_inputs = self.model.tokenizer(
+            prompts,
+            padding="max_length",
+            max_length=self.model.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids.to(self.device)
+        attention_mask = text_inputs.attention_mask.to(self.device)
+
+        self.model.text_encoder.eval()
+        with torch.no_grad():
+            prompt_embeds = self.model.text_encoder(text_input_ids, attention_mask=attention_mask)[0]
+
+        if negative and attention_mask is not None:  # set the masked tokens to the null embed
+            prompt_embeds = torch.where(attention_mask.to(torch.bool).unsqueeze(2), prompt_embeds, 0.0)
+
+        prompt_embeds = self.model.projection_model(text_hidden_states=prompt_embeds).text_hidden_states
+
+        if attention_mask is None:
+            raise gr.Error("Shouldn't reach here. Please raise an issue if you do.")
+            """prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+            if attention_mask is not None and negative_attention_mask is None:
+                negative_attention_mask = torch.ones_like(attention_mask)
+            elif attention_mask is None and negative_attention_mask is not None:
+                attention_mask = torch.ones_like(negative_attention_mask)"""
+
+        if prompts == [""]:  # empty
+            return torch.zeros_like(prompt_embeds, device=prompt_embeds.device), None, None
+
+        prompt_embeds = prompt_embeds * attention_mask.unsqueeze(-1).to(prompt_embeds.dtype)
+        prompt_embeds = prompt_embeds * attention_mask.unsqueeze(-1).to(prompt_embeds.dtype)
+        return prompt_embeds, None, attention_mask
+
+    def get_fn_STFT(self) -> torch.nn.Module:
+        from audioldm.audio import TacotronSTFT
+        return TacotronSTFT(
+            filter_length=1024,
+            hop_length=160,
+            win_length=1024,
+            n_mel_channels=64,
+            sampling_rate=44100,
+            mel_fmin=0,
+            mel_fmax=22050,
+        )
+
+    def vae_encode(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.unsqueeze(0)
+
+        audio_vae_length = int(self.model.transformer.config.sample_size * self.model.vae.hop_length)
+        audio_shape = (1, self.model.vae.config.audio_channels, audio_vae_length)
+
+        # check num_channels
+        if x.shape[1] == 1 and self.model.vae.config.audio_channels == 2:
+            x = x.repeat(1, 2, 1)
+
+        audio_length = x.shape[-1]
+        audio = x.new_zeros(audio_shape)
+        audio[:, :, : min(audio_length, audio_vae_length)] = x[:, :, :audio_vae_length]
+
+        encoded_audio = self.model.vae.encode(audio.to(self.device)).latent_dist
+        encoded_audio = encoded_audio.sample()
+        return encoded_audio
+
+    def vae_decode(self, x: torch.Tensor) -> torch.Tensor:
+        torch.cuda.empty_cache()
+        # return self.model.vae.decode(1 / self.model.vae.config.scaling_factor * x).sample
+        aud = self.model.vae.decode(x).sample
+        return aud[:, :, self.waveform_start:self.waveform_end]
+
+    def setup_extra_inputs(self, x: torch.Tensor, init_timestep: torch.Tensor,
+                           extra_info: Optional[Any] = None,
+                           audio_start_in_s: float = 0, audio_end_in_s: Optional[float] = None,
+                           save_compute: bool = False) -> None:
+        max_audio_length_in_s = self.model.transformer.config.sample_size * self.model.vae.hop_length / \
+            self.model.vae.config.sampling_rate
+        if audio_end_in_s is None:
+            audio_end_in_s = max_audio_length_in_s
+
+        if audio_end_in_s - audio_start_in_s > max_audio_length_in_s:
+            raise ValueError(
+                f"The total audio length requested ({audio_end_in_s-audio_start_in_s}s) is longer "
+                f"than the model maximum possible length ({max_audio_length_in_s}). "
+                f"Make sure that 'audio_end_in_s-audio_start_in_s<={max_audio_length_in_s}'."
+            )
+
+        self.waveform_start = int(audio_start_in_s * self.model.vae.config.sampling_rate)
+        self.waveform_end = int(audio_end_in_s * self.model.vae.config.sampling_rate)
+
+        self.seconds_start_hidden_states, self.seconds_end_hidden_states = self.model.encode_duration(
+            audio_start_in_s, audio_end_in_s, self.device, False, 1)
+
+        if save_compute:
+            self.seconds_start_hidden_states = torch.cat([self.seconds_start_hidden_states, self.seconds_start_hidden_states], dim=0)
+            self.seconds_end_hidden_states = torch.cat([self.seconds_end_hidden_states, self.seconds_end_hidden_states], dim=0)
+
+        self.audio_duration_embeds = torch.cat([self.seconds_start_hidden_states,
+                                                self.seconds_end_hidden_states], dim=2)
+
+        # 7. Prepare rotary positional embedding
+        self.rotary_embedding = get_1d_rotary_pos_embed(
+            self.model.rotary_embed_dim,
+            x.shape[2] + self.audio_duration_embeds.shape[1],
+            use_real=True,
+            repeat_interleave_real=False,
+        )
+
+        self.model.scheduler._init_step_index(init_timestep)
+
+        # fix lower_order_nums for the reverse step - Option 1: only start from first order
+        # self.model.scheduler.lower_order_nums = 0
+        # self.model.scheduler.model_outputs = [None] * self.model.scheduler.config.solver_order
+        # fix lower_order_nums for the reverse step - Option 2: start from the correct order with history
+        t_to_idx = {float(v): k for k, v in enumerate(self.model.scheduler.timesteps)}
+        idx = len(self.model.scheduler.timesteps) - t_to_idx[float(init_timestep)] - 1
+        self.model.scheduler.model_outputs = [None, extra_info[idx] if extra_info is not None else None]
+        self.model.scheduler.lower_order_nums = min(self.model.scheduler.step_index,
+                                                    self.model.scheduler.config.solver_order)
+
+        # if rand check:
+        #     x *= self.model.scheduler.init_noise_sigma
+        # return x
+
+    def sample_xts_from_x0(self, x0: torch.Tensor, num_inference_steps: int = 50) -> torch.Tensor:
+        """
+        Samples from P(x_1:T|x_0)
+        """
+
+        sigmas = self.model.scheduler.sigmas
+        shapes = self.get_noise_shape(x0, num_inference_steps + 1)
+        xts = torch.zeros(shapes).to(x0.device)
+        xts[0] = x0
+
+        timesteps = self.model.scheduler.timesteps.to(self.device)
+        t_to_idx = {float(v): k for k, v in enumerate(timesteps)}
+        for t in reversed(timesteps):
+            # idx = t_to_idx[int(t)]
+            idx = num_inference_steps - t_to_idx[float(t)]
+            n = torch.randn_like(x0)
+            xts[idx] = x0 + n * sigmas[t_to_idx[float(t)]]
+        return xts
+
+    def get_zs_from_xts(self, xt: torch.Tensor, xtm1: torch.Tensor, data_pred: torch.Tensor,
+                        t: torch.Tensor, numerical_fix: bool = True, first_order: bool = False, **kwargs
+                        ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        # pred of x0
+        sigmas = self.model.scheduler.sigmas
+        timesteps = self.model.scheduler.timesteps
+        solver_order = self.model.scheduler.config.solver_order
+
+        if self.model.scheduler.step_index is None:
+            self.model.scheduler._init_step_index(t)
+        curr_step_index = self.model.scheduler.step_index
+
+        # Improve numerical stability for small number of steps
+        lower_order_final = (curr_step_index == len(timesteps) - 1) and (
+            self.model.scheduler.config.euler_at_final
+            or (self.model.scheduler.config.lower_order_final and len(timesteps) < 15)
+            or self.model.scheduler.config.final_sigmas_type == "zero")
+        lower_order_second = ((curr_step_index == len(timesteps) - 2) and
+                              self.model.scheduler.config.lower_order_final and len(timesteps) < 15)
+
+        data_pred = self.model.scheduler.convert_model_output(data_pred, sample=xt)
+        for i in range(solver_order - 1):
+            self.model.scheduler.model_outputs[i] = self.model.scheduler.model_outputs[i + 1]
+        self.model.scheduler.model_outputs[-1] = data_pred
+
+        # instead of brownian noise, here we calculate the noise ourselves
+        if (curr_step_index == len(timesteps) - 1) and self.model.scheduler.config.final_sigmas_type == "zero":
+            z = torch.zeros_like(xt)
+        elif first_order or solver_order == 1 or self.model.scheduler.lower_order_nums < 1 or lower_order_final:
+            sigma_t, sigma_s = sigmas[curr_step_index + 1], sigmas[curr_step_index]
+            h = torch.log(sigma_s) - torch.log(sigma_t)
+            z = (xtm1 - (sigma_t / sigma_s * torch.exp(-h)) * xt - (1 - torch.exp(-2.0 * h)) * data_pred) \
+                / (sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)))
+        elif solver_order == 2 or self.model.scheduler.lower_order_nums < 2 or lower_order_second:
+            sigma_t = sigmas[curr_step_index + 1]
+            sigma_s0 = sigmas[curr_step_index]
+            sigma_s1 = sigmas[curr_step_index - 1]
+            m0, m1 = self.model.scheduler.model_outputs[-1], self.model.scheduler.model_outputs[-2]
+            h, h_0 = torch.log(sigma_s0) - torch.log(sigma_t), torch.log(sigma_s1) - torch.log(sigma_s0)
+            r0 = h_0 / h
+            D0, D1 = m0, (1.0 / r0) * (m0 - m1)
+
+            # sde-dpmsolver++
+            z = (xtm1 - (sigma_t / sigma_s0 * torch.exp(-h)) * xt
+                 - (1 - torch.exp(-2.0 * h)) * D0
+                 - 0.5 * (1 - torch.exp(-2.0 * h)) * D1) \
+                / (sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)))
+
+        # correction to avoid error accumulation
+        if numerical_fix:
+            if first_order or solver_order == 1 or self.model.scheduler.lower_order_nums < 1 or lower_order_final:
+                xtm1 = self.model.scheduler.dpm_solver_first_order_update(data_pred, sample=xt, noise=z)
+            elif solver_order == 2 or self.model.scheduler.lower_order_nums < 2 or lower_order_second:
+                xtm1 = self.model.scheduler.multistep_dpm_solver_second_order_update(
+                    self.model.scheduler.model_outputs, sample=xt, noise=z)
+            # If not perfect recon - maybe TODO fix self.model.scheduler.model_outputs as well?
+
+        if self.model.scheduler.lower_order_nums < solver_order:
+            self.model.scheduler.lower_order_nums += 1
+        # upon completion increase step index by one
+        self.model.scheduler._step_index += 1
+
+        return z, xtm1, self.model.scheduler.model_outputs[-2]
+
+    def get_sr(self) -> int:
+        return self.model.vae.config.sampling_rate
+
+    def get_noise_shape(self, x0: torch.Tensor, num_steps: int) -> Tuple[int, int, int]:
+        variance_noise_shape = (num_steps,
+                                self.model.transformer.config.in_channels,
+                                int(self.model.transformer.config.sample_size))
+        return variance_noise_shape
+
+    def reverse_step_with_custom_noise(self, model_output: torch.Tensor, timestep: torch.Tensor, sample: torch.Tensor,
+                                       variance_noise: Optional[torch.Tensor] = None,
+                                       first_order: bool = False, **kwargs
+                                       ) -> torch.Tensor:
+        if self.model.scheduler.step_index is None:
+            self.model.scheduler._init_step_index(timestep)
+
+        # Improve numerical stability for small number of steps
+        lower_order_final = (self.model.scheduler.step_index == len(self.model.scheduler.timesteps) - 1) and (
+            self.model.scheduler.config.euler_at_final
+            or (self.model.scheduler.config.lower_order_final and len(self.model.scheduler.timesteps) < 15)
+            or self.model.scheduler.config.final_sigmas_type == "zero"
+        )
+        lower_order_second = (
+            (self.model.scheduler.step_index == len(self.model.scheduler.timesteps) - 2) and
+            self.model.scheduler.config.lower_order_final and len(self.model.scheduler.timesteps) < 15
+        )
+
+        model_output = self.model.scheduler.convert_model_output(model_output, sample=sample)
+        for i in range(self.model.scheduler.config.solver_order - 1):
+            self.model.scheduler.model_outputs[i] = self.model.scheduler.model_outputs[i + 1]
+        self.model.scheduler.model_outputs[-1] = model_output
+
+        if variance_noise is None:
+            if self.model.scheduler.noise_sampler is None:
+                self.model.scheduler.noise_sampler = BrownianTreeNoiseSampler(
+                    model_output, sigma_min=self.model.scheduler.config.sigma_min,
+                    sigma_max=self.model.scheduler.config.sigma_max, seed=None)
+            variance_noise = self.model.scheduler.noise_sampler(
+                self.model.scheduler.sigmas[self.model.scheduler.step_index],
+                self.model.scheduler.sigmas[self.model.scheduler.step_index + 1]).to(model_output.device)
+
+        if first_order or self.model.scheduler.config.solver_order == 1 or \
+                self.model.scheduler.lower_order_nums < 1 or lower_order_final:
+            prev_sample = self.model.scheduler.dpm_solver_first_order_update(
+                model_output, sample=sample, noise=variance_noise)
+        elif self.model.scheduler.config.solver_order == 2 or \
+                self.model.scheduler.lower_order_nums < 2 or lower_order_second:
+            prev_sample = self.model.scheduler.multistep_dpm_solver_second_order_update(
+                self.model.scheduler.model_outputs, sample=sample, noise=variance_noise)
+
+        if self.model.scheduler.lower_order_nums < self.model.scheduler.config.solver_order:
+            self.model.scheduler.lower_order_nums += 1
+
+        # upon completion increase step index by one
+        self.model.scheduler._step_index += 1
+
+        return prev_sample
+
+    def unet_forward(self,
+                     sample: torch.FloatTensor,
+                     timestep: Union[torch.Tensor, float, int],
+                     encoder_hidden_states: torch.Tensor,
+                     encoder_attention_mask: Optional[torch.Tensor] = None,
+                     return_dict: bool = True,
+                     **kwargs) -> Tuple:
+
+        # Create text_audio_duration_embeds and audio_duration_embeds
+        embeds = torch.cat([encoder_hidden_states, self.seconds_start_hidden_states, self.seconds_end_hidden_states],
+                           dim=1)
+        if encoder_attention_mask is None:
+            # handle the batched case
+            if embeds.shape[0] > 1:
+                embeds[0] = torch.zeros_like(embeds[0], device=embeds.device)
+            else:
+                embeds = torch.zeros_like(embeds, device=embeds.device)
+
+        noise_pred = self.model.transformer(sample,
+                                            timestep.unsqueeze(0),
+                                            encoder_hidden_states=embeds,
+                                            global_hidden_states=self.audio_duration_embeds,
+                                            rotary_embedding=self.rotary_embedding)
+
+        if not return_dict:
+            return (noise_pred.sample,)
+
+        return noise_pred, None, None
 
 
-def load_model(model_id, device, double_precision=False):
-    ldm_stable = AudioLDM2Wrapper(model_id=model_id, device=device, double_precision=double_precision)
+def load_model(model_id: str, device: torch.device,
+               double_precision: bool = False, token: Optional[str] = None) -> PipelineWrapper:
+    if 'audioldm2' in model_id:
+        ldm_stable = AudioLDM2Wrapper(model_id=model_id, device=device, double_precision=double_precision, token=token)
+    elif 'stable-audio' in model_id:
+        ldm_stable = StableAudWrapper(model_id=model_id, device=device, double_precision=double_precision, token=token)
     ldm_stable.load_scheduler()
     torch.cuda.empty_cache()
     return ldm_stable

requirements.txt

@@ -1,8 +1,9 @@
-torch
-numpy<2
+torch>2.2.0
+numpy<2.0.0
 torchaudio
 diffusers
 accelerate
+torchsde
 transformers
 tqdm
 soundfile

utils.py

@@ -2,8 +2,11 @@ import numpy as np
 import torch
 from typing import Optional, List, Tuple, NamedTuple, Union
 from models import PipelineWrapper
+import torchaudio
 from audioldm.utils import get_duration
 
+MAX_DURATION = 30
+
 
 class PromptEmbeddings(NamedTuple):
     embedding_hidden_states: torch.Tensor
@@ -11,26 +14,57 @@ class PromptEmbeddings(NamedTuple):
     boolean_prompt_mask: torch.Tensor
 
 
-def load_audio(audio_path: Union[str, np.array], fn_STFT, left: int = 0, right: int = 0, device: Optional[torch.device] = None
-               ) -> torch.tensor:
-    if type(audio_path) is str:
-        import audioldm
-        import audioldm.audio
+def load_audio(audio_path: Union[str, np.array], fn_STFT, left: int = 0, right: int = 0,
+               device: Optional[torch.device] = None,
+               return_wav: bool = False, stft: bool = False, model_sr: Optional[int] = None) -> torch.Tensor:
+    if stft:  # AudioLDM/tango loading to spectrogram
+        if type(audio_path) is str:
+            import audioldm
+            import audioldm.audio
 
-        duration = min(get_duration(audio_path), 30)
+            duration = get_duration(audio_path)
+            if MAX_DURATION is not None:
+                duration = min(duration, MAX_DURATION)
 
-        mel, _, _ = audioldm.audio.wav_to_fbank(audio_path, target_length=int(duration * 102.4), fn_STFT=fn_STFT)
-        mel = mel.unsqueeze(0)
-    else:
-        mel = audio_path
+            mel, _, wav = audioldm.audio.wav_to_fbank(audio_path, target_length=int(duration * 102.4), fn_STFT=fn_STFT)
+            mel = mel.unsqueeze(0)
+        else:
+            mel = audio_path
 
-    c, h, w = mel.shape
-    left = min(left, w-1)
-    right = min(right, w - left - 1)
-    mel = mel[:, :, left:w-right]
-    mel = mel.unsqueeze(0).to(device)
+        c, h, w = mel.shape
+        left = min(left, w-1)
+        right = min(right, w - left - 1)
+        mel = mel[:, :, left:w-right]
+        mel = mel.unsqueeze(0).to(device)
 
-    return mel
+        if return_wav:
+            return mel, 16000, duration, wav
+
+        return mel, model_sr, duration
+    else:
+        waveform, sr = torchaudio.load(audio_path)
+        if sr != model_sr:
+            waveform = torchaudio.functional.resample(waveform, orig_freq=sr, new_freq=model_sr)
+        # waveform = waveform.numpy()[0, ...]
+
+        def normalize_wav(waveform):
+            waveform = waveform - torch.mean(waveform)
+            waveform = waveform / (torch.max(torch.abs(waveform)) + 1e-8)
+            return waveform * 0.5
+
+        waveform = normalize_wav(waveform)
+        # waveform = waveform[None, ...]
+        # waveform = pad_wav(waveform, segment_length)
+
+        # waveform = waveform[0, ...]
+        waveform = torch.FloatTensor(waveform)
+        if MAX_DURATION is not None:
+            duration = min(waveform.shape[-1] / model_sr, MAX_DURATION)
+            waveform = waveform[:, :int(duration * model_sr)]
+
+        # cut waveform
+        duration = waveform.shape[-1] / model_sr
+        return waveform, model_sr, duration
 
 
 def get_height_of_spectrogram(length: int, ldm_stable: PipelineWrapper) -> int: