Update app.py

app.py

@@ -1,9 +1,11 @@
 import os
 import gradio as gr
 
-from haystack.nodes import TransformersImageToText
-from haystack.nodes import PromptNode, PromptTemplate
+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 📸
@@ -15,31 +17,35 @@ description = """
 
 `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](https://github.com/deepset-ai/haystack) 💙
+Built by [Bilge Yucel](https://twitter.com/bilgeycl) using [Haystack 2.0](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="""
+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: {documents};
+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()
-    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]
+    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)