File size: 2,831 Bytes
aac5437
95efa40
 
b8462d5
 
95efa40
b8462d5
 
95efa40
8334aa7
b96a262
 
cfadd82
dc74825
b96a262
 
 
dc74825
cfadd82
b8462d5
8334aa7
 
b8462d5
8606aa2
6cd4cbd
8334aa7
40673d5
418a286
b8462d5
95efa40
b96a262
95efa40
40673d5
b8462d5
 
 
 
40673d5
6bf02ba
b8462d5
 
 
 
40673d5
b8462d5
 
40673d5
 
95efa40
b96a262
4d2d71e
cfadd82
 
418a286
8520592
b96a262
95efa40
cfadd82
85f65db
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
54
55
56
57
58
59
60
61
62
import os
import gradio as gr

from haystack.components.generators import HuggingFaceTGIGenerator
from haystack.components.builders.prompt_builder import PromptBuilder
from haystack import Pipeline
from haystack.utils import Secret
from image_captioner import ImageCaptioner

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

* Upload your photo or select one from the examples
* Choose your model
* ✨ Captionate! ✨ 

`mistralai/Mistral-7B-Instruct-v0.2` performs best, but try different models to see how they react to the same prompt.

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

prompt_template = """
You will receive a descriptive text of a photo.
Try to generate a nice Instagram caption with a phrase rhyming with the text. Include emojis in the caption.
                                                                                                                                   
Descriptive text: {{caption}};
Instagram Caption:
"""

hf_api_key = os.environ["HF_API_KEY"]

def generate_caption(image_file_path, model_name):
    image_to_text = ImageCaptioner(
        model_name="nlpconnect/vit-gpt2-image-captioning",
    )
    prompt_builder = PromptBuilder(template=prompt_template)
    generator = HuggingFaceTGIGenerator(model=model_name, token=Secret.from_token(hf_api_key), generation_kwargs={"max_new_tokens":50})
    captioning_pipeline = Pipeline()
    captioning_pipeline.add_component("image_to_text", image_to_text)
    captioning_pipeline.add_component("prompt_builder", prompt_builder)
    captioning_pipeline.add_component("generator", generator)
    
    captioning_pipeline.connect("image_to_text.caption", "prompt_builder.caption")
    captioning_pipeline.connect("prompt_builder", "generator")

    result = captioning_pipeline.run({"image_to_text":{"image_file_path":image_file_path}})
    return result["generator"]["replies"][0]

with gr.Blocks(theme="soft") as demo:
    gr.Markdown(value=description)
    with gr.Row():
        image = gr.Image(type="filepath")
        with gr.Column():
            model_name = gr.Dropdown(["mistralai/Mistral-7B-Instruct-v0.2","OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5", "tiiuae/falcon-7b-instruct", "tiiuae/falcon-7b", "HuggingFaceH4/starchat-beta", "bigscience/bloom", "google/flan-t5-xxl"], value="mistralai/Mistral-7B-Instruct-v0.2", label="Choose your model!")
            gr.Examples(["./whale.png", "./rainbow.jpeg", "./selfie.png"], inputs=image, label="Click on any example") 
    submit_btn = gr.Button("✨ Captionate ✨")
    caption = gr.Textbox(label="Caption", show_copy_button=True)
    submit_btn.click(fn=generate_caption, inputs=[image, model_name], outputs=[caption])

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