Update app.py

app.py

@@ -21,7 +21,7 @@ tokenizer = AutoTokenizer.from_pretrained(decoder_checkpoint)
 model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint)
 
 
-def generate_story(image, theme, genre):
+def generate_story(image, theme, genre, word_count):
     try:
         # Preprocess the image
         image = image.convert('RGB')
@@ -34,7 +34,7 @@ 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 within 100 words about {caption_text}."
+        story_prompt = f"Write an interesting {theme} story in the {genre} genre. The story should be within {word_count} words about {caption_text}."
 
         llm = OpenAI(model_name="gpt-3.5-turbo-instruct", openai_api_key=openai_api_key)
         story = llm.invoke(story_prompt)
@@ -51,13 +51,14 @@ 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=10)
 examples = [["example1.jpg"], ["example2.jpg"]]  # List of example image paths
+word_count_slider = gr.Slider(minimum=50, maximum=200, default=100, label="Word Count")
 
 
 gr.Interface(
     fn=generate_story,
-    inputs=[input_image, input_theme, input_genre],
+    inputs=[input_image, input_theme, input_genre, word_count_slider],
     theme='freddyaboulton/dracula_revamped',
     outputs=[output_caption, output_text],
     examples = examples,