init!

app.py

@@ -1,40 +1,52 @@
-# test.py
 import torch
 from PIL import Image
 from transformers import AutoModel, AutoTokenizer
+import gradio as gr
 
-model = AutoModel.from_pretrained('openbmb/MiniCPM-Llama3-V-2_5', trust_remote_code=True, torch_dtype=torch.float16)
-model = model.to(device='cuda')
+# Load a smaller model and tokenizer
+model_name = 'google/vit-base-patch16-224' # Example of a smaller model, adjust as needed
 
-tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-Llama3-V-2_5', trust_remote_code=True)
-model.eval()
+try:
+ model = AutoModel.from_pretrained(model_name, torch_dtype=torch.float16)
+ model = model.to(device='cuda' if torch.cuda.is_available() else 'cpu')
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
+ model.eval()
+except Exception as e:
+ print(f"Error loading model or tokenizer: {e}")
+ exit()
 
-image = Image.open('xx.jpg').convert('RGB')
-question = 'What is in the image?'
-msgs = [{'role': 'user', 'content': question}]
+def process_image(image, question):
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-res = model.chat(
- image=image,
- msgs=msgs,
- tokenizer=tokenizer,
- sampling=True, # if sampling=False, beam_search will be used by default
- temperature=0.7,
- # system_prompt='' # pass system_prompt if needed
-)
-print(res)
+ # Convert Gradio image to PIL Image
+ image = Image.fromarray(image).convert('RGB')
+ 
+ # Create message list
+ msgs = [{'role': 'user', 'content': question}]
+ 
+ # Perform inference
+ try:
+ with torch.no_grad():
+ res = model.chat(
+ image=image,
+ msgs=msgs,
+ tokenizer=tokenizer,
+ sampling=True, # if sampling=False, beam_search will be used by default
+ temperature=0.7,
+ stream=False # Set to False for non-streaming output
+ )
+ return res
+ except Exception as e:
+ return f"Error during model inference: {e}"
 
-## if you want to use streaming, please make sure sampling=True and stream=True
-## the model.chat will return a generator
-res = model.chat(
- image=image,
- msgs=msgs,
- tokenizer=tokenizer,
- sampling=True,
- temperature=0.7,
- stream=True
+# Define the Gradio interface
+interface = gr.Interface(
+ fn=process_image,
+ inputs=[gr.inputs.Image(type='numpy'), gr.inputs.Textbox(label="Question")],
+ outputs="text",
+ title="Image Question Answering",
+ description="Upload an image and ask a question about it. The model will provide an answer."
 )
 
-generated_text = ""
-for new_text in res:
- generated_text += new_text
- print(new_text, flush=True, end='')
+# Launch the Gradio app
+interface.launch()