Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import spaces
 import os
 import uuid
@@ -14,7 +13,6 @@ os.environ["SAFETENSORS_FAST_GPU"] = "1"
 os.putenv('HF_HUB_ENABLE_HF_TRANSFER','1')
 
 import torch
-
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
 torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
@@ -29,44 +27,28 @@ import torchaudio
 from einops import rearrange
 import gradio as gr
 
-# Importing the model-related functions
 from stable_audio_tools import get_pretrained_model
 from stable_audio_tools.inference.generation import generate_diffusion_cond
 
 model, model_config = get_pretrained_model("ford442/stable-audio-open-1.0")
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {device}")
-    
 model.to(device,torch.bfloat16)
 
-# Function to set up, generate, and process the audio
-@spaces.GPU(duration=60)  # Allocate GPU only when this function is called
+@spaces.GPU(duration=60)
 def generate_audio(prompt, seconds_total=30, steps=100, cfg_scale=7):
     print(f"Prompt received: {prompt}")
     print(f"Settings: Duration={seconds_total}s, Steps={steps}, CFG Scale={cfg_scale}")
-
-    # Fetch the Hugging Face token from the environment variable
-    hf_token = os.getenv('HF_TOKEN')
-    print(f"Hugging Face token: {hf_token}")
-
-    # Use pre-loaded model and configuration
     sample_rate = model_config["sample_rate"]
     sample_size = model_config["sample_size"]
-
     print(f"Sample rate: {sample_rate}, Sample size: {sample_size}")
-
     print("Model moved to device.")
-
-    # Set up text and timing conditioning
     conditioning = [{
         "prompt": prompt,
         "seconds_start": 0,
         "seconds_total": seconds_total
     }]
     print(f"Conditioning: {conditioning}")
-
-    # Generate stereo audio
     print("Generating audio...")
     output = generate_diffusion_cond(
         model,
@@ -80,27 +62,22 @@ def generate_audio(prompt, seconds_total=30, steps=100, cfg_scale=7):
         device=device
     )
     print("Audio generated.")
-
-    # Rearrange audio batch to a single sequence
     output = rearrange(output, "b d n -> d (b n)")
-    print("Audio rearranged.")
-
     # Peak normalize, clip, convert to int16
     output = output.to(torch.float32).div(torch.max(torch.abs(output))).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
-    print("Audio normalized and converted.")
-
-    # Generate a unique filename for the output
     unique_filename = f"output_{uuid.uuid4().hex}.wav"
     print(f"Saving audio to file: {unique_filename}")
-
-    # Save to file
-    torchaudio.save(unique_filename, output, sample_rate)
+    torchaudio.save(
+    unique_filename,
+    output,
+    sample_rate,
+    format="mp3",
+    encoding="MP3",
+    bits_per_sample=320
+    )
     print(f"Audio saved: {unique_filename}")
-
-    # Return the path to the generated audio file
     return unique_filename
 
-# Setting up the Gradio Interface
 interface = gr.Interface(
     fn=generate_audio,
     inputs=[
@@ -115,14 +92,12 @@ interface = gr.Interface(
     examples=[
     [
         "Create a serene soundscape of a quiet beach at sunset.",  # Text prompt
- 
-        45,  # Duration in Seconds
+         45,  # Duration in Seconds
         100,  # Number of Diffusion Steps
         10,  # CFG Scale
     ],
     [
         "Generate an energetic and bustling city street scene with distant traffic and close conversations.",  # Text prompt
-        
         30,  # Duration in Seconds
         120,  # Number of Diffusion Steps
         5,  # CFG Scale
@@ -146,7 +121,6 @@ interface = gr.Interface(
         25,  # Duration in Seconds
         90,  # Number of Diffusion Steps
         6,  # CFG Scale
-  
     ],
     ["Rock beat played in a treated studio, session drumming on an acoustic kit.",
         30,  # Duration in Seconds
@@ -156,5 +130,4 @@ interface = gr.Interface(
     ]
 ])
 
-# Launch the Interface
 interface.launch()