Update app.py

app.py

@@ -1,57 +1,44 @@
 import gradio as gr
-import spaces
 from transformers import WhisperProcessor, WhisperForConditionalGeneration
 import torch
 import os
 import soundfile as sf
 from scipy.signal import resample
 
-print(f"Is CUDA available: {torch.cuda.is_available()}")
-if torch.cuda.is_available():
-    print(f"Using CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
-else:
-    print("No GPU detected, defaulting to CPU.")
-
 # Define the model ID
 MODEL_ID = "WMRNORDIC/whisper-swedish-telephonic"
 
-# Load token from environment variables
+# 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 is not set. Please set it in the environment variables or Space settings.")
-
-# GPU Initialization
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Running on device: {DEVICE}")
+    raise ValueError("HF_API_TOKEN not found in environment variables. Please set it in the Space settings.")
 
-# Load model and processor with ZeroGPU integration
-@spaces.GPU
+# Function to initialize the model and processor lazily
 def initialize_model():
-    try:
-        print("Loading model and processor...")
-        processor = WhisperProcessor.from_pretrained(MODEL_ID, use_auth_token=HF_API_TOKEN)
-        model = WhisperForConditionalGeneration.from_pretrained(MODEL_ID, use_auth_token=HF_API_TOKEN)
-        model = model.to(DEVICE)  # Move model to GPU or CPU
-        print("Model loaded successfully.")
-        return processor, model
-    except Exception as e:
-        print(f"Error loading model or processor: {e}")
-        raise e
-
-processor, model = initialize_model()
+    # This function will be executed only when Gradio is processing a request
+    print("Loading 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")  # Ensure GPU is used if available
+    print("Model loaded successfully.")
+    return processor, model
 
 # Function to resample audio to 16kHz
 def resample_audio(audio_data, original_rate, target_rate=16000):
     if original_rate != target_rate:
-        print(f"Resampling audio from {original_rate}Hz to {target_rate}Hz...")
         num_samples = int(len(audio_data) * target_rate / original_rate)
         return resample(audio_data, num_samples)
     return audio_data
 
-# Transcription function with GPU allocation
-@spaces.GPU
+# Transcription function
 def transcribe_audio(audio):
     try:
+        # Lazy-load the model and processor inside the request handler
+        global processor, model
+        if 'processor' not in globals() or 'model' not in globals():
+            processor, model = initialize_model()
+        
+        # Handle microphone input or uploaded file
         if isinstance(audio, tuple):  # Microphone input
             audio_data = audio[1]
             sample_rate = audio[0]
@@ -61,8 +48,8 @@ def transcribe_audio(audio):
             audio_data = resample_audio(audio_data, sample_rate)
 
         # Preprocess and perform inference
-        input_features = processor(audio_data, return_tensors="pt", sampling_rate=16000).input_features
-        input_features = input_features.to(DEVICE)  # Move input to GPU or CPU
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        input_features = processor(audio_data, return_tensors="pt", sampling_rate=16000).input_features.to(device)
         with torch.no_grad():
             predicted_ids = model.generate(input_features)