Update app.py

app.py

@@ -12,7 +12,7 @@ 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)
+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)