Khadaffe Abubakar Sulaiman commited on
Commit
efb087f
Β·
1 Parent(s): b1a2052

feat: add llava script

Browse files
Files changed (2) hide show
  1. app.py +27 -0
  2. requirements.txt +5 -0
app.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoProcessor, LlavaForConditionalGeneration
3
+ import torch
4
+ from PIL import Image
5
+
6
+ # Load model and processor
7
+ model_id = "llava-hf/llava-1.5-7b-hf"
8
+ processor = AutoProcessor.from_pretrained(model_id)
9
+ model = LlavaForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
10
+
11
+ # Inference function
12
+ def chat_with_image(image, prompt):
13
+ inputs = processor(prompt, image, return_tensors="pt").to("cuda", torch.float16)
14
+ output = model.generate(**inputs, max_new_tokens=100)
15
+ answer = processor.decode(output[0], skip_special_tokens=True)
16
+ return answer
17
+
18
+ # Gradio interface
19
+ demo = gr.Interface(
20
+ fn=chat_with_image,
21
+ inputs=[gr.Image(type="pil"), gr.Textbox(label="Prompt")],
22
+ outputs=gr.Textbox(label="Answer"),
23
+ title="LLaVA 1.5 - Image Chatbot",
24
+ description="Ask questions about an image using LLaVA 1.5 (7B)"
25
+ )
26
+
27
+ demo.launch()
requirements.txt ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ torch
2
+ transformers
3
+ gradio
4
+ accelerate
5
+ Pillow