Update app.py

app.py

@@ -63,12 +63,12 @@ def generate_story(image, theme, genre, word_count):
         
         
         # Generate story based on the caption
-        story_prompt = f"Write an interesting {theme} story in the {genre} genre. The story should be within {word_count} words about {caption_text}."
+        story_prompt = f"Write an interesting {theme} story in the {genre} genre about {caption_text}. The story should be within {word_count} words."
 
         llm = OpenAI(model_name="gpt-3.5-turbo-instruct", openai_api_key=openai_api_key, max_tokens=1000)
         story = llm.invoke(story_prompt)
         
-        return story_prompt, story
+        return caption_text, story
     except Exception as e:
         return f"An error occurred during inference: {str(e)}"
     
@@ -77,15 +77,21 @@ def generate_story(image, theme, genre, word_count):
 
 '''
 from transformers import pipeline, AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecoderModel 
+
 # Load text generation model
+
 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"
 model_checkpoint = "nlpconnect/vit-gpt2-image-captioning"
+
 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, word_count):
     try:
         # Preprocess the image
@@ -103,6 +109,7 @@ def generate_story(image, theme, genre, word_count):
         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)}"
 '''