Update brain.py

brain.py

@@ -1,6 +1,5 @@
 from dotenv import load_dotenv, find_dotenv
 from transformers import pipeline
-from transformers import pipeline
 from langchain_huggingface import HuggingFaceEndpoint
 from langchain_core.prompts import PromptTemplate
 load_dotenv(find_dotenv())
@@ -44,9 +43,10 @@ def generate_story(scenario, repo_id="mistralai/Mistral-7B-Instruct-v0.2"):
     """
     prompt = PromptTemplate.from_template(prompt_template)
     story = prompt | llm
-    return story.invoke(input={"segmentation_results": scenario})
+    return story.invoke(input={"scenario": scenario})
 
 
 if __name__ == "__main__":
     my_story = generate_story(img2txt("image.jpg"))
     print(my_story)
+