Update app.py

app.py

@@ -20,9 +20,8 @@ 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(
@@ -31,7 +30,7 @@ image_model = genai.GenerativeModel(
 )
 hf_client = InferenceClient(token=hf_token)
 
-# -------------------- Utility Functions --------------------
+# -------------- Utility Functions --------------
 
 def generate_image_from_text(prompt):
     try:
@@ -96,25 +95,18 @@ def generate_pdf(images, explanations):
         doc.build(story)
         return tmp.name
 
-# -------------------- Core Logic --------------------
+# ------------------ Core Scene Logic ------------------
 
 def generate_scene(num_scenes, theme, char_count, character_names, dialogue,
-                   images, summaries, explanations,
-                   char_styles, char_moods, bg_style, dialogue_speaker):
+                   images, summaries, explanations, current_scene_index):
     
-    if len(images) >= int(num_scenes):
+    if current_scene_index >= int(num_scenes):
         return gr.update(), gr.update(), gr.update(), images, summaries, explanations, \
-               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}."
-    )
+               "✅ All scenes have been generated.", images[-1], current_scene_index
 
+    prompt = f"Scene {current_scene_index + 1} set in {theme}. Number of characters: {char_count}. Characters: {character_names}."
     if 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.'
+        prompt += f' Include this dialogue: "{dialogue}".'
 
     image = generate_image_from_text(prompt)
     summary = summarize_scene(prompt)
@@ -124,29 +116,14 @@ def generate_scene(num_scenes, theme, char_count, character_names, dialogue,
     summaries.append(summary)
     explanations.append(explanation)
 
-    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)
-
-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
+    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
 
 def ai_write_scene(theme, total_scenes, scene_summaries):
     try:
         scene_index = len(scene_summaries) + 1
         if scene_index > int(total_scenes):
-            return "", "", "", "", "", "Fantasy", f"✅ All {total_scenes} scenes completed."
+            return "", "", 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"""
@@ -157,60 +134,60 @@ 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.
+1. Character names,
+2. One line of dialogue
 """
+
         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."
+        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."
     except Exception as e:
         print("AI write error:", e)
-        return "", "", "", "", "", "Fantasy", "⚠️ AI scene generation failed."
+        return "", "", "⚠️ AI scene generation failed."
+
+def finalize_story(images, explanations):
+    if not images or not explanations:
+        return None, None
 
+    pdf = generate_pdf(images, explanations)
 
-# -------------------- UI --------------------
+    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 ------------------
 
 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", 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")
+        character_names = gr.Textbox(label="🧙‍♂️ Character Names")
+        dialogue = gr.Textbox(label="💬 Dialogue (optional)")
 
     with gr.Row():
         generate_btn = gr.Button("➕ Generate This Scene")
         ai_coauthor_btn = gr.Button("🤖 Let AI Write This One")
 
+    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")
@@ -220,32 +197,57 @@ 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")
 
-    # Persistent session states
+    # 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
+        ],
+        outputs=[
+            image_output, summary_output, explanation_output,
             scene_images, scene_summaries, scene_explanations,
-            char_styles, char_moods, bg_style, dialogue_speaker
+            status, last_image_display, current_scene_index
+        ]
+    )
+
+    recreate_btn.click(
+        fn=generate_scene,
+        inputs=[
+            scene_total, theme, char_count, character_names, dialogue,
+            scene_images, scene_summaries, scene_explanations, current_scene_index
         ],
         outputs=[
             image_output, summary_output, explanation_output,
             scene_images, scene_summaries, scene_explanations,
-            status, done_btn
+            status, last_image_display, current_scene_index
+        ]
+    )
+
+    next_scene_btn.click(
+        fn=generate_scene,
+        inputs=[
+            scene_total, theme, char_count, character_names, dialogue,
+            scene_images, scene_summaries, scene_explanations, current_scene_index
+        ],
+        outputs=[
+            image_output, summary_output, explanation_output,
+            scene_images, scene_summaries, scene_explanations,
+            status, last_image_display, current_scene_index
         ]
     )
 
     ai_coauthor_btn.click(
         fn=ai_write_scene,
         inputs=[theme, scene_total, scene_summaries],
-        outputs=[character_names, dialogue, dialogue_speaker, char_styles, char_moods, bg_style, status]
+        outputs=[character_names, dialogue, status]
     )
 
-
     done_btn.click(
         fn=finalize_story,
         inputs=[scene_images, scene_explanations],