Update app.py

app.py

@@ -23,12 +23,15 @@ def infer(prompt, progress=gr.Progress(track_tqdm=True)):
 
     return f"{prompt}.wav"
 
-def infer_img2img(prompt, audio_path):
+def infer_img2img(prompt, audio_path, progress=gr.Progress(track_tqdm=True)):
 
     pretrained_model_name_or_path = "auffusion/auffusion-full-no-adapter"
     dtype = torch.float16
     device = "cuda"
     
+    if not os.path.isdir(pretrained_model_name_or_path):
+        pretrained_model_name_or_path = snapshot_download(pretrained_model_name_or_path)     
+    
     vocoder = Generator.from_pretrained(pretrained_model_name_or_path, subfolder="vocoder")
     vocoder = vocoder.to(device=device, dtype=dtype)