Spaces:
Runtime error
Runtime error
Commit
·
86573d5
1
Parent(s):
364461b
Update app.py
Browse files
app.py
CHANGED
@@ -10,20 +10,26 @@ CAPTION_MODELS = {
|
|
10 |
}
|
11 |
|
12 |
# Simple caption creation
|
13 |
-
def caption_image(model_choice, image_input):
|
|
|
|
|
|
|
|
|
|
|
14 |
captioner = pipeline(task="image-to-text",
|
15 |
model=CAPTION_MODELS[model_choice],
|
16 |
max_new_tokens=30,
|
17 |
device_map="cpu", use_fast=True
|
18 |
)
|
19 |
-
caption = captioner(
|
20 |
return str(caption).strip()
|
21 |
|
22 |
-
def launch(model_choice, image_input):
|
23 |
-
return caption_image(model_choice, image_input)
|
24 |
|
25 |
-
model_dropdown = gr.Dropdown(choices=list(CAPTION_MODELS.keys()), label='Model
|
26 |
-
image_input = gr.Image(type="pil", label="Input Image
|
|
|
27 |
|
28 |
-
iface = gr.Interface(launch, inputs=[model_dropdown, image_input], outputs="text")
|
29 |
iface.launch()
|
|
|
10 |
}
|
11 |
|
12 |
# Simple caption creation
|
13 |
+
def caption_image(model_choice, image_input, url_input):
|
14 |
+
if image_input is not None:
|
15 |
+
input_data = image_input
|
16 |
+
else:
|
17 |
+
input_data = url_input
|
18 |
+
|
19 |
captioner = pipeline(task="image-to-text",
|
20 |
model=CAPTION_MODELS[model_choice],
|
21 |
max_new_tokens=30,
|
22 |
device_map="cpu", use_fast=True
|
23 |
)
|
24 |
+
caption = captioner(input_data)[0]['generated_text']
|
25 |
return str(caption).strip()
|
26 |
|
27 |
+
def launch(model_choice, image_input, url_input):
|
28 |
+
return caption_image(model_choice, image_input, url_input)
|
29 |
|
30 |
+
model_dropdown = gr.Dropdown(choices=list(CAPTION_MODELS.keys()), label='Select Caption Model')
|
31 |
+
image_input = gr.Image(type="pil", label="Input Image")
|
32 |
+
url_input = gr.Text(label="Input URL")
|
33 |
|
34 |
+
iface = gr.Interface(launch, inputs=[model_dropdown, image_input, url_input], outputs="text")
|
35 |
iface.launch()
|