Spaces:
Running
Running
| 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 pics! | |
| * Upload your photo or select one from examples | |
| * Choose your model | |
| * ✨ Captionate! ✨ | |
| `OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5` and `tiiuae/falcon-7b-instruct` perform the best but try out different models to see how they react to the same prompt. | |
| 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 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: {documents}; | |
| Instagram Caption: | |
| """) | |
| hf_api_key = os.environ["HF_API_KEY"] | |
| def generate_caption(image_file_paths, model_name): | |
| captioning_pipeline = Pipeline() | |
| prompt_node = PromptNode(model_name_or_path=model_name, api_key=hf_api_key, default_prompt_template=prompt_template, model_kwargs={"trust_remote_code":True}) | |
| 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"]) | |
| caption = captioning_pipeline.run(file_paths=[image_file_paths]) | |
| return caption["results"][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-v0.1","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-v0.1", 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() |