Update app.py

app.py

@@ -1,51 +1,33 @@
-# Import spaces first to avoid CUDA initialization issues
 import spaces
-
-# Then import other libraries
+import gradio as gr
 import torch
-import librosa
+import os
 from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration, AutoModelForCausalLM, AutoProcessor
 from gtts import gTTS
-import gradio as gr
-from PIL import Image
-import os
 from langdetect import detect
-import subprocess
 
-print("Using GPU for operations when available")
+# Disable CUDA initialization at import
+os.environ['CUDA_VISIBLE_DEVICES'] = ''
+torch.set_grad_enabled(False)
 
-# Install flash-attn
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-
-# Function to safely load pipeline within a GPU-decorated function
-@spaces.GPU
-def load_pipeline(model_name, **kwargs):
-    try:
-        device = 0 if torch.cuda.is_available() else "cpu"
-        return pipeline(model=model_name, device=device, **kwargs)
-    except Exception as e:
-        print(f"Error loading {model_name} pipeline: {e}")
-        return None
+print("CUDA initialization disabled at import")
 
-# Load Whisper model for speech recognition within a GPU-decorated function
 @spaces.GPU
 def load_whisper():
     try:
-        device = 0 if torch.cuda.is_available() else "cpu"
         processor = WhisperProcessor.from_pretrained("openai/whisper-small")
-        model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small").to(device)
+        model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
         return processor, model
     except Exception as e:
         print(f"Error loading Whisper model: {e}")
         return None, None
 
-# Load vision model within a GPU-decorated function
 @spaces.GPU
 def load_vision_model():
     try:
         model_id = "microsoft/Phi-3.5-vision-instruct"
         model = AutoModelForCausalLM.from_pretrained(
-            model_id, trust_remote_code=True, torch_dtype=torch.float16, use_flash_attention_2=False
+            model_id, trust_remote_code=True, torch_dtype=torch.float16
         )
         processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True, num_crops=16)
         return model, processor
@@ -53,43 +35,32 @@ def load_vision_model():
         print(f"Error loading vision model: {e}")
         return None, None
 
-# Load sarvam-2b for text generation within a GPU-decorated function
 @spaces.GPU
 def load_sarvam():
-    return load_pipeline('sarvamai/sarvam-2b-v0.5')
-
-# Load all models
-whisper_processor, whisper_model = load_whisper()
-vision_model, vision_processor = load_vision_model()
-sarvam_pipe = load_sarvam()
+    try:
+        return pipeline('sarvamai/sarvam-2b-v0.5')
+    except Exception as e:
+        print(f"Error loading Sarvam model: {e}")
+        return None
 
 @spaces.GPU
-def process_audio_input(audio):
-    if whisper_processor is None or whisper_model is None:
-        return "Error: Speech recognition model is not available. Please type your message instead."
-    
+def process_audio(audio_path, whisper_processor, whisper_model):
+    import librosa
     try:
-        audio, sr = librosa.load(audio, sr=16000)
-        input_features = whisper_processor(audio, sampling_rate=sr, return_tensors="pt").input_features.to(whisper_model.device)
+        audio, sr = librosa.load(audio_path, sr=16000)
+        input_features = whisper_processor(audio, sampling_rate=sr, return_tensors="pt").input_features
         predicted_ids = whisper_model.generate(input_features)
         transcription = whisper_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
         return transcription
     except Exception as e:
-        return f"Error processing audio: {str(e)}. Please type your message instead."
+        return f"Error processing audio: {str(e)}"
 
 @spaces.GPU
-def process_image_input(image, text_prompt):
-    if vision_model is None or vision_processor is None:
-        return "Error: Vision model is not available."
-    
+def process_image(image, text_prompt, vision_model, vision_processor):
     try:
-        messages = [
-            {"role": "user", "content": f"{text_prompt}\n<|image_1|>"},
-        ]
-        prompt = vision_processor.tokenizer.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
-        )
-        inputs = vision_processor(prompt, image, return_tensors="pt").to(vision_model.device)
+        messages = [{"role": "user", "content": f"{text_prompt}\n<|image_1|>"}]
+        prompt = vision_processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = vision_processor(prompt, image, return_tensors="pt")
         generate_ids = vision_model.generate(**inputs, max_new_tokens=1000, temperature=0.2, do_sample=True)
         generate_ids = generate_ids[:, inputs['input_ids'].shape[1]:]
         response = vision_processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
@@ -97,10 +68,8 @@ def process_image_input(image, text_prompt):
     except Exception as e:
         return f"Error processing image: {str(e)}"
 
-def generate_response(transcription):
-    if sarvam_pipe is None:
-        return "Error: Text generation model is not available."
-    
+@spaces.GPU
+def generate_response(transcription, sarvam_pipe):
     try:
         response = sarvam_pipe(transcription, max_length=100, num_return_sequences=1)[0]['generated_text']
         return response
@@ -119,17 +88,21 @@ def text_to_speech(text, lang='hi'):
 @spaces.GPU
 def indic_vision_assistant(input_type, audio_input, text_input, image_input):
     try:
+        whisper_processor, whisper_model = load_whisper()
+        vision_model, vision_processor = load_vision_model()
+        sarvam_pipe = load_sarvam()
+
         if input_type == "audio" and audio_input is not None:
-            transcription = process_audio_input(audio_input)
+            transcription = process_audio(audio_input, whisper_processor, whisper_model)
         elif input_type == "text" and text_input:
             transcription = text_input
         elif input_type == "image" and image_input is not None:
             text_prompt = text_input if text_input else "Describe this image in detail."
-            transcription = process_image_input(image_input, text_prompt)
+            transcription = process_image(image_input, text_prompt, vision_model, vision_processor)
         else:
             return "Please provide either audio, text, or image input.", "No input provided.", None
 
-        response = generate_response(transcription)
+        response = generate_response(transcription, sarvam_pipe)
         lang = detect(response)
         audio_response = text_to_speech(response, lang)