Spaces:
Running
Running
File size: 1,744 Bytes
95efa40 05aaa8f 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 |
import gradio as gr
from haystack.nodes import TransformersImageToText
from haystack.nodes import PromptNode, PromptTemplate
from haystack import Pipeline
image_to_text = TransformersImageToText(
model_name_or_path="nlpconnect/vit-gpt2-image-captioning",
use_gpu=True,
batch_size=16,
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 that has a phrase rhyming with the text
Describing text:{documents};
Caption:
""")
# prompt_node = PromptNode(model_name_or_path="gpt-3.5-turbo", api_key=api_key, default_prompt_template=pt)
prompt_node = PromptNode(model_name_or_path="google/flan-t5-large", default_prompt_template=prompt_template)
# prompt_node = PromptNode(model_name_or_path="tiiuae/falcon-7b-instruct", api_key=hf_api_key, default_prompt_template=pt, 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):
print(image_file_paths)
# documents = image_to_text.generate_captions(image_file_paths=[image_file_paths])
# print(documents[0].content)
caption = captioning_pipeline.run(file_paths=[image_file_paths])
print(caption)
return caption["results"][0]
with gr.Blocks(theme="soft") as demo:
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() |