testing / app.py
ArnavLatiyan's picture
Update app.py
eaee34d verified
raw
history blame contribute delete
833 Bytes
from transformers import AutoProcessor
from unsloth import UnslothModel
from PIL import Image
import gradio as gr
import torch
model, tokenizer = UnslothModel.from_pretrained(
model_name="unsloth/Llama-3.2-11B-Vision-bnb-4bit",
adapter_path="ArnavLatiyan/my-lora-leaf-model",
load_in_4bit=True
)
processor = AutoProcessor.from_pretrained("unsloth/llama-3.2-vision-11b")
def predict(image):
inputs = processor(images=image, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
prompt = tokenizer("What disease is this?", return_tensors="pt").to(inputs["pixel_values"].device)
inputs.update(prompt)
outputs = model.generate(**inputs, max_new_tokens=50)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
gr.Interface(fn=predict, inputs="image", outputs="text").launch()