allow user to set steps, pick scheduler, and make "gradio app.py" work

app.py

@@ -6,11 +6,25 @@ from modules.AudioToken.embedder import FGAEmbedder
 from diffusers import AutoencoderKL, UNet2DConditionModel
 from diffusers.models.attention_processor import LoRAAttnProcessor
 from diffusers import StableDiffusionPipeline
+from diffusers import (
+    DDPMScheduler,
+    DDIMScheduler,
+    PNDMScheduler,
+    LMSDiscreteScheduler,
+    EulerDiscreteScheduler,
+    EulerAncestralDiscreteScheduler,
+    DPMSolverMultistepScheduler,
+    DPMSolverSinglestepScheduler,
+    DEISMultistepScheduler,
+    UniPCMultistepScheduler,
+    HeunDiscreteScheduler,
+    KDPM2AncestralDiscreteScheduler,
+    KDPM2DiscreteScheduler,
+)
 import numpy as np
 import gradio as gr
 from scipy import signal
 
-
 class AudioTokenWrapper(torch.nn.Module):
     """Simple wrapper module for Stable Diffusion that holds all the models together"""
 
@@ -22,17 +36,33 @@ class AudioTokenWrapper(torch.nn.Module):
 
         super().__init__()
         # Load scheduler and models
+        self.ddpm = DDPMScheduler.from_pretrained(repo_id, subfolder="scheduler")
+        self.ddim = DDIMScheduler.from_pretrained(repo_id, subfolder="scheduler")
+        self.pndm = PNDMScheduler.from_pretrained(repo_id, subfolder="scheduler")
+        self.lms = LMSDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
+        self.euler_anc = EulerAncestralDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
+        self.euler = EulerDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
+        self.dpm = DPMSolverMultistepScheduler.from_pretrained(repo_id, subfolder="scheduler")
+        self.dpms = DPMSolverSinglestepScheduler.from_pretrained(repo_id, subfolder="scheduler")
+        self.deis = DEISMultistepScheduler.from_pretrained(repo_id, subfolder="scheduler")
+        self.unipc = UniPCMultistepScheduler.from_pretrained(repo_id, subfolder="scheduler")
+        self.heun = HeunDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
+        self.kdpm2_anc = KDPM2AncestralDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
+        self.kdpm2 = KDPM2DiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
+
+
+
         self.tokenizer = CLIPTokenizer.from_pretrained(
-            "CompVis/stable-diffusion-v1-4", subfolder="tokenizer"
+            repo_id, subfolder="tokenizer"
         )
         self.text_encoder = CLIPTextModel.from_pretrained(
-            "CompVis/stable-diffusion-v1-4", subfolder="text_encoder", revision=None
+            repo_id, subfolder="text_encoder", revision=None
         )
         self.unet = UNet2DConditionModel.from_pretrained(
-            "CompVis/stable-diffusion-v1-4", subfolder="unet", revision=None
+            repo_id, subfolder="unet", revision=None
         )
         self.vae = AutoencoderKL.from_pretrained(
-            "CompVis/stable-diffusion-v1-4", subfolder="vae", revision=None
+            repo_id, subfolder="vae", revision=None
         )
 
         checkpoint = torch.load(
@@ -90,11 +120,39 @@ class AudioTokenWrapper(torch.nn.Module):
         self.text_encoder.resize_token_embeddings(len(self.tokenizer))
 
 
-def greet(audio):
+def greet(audio, steps=25, scheduler="ddpm"):
     sample_rate, audio = audio
     audio = audio.astype(np.float32, order='C') / 32768.0
     desired_sample_rate = 16000
 
+    match scheduler:
+        case "ddpm":
+            use_sched = model.ddpm
+        case "ddim":
+            use_sched = model.ddim
+        case "pndm":
+            use_sched = model.pndm
+        case "lms":
+            use_sched = model.lms
+        case "euler_anc":
+            use_sched = model.euler_anc
+        case "euler":
+            use_sched = model.euler
+        case "dpm":
+            use_sched = model.dpm
+        case "dpms":
+            use_sched = model.dpms
+        case "deis":
+            use_sched = model.deis
+        case "unipc":
+            use_sched = model.unipc
+        case "heun":
+            use_sched = model.heun
+        case "kdpm2_anc":
+            use_sched = model.kdpm2_anc
+        case "kdpm2":
+            use_sched = model.kdpm2
+
     if audio.ndim == 2:
         audio = audio.sum(axis=1) / 2
 
@@ -114,49 +172,61 @@ def greet(audio):
     audio_values = torch.unsqueeze(torch.tensor(audio), dim=0).to(device).to(dtype=weight_dtype)
     if audio_values.ndim == 1:
         audio_values = torch.unsqueeze(audio_values, dim=0)
-    aud_features = model.aud_encoder.extract_features(audio_values)[1]
-    audio_token = model.embedder(aud_features)
+    with torch.no_grad():
+        torch.cuda.empty_cache()
+        aud_features = model.aud_encoder.extract_features(audio_values)[1]
+        audio_token = model.embedder(aud_features)
 
     token_embeds = model.text_encoder.get_input_embeddings().weight.data
     token_embeds[model.placeholder_token_id] = audio_token.clone()
-
+    g_gpu = torch.Generator(device='cuda')
+    g_gpu.manual_seed(23029249075547) # no reason this can't be input by the user!
     pipeline = StableDiffusionPipeline.from_pretrained(
-        "CompVis/stable-diffusion-v1-4",
+        "philz1337/reliberate",
         tokenizer=model.tokenizer,
         text_encoder=model.text_encoder,
         vae=model.vae,
         unet=model.unet,
+	scheduler=use_sched,
+        safety_checker=None,
     ).to(device)
-    image = pipeline(prompt, num_inference_steps=40, guidance_scale=7.5).images[0]
+    pipeline.enable_attention_slicing()
+    # pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config)
+    # pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
+    print(f"taking {steps} steps using the {scheduler} scheduler")
+    image = pipeline(prompt, num_inference_steps=steps, guidance_scale=8.5, generator=g_gpu).images[0]
     return image
 
 
-if __name__ == "__main__":
-
-    lora = False
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    model = AudioTokenWrapper(lora, device)
-
-    description = """<p>
-    This is a demo of <a href='https://pages.cs.huji.ac.il/adiyoss-lab/AudioToken' target='_blank'>AudioToken: Adaptation of Text-Conditioned Diffusion Models for Audio-to-Image Generation</a>.<br><br>
-    A novel method utilizing latent diffusion models trained for text-to-image-generation to generate images conditioned on audio recordings. Using a pre-trained audio encoding model, the proposed method encodes audio into a new token, which can be considered as an adaptation layer between the audio and text representations.<br><br>
-    For more information, please see the original <a href='https://arxiv.org/abs/2305.13050' target='_blank'>paper</a> and <a href='https://github.com/guyyariv/AudioToken' target='_blank'>repo</a>.
-    </p>"""
-
-    examples = [
-        # ["assets/train.wav"],
-        ["assets/dog barking.wav"],
-        ["assets/airplane taking off.wav"],
-        # ["assets/electric guitar.wav"],
-        # ["assets/female sings.wav"],
-    ]
-
-    demo = gr.Interface(
-        fn=greet,
-        inputs="audio",
-        outputs="image",
-        title='AudioToken',
-        description=description,
-        examples=examples
-    )
-    demo.launch()
+lora = False
+repo_id = "philz1337/reliberate"
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+model = AudioTokenWrapper(lora, device)
+model = model.to(device)
+description = """<p>
+This is a demo of <a href='https://pages.cs.huji.ac.il/adiyoss-lab/AudioToken' target='_blank'>AudioToken: Adaptation of Text-Conditioned Diffusion Models for Audio-to-Image Generation</a>.<br><br>
+A novel method utilizing latent diffusion models trained for text-to-image-generation to generate images conditioned on audio recordings. Using a pre-trained audio encoding model, the proposed method encodes audio into a new token, which can be considered as an adaptation layer between the audio and text representations.<br><br>
+For more information, please see the original <a href='https://arxiv.org/abs/2305.13050' target='_blank'>paper</a> and <a href='https://github.com/guyyariv/AudioToken' target='_blank'>repo</a>.
+</p>"""
+
+examples = [
+    # ["assets/train.wav"],
+    # ["assets/dog barking.wav"],
+    # ["assets/airplane taking off.wav"],
+    # ["assets/electric guitar.wav"],
+    # ["assets/female sings.wav"],
+]
+
+my_demo = gr.Interface(
+    fn=greet,
+    inputs=[
+	"audio",
+	gr.Slider(value=25,step=1,label="diffusion steps"),
+	gr.Dropdown(choices=["ddim","ddpm","pndm","lms","euler_anc","euler","dpm","dpms","deis","unipc","heun","kdpm2_anc","kdpm2"],value="unipc"),
+    ],
+    outputs="image",
+    title='AudioToken',
+    description=description,
+    examples=examples
+)
+my_demo.launch()