Spestly commited on
Commit
a05fede
·
verified ·
1 Parent(s): b6b0e21

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +126 -0
app.py ADDED
@@ -0,0 +1,126 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from huggingface_hub import InferenceClient
2
+ import gradio as gr
3
+ import base64
4
+ from PIL import Image
5
+ import io
6
+
7
+ client = InferenceClient(
8
+ provider="cohere",
9
+ api_key="HF_TOKEN", # Will add soon as soon as I get the auth working
10
+ )
11
+
12
+ def image_to_data_url(image_path):
13
+ if image_path is None:
14
+ return None
15
+ with Image.open(image_path) as img:
16
+ buffered = io.BytesIO()
17
+ img.save(buffered, format=img.format)
18
+ img_str = base64.b64encode(buffered.getvalue()).decode()
19
+ return f"data:image/{img.format.lower()};base64,{img_str}"
20
+
21
+ def process_input(image, image_url, prompt, model):
22
+ image_data = None
23
+ if image is not None:
24
+ image_data = image_to_data_url(image)
25
+ elif image_url:
26
+ image_data = image_url
27
+
28
+ if not image_data:
29
+ raise gr.Error("Please provide either an image upload or image URL")
30
+
31
+ messages = [
32
+ {
33
+ "role": "user",
34
+ "content": [
35
+ {"type": "text", "text": prompt},
36
+ {"type": "image_url", "image_url": {"url": image_data}}
37
+ ]
38
+ }
39
+ ]
40
+
41
+ stream = client.chat.completions.create(
42
+ model=model,
43
+ messages=messages,
44
+ max_tokens=512,
45
+ stream=True,
46
+ )
47
+
48
+ full_response = ""
49
+ for chunk in stream:
50
+ content = chunk.choices[0].delta.content or ""
51
+ full_response += content
52
+ yield full_response
53
+
54
+ models = [
55
+ "CohereLabs/aya-vision-32b",
56
+ "CohereLabs/aya-vision-8b",
57
+ ]
58
+
59
+ with gr.Blocks() as demo:
60
+ gr.Markdown("# Cohere Aya Vision model UI")
61
+
62
+ with gr.Row():
63
+ with gr.Column():
64
+ model_choice = gr.Dropdown(
65
+ label="Select Model",
66
+ choices=models,
67
+ value=models[0],
68
+ interactive=True
69
+ )
70
+
71
+ with gr.Tab("Upload Image"):
72
+ image_input = gr.Image(
73
+ label="Upload Image",
74
+ type="filepath",
75
+ sources=["upload"]
76
+ )
77
+ with gr.Tab("Image URL"):
78
+ image_url = gr.Textbox(
79
+ label="Image URL",
80
+ placeholder="Paste image URL here...",
81
+ value=""
82
+ )
83
+
84
+ prompt = gr.Textbox(
85
+ label="Prompt",
86
+ value="Describe this image in one sentence.",
87
+ interactive=True
88
+ )
89
+ submit_btn = gr.Button("Generate", variant="primary")
90
+
91
+ with gr.Column():
92
+ output = gr.Textbox(
93
+ label="Model Response",
94
+ interactive=False,
95
+ lines=10,
96
+ autoscroll=True
97
+ )
98
+
99
+ submit_btn.click(
100
+ fn=process_input,
101
+ inputs=[image_input, image_url, prompt, model_choice],
102
+ outputs=output,
103
+ concurrency_limit=None
104
+ )
105
+
106
+ gr.Examples(
107
+ examples=[
108
+ [
109
+ None,
110
+ "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg",
111
+ "Describe this image in one sentence.",
112
+ models[0]
113
+ ],
114
+ [
115
+ None,
116
+ "https://upload.wikimedia.org/wikipedia/commons/thumb/4/4d/Cat_November_2010-1a.jpg/1200px-Cat_November_2010-1a.jpg",
117
+ "What is the main subject of this image?",
118
+ models[1]
119
+ ]
120
+ ],
121
+ inputs=[image_input, image_url, prompt, model_choice],
122
+ label="Example Inputs"
123
+ )
124
+
125
+ if __name__ == "__main__":
126
+ demo.queue().launch()