bilgeyucel commited on
Commit
40673d5
1 Parent(s): 4397b1e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -6
app.py CHANGED
@@ -24,28 +24,28 @@ prompt_template = """
24
  You will receive a descriptive text of a photo.
25
  Try to generate a nice Instagram caption with a phrase rhyming with the text. Include emojis in the caption.
26
 
27
- Descriptive text: {{captions[0]}};
28
  Instagram Caption:
29
  """
30
 
31
  hf_api_key = os.environ["HF_API_KEY"]
32
 
33
- def generate_caption(image_file_paths, model_name):
34
  image_to_text = ImageCaptioner(
35
  model_name="nlpconnect/vit-gpt2-image-captioning",
36
  )
37
  prompt_builder = PromptBuilder(template=prompt_template)
38
- generator = HuggingFaceTGIGenerator(model=model_name, token=Secret.from_token(hf_api_key))
39
  captioning_pipeline = Pipeline()
40
  captioning_pipeline.add_component("image_to_text", image_to_text)
41
  captioning_pipeline.add_component("prompt_builder", prompt_builder)
42
  captioning_pipeline.add_component("generator", generator)
43
 
44
- captioning_pipeline.connect("image_to_text.captions", "prompt_builder.captions")
45
  captioning_pipeline.connect("prompt_builder", "generator")
46
 
47
- result = captioning_pipeline.run({"image_to_text":{"image_file_paths":image_file_paths}})
48
- return result["generator"][0]
49
 
50
  with gr.Blocks(theme="soft") as demo:
51
  gr.Markdown(value=description)
 
24
  You will receive a descriptive text of a photo.
25
  Try to generate a nice Instagram caption with a phrase rhyming with the text. Include emojis in the caption.
26
 
27
+ Descriptive text: {{caption}};
28
  Instagram Caption:
29
  """
30
 
31
  hf_api_key = os.environ["HF_API_KEY"]
32
 
33
+ def generate_caption(image_file_path, model_name):
34
  image_to_text = ImageCaptioner(
35
  model_name="nlpconnect/vit-gpt2-image-captioning",
36
  )
37
  prompt_builder = PromptBuilder(template=prompt_template)
38
+ generator = HuggingFaceTGIGenerator(model=model_name, token=Secret.from_token(hf_api_key), generation_kwargs={"max_new_tokens":50})
39
  captioning_pipeline = Pipeline()
40
  captioning_pipeline.add_component("image_to_text", image_to_text)
41
  captioning_pipeline.add_component("prompt_builder", prompt_builder)
42
  captioning_pipeline.add_component("generator", generator)
43
 
44
+ captioning_pipeline.connect("image_to_text.caption", "prompt_builder.caption")
45
  captioning_pipeline.connect("prompt_builder", "generator")
46
 
47
+ result = captioning_pipeline.run({"image_to_text":{"image_file_path":image_file_path}})
48
+ return result["generator"]["replies"][0]
49
 
50
  with gr.Blocks(theme="soft") as demo:
51
  gr.Markdown(value=description)