Update app.py

app.py

@@ -24,28 +24,28 @@ 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]}};
+Descriptive text: {{caption}};
 Instagram Caption:
 """
 
 hf_api_key = os.environ["HF_API_KEY"]
 
-def generate_caption(image_file_paths, model_name):
+def generate_caption(image_file_path, 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))
+    generator = HuggingFaceTGIGenerator(model=model_name, token=Secret.from_token(hf_api_key), generation_kwargs={"max_new_tokens":50})
     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("image_to_text.caption", "prompt_builder.caption")
     captioning_pipeline.connect("prompt_builder", "generator")
 
-    result = captioning_pipeline.run({"image_to_text":{"image_file_paths":image_file_paths}})
-    return result["generator"][0]
+    result = captioning_pipeline.run({"image_to_text":{"image_file_path":image_file_path}})
+    return result["generator"]["replies"][0]
 
 with gr.Blocks(theme="soft") as demo:
     gr.Markdown(value=description)