Mediocreatmybest commited on
Commit
86573d5
·
1 Parent(s): 364461b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +13 -7
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(image_input)[0]['generated_text']
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 Choice')
26
- image_input = gr.Image(type="pil", label="Input Image or URL")
 
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()