Update app.py

app.py

@@ -20,8 +20,9 @@ api_key = os.getenv("GOOGLE_API_KEY")
 hf_token = os.getenv("HF_token")
 
 if not api_key or not hf_token:
-    raise EnvironmentError("Missing API keys. Check GOOGLE_API_KEY and HF_TOKEN in .env.")
+    raise EnvironmentError("Missing API keys. Check GOOGLE_API_KEY and HF_token in .env.")
 
+# Configure models
 genai.configure(api_key=api_key)
 text_model = genai.GenerativeModel("gemini-1.5-flash")
 image_model = genai.GenerativeModel(
@@ -30,7 +31,7 @@ image_model = genai.GenerativeModel(
 )
 hf_client = InferenceClient(token=hf_token)
 
-# -------------- Utility Functions --------------
+# -------------------- Utility Functions --------------------
 
 def generate_image_from_text(prompt):
     try:
@@ -95,18 +96,25 @@ def generate_pdf(images, explanations):
         doc.build(story)
         return tmp.name
 
-# ------------------ Core Scene Logic ------------------
+# -------------------- Core Logic --------------------
 
 def generate_scene(num_scenes, theme, char_count, character_names, dialogue,
-                   images, summaries, explanations, current_scene_index):
+                   images, summaries, explanations,
+                   char_styles, char_moods, bg_style, dialogue_speaker):
     
-    if current_scene_index >= int(num_scenes):
+    if len(images) >= int(num_scenes):
         return gr.update(), gr.update(), gr.update(), images, summaries, explanations, \
-               "✅ All scenes have been generated.", images[-1], current_scene_index
+               f"✅ All {num_scenes} scenes have been generated.", gr.update(visible=True)
+
+    prompt = (
+        f"Scene {len(images)+1} set in {theme}. Number of characters: {char_count}. "
+        f"Characters: {character_names}. Outfits: {char_styles}. Mood: {char_moods}. "
+        f"Style: {bg_style}."
+    )
 
-    prompt = f"Scene {current_scene_index + 1} set in {theme}. Number of characters: {char_count}. Characters: {character_names}."
     if dialogue:
-        prompt += f' Include this dialogue: "{dialogue}".'
+        prompt += f' The character "{dialogue_speaker}" says: "{dialogue}". '
+        prompt += f' Please visualize this as a speech bubble above {dialogue_speaker}, like in a cartoon.'
 
     image = generate_image_from_text(prompt)
     summary = summarize_scene(prompt)
@@ -116,14 +124,29 @@ def generate_scene(num_scenes, theme, char_count, character_names, dialogue,
     summaries.append(summary)
     explanations.append(explanation)
 
-    status = f"✅ Scene {current_scene_index + 1} of {int(num_scenes)} generated."
-    return image, summary, explanation, images, summaries, explanations, status, image, current_scene_index + 1
+    status = f"✅ Scene {len(images)} of {int(num_scenes)} generated."
+    done_visible = len(images) == int(num_scenes)
+
+    return image, summary, explanation, images, summaries, explanations, status, gr.update(visible=done_visible), image
+
+def finalize_story(images, explanations):
+    if not images or not explanations:
+        return None, None
+
+    pdf = generate_pdf(images, explanations)
+
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".txt", mode="w", encoding="utf-8") as txt:
+        for i, exp in enumerate(explanations):
+            txt.write(f"Scene {i+1}:\n{exp}\n\n")
+        txt_path = txt.name
+
+    return txt_path, pdf
 
 def ai_write_scene(theme, total_scenes, scene_summaries):
     try:
         scene_index = len(scene_summaries) + 1
         if scene_index > int(total_scenes):
-            return "", "", f"✅ All {total_scenes} scenes completed."
+            return "", "", "", "", "", "Fantasy", f"✅ All {total_scenes} scenes completed."
 
         story_so_far = "\n".join([f"Scene {i+1}: {s}" for i, s in enumerate(scene_summaries)]) if scene_summaries else ""
         prompt = f"""
@@ -134,60 +157,64 @@ Continue the story logically based on previous scenes (if any):
 {story_so_far}
 
 Return:
-1. Character names,
-2. One line of dialogue
+1. Character names (existing or new),
+2. A single dialogue line,
+3. The speaker of that dialogue,
+4. Outfits worn,
+5. Characters' emotional moods,
+6. Background style (choose from: Realistic, Cartoon, Fantasy, Dark Fantasy).
+Only provide the raw values, no headers.
 """
-
         response = text_model.generate_content(prompt).text.strip()
-        lines = response.split("\n")
-        char_names = lines[0] if lines else ""
-        dialogue = lines[1] if len(lines) > 1 else ""
-
-        return char_names, dialogue, f"📝 Scene {scene_index} prompt generated."
+        parts = [line.strip() for line in response.split("\n") if line.strip()]
+        char_names = parts[0] if len(parts) > 0 else ""
+        dialogue = parts[1] if len(parts) > 1 else ""
+        speaker = parts[2] if len(parts) > 2 else ""
+        outfits = parts[3] if len(parts) > 3 else ""
+        moods = parts[4] if len(parts) > 4 else ""
+        bg_style = parts[5] if len(parts) > 5 else "Fantasy"
+
+        return char_names, dialogue, speaker, outfits, moods, bg_style, f"📝 Scene {scene_index} ready to generate."
     except Exception as e:
         print("AI write error:", e)
-        return "", "", "⚠️ AI scene generation failed."
+        return "", "", "", "", "", "Fantasy", "⚠️ AI scene generation failed."
 
-def finalize_story(images, explanations):
-    if not images or not explanations:
-        return None, None
-
-    pdf = generate_pdf(images, explanations)
-
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".txt", mode="w", encoding="utf-8") as txt:
-        for i, exp in enumerate(explanations):
-            txt.write(f"Scene {i+1}:\n{exp}\n\n")
-        txt_path = txt.name
 
-    return txt_path, pdf
-
-# ------------------ UI ------------------
+# -------------------- UI --------------------
 
 with gr.Blocks(title="Scene-by-Scene Story Generator") as demo:
     gr.Markdown("## 🎬 AI Scene-by-Scene Story Creator")
+    gr.Markdown("Describe your story one scene at a time, with AI-generated images, summaries, and explanations.")
 
     with gr.Accordion("🧩 Story Setup", open=True):
         scene_total = gr.Number(label="🔢 Number of Scenes", precision=0, value=3)
         theme = gr.Textbox(label="🌍 Global Theme", placeholder="e.g. A magical forest full of glowing creatures")
 
+    gr.Markdown("### ✨ Describe Your Next Scene")
+
     with gr.Row():
         char_count = gr.Number(label="👥 Number of Characters", precision=0, value=2)
-        character_names = gr.Textbox(label="🧙‍♂️ Character Names")
-        dialogue = gr.Textbox(label="💬 Dialogue (optional)")
+        character_names = gr.Textbox(label="🧙‍♂️ Character Names", placeholder="e.g. Elora, Bramble the Bear")
+        dialogue = gr.Textbox(label="💬 Dialogue (optional)", placeholder="e.g. 'Protect the forest!'")
+
+    with gr.Row():
+        dialogue_speaker = gr.Textbox(label="🗣️ Who says the dialogue?", placeholder="e.g. Bramble")
+        char_styles = gr.Textbox(label="🧥 Outfit Descriptions", placeholder="e.g. Elora wears a leafy cloak, Bramble has a warrior vest")
+        char_moods = gr.Textbox(label="😠 Character Moods", placeholder="e.g. Elora is cautious, Bramble is brave")
+        bg_style = gr.Dropdown(label="🎨 Background Style", choices=["Realistic", "Cartoon", "Fantasy", "Dark Fantasy"], value="Fantasy")
 
     with gr.Row():
         generate_btn = gr.Button("➕ Generate This Scene")
         ai_coauthor_btn = gr.Button("🤖 Let AI Write This One")
+        recreate_btn = gr.Button("🔄 Recreate This Scene")
+        next_btn = gr.Button("⏭️ Next Scene")
+    last_image_display = gr.Image(label="📸 Last Scene Image", type="pil")
 
-    with gr.Row():
-        recreate_btn = gr.Button("🔄 Recreate Scene")
-        next_scene_btn = gr.Button("⏭️ Next Scene")
 
     status = gr.Markdown()
     image_output = gr.Image(label="🖼️ Scene Image", type="pil")
     summary_output = gr.Markdown(label="📝 Scene Summary")
     explanation_output = gr.Textbox(label="📖 Scene Explanation", lines=6)
-    last_image_display = gr.Image(label="🖼️ Last Scene Image")
 
     with gr.Row():
         tts_btn = gr.Button("🔊 Read Aloud")
@@ -197,57 +224,58 @@ with gr.Blocks(title="Scene-by-Scene Story Generator") as demo:
     txt_file = gr.File(label="📄 Explanations (.txt)")
     pdf_file = gr.File(label="📘 Scene PDF")
 
-    # States
+    # Persistent session states
     scene_images = gr.State([])
     scene_explanations = gr.State([])
     scene_summaries = gr.State([])
-    current_scene_index = gr.State(0)
 
     generate_btn.click(
         fn=generate_scene,
         inputs=[
             scene_total, theme, char_count, character_names, dialogue,
-            scene_images, scene_summaries, scene_explanations, current_scene_index
+            scene_images, scene_summaries, scene_explanations
         ],
         outputs=[
             image_output, summary_output, explanation_output,
             scene_images, scene_summaries, scene_explanations,
-            status, last_image_display, current_scene_index
+            status, done_btn, last_image_display
         ]
     )
 
     recreate_btn.click(
-        fn=generate_scene,
+        fn=lambda *args: generate_scene(*args[:-1], args[-3][:-1], args[-2][:-1], args[-1][:-1]),
         inputs=[
             scene_total, theme, char_count, character_names, dialogue,
-            scene_images, scene_summaries, scene_explanations, current_scene_index
+            scene_images, scene_summaries, scene_explanations
         ],
         outputs=[
             image_output, summary_output, explanation_output,
             scene_images, scene_summaries, scene_explanations,
-            status, last_image_display, current_scene_index
+            status, done_btn, last_image_display
         ]
     )
 
-    next_scene_btn.click(
+    next_btn.click(
         fn=generate_scene,
         inputs=[
             scene_total, theme, char_count, character_names, dialogue,
-            scene_images, scene_summaries, scene_explanations, current_scene_index
+            scene_images, scene_summaries, scene_explanations
         ],
         outputs=[
             image_output, summary_output, explanation_output,
             scene_images, scene_summaries, scene_explanations,
-            status, last_image_display, current_scene_index
+            status, done_btn, last_image_display
         ]
     )
 
+
     ai_coauthor_btn.click(
         fn=ai_write_scene,
         inputs=[theme, scene_total, scene_summaries],
-        outputs=[character_names, dialogue, status]
+        outputs=[character_names, dialogue, dialogue_speaker, char_styles, char_moods, bg_style, status]
     )
 
+
     done_btn.click(
         fn=finalize_story,
         inputs=[scene_images, scene_explanations],