File size: 1,892 Bytes
aac5437
95efa40
 
 
 
 
 
8334aa7
 
 
 
afc9aca
8334aa7
 
95efa40
 
 
 
 
 
 
8334aa7
95efa40
8334aa7
 
 
 
 
95efa40
 
 
 
aac5437
 
95efa40
 
 
 
 
 
05aaa8f
 
95efa40
 
 
4d2d71e
95efa40
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
import os
import gradio as gr

from haystack.nodes import TransformersImageToText
from haystack.nodes import PromptNode, PromptTemplate
from haystack import Pipeline

description = """
# Captionate ✨ 📸
## Create Instagram captions for your insta pics! 

Built by [Bilge Yucel](https://twitter.com/bilgeycl) using [Haystack](https://github.com/deepset-ai/haystack)💙
"""

image_to_text = TransformersImageToText(
  model_name_or_path="nlpconnect/vit-gpt2-image-captioning",
  progress_bar=True
)

prompt_template = PromptTemplate(prompt="""
You will receive a describing text of a photo.
Try to come up with a nice Instagram caption. 

Requirements for the caption: 
* Must rhyme with the describing text
* Should be at least 10 words
* Needs to include one emoji and suitable hastags
                                                                                                                                   
Describing text:{documents};
Caption:
""")

hf_api_key = os.environ["HF_API_KEY"]
prompt_node = PromptNode(model_name_or_path="tiiuae/falcon-7b-instruct", api_key=hf_api_key, default_prompt_template=prompt_template, model_kwargs={"trust_remote_code":True})

captioning_pipeline = Pipeline()
captioning_pipeline.add_node(component=image_to_text, name="image_to_text", inputs=["File"])
captioning_pipeline.add_node(component=prompt_node, name="prompt_node", inputs=["image_to_text"])

def generate_caption(image_file_paths):
    caption = captioning_pipeline.run(file_paths=[image_file_paths])
    print(caption)
    return caption["results"][0]

with gr.Blocks(theme="soft") as demo:
    gr.Markdown(value=description)
    image = gr.Image(type="filepath")
    submit_btn = gr.Button("✨ Captionate ✨")
    caption = gr.Textbox(label="Caption")
    submit_btn.click(fn=generate_caption, inputs=[image], outputs=[caption])

if __name__ == "__main__":
    demo.launch()