Update app.py

app.py

@@ -1,9 +1,16 @@
 import gradio as gr
+# from langchain.llms import OpenAI
+from langchain_openai import OpenAI
+from transformers import pipeline
 from transformers import AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecoderModel 
-import os
 
+import os
 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="distilbert/distilgpt2")
+
 # Load image captioning model
 encoder_checkpoint = "nlpconnect/vit-gpt2-image-captioning"
 decoder_checkpoint = "nlpconnect/vit-gpt2-image-captioning"
@@ -13,6 +20,7 @@ feature_extractor = ViTFeatureExtractor.from_pretrained(encoder_checkpoint)
 tokenizer = AutoTokenizer.from_pretrained(decoder_checkpoint)
 model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint)
 
+
 def generate_story(image, theme, genre):
     try:
         # Preprocess the image
@@ -28,23 +36,23 @@ def generate_story(image, theme, genre):
         # Generate story based on the caption
         story_prompt = f"Write an interesting {theme} story in the {genre} genre. The story should be within 100 words about {caption_text}."
 
-        # Assume OpenAI class/functionality here
-        # llm = OpenAI(model_name="gpt-3.5-turbo-instruct", openai_api_key=openai_api_key)
-        # story = llm.invoke(story_prompt)
-        
-        # Placeholder for story generation
-        story = "Generated story placeholder"
+        llm = OpenAI(model_name="gpt-3.5-turbo-instruct", openai_api_key=openai_api_key)
+        story = llm.invoke(story_prompt)
+        # story = text_generation_model(story_prompt, max_length=150)[0]["generated_text"]
         
         return caption_text, story
     except Exception as e:
         return f"An error occurred during inference: {str(e)}"
 
+
+
 # Gradio interface
-input_image = gr.Image(label="Select Image", type="pil")
+input_image = gr.Image(label="Select Image",type="pil")
 input_theme = gr.Dropdown(["Love and Loss", "Identity and Self-Discovery", "Power and Corruption", "Redemption and Forgiveness", "Survival and Resilience", "Nature and the Environment", "Justice and Injustice", "Friendship and Loyalty", "Hope and Despair"], label="Input Theme")
 input_genre = gr.Dropdown(["Fantasy", "Science Fiction", "Poetry", "Mystery/Thriller", "Romance", "Historical Fiction", "Horror", "Adventure", "Drama", "Comedy"], label="Input Genre")
 output_caption = gr.Textbox(label="Image Caption", lines=2)
-output_text = gr.Textbox(label="Generated Story", lines=8)
+output_text = gr.Textbox(label="Generated Story",lines=8)
+
 
 gr.Interface(
     fn=generate_story,
@@ -52,4 +60,4 @@ gr.Interface(
     outputs=[output_caption, output_text],
     title="Image to Story Generator",
     description="Generate a story from an image taking theme and genre as input. It leverages image captioning and text generation models.",
-).launch()
+).launch()