Update app.py

app.py

@@ -1,19 +1,17 @@
-import torch
 import os
+import torch
 import soundfile as sf
-from diffusers import StableAudioPipeline
 from huggingface_hub import login
+from diffusers import StableAudioPipeline
 import gradio as gr
 
-# Retrieve the token securely from environment variables
+# Load Hugging Face token securely
 HUGGINGFACE_TOKEN = os.getenv("HF_TOKEN")
 if HUGGINGFACE_TOKEN is None:
-    raise ValueError("Missing Hugging Face API token. Set 'HF_TOKEN' in Secrets.")
-
-# Authenticate with Hugging Face Hub
+    raise ValueError("Missing Hugging Face token. Please set it in Spaces Secrets.")
 login(HUGGINGFACE_TOKEN)
 
-# Set up the device
+# Set device for PyTorch
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if device == "cuda" else torch.float32
 
@@ -24,22 +22,21 @@ pipe = StableAudioPipeline.from_pretrained(
 )
 pipe = pipe.to(device)
 
-# Audio generation function
-def generate_audio(prompt, negative_prompt, duration, seed):
+# Function to generate audio
+def generate_audio(prompt, negative_prompt, duration, diffusion_steps, seed):
     generator = torch.Generator(device).manual_seed(seed)
     audio_output = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
-        num_inference_steps=50,
+        num_inference_steps=int(diffusion_steps),  # Number of diffusion steps
         audio_end_in_s=duration,
         num_waveforms_per_prompt=1,
         generator=generator
     ).audios
-    
-    # Save the generated audio
     output_audio = audio_output[0].T.float().cpu().numpy()
-    sf.write("output.wav", output_audio, pipe.vae.sampling_rate)
-    return "output.wav"
+    output_file = "output.wav"
+    sf.write(output_file, output_audio, pipe.vae.sampling_rate)
+    return output_file
 
 # Gradio UI
 with gr.Blocks() as demo:
@@ -69,4 +66,3 @@ with gr.Blocks() as demo:
 
 # Launch the app
 demo.launch()
-