Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces  # Required for ZeroGPU compliance
 import gradio as gr
 from transformers import WhisperProcessor, WhisperForConditionalGeneration
 import torch
@@ -5,52 +6,53 @@ import os
 import soundfile as sf
 from scipy.signal import resample
 
-# Define the model ID
+# Model ID and Hugging Face Token
 MODEL_ID = "WMRNORDIC/whisper-swedish-telephonic"
-
-# Load the Hugging Face token from the environment
 HF_API_TOKEN = os.getenv("HF_API_TOKEN")
+
 if not HF_API_TOKEN:
-    raise ValueError("HF_API_TOKEN not found in environment variables. Please set it in the Space settings.")
+    raise ValueError("HF_API_TOKEN not found. Set it in the environment variables.")
 
 # Sample file path
 SAMPLE_FILE_PATH = "trimmed_resampled_audio.wav"  # Update this path if necessary
 
-# Function to initialize the model and processor lazily
+
+@spaces.GPU
 def initialize_model():
-    print("Loading model and processor...")
+    """Lazy initialization of model and processor with GPU allocation."""
+    print("Initializing model and processor...")
     processor = WhisperProcessor.from_pretrained(MODEL_ID, token=HF_API_TOKEN)
     model = WhisperForConditionalGeneration.from_pretrained(MODEL_ID, token=HF_API_TOKEN)
-    model = model.to("cuda" if torch.cuda.is_available() else "cpu")  # Use GPU if available
-    print("Model loaded successfully.")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = model.to(device)
+    print(f"Model loaded on device: {device}")
     return processor, model
 
-# Function to resample audio to 16kHz
-def resample_audio(audio_data, original_rate, target_rate=16000):
-    if original_rate != target_rate:
-        num_samples = int(len(audio_data) * target_rate / original_rate)
-        return resample(audio_data, num_samples)
-    return audio_data
-
-# Transcription function
+@spaces.GPU
 def transcribe_audio(audio):
+    """Transcription logic with ZeroGPU compliance."""
     try:
+        # Lazy-load model and processor
         global processor, model
         if 'processor' not in globals() or 'model' not in globals():
             processor, model = initialize_model()
-        
-        # Handle microphone input or uploaded file
+
+        # Handle audio input
         if isinstance(audio, tuple):  # Microphone input
-            audio_data = audio[1]
-            sample_rate = audio[0]
-            audio_data = resample_audio(audio_data, sample_rate)
+            audio_data, sample_rate = audio[1], audio[0]
         else:  # Uploaded file
             audio_data, sample_rate = sf.read(audio)
-            audio_data = resample_audio(audio_data, sample_rate)
 
-        # Preprocess and perform inference
+        # Resample to 16kHz
+        if sample_rate != 16000:
+            num_samples = int(len(audio_data) * 16000 / sample_rate)
+            audio_data = resample(audio_data, num_samples)
+
+        # Prepare inputs for the model
         device = "cuda" if torch.cuda.is_available() else "cpu"
         input_features = processor(audio_data, return_tensors="pt", sampling_rate=16000).input_features.to(device)
+
+        # Generate transcription
         with torch.no_grad():
             predicted_ids = model.generate(input_features)
 
@@ -61,7 +63,7 @@ def transcribe_audio(audio):
     except Exception as e:
         return f"Error during transcription: {str(e)}"
 
-# Gradio interface
+# Gradio Interface
 def create_demo():
     """Set up the Gradio app."""
     with gr.Blocks() as demo:
@@ -79,6 +81,7 @@ def create_demo():
         audio_input.change(transcribe_audio, inputs=audio_input, outputs=transcription_output)
     return demo
 
+
 # Initialize Gradio app
 demo = create_demo()