Update app.py

app.py

@@ -16,8 +16,7 @@ tokenizer = AutoTokenizer.from_pretrained(decoder_checkpoint)
 model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint)
 
 
-
-def generate_story(image, genre, style):
+def generate_story(image, theme, genre, style):
     try:
         # Preprocess the image
         image = image.convert('RGB')
@@ -30,7 +29,8 @@ def generate_story(image, genre, style):
         caption_text = tokenizer.batch_decode(caption_ids, skip_special_tokens=True)[0]
         
         # Generate story based on the caption
-        story_prompt = f"{genre.capitalize()} story in {style.capitalize()} style about: {caption_text}."
+        story_prompt = f"Write an interesting {theme} story in the {genre} genre, written in {style} style. The story should be about {caption_text}."
+
         story = text_generation_model(story_prompt, max_length=150)[0]["generated_text"]
         
         return story
@@ -41,14 +41,15 @@ def generate_story(image, genre, style):
 
 # Gradio interface
 input_image = gr.Image(label="Select Image",type="pil")
-input_genre = gr.Dropdown(["Hindi", "Spanish", "Portuguese", "French", "German", "Italian", "Russian", "Japanese"], label="Input Genre")
-input_style = gr.Dropdown(["Hindi", "Spanish", "Portuguese", "French", "German", "Italian", "Russian", "Japanese"], label="Input Style")
+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", "Mystery/Thriller", "Romance", "Historical Fiction", "Horror", "Adventure", "Drama", "Comedy"], label="Input Genre")
+input_style = gr.Dropdown(["Classic Literature", "Contemporary Fiction", "Poetic/Prose", "Experimental", "Dialogue-driven", "Descriptive", "Epistolary", "Minimalist", "Surreal/Fantastical"], label="Input Style")
 output_text = gr.Textbox(label="Generated Story",lines=8)
 
 
 gr.Interface(
     fn=generate_story,
-    inputs=[input_image, input_genre, input_style],
+    inputs=[input_image, input_theme, input_genre, input_style],
     outputs=output_text,
     title="Image to Story Generator",
     description="Generate a story from an image taking genre and style as input.",