Update app.py

app.py

@@ -1,84 +1,157 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoProcessor, pipeline
+import spaces
 from PIL import Image
+import os
 import torch
-import warnings
+from transformers import AutoModelForCausalLM, AutoProcessor
 
-# 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)
 
-# 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"
+@spaces.GPU(duration=120)
+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
     )
-
-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)
     
-    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:"
+    # Process the input
+    inputs = processor(prompt, image, return_tensors="pt").to("cuda")
     
-    response = llama_pipeline(prompt, max_new_tokens=100, do_sample=True, temperature=0.7)[0]['generated_text']
-    return response.split("AI:")[-1].strip()
-
-def chatbot(image, query, history):
-    if image is not None:
-        response = analyze_image(Image.fromarray(image), query)
-    else:
-        response = generate_text(query, history)
+    # 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)
     
-    history.append((query, response))
-    return "", history, history
-
-with gr.Blocks() as demo:
-    gr.Markdown("# Multi-Modal AI Assistant")
-    
-    with gr.Row():
-        image_input = gr.Image(type="numpy", label="Upload an image (optional)")
-        chat_history = gr.Chatbot(label="Chat History")
+    # 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]
     
-    query_input = gr.Textbox(label="Ask a question or enter a prompt")
-    submit_button = gr.Button("Submit")
+    # Move model back to CPU to free up GPU memory
+    model.to('cpu')
+    return response
+
+# 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>
+"""
+
+# 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>
+    """)
     
-    state = gr.State([])
+    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)
     
-    submit_button.click(
-        chatbot,
-        inputs=[image_input, query_input, state],
-        outputs=[query_input, chat_history, state]
+    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,
     )
+    
+    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)
 
-demo.launch()
+# Launch the app
+iface.launch()