Update app.py

app.py

@@ -1,348 +1,124 @@
-# Import spaces first to avoid CUDA initialization issues
-import spaces
-
-# Then import other libraries
+import os
 import torch
-import librosa
-from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration, AutoModelForCausalLM, AutoProcessor
-from gtts import gTTS
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, TextIteratorStreamer, BitsAndBytesConfig
 import gradio as gr
+from threading import Thread
 from PIL import Image
-import os
-import base64
-from io import BytesIO
-
-import io
 import subprocess
-from langdetect import detect
+import spaces  # Add this import
 
-print("Using GPU for operations when available")
-
-# Install flash-attn
+# Install flash-attention
 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
-
-# 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)
-        return processor, model
-    except Exception as e:
-        print(f"Error loading Whisper 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 Phi-3.5-vision-instruct model
-@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
-        )
-        processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True, num_crops=16)
-        return model, processor
-    except Exception as e:
-        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 Phi-3.5-vision-instruct model
+# Constants
+TITLE = "<h1><center>Phi 3.5 Multimodal (Text + Vision)</center></h1>"
+DESCRIPTION = "# Phi-3.5 Multimodal Demo (Text + Vision)"
+
+# Model configurations
+TEXT_MODEL_ID = "microsoft/Phi-3.5-mini-instruct"
+VISION_MODEL_ID = "microsoft/Phi-3.5-vision-instruct"
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Quantization config for text model
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4"
+)
+
+# Load models and tokenizers
+text_tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_ID)
+text_model = AutoModelForCausalLM.from_pretrained(
+    TEXT_MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    quantization_config=quantization_config
+)
+
+vision_model = AutoModelForCausalLM.from_pretrained(
+    VISION_MODEL_ID, 
+    trust_remote_code=True, 
+    torch_dtype="auto", 
+    attn_implementation="flash_attention_2"
+).to(device).eval()
+
+vision_processor = AutoProcessor.from_pretrained(VISION_MODEL_ID, trust_remote_code=True)
+
+# Helper functions
 @spaces.GPU
-def load_vision_model():
-    try:
-        print("Starting to load vision model...")
-        model_id = "microsoft/Phi-3.5-vision-instruct"
-        print(f"Loading model from {model_id}")
-        
-        # Check for CUDA availability
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"Using device: {device}")
-        
-        # Load model with potential memory optimization
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id, 
-            trust_remote_code=True, 
-            torch_dtype=torch.float16, 
-            use_flash_attention_2=True,  # Enable if supported
-            device_map="auto",  # Automatically manage model placement
-            low_cpu_mem_usage=True
-        )
-        print("Model loaded successfully")
-        
-        print("Loading processor...")
-        processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True, num_crops=16)
-        print("Processor loaded successfully")
-        
-        return model, processor
-    except ImportError as e:
-        print(f"Error importing required modules: {str(e)}")
-        print("Please ensure all required dependencies are installed.")
-    except RuntimeError as e:
-        print(f"Runtime error (possibly CUDA out of memory): {str(e)}")
-        print("Consider using a smaller model or enabling GPU offloading.")
-    except Exception as e:
-        print(f"Unexpected error in loading vision model: {str(e)}")
-    
-    return None, None
+def stream_text_chat(message, history, system_prompt, temperature=0.8, max_new_tokens=1024, top_p=1.0, top_k=20):
+    conversation = [{"role": "system", "content": system_prompt}]
+    for prompt, answer in history:
+        conversation.extend([
+            {"role": "user", "content": prompt},
+            {"role": "assistant", "content": answer},
+        ])
+    conversation.append({"role": "user", "content": message})
+
+    input_ids = text_tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(text_model.device)
+    streamer = TextIteratorStreamer(text_tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+
+    generate_kwargs = dict(
+        input_ids=input_ids,
+        max_new_tokens=max_new_tokens,
+        do_sample=temperature > 0,
+        top_p=top_p,
+        top_k=top_k,
+        temperature=temperature,
+        eos_token_id=[128001, 128008, 128009],
+        streamer=streamer,
+    )
 
+    with torch.no_grad():
+        thread = Thread(target=text_model.generate, kwargs=generate_kwargs)
+        thread.start()
 
-# Process audio input within a GPU-decorated function
-@spaces.GPU
-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."
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text
+        yield history + [[message, buffer]]
 
-# Updated process_image_input function
-@spaces.GPU
-@spaces.GPU
-def process_image_input(image, text_prompt, vision_model, processor):
-    if vision_model is None or processor is None:
-        return "Error: Vision model is not available."
+@spaces.GPU  # Add this decorator
+def process_vision_query(image, text_input):
+    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(device)
     
-    try:
-        # Convert image to base64
-        if isinstance(image, Image.Image):
-            buffered = BytesIO()
-            image.save(buffered, format="PNG")
-            img_str = base64.b64encode(buffered.getvalue()).decode()
-        else:
-            # If it's not a PIL Image, assume it's a file path
-            with open(image, "rb") as image_file:
-                img_str = base64.b64encode(image_file.read()).decode()
-        
-        # Format the input with image tag
-        formatted_prompt = f"{text_prompt}\n<image>data:image/png;base64,{img_str}</image>"
-        
-        # Process the formatted prompt
-        inputs = processor(text=formatted_prompt, return_tensors="pt").to(vision_model.device)
-        
-        # Generate text
-        with torch.no_grad():
-            outputs = vision_model.generate(
-                **inputs,
-                max_new_tokens=100,
-                do_sample=True,
-                top_k=50,
-                top_p=0.95,
-                num_return_sequences=1
-            )
-        
-        generated_text = processor.batch_decode(outputs, skip_special_tokens=True)[0]
-        return generated_text
-    except Exception as e:
-        return f"Error processing image: {str(e)}"
-
-# Generate response within a GPU-decorated function
-@spaces.GPU
-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)}"
-
-# Text-to-speech function
-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']:
-            # You might want to use a different TTS library here
-            # For example, you could use the Google Cloud Text-to-Speech API
-            # or a specialized Indic language TTS library
-            
-            # This is a placeholder for a better Indic TTS solution
-            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
-
-# Improved language detection function
-def detect_language(text):
-    lang_codes = {
-        'bn': 'Bengali', 'gu': 'Gujarati', 'hi': 'Hindi', 'kn': 'Kannada',
-        'ml': 'Malayalam', 'mr': 'Marathi', 'or': 'Oriya', 'pa': 'Punjabi',
-        'ta': 'Tamil', 'te': 'Telugu', 'en': 'English'
-    }
+    with torch.no_grad():
+        generate_ids = vision_model.generate(
+            **inputs, 
+            max_new_tokens=1000, 
+            eos_token_id=vision_processor.tokenizer.eos_token_id
+        )
     
-    try:
-        detected_lang = detect(text)
-        return detected_lang if detected_lang in lang_codes else 'en'
-    except:
-        # Fallback to simple script-based detection
-        for code, lang in lang_codes.items():
-            if any(ord(char) >= 0x0900 and ord(char) <= 0x097F for char in text):  # Devanagari script
-                return 'hi'
-        return 'en'  # Default to English if no Indic script is detected
-
-@spaces.GPU
-def indic_vision_assistant(input_type, audio_input, text_input, image_input):
-    try:
-        whisper_processor, whisper_model = load_whisper()
-        sarvam_pipe = load_sarvam()
-        vision_model, 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:
-            # Use a default prompt if no text input is provided
-            text_prompt = text_input if text_input else "Describe this image in detail."
-            transcription = process_image_input(image_input, text_prompt, vision_model, processor)
-        else:
-            return "Please provide either audio, text, or image input.", "No input provided.", None
-
-        response = generate_response(transcription, sarvam_pipe)
-        lang = detect_language(response)
-        audio_response = text_to_speech(response, lang)
-        
-        return transcription, response, audio_response
-    except Exception as e:
-        error_message = f"An error occurred: {str(e)}"
-        return error_message, error_message, None
-
+    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
 
 # Custom CSS
 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 {@spaces.GPU
-def indic_vision_assistant(input_type, audio_input, text_input, image_input):
-    try:
-        whisper_processor, whisper_model = load_whisper()
-        sarvam_pipe = load_sarvam()
-        vision_model, 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:
-            # Use a default prompt if no text input is provided
-            text_prompt = text_input if text_input else "Describe this image in detail."
-            transcription = process_image_input(image_input, text_prompt, vision_model, processor)
-        else:
-            return "Please provide either audio, text, or image input.", "No input provided.", None
-
-        response = generate_response(transcription, sarvam_pipe)
-        lang = detect_language(response)
-        audio_response = text_to_speech(response, lang)
-        
-        return transcription, response, audio_response
-    except Exception as e:
-        error_message = f"An error occurred: {str(e)}"
-        return error_message, error_message, None
-
-    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;
-}
+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
 custom_header = """
 <div id="custom-header">
-    <h1>
-        <span class="blue">Hello,</span>
-        <span class="pink">User</span>
-    </h1>
-    <h2>How can I help you today?</h2>
+    <h1><span class="blue">Phi 3.5</span> <span class="pink">Multimodal Assistant</span></h1>
+    <h2>Text and Vision AI at Your Service</h2>
 </div>
 """
 
@@ -350,28 +126,25 @@ custom_header = """
 custom_suggestions = """
 <div class="suggestions">
     <div class="suggestion">
-        <span class="suggestion-icon">🎤</span>
-        <p>Speak in any Indic language</p>
-    </div>
-    <div class="suggestion">
-        <span class="suggestion-icon">⌨️</span>
-        <p>Type in any Indic language</p>
+        <span class="suggestion-icon">💬</span>
+        <p>Chat with the Text Model</p>
     </div>
     <div class="suggestion">
         <span class="suggestion-icon">🖼️</span>
-        <p>Upload an image for analysis</p>
+        <p>Analyze Images with Vision Model</p>
     </div>
     <div class="suggestion">
         <span class="suggestion-icon">🤖</span>
         <p>Get AI-generated responses</p>
     </div>
     <div class="suggestion">
-        <span class="suggestion-icon">🔊</span>
-        <p>Listen to audio responses</p>
+        <span class="suggestion-icon">🔍</span>
+        <p>Explore advanced options</p>
     </div>
 </div>
 """
-# Update the Gradio interface to allow text input for image processing
+
+# Gradio interface
 with gr.Blocks(css=custom_css, theme=gr.themes.Base().set(
     body_background_fill="#0b0f19",
     body_text_color="#e2e8f0",
@@ -380,30 +153,38 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Base().set(
     button_primary_text_color="white",
     block_title_text_color="#94a3b8",
     block_label_text_color="#94a3b8",
-)) as iface:
+)) as demo:
     gr.HTML(custom_header)
     gr.HTML(custom_suggestions)
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### Indic Vision 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 prompt")
-    image_input = gr.Image(type="pil", label="Upload an image (if image input selected)")
-    
-    submit_btn = gr.Button("Submit")
-    
-    output_transcription = gr.Textbox(label="Transcription/Input")
-    output_response = gr.Textbox(label="Generated Response")
-    output_audio = gr.Audio(label="Audio Response")
-    
-    submit_btn.click(
-        fn=indic_vision_assistant,
-        inputs=[input_type, audio_input, text_input, image_input],
-        outputs=[output_transcription, output_response, output_audio]
-    )
-    gr.HTML("<footer>Powered by Indic Language AI with Vision Capabilities</footer>")
-# Launch the app
-iface.launch()
+
+    with gr.Tab("Text Model (Phi-3.5-mini)"):
+        chatbot = gr.Chatbot(height=400)
+        msg = gr.Textbox(label="Message", placeholder="Type your message here...")
+        with gr.Accordion("Advanced Options", open=False):
+            system_prompt = gr.Textbox(value="You are a helpful assistant", label="System Prompt")
+            temperature = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.8, label="Temperature")
+            max_new_tokens = gr.Slider(minimum=128, maximum=8192, step=1, value=1024, label="Max new tokens")
+            top_p = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=1.0, label="top_p")
+            top_k = gr.Slider(minimum=1, maximum=20, step=1, value=20, label="top_k")
+        
+        submit_btn = gr.Button("Submit", variant="primary")
+        clear_btn = gr.Button("Clear Chat", variant="secondary")
+
+        submit_btn.click(stream_text_chat, [msg, chatbot, system_prompt, temperature, max_new_tokens, top_p, top_k], [chatbot])
+        clear_btn.click(lambda: None, None, chatbot, queue=False)
+
+    with gr.Tab("Vision Model (Phi-3.5-vision)"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                vision_input_img = gr.Image(label="Upload an Image", type="pil")
+                vision_text_input = gr.Textbox(label="Ask a question about the image", placeholder="What do you see in this image?")
+                vision_submit_btn = gr.Button("Analyze Image", variant="primary")
+            with gr.Column(scale=1):
+                vision_output_text = gr.Textbox(label="AI Analysis", lines=10)
+        
+        vision_submit_btn.click(process_vision_query, [vision_input_img, vision_text_input], [vision_output_text])
+
+    gr.HTML("<footer>Powered by Phi 3.5 Multimodal AI</footer>")
+
+if __name__ == "__main__":
+    demo.launch()