Update app.py

app.py

@@ -142,31 +142,43 @@ def finalize_story(images, explanations):
 
     return txt_path, pdf
 
-def ai_write_scene(theme):
+def ai_write_scene(theme, total_scenes, scene_summaries):
     try:
-        prompt = f"""You're helping a writer generate a new scene for a story set in: "{theme}".
-Please respond with:
-1. Character names (2-3 max),
-2. A short dialogue line,
-3. Who says the dialogue,
-4. Their outfits,
-5. Their mood,
-6. Background style (one of: Realistic, Cartoon, Fantasy, Dark Fantasy)."""
-
-        response = text_model.generate_content(prompt).text
-        lines = response.splitlines()
-
-        char_names = next((l.split(":", 1)[1].strip() for l in lines if "Character" in l), "")
-        dialogue = next((l.split(":", 1)[1].strip() for l in lines if "dialogue" in l.lower()), "")
-        speaker = next((l.split(":", 1)[1].strip() for l in lines if "says" in l.lower()), "")
-        outfits = next((l.split(":", 1)[1].strip() for l in lines if "outfits" in l.lower()), "")
-        moods = next((l.split(":", 1)[1].strip() for l in lines if "mood" in l.lower()), "")
-        bg_style = next((l.split(":", 1)[1].strip() for l in lines if "background" in l.lower()), "Fantasy")
-
-        return char_names, dialogue, speaker, outfits, moods, bg_style
+        scene_index = len(scene_summaries) + 1
+        if scene_index > int(total_scenes):
+            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"""
+You are co-writing a story set in this theme: "{theme}".
+Generate Scene {scene_index} of {total_scenes}.
+Continue the story logically based on previous scenes (if any):
+
+{story_so_far}
+
+Return:
+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()
+        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 "", "", "", "", "", "Fantasy"
+        return "", "", "", "", "", "Fantasy", "⚠️ AI scene generation failed."
+
 
 # -------------------- UI --------------------
 
@@ -229,10 +241,11 @@ with gr.Blocks(title="Scene-by-Scene Story Generator") as demo:
 
     ai_coauthor_btn.click(
         fn=ai_write_scene,
-        inputs=[theme],
-        outputs=[character_names, dialogue, dialogue_speaker, char_styles, char_moods, bg_style]
+        inputs=[theme, scene_total, scene_summaries],
+        outputs=[character_names, dialogue, dialogue_speaker, char_styles, char_moods, bg_style, status]
     )
 
+
     done_btn.click(
         fn=finalize_story,
         inputs=[scene_images, scene_explanations],