Update app.py

app.py

@@ -6,13 +6,13 @@ from tqdm import tqdm
 from huggingface_hub import snapshot_download
 from models import AudioDiffusion, DDPMScheduler
 from audioldm.audio.stft import TacotronSTFT
-from audioldm.variational_autoencoder import AutoencoderKL
+#from audioldm.variational_autoencoder import AutoencoderKL
 from pydub import AudioSegment
 from gradio import Markdown
 import spaces
 
 import torch
-#from diffusers.models.autoencoder_kl import AutoencoderKL
+from diffusers.models.autoencoder_kl import AutoencoderKL
 from diffusers.models.unet_2d_condition import UNet2DConditionModel
 from diffusers import DiffusionPipeline,AudioPipelineOutput
 from transformers import CLIPTextModel, T5EncoderModel, AutoModel, T5Tokenizer, T5TokenizerFast
@@ -239,21 +239,21 @@ class Tango:
 
 tango = Tango(device="cpu")
     
-#pipe = Tango2Pipeline(vae=tango.vae,
-#                      text_encoder=tango.model.text_encoder,
-#                      tokenizer=tango.model.tokenizer,
-#                      unet=tango.model.unet,
-#                      scheduler=tango.scheduler
-#                      )
-#pipe.to(device)
-tango.vae.to(device_type)
-tango.stft.to(device_type)
-tango.model.to(device_type)
+pipe = Tango2Pipeline(vae=tango.vae,
+                      text_encoder=tango.model.text_encoder,
+                      tokenizer=tango.model.tokenizer,
+                      unet=tango.model.unet,
+                      scheduler=tango.scheduler
+                      )
+pipe.to(device_type)
+#tango.vae.to(device_type)
+#tango.stft.to(device_type)
+#tango.model.to(device_type)
     
 @spaces.GPU(duration=60)
 def gradio_generate(prompt, output_format, steps, guidance):
-    #output_wave = pipe(prompt,steps,guidance) ## Using pipeliine automatically uses flash attention for torch2.0 above
-    output_wave = tango.generate(prompt, steps, guidance)
+    output_wave = pipe(prompt,steps,guidance) ## Using pipeliine automatically uses flash attention for torch2.0 above
+    #output_wave = tango.generate(prompt, steps, guidance)
     # output_filename = f"{prompt.replace(' ', '_')}_{steps}_{guidance}"[:250] + ".wav"
     output_filename = "temp.wav"
     wavio.write(output_filename, output_wave, rate=16000, sampwidth=2)