Spaces:
Running
Running
Commit
·
95efa40
1
Parent(s):
70c8cdb
Add the initial version
Browse files- app.py +46 -0
- requirements.txt +2 -0
app.py
ADDED
|
@@ -0,0 +1,46 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import gradio as gr
|
| 2 |
+
|
| 3 |
+
from haystack.nodes import TransformersImageToText
|
| 4 |
+
from haystack.nodes import PromptNode, PromptTemplate
|
| 5 |
+
from haystack import Pipeline
|
| 6 |
+
|
| 7 |
+
image_to_text = TransformersImageToText(
|
| 8 |
+
model_name_or_path="nlpconnect/vit-gpt2-image-captioning",
|
| 9 |
+
use_gpu=True,
|
| 10 |
+
batch_size=16,
|
| 11 |
+
progress_bar=True
|
| 12 |
+
)
|
| 13 |
+
|
| 14 |
+
prompt_template = PromptTemplate(prompt="""
|
| 15 |
+
You will receive a describing text of a photo.
|
| 16 |
+
Try to come up with a nice Instagram caption that has a phrase rhyming with the text
|
| 17 |
+
|
| 18 |
+
Describing text:{documents};
|
| 19 |
+
Caption:
|
| 20 |
+
""")
|
| 21 |
+
|
| 22 |
+
# prompt_node = PromptNode(model_name_or_path="gpt-3.5-turbo", api_key=api_key, default_prompt_template=pt)
|
| 23 |
+
|
| 24 |
+
prompt_node = PromptNode(model_name_or_path="google/flan-t5-large", default_prompt_template=prompt_template)
|
| 25 |
+
# 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})
|
| 26 |
+
|
| 27 |
+
captioning_pipeline = Pipeline()
|
| 28 |
+
|
| 29 |
+
captioning_pipeline.add_node(component=image_to_text, name="image_to_text", inputs=["File"])
|
| 30 |
+
captioning_pipeline.add_node(component=prompt_node, name="prompt_node", inputs=["image_to_text"])
|
| 31 |
+
|
| 32 |
+
def generate_caption(image_file_paths):
|
| 33 |
+
print(image_file_paths)
|
| 34 |
+
# documents = image_to_text.generate_captions(image_file_paths=[image_file_paths])
|
| 35 |
+
# print(documents[0].content)
|
| 36 |
+
caption = captioning_pipeline.run(file_paths=image_file_paths)
|
| 37 |
+
return caption["results"][0]
|
| 38 |
+
|
| 39 |
+
with gr.Blocks(theme="soft") as demo:
|
| 40 |
+
image = gr.Image(type="filepath")
|
| 41 |
+
submit_btn = gr.Button("✨ Captionate ✨")
|
| 42 |
+
caption = gr.Textbox(label="Caption")
|
| 43 |
+
submit_btn.click(fn=generate_caption, inputs=[image], outputs=[caption])
|
| 44 |
+
|
| 45 |
+
if __name__ == "__main__":
|
| 46 |
+
demo.launch()
|
requirements.txt
ADDED
|
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
|
|
|
| 1 |
+
farm-haystack==1.19.0
|
| 2 |
+
gradio==3.39.0
|