Update app.py

app.py

@@ -1,13 +1,13 @@
 import gradio as gr
+from langchain_openai import OpenAI
 from transformers import pipeline
 from transformers import AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecoderModel 
 
 import os
-api_key = os.getenv("OPENAI_API_KEY")
-client = OpenAI(api_key=api_key)
+openai_api_key = os.getenv("OPENAI_API_KEY")
 
 # Load text generation model
-text_generation_model = pipeline("text-generation", model="openai-community/gpt2-large")
+# text_generation_model = pipeline("text-generation", model="openai-community/gpt2-large")
 # text_generation_model = pipeline("text-generation", model="distilbert/distilgpt2")
 
 # Load image captioning model
@@ -33,9 +33,10 @@ def generate_story(image, theme, genre):
         caption_text = tokenizer.batch_decode(caption_ids, skip_special_tokens=True)[0]
         
         # Generate story based on the caption
-        story_prompt = f"Write an interesting {theme} story in the {genre} genre. The story should be about {caption_text}."
+        story_prompt = f"Write an interesting {theme} story in the {genre} genre. The story should be within 100 words about {caption_text}."
 
-        story = text_generation_model(story_prompt, max_length=150)[0]["generated_text"]
+        llm = OpenAI(model_name="gpt-3.5-turbo-instruct", openai_api_key=openai_api_key)
+        story = llm.invoke(story_prompt)
         
         return story
     except Exception as e: