Update create_wav.py

create_wav.py

@@ -1,20 +1,40 @@
 import scipy
 import torch
 from diffusers import AudioLDMPipeline
+from diffusers import AudioLDM2Pipeline
 
-repo_id = "cvssp/audioldm-s-full-v2"
-pipe = AudioLDMPipeline.from_pretrained(repo_id, torch_dtype=torch.float16)
-pipe = pipe.to("cpu")
+def processing(prompt:str, name_file: str, model:str):
 
-def create(prompt:str, name_file: str):
-  prompt = prompt
-  audio = pipe(prompt, num_inference_steps=10, audio_length_in_s=5.0).audios[0]
+    if model == "good":
+      repo_id = "cvssp/audioldm-s-full-v2"
+      pipe = AudioLDMPipeline.from_pretrained(repo_id, torch_dtype=torch.float16)
+      pipe = pipe.to("cuda")
 
-  # save the audio sample as a .wav file
-  scipy.io.wavfile.write(f'{name_file}.wav', rate=16000, data=audio)
+      prompt = prompt
+      audio = pipe(prompt, num_inference_steps=100, audio_length_in_s=5).audios[0]
 
-  # Assuming your file is named 'audio_file.mp3'
-  audio_file = f'{name_file}.wav'
+      scipy.io.wavfile.write(f'{name_file}.wav', rate=16000, data=audio)
 
-  # Play the audio file
-  return audio_file
+    elif model == "better":
+      repo_id = "cvssp/audioldm2"
+      pipe = AudioLDM2Pipeline.from_pretrained(repo_id, torch_dtype=torch.float16)
+      pipe = pipe.to("cuda")
+
+      prompt = prompt
+      negative_prompt = "Low quality."
+
+      generator = torch.Generator("cuda").manual_seed(0)
+
+      audio = pipe(
+          prompt,
+          negative_prompt=negative_prompt,
+          num_inference_steps=200,
+          audio_length_in_s=10.0,
+          num_waveforms_per_prompt=3,
+          generator=generator
+          ).audios
+
+      scipy.io.wavfile.write(f'{name_file}.wav', rate=16000, data=audio[0])
+
+    audio_file = f'{name_file}.wav'
+    return audio_file