Update app.py

app.py

@@ -4,17 +4,17 @@ from PIL import Image
 from io import BytesIO
 from gradio_client import Client
 import json
-import os
-def generate(outline, characters, settings ):
-    prompt = f"""Hello! I would like to request a 4-paragraph and 700-word per paragraph story and a cover image prompt for sd3 in JSON format described later in the prompt with the following detailed outline:\n\n{outline}\n\nCharacters: {characters}\n\nSettings: {settings}\n\nPlease generate the story with the following detailed JSON format: p1, p2, p3, p4: Keys for story paragraphs; title: Key for story title; prompt: "Mother Nature, a gentle figure surrounded by flowers, stands over a steaming cauldron in a vibrant forest, stirring its contents with a wooden spoon". Please do not include any other text in the output. Thank you. Only the JSON is needed or it will break the whole system and make us lose 10 million dollars. Please don't say 'Full response: Here is the requested output in JSON format:' or 'Here is the full response.' Only JSON. If you give plain text, it will not work and count as an error and we will lose customers. Please do not give text. You are not ChatGPT. Don't say 'Here is the full JSON.' You are not an assistant; you are used by an AI. Thank you.\n\n"""
+
+def generate(outline, characters, settings, num_images):
+    # Create a dynamic prompt with multiple image prompts based on the number of images
+    prompt = f"""Hello! I would like to request a 4-paragraph story (700 words per paragraph) and a corresponding image prompt for each paragraph in JSON format described later in the prompt with the following detailed outline:\n\n{outline}\n\nCharacters: {characters}\n\nSettings: {settings}\n\nPlease generate the story with the following detailed JSON format:\n\np1, p2, p3, p4: Keys for story paragraphs; title: Key for story title; {', '.join([f'prompt{i+1}' for i in range(num_images)])}: Keys for image prompts for each paragraph (or more if additional images are needed). Please do not include any other text in the output. Only JSON is needed or it will break the system and cause loss. Please don't say 'Here is the requested output in JSON format' or anything similar. Thank you."""
 
     client = Client("Be-Bo/llama-3-chatbot_70b")
     hikaye = client.predict(
         message=prompt,
         api_name="/chat"
     )
-    
-    # Debug print to check hikaye
+
     print("Debug: Generated hikaye:", hikaye)
     
     return hikaye
@@ -39,31 +39,36 @@ def cover(prompts):
 
     return images
 
-def parse_story_response(response):
+def parse_story_response(response, num_images):
     print("Debug: Raw response from API:", response)
-    
+
     if not response:
         print("Debug: Response is empty or None.")
-        return None, None, None, None, None, None
-    
+        return None, None, None, None, []
+
     title = response.get('title', '')
     p1 = response.get('p1', '')
     p2 = response.get('p2', '')
     p3 = response.get('p3', '')
     p4 = response.get('p4', '')
-    prompt = response.get('prompt', '')
+
+    prompts = [response.get(f'prompt{i+1}', '') for i in range(num_images)]
     
-    return title, p1, p2, p3, p4, prompt
+    return title, p1, p2, p3, p4, prompts
 
+# Streamlit UI
 st.title('Story AI By Ozi')
 
 characters = st.text_area(label="Characters")
 outline = st.text_area(label="Story Outline")
 settings = st.text_area(label="Setting")
 
+# Number chooser for images (up to 5)
+num_images = st.slider('Number of Images to Generate', 1, 5, 4)
+
 if st.button(label="Generate"):
     with st.spinner('Generating story and cover images...'):
-        hikaye = generate(outline, characters, settings)
+        hikaye = generate(outline, characters, settings, num_images)
         print("Debug: Story generation response:", hikaye)
 
         if hikaye:
@@ -73,19 +78,16 @@ if st.button(label="Generate"):
                 st.error(f"Failed to parse JSON response: {e}")
                 st.stop()
 
-            title, p1, p2, p3, p4, prmt = parse_story_response(hikaye_json)
+            title, p1, p2, p3, p4, prompts = parse_story_response(hikaye_json, num_images)
 
             if title and p1 and p2 and p3 and p4:
                 st.markdown(f'### {title}')
 
-                # Prepare prompts for each paragraph
-                prompts = [prmt, f"Image for paragraph 1: {p1}", f"Image for paragraph 2: {p2}", f"Image for paragraph 3: {p3}", f"Image for paragraph 4: {p4}"]
-
                 # Generate and display images
                 images = cover(prompts)
                 for i, image in enumerate(images):
                     if image:
-                        st.image(image, caption=prompts[i])
+                        st.image(image, caption=f"Image {i+1}: {prompts[i]}")
                     else:
                         st.error(f"Failed to generate image {i+1}.")