import gradio as gr
import spaces
from PIL import Image
import requests
from transformers import AutoModelForCausalLM, AutoProcessor
import torch
import subprocess
from io import BytesIO
# Install flash-attn
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
# 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
)
# 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]
# Move model back to CPU to free up GPU memory
model.to('cpu')
return response
# Function to load image from URL
def load_image_from_url(url):
response = requests.get(url)
img = Image.open(BytesIO(response.content))
return img
# Custom CSS
custom_css = """
"""
# Custom HTML
custom_html = """
"""
# Create the Gradio interface
with gr.Blocks(css=custom_css) as iface:
gr.HTML(custom_html)
with gr.Row():
with gr.Column(scale=1):
input_image = gr.Image(type="pil", label="Upload Math Problem Image")
submit_btn = gr.Button("Solve Problem")
with gr.Column(scale=1):
output_text = gr.Textbox(label="Solution", lines=10)
submit_btn.click(fn=solve_math_problem, inputs=input_image, outputs=output_text)
gr.Examples(
examples=[
"https://i.imgur.com/2Gwd3bN.jpg", # Replace with actual URLs of math problem images
"https://i.imgur.com/wPw5YtB.jpg"
],
inputs=input_image,
outputs=output_text,
fn=lambda url: solve_math_problem(load_image_from_url(url)),
cache_examples=True,
)
# Launch the app
iface.launch()