Update app.py

app.py

@@ -1,15 +1,8 @@
 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
-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")
+openai_api_key = os.getenv("OPENAI_API_KEY")
 
 # Load image captioning model
 encoder_checkpoint = "nlpconnect/vit-gpt2-image-captioning"
@@ -20,7 +13,6 @@ 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
@@ -36,27 +28,28 @@ 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}."
 
-        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"]
+        # 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)
         
-        return story
+        # Placeholder for story generation
+        story = "Generated story placeholder"
+        
+        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_text = gr.Textbox(label="Generated Story",lines=8)
-
+output_caption = gr.Textbox(label="Image Caption", lines=2)
+output_text = gr.Textbox(label="Generated Story", lines=8)
 
 gr.Interface(
     fn=generate_story,
     inputs=[input_image, input_theme, input_genre],
-    outputs=output_text,
+    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()