Update app.py

app.py

@@ -36,7 +36,7 @@ print(f"Using device: {device}")
 model.to(device,torch.float32)
 
 @spaces.GPU(duration=60)
-def generate_audio(prompt, seconds_total=30, steps=100, cfg_scale=7):
+def generate_audio(prompt, seconds_total=30, steps=100, cfg_scale=7, use_bfloat=False):
     print(f"Prompt received: {prompt}")
     print(f"Settings: Duration={seconds_total}s, Steps={steps}, CFG Scale={cfg_scale}")
     sample_rate = model_config["sample_rate"]
@@ -50,6 +50,8 @@ def generate_audio(prompt, seconds_total=30, steps=100, cfg_scale=7):
     }]
     print(f"Conditioning: {conditioning}")
     print("Generating audio...")
+    if use_bfloat==True:
+        model.to(torch.bfloat16)
     output = generate_diffusion_cond(
         model,
         steps=steps,
@@ -65,7 +67,7 @@ def generate_audio(prompt, seconds_total=30, steps=100, cfg_scale=7):
     output = rearrange(output, "b d n -> d (b n)")
     # 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()
-    unique_filename = f"output_{uuid.uuid4().hex}.wav"
+    unique_filename = f"output_{uuid.uuid4().hex}.mp3"
     print(f"Saving audio to file: {unique_filename}")
     torchaudio.save(
     unique_filename,
@@ -78,13 +80,15 @@ def generate_audio(prompt, seconds_total=30, steps=100, cfg_scale=7):
     print(f"Audio saved: {unique_filename}")
     return unique_filename
 
+
 interface = gr.Interface(
     fn=generate_audio,
     inputs=[
         gr.Textbox(label="Prompt", placeholder="Enter your text prompt here"),
         gr.Slider(0, 420, value=30, label="Duration in Seconds"),
         gr.Slider(10, 420, value=100, step=10, label="Number of Diffusion Steps"),
-        gr.Slider(1.0, 32.0, value=7.0, step=0.1, label="CFG Scale")
+        gr.Slider(1.0, 32.0, value=7.0, step=0.1, label="CFG Scale"),
+        gr.Checkbox(value=False, label="Use Brainfloat")
     ],
     outputs=gr.Audio(type="filepath", label="Generated Audio"),
     title="Stable Audio Generator",