Spaces:
Running
Running
File size: 2,788 Bytes
aac5437 95efa40 a5b68e6 95efa40 a5b68e6 95efa40 8334aa7 b96a262 cfadd82 dc74825 b96a262 dc74825 cfadd82 a5b68e6 8334aa7 a5b68e6 8606aa2 6cd4cbd 8334aa7 a5b68e6 418a286 a5b68e6 95efa40 b96a262 95efa40 85f65db a5b68e6 6bf02ba a5b68e6 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: {{captions[0]}};
Instagram Caption:
"""
hf_api_key = os.environ["HF_API_KEY"]
def generate_caption(image_file_paths, 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))
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.captions", "prompt_builder.captions")
captioning_pipeline.connect("prompt_builder", "generator")
result = captioning_pipeline.run({"image_to_text":{"image_file_paths":image_file_paths}})
return result["generator"][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() |