import torch import librosa from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration, AutoModelForCausalLM, AutoProcessor from gtts import gTTS import gradio as gr from PIL import Image import subprocess print("Using GPU for operations when available") # 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 def load_pipeline(model_name, **kwargs): try: device = "cuda" 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 # Load Whisper model for speech recognition def load_whisper(): try: device = "cuda" if torch.cuda.is_available() else "cpu" processor = WhisperProcessor.from_pretrained("openai/whisper-small") model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small").to(device) return processor, model except Exception as e: print(f"Error loading Whisper model: {e}") return None, None # Load sarvam-2b for text generation def load_sarvam(): return load_pipeline('sarvamai/sarvam-2b-v0.5') # Load vision model def load_vision_model(): model_id = "microsoft/Phi-3.5-vision-instruct" device = "cuda" if torch.cuda.is_available() else "cpu" model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, torch_dtype="auto", attn_implementation="flash_attention_2").to(device).eval() processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) return model, processor # Process audio input def process_audio_input(audio, whisper_processor, whisper_model): if whisper_processor is None or whisper_model is None: return "Error: Speech recognition model is not available. Please type your message instead." 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) 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." # Generate response def text_to_speech(text, lang='hi'): try: # Use a better TTS engine for Indic languages if lang in ['hi', 'bn', 'gu', 'kn', 'ml', 'mr', 'or', 'pa', 'ta', 'te']: tts = gTTS(text=text, lang=lang, tld='co.in') # Use Indian TLD else: tts = gTTS(text=text, lang=lang) tts.save("response.mp3") return "response.mp3" except Exception as e: print(f"Error in text-to-speech: {str(e)}") return None # Detect language (placeholder function, replace with actual implementation) def detect_language(text): # Implement language detection logic here return 'en' # Default to English for now def generate_response(transcription, sarvam_pipe): if sarvam_pipe is None: return "Error: Text generation model is not available." try: # Generate response using the sarvam-2b model response = sarvam_pipe(transcription, max_length=100, num_return_sequences=1)[0]['generated_text'] return response except Exception as e: return f"Error generating response: {str(e)}" def process_image(image, text_input, vision_model, vision_processor): try: prompt = f"<|user|>\n<|image_1|>\n{text_input}<|end|>\n<|assistant|>\n" image = Image.fromarray(image).convert("RGB") inputs = vision_processor(prompt, image, return_tensors="pt").to(vision_model.device) generate_ids = vision_model.generate(**inputs, max_new_tokens=1000, eos_token_id=vision_processor.tokenizer.eos_token_id) 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] return response except Exception as e: return f"Error processing image: {str(e)}" def multimodal_assistant(input_type, audio_input, text_input, image_input): try: # Load models whisper_processor, whisper_model = load_whisper() sarvam_pipe = load_sarvam() vision_model, vision_processor = load_vision_model() if input_type == "audio" and audio_input is not None: transcription = process_audio_input(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: return process_image(image_input, text_input, vision_model, vision_processor), None else: return "Please provide either audio, text, or image input.", None response = generate_response(transcription, sarvam_pipe) lang = detect_language(response) audio_response = text_to_speech(response, lang) return response, audio_response except Exception as e: error_message = f"An error occurred: {str(e)}" return error_message, None # Custom CSS (you can keep your existing custom CSS here) custom_css = """ body { background-color: #0b0f19; color: #e2e8f0; font-family: 'Arial', sans-serif; } #custom-header { text-align: center; padding: 20px 0; background-color: #1a202c; margin-bottom: 20px; border-radius: 10px; } #custom-header h1 { font-size: 2.5rem; margin-bottom: 0.5rem; } #custom-header h1 .blue { color: #60a5fa; } #custom-header h1 .pink { color: #f472b6; } #custom-header h2 { font-size: 1.5rem; color: #94a3b8; } .suggestions { display: flex; justify-content: center; flex-wrap: wrap; gap: 1rem; margin: 20px 0; } .suggestion { background-color: #1e293b; border-radius: 0.5rem; padding: 1rem; display: flex; align-items: center; transition: transform 0.3s ease; width: 200px; } .suggestion:hover { transform: translateY(-5px); } .suggestion-icon { font-size: 1.5rem; margin-right: 1rem; background-color: #2d3748; padding: 0.5rem; border-radius: 50%; } .gradio-container { max-width: 100% !important; } #component-0, #component-1, #component-2 { max-width: 100% !important; } footer { text-align: center; margin-top: 2rem; color: #64748b; } """ # Custom HTML for the header (you can keep your existing custom header here) custom_header = """

Multimodal Indic Assistant

How can I help you today?

""" # Custom HTML for suggestions custom_suggestions = """
🎤

Speak in any Indic language

⌨️

Type in any Indic language

📷

Upload an image for analysis

🤖

Get AI-generated responses

🔊

Listen to audio responses

""" # Create Gradio interface with gr.Blocks(css=custom_css, theme=gr.themes.Base().set( body_background_fill="#0b0f19", body_text_color="#e2e8f0", button_primary_background_fill="#3b82f6", button_primary_background_fill_hover="#2563eb", button_primary_text_color="white", block_title_text_color="#94a3b8", block_label_text_color="#94a3b8", )) as iface: gr.HTML(custom_header) gr.HTML(custom_suggestions) with gr.Row(): with gr.Column(scale=1): gr.Markdown("### Multimodal Indic Assistant") input_type = gr.Radio(["audio", "text", "image"], label="Input Type", value="audio") audio_input = gr.Audio(type="filepath", label="Speak (if audio input selected)") text_input = gr.Textbox(label="Type your message or image question") image_input = gr.Image(label="Upload an image (if image input selected)") submit_btn = gr.Button("Submit") output_response = gr.Textbox(label="Generated Response") output_audio = gr.Audio(label="Audio Response") submit_btn.click( fn=multimodal_assistant, inputs=[input_type, audio_input, text_input, image_input], outputs=[output_response, output_audio] ) gr.HTML("") # Launch the app iface.launch()