Update app.py

app.py

@@ -1,197 +1,84 @@
 import gradio as gr
-import spaces
+from transformers import AutoModelForCausalLM, AutoProcessor, pipeline
 from PIL import Image
-import requests
-from transformers import AutoModelForCausalLM, AutoProcessor
 import torch
-import subprocess
-from io import BytesIO
-import os
+import warnings
 
-# Install flash-attn
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+# Suppress warnings
+warnings.filterwarnings("ignore")
 
-# Load the model and processor
-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,  # Explicitly disable Flash Attention 2
-)
-processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True, num_crops=16)
-
-@spaces.GPU(duration=120)  # Adjust the duration as needed
-def solve_math_problem(image):
-    # Move model to GPU for this function call
-    model.to('cuda')
-    
-    # Prepare the input
-    messages = [
-        {"role": "user", "content": "<|image_1|>\nSolve this math problem step by step. Explain your reasoning clearly."},
-    ]
-    prompt = processor.tokenizer.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
+# Load Phi-3.5-vision model
+phi_model_id = "microsoft/Phi-3.5-vision-instruct"
+try:
+    phi_model = AutoModelForCausalLM.from_pretrained(
+        phi_model_id, 
+        device_map="auto", 
+        trust_remote_code=True,
+        torch_dtype=torch.float16,  # Use float16 to reduce memory usage
+        _attn_implementation="eager"  # Fall back to eager implementation if flash attention is not available
     )
+except ImportError:
+    print("FlashAttention not available, falling back to eager implementation.")
+    phi_model = AutoModelForCausalLM.from_pretrained(
+        phi_model_id, 
+        device_map="auto", 
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        _attn_implementation="eager"
+    )
+
+phi_processor = AutoProcessor.from_pretrained(phi_model_id, trust_remote_code=True)
+
+# Load Llama 3.1 model
+llama_model_id = "meta-llama/Llama-3.1-8B"
+try:
+    llama_pipeline = pipeline("text-generation", model=llama_model_id, device_map="auto", torch_dtype=torch.float16)
+except Exception as e:
+    print(f"Error loading Llama 3.1 model: {e}")
+    print("Falling back to a smaller, open-source model.")
+    llama_model_id = "gpt2"  # Fallback to a smaller, open-source model
+    llama_pipeline = pipeline("text-generation", model=llama_model_id, device_map="auto")
+
+def analyze_image(image, query):
+    prompt = f"<|user|>\n<|image_1|>\n{query}<|end|>\n<|assistant|>\n"
+    inputs = phi_processor(prompt, images=image, return_tensors="pt").to(phi_model.device)
     
-    # Process the input
-    inputs = processor(prompt, image, return_tensors="pt").to("cuda")
-    
-    # Generate the response
-    generation_args = {
-        "max_new_tokens": 1000,
-        "temperature": 0.2,
-        "do_sample": True,
-    }
-    generate_ids = model.generate(**inputs, eos_token_id=processor.tokenizer.eos_token_id, **generation_args)
-    
-    # Decode the response
-    generate_ids = generate_ids[:, inputs['input_ids'].shape[1]:]
-    response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    with torch.no_grad():
+        output = phi_model.generate(**inputs, max_new_tokens=100)
+    return phi_processor.decode(output[0], skip_special_tokens=True)
+
+def generate_text(query, history):
+    context = "\n".join([f"{h[0]}\n{h[1]}" for h in history])
+    prompt = f"{context}\nHuman: {query}\nAI:"
     
-    # Move model back to CPU to free up GPU memory
-    model.to('cpu')
-    return response
+    response = llama_pipeline(prompt, max_new_tokens=100, do_sample=True, temperature=0.7)[0]['generated_text']
+    return response.split("AI:")[-1].strip()
 
-def load_image_from_file(file_path):
-    if os.path.exists(file_path):
-        return Image.open(file_path)
+def chatbot(image, query, history):
+    if image is not None:
+        response = analyze_image(Image.fromarray(image), query)
     else:
-        raise FileNotFoundError(f"Image file not found: {file_path}")
-# Custom CSS
-# Custom CSS
-custom_css = """
-<style>
-    body {
-        font-family: 'Arial', sans-serif;
-        background-color: #f0f3f7;
-        margin: 0;
-        padding: 0;
-    }
-    .container {
-        max-width: 1200px;
-        margin: 0 auto;
-        padding: 20px;
-    }
-    .header {
-        background-color: #2c3e50;
-        color: white;
-        padding: 20px 0;
-        text-align: center;
-    }
-    .header h1 {
-        margin: 0;
-        font-size: 2.5em;
-    }
-    .main-content {
-        display: flex;
-        justify-content: space-between;
-        margin-top: 30px;
-    }
-    .input-section, .output-section {
-        width: 48%;
-        background-color: white;
-        border-radius: 8px;
-        padding: 20px;
-        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
-    }
-    .gr-button {
-        background-color: #27ae60;
-        color: white;
-        border: none;
-        padding: 10px 20px;
-        border-radius: 5px;
-        cursor: pointer;
-        transition: background-color 0.3s;
-    }
-    .gr-button:hover {
-        background-color: #2ecc71;
-    }
-    .examples-section {
-        margin-top: 30px;
-        background-color: white;
-        border-radius: 8px;
-        padding: 20px;
-        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
-    }
-    .examples-section h3 {
-        margin-top: 0;
-        color: #2c3e50;
-    }
-    .footer {
-        text-align: center;
-        margin-top: 30px;
-        color: #7f8c8d;
-    }
-</style>
-"""
-
-# Custom HTML
-custom_html = """
-<div class="container">
-    <div class="header">
-        <h1>AI Math Equation Solver</h1>
-        <p>Upload an image of a math problem, and our AI will solve it step by step!</p>
-    </div>
-    <div class="main-content">
-        <div class="input-section">
-            <h2>Upload Your Math Problem</h2>
-            {input_image}
-            {submit_btn}
-        </div>
-        <div class="output-section">
-            <h2>Solution</h2>
-            {output_text}
-        </div>
-    </div>
-    <div class="examples-section">
-        <h3>Try These Examples</h3>
-        {examples}
-    </div>
-    <div class="footer">
-        <p>Powered by Gradio and AI - Created for educational purposes</p>
-    </div>
-</div>
-"""
+        response = generate_text(query, history)
+    
+    history.append((query, response))
+    return "", history, history
 
-# Create the Gradio interface
-with gr.Blocks(css=custom_css) as iface:
-    gr.HTML("""
-        <div class="header">
-            <h1>AI Math Equation Solver</h1>
-            <p>Upload an image of a math problem, and our AI will solve it step by step!</p>
-        </div>
-    """)
+with gr.Blocks() as demo:
+    gr.Markdown("# Multi-Modal AI Assistant")
     
-    with gr.Row(equal_height=True):
-        with gr.Column():
-            gr.HTML("<h2>Upload Your Math Problem</h2>")
-            input_image = gr.Image(type="pil", label="Upload Math Problem Image")
-            submit_btn = gr.Button("Solve Problem", elem_classes=["gr-button"])
-        
-        with gr.Column():
-            gr.HTML("<h2>Solution</h2>")
-            output_text = gr.Textbox(label="Step-by-step Solution", lines=10)
+    with gr.Row():
+        image_input = gr.Image(type="numpy", label="Upload an image (optional)")
+        chat_history = gr.Chatbot(label="Chat History")
     
-    gr.HTML("<h3>Try These Examples</h3>")
-    examples = gr.Examples(
-        examples=[
-            os.path.join(os.path.dirname(__file__), "eqn1.png"),
-            os.path.join(os.path.dirname(__file__), "eqn2.png")
-        ],
-        inputs=input_image,
-        outputs=output_text,
-        fn=solve_math_problem,
-        cache_examples=True,
-    )
+    query_input = gr.Textbox(label="Ask a question or enter a prompt")
+    submit_button = gr.Button("Submit")
     
-    gr.HTML("""
-        <div class="footer">
-            <p>Powered by Gradio and AI - Created for educational purposes</p>
-        </div>
-    """)
-
-    submit_btn.click(fn=solve_math_problem, inputs=input_image, outputs=output_text)
+    state = gr.State([])
+    
+    submit_button.click(
+        chatbot,
+        inputs=[image_input, query_input, state],
+        outputs=[query_input, chat_history, state]
+    )
 
-# Launch the app
-iface.launch()
+demo.launch()