Update app.py

app.py

@@ -1,10 +1,35 @@
+
+import spaces
+import os
+import uuid
+
+os.putenv('PYTORCH_NVML_BASED_CUDA_CHECK','1')
+os.putenv('TORCH_LINALG_PREFER_CUSOLVER','1')
+alloc_conf_parts = [
+    'expandable_segments:True',
+    'pinned_use_background_threads:True'  # Specific to pinned memory.
+]
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = ','.join(alloc_conf_parts)
+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
+torch.backends.cudnn.allow_tf32 = False
+torch.backends.cudnn.deterministic = False
+torch.backends.cudnn.benchmark = False
+torch.backends.cuda.preferred_blas_library="cublas"
+torch.backends.cuda.preferred_linalg_library="cusolver"
+torch.set_float32_matmul_precision("highest")
+
+
 import torchaudio
 from einops import rearrange
 import gradio as gr
-import spaces
-import os
-import uuid
 
 # Importing the model-related functions
 from stable_audio_tools import get_pretrained_model
@@ -18,7 +43,7 @@ def load_model():
     return model, model_config
 
 # Function to set up, generate, and process the audio
-@spaces.GPU(duration=120)  # Allocate GPU only when this function is called
+@spaces.GPU(duration=60)  # Allocate GPU only when this function is called
 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}")
@@ -37,7 +62,7 @@ def generate_audio(prompt, seconds_total=30, steps=100, cfg_scale=7):
 
     print(f"Sample rate: {sample_rate}, Sample size: {sample_size}")
 
-    model = model.to(device)
+    model = model.to(device,torch.bfloat16)
     print("Model moved to device.")
 
     # Set up text and timing conditioning