Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py - Final Version with Enhanced Animations
+# app.py - Final Version with AI-Powered Animations
 import streamlit as st
 import os
 import time
@@ -15,6 +15,10 @@ import matplotlib.pyplot as plt
 import requests
 from io import BytesIO
 import json
+import torch
+from diffusers import DiffusionPipeline, StableDiffusionPipeline
+import torch
+from transformers import pipeline
 
 # Configure Streamlit page
 st.set_page_config(
@@ -155,6 +159,14 @@ st.markdown("""
         border-radius: 10px;
         box-shadow: 0 8px 16px rgba(0,0,0,0.2);
     }
+    
+    .ai-animation {
+        border: 3px solid var(--accent);
+        border-radius: 15px;
+        padding: 15px;
+        background: white;
+        margin: 20px 0;
+    }
 </style>
 """, unsafe_allow_html=True)
 
@@ -197,143 +209,37 @@ CONCEPTS = {
     }
 }
 
-# Character database with enhanced visuals
-CHARACTERS = [
-    {"name": "rabbit", "emoji": "🐰", "color": "#FFB6C1", "speed": 1.5},
-    {"name": "dragon", "emoji": "🐉", "color": "#FF6347", "speed": 0.8},
-    {"name": "cat", "emoji": "🐱", "color": "#DDA0DD", "speed": 2.0},
-    {"name": "dog", "emoji": "🐶", "color": "#FFD700", "speed": 2.5},
-    {"name": "knight", "emoji": "🤺", "color": "#87CEEB", "speed": 1.2},
-    {"name": "wizard", "emoji": "🧙", "color": "#98FB98", "speed": 1.0},
-    {"name": "scientist", "emoji": "🔬", "color": "#20B2AA", "speed": 1.3},
-    {"name": "pirate", "emoji": "🏴‍☠️", "color": "#FFA500", "speed": 1.4}
-]
-
-# Object database for stories
-STORY_OBJECTS = [
-    {"name": "carrot", "emoji": "🥕", "color": "#FF8C00"},
-    {"name": "treasure", "emoji": "💰", "color": "#FFD700"},
-    {"name": "castle", "emoji": "🏰", "color": "#A9A9A9"},
-    {"name": "flower", "emoji": "🌸", "color": "#FF69B4"},
-    {"name": "book", "emoji": "📚", "color": "#8B4513"},
-    {"name": "star", "emoji": "⭐", "color": "#FFFF00"},
-    {"name": "key", "emoji": "🔑", "color": "#DAA520"},
-    {"name": "apple", "emoji": "🍎", "color": "#FF0000"}
-]
-
-# Animation templates with visual cues
-ANIMATION_TEMPLATES = {
-    "loop": {
-        "description": "Character moving to a target multiple times",
-        "visual_cue": "🔄 Repeating action multiple times",
-        "code": """
-# Loop Animation
-import matplotlib.pyplot as plt
-import numpy as np
-
-fig, ax = plt.subplots(figsize=(10, 6))
-ax.set_title('Your Story: {story}')
-ax.set_xlim(0, 10)
-ax.set_ylim(0, 10)
-
-character = ax.text(1, 5, '{emoji}', fontsize=48, color='{color}')
-target = ax.text(9, 5, '{target_emoji}', fontsize=48)
-info = ax.text(5, 8, 'Loop: Repeating {count} times', 
-               fontsize=16, ha='center', color='{concept_color}')
-
-for i in range({count}):
-    # Update position
-    for pos in np.linspace(1, 9, 20):
-        character.set_position((pos, 5))
-        plt.pause(0.05)
-    
-    # Show loop counter
-    info.set_text(f'Loop: Completed {{i+1}}/{count}')
-    plt.pause(0.2)
-    
-    # Return to start
-    for pos in np.linspace(9, 1, 20):
-        character.set_position((pos, 5))
-        plt.pause(0.05)
-plt.show()
-        """
-    },
-    "conditional": {
-        "description": "Character making a decision based on a condition",
-        "visual_cue": "❓ Making a decision based on a condition",
-        "code": """
-# Conditional Animation
-import matplotlib.pyplot as plt
-import numpy as np
-import random
-
-fig, ax = plt.subplots(figsize=(10, 6))
-ax.set_title('Your Story: {story}')
-ax.set_xlim(0, 10)
-ax.set_ylim(0, 10)
-
-character = ax.text(5, 5, '{emoji}', fontsize=48, color='{color}')
-condition = random.choice([True, False])
-info = ax.text(5, 8, 'Condition: {condition_desc}', 
-               fontsize=16, ha='center', color='{concept_color}')
-
-if condition:
-    # Path for True condition
-    decision = ax.text(8, 7, '✅ True Path', fontsize=20, color='green')
-    path = np.linspace(5, 8, 20)
-else:
-    # Path for False condition
-    decision = ax.text(2, 7, '❌ False Path', fontsize=20, color='red')
-    path = np.linspace(5, 2, 20)
-
-for pos in path:
-    character.set_position((pos, 5))
-    plt.pause(0.05)
-plt.show()
-        """
-    },
-    "function": {
-        "description": "Character performing an action multiple times",
-        "visual_cue": "✨ Performing action with a function",
-        "code": """
-# Function Animation
-import matplotlib.pyplot as plt
-import numpy as np
-
-fig, ax = plt.subplots(figsize=(10, 6))
-ax.set_title('Your Story: {story}')
-ax.set_xlim(0, 10)
-ax.set_ylim(0, 10)
-
-character = ax.text(5, 5, '{emoji}', fontsize=48, color='{color}')
-info = ax.text(5, 8, 'Function: Performing action {count} times', 
-               fontsize=16, ha='center', color='{concept_color}')
-
-def perform_action(position):
-    # Action animation
-    character.set_fontsize(60)
-    plt.pause(0.1)
-    character.set_fontsize(48)
-    plt.pause(0.1)
-    # Move to new position
-    character.set_position((position, 5))
-
-for i in range({count}):
-    perform_action(5 + i)
-    info.set_text(f'Function: Action {{i+1}}/{count}')
-    plt.pause(0.3)
-plt.show()
-        """
-    }
-}
-
-# Animation examples for different concepts
+# Pre-generated animation examples
 ANIMATION_EXAMPLES = {
     "loop": "https://i.imgur.com/7zQY1eE.gif",
     "conditional": "https://i.imgur.com/5X8jYAy.gif",
     "function": "https://i.imgur.com/9zJkQ7P.gif"
 }
 
+# Initialize AI models
+@st.cache_resource
+def load_models():
+    """Load AI models for animation generation"""
+    models = {}
+    
+    try:
+        # Text-to-image model for scene generation
+        models['text_to_image'] = DiffusionPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-2-1", 
+            torch_dtype=torch.float16
+        )
+    except Exception as e:
+        st.error(f"Could not load text-to-image model: {str(e)}")
+        models['text_to_image'] = None
+    
+    try:
+        # Text-to-speech model
+        models['text_to_speech'] = pipeline("text-to-speech", model="suno/bark-small")
+    except:
+        models['text_to_speech'] = None
+    
+    return models
+
 def analyze_story(story):
     """Analyze story and identify programming concepts"""
     story_lower = story.lower()
@@ -368,172 +274,163 @@ def extract_count_from_story(story):
             return min(int(word), 10)
     return 3  # Default value
 
-def extract_story_elements(story):
-    """Extract characters and objects from the story"""
-    # Choose a random character
-    character = random.choice(CHARACTERS)
-    
-    # Choose a random object
-    story_object = random.choice(STORY_OBJECTS)
-    
-    # Try to find character in story
-    for char in CHARACTERS:
-        if char["name"] in story.lower():
-            character = char
-            break
-    
-    # Try to find object in story
-    for obj in STORY_OBJECTS:
-        if obj["name"] in story.lower():
-            story_object = obj
-            break
-    
-    return character, story_object
-
-def create_visualization(story, concepts):
-    """Create a visualization image based on the story and concepts"""
+def create_story_scene(story, concept, models):
+    """Create a story scene using AI text-to-image model"""
     try:
-        # Get story elements
-        character, story_object = extract_story_elements(story)
-        count = extract_count_from_story(story)
-        concept = concepts[0] if concepts else "variable"
-        concept_color = CONCEPTS[concept]["color"]
-        
-        # Create figure
-        fig, ax = plt.subplots(figsize=(10, 6))
-        ax.set_facecolor('#f0f8ff')
-        ax.set_xlim(0, 10)
-        ax.set_ylim(0, 10)
-        ax.axis('off')
-        
-        # Add title
-        ax.text(5, 9, '✨ Your Story Visualization ✨', 
-                fontsize=20, ha='center', color='purple', fontweight='bold')
+        # Create a prompt based on the story and concept
+        style = "cartoon style, bright colors, children's book illustration"
+        prompt = f"{story}. {CONCEPTS[concept]['name']} concept. {style}"
         
-        # Add story text
-        story_text = story[:80] + ('...' if len(story) > 80 else '')
-        ax.text(5, 8, f'"{story_text}"', 
-                fontsize=14, ha='center', color='#333')
+        # Generate image
+        image = models['text_to_image'](prompt).images[0]
         
-        # Add character
-        ax.text(3, 5, character["emoji"], 
-                fontsize=100, ha='center', color=character["color"])
-        
-        # Add object
-        ax.text(7, 5, story_object["emoji"], 
-                fontsize=100, ha='center', color=story_object["color"])
+        # Convert to bytes
+        buf = io.BytesIO()
+        image.save(buf, format='PNG')
+        buf.seek(0)
         
-        # Add connection arrow
-        ax.annotate('', xy=(7, 5), xytext=(3, 5),
-                    arrowprops=dict(arrowstyle='->', lw=3, color=concept_color))
+        return buf
+    
+    except Exception as e:
+        st.error(f"AI scene generation error: {str(e)}")
+        return None
+
+def create_animation_preview(story, concept):
+    """Create an animation preview for the concept"""
+    try:
+        # Return a sample GIF for the concept
+        return ANIMATION_EXAMPLES.get(concept, "https://i.imgur.com/7zQY1eE.gif")
+    except:
+        return "https://i.imgur.com/7zQY1eE.gif"
+
+def text_to_speech_custom(text, models, filename="story_audio.wav"):
+    """Convert text to speech using AI model"""
+    try:
+        if models['text_to_speech']:
+            # Generate audio
+            audio = models['text_to_speech'](text)
+            
+            # Save to file
+            with open(filename, "wb") as f:
+                f.write(audio["audio"])
+            
+            return filename
         
-        # Add concept visualization
-        concept_text = CONCEPTS[concept]["emoji"] + " " + CONCEPTS[concept]["name"]
-        ax.text(5, 3, concept_text, 
-                fontsize=24, ha='center', color=concept_color, fontweight='bold')
+        # Fallback to gTTS
+        tts = gTTS(text=text, lang='en')
+        tts.save(filename)
+        return filename
+    except Exception as e:
+        st.error(f"Audio creation error: {str(e)}")
+        return None
+
+def autoplay_audio(file_path):
+    """Autoplay audio in Streamlit"""
+    with open(file_path, "rb") as f:
+        data = f.read()
+        b64 = base64.b64encode(data).decode()
+        md = f"""
+            <audio autoplay>
+            <source src="data:audio/wav;base64,{b64}" type="audio/wav">
+            </audio>
+            """
+        st.markdown(md, unsafe_allow_html=True)
+
+def generate_animation_code(story, concept):
+    """Generate Python animation code based on story"""
+    try:
+        # Sample code based on concept
+        if concept == "loop":
+            code = f"""
+# Loop Animation for: {story[:30]}...
+import time
+
+def animate_story():
+    actions = {extract_count_from_story(story)}
+    
+    for i in range(actions):
+        print(f"Action {{i+1}}: {story[:20]}...")
+        # Animation code would go here
+        time.sleep(0.5)
         
-        # Add action text
-        action = f"Action: {character['name'].title()} moves toward {story_object['name']}"
-        ax.text(5, 2.5, action, 
-                fontsize=16, ha='center', color='#333')
+animate_story()
+            """
+            description = "Looping through actions multiple times"
+            visual_cue = "Repeating actions in a loop"
         
-        # Add footer
-        ax.text(5, 1, "Created with StoryCoder", 
-                fontsize=12, ha='center', style='italic', color='gray')
+        elif concept == "conditional":
+            code = f"""
+# Conditional Animation for: {story[:30]}...
+import random
+
+condition = random.choice([True, False])
+
+if condition:
+    print("Condition is True: {story[:20]}...")
+    # True branch animation
+else:
+    print("Condition is False: {story[:20]}...")
+    # False branch animation
+            """
+            description = "Making decisions based on conditions"
+            visual_cue = "Branching paths based on conditions"
+        
+        else:  # function
+            code = f"""
+# Function Animation for: {story[:30]}...
+def perform_action():
+    print("Performing action: {story[:20]}...")
+    # Action animation code
+
+# Call the function multiple times
+for _ in range({extract_count_from_story(story)}):
+    perform_action()
+            """
+            description = "Reusing code with functions"
+            visual_cue = "Calling a reusable function"
         
-        # Save to buffer
-        buf = io.BytesIO()
-        plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0.5, dpi=150)
-        buf.seek(0)
-        plt.close()
-        return buf
+        return code, description, visual_cue
     
     except Exception as e:
-        st.error(f"Visualization error: {str(e)}")
-        return None
+        return f"# Error generating code\nprint('Could not generate code: {str(e)}')", "", ""
 
-def create_interactive_animation(story, concepts):
+def create_interactive_animation(story, concept):
     """Create an interactive animation using Plotly"""
     try:
-        # Get story elements
-        character, story_object = extract_story_elements(story)
-        count = extract_count_from_story(story)
-        concept = concepts[0] if concepts else "variable"
-        concept_color = CONCEPTS[concept]["color"]
-        
         # Create animation data
         frames = []
+        count = extract_count_from_story(story)
+        
         for i in range(count):
-            # Create a path for the character
-            path_length = 20
-            for step in range(path_length):
-                progress = (i * path_length + step) / (count * path_length)
-                frames.append({
-                    "frame": i * path_length + step,
-                    "x": 3 + 4 * progress,
-                    "y": 5 + np.sin(progress * np.pi * 2) * 2,  # Wave motion
-                    "size": 20 + np.sin(step * 0.3) * 10,  # Pulsing effect
-                    "character": character["emoji"],
-                    "color": character["color"],
-                    "info": f"{character['name'].title()} moving to {story_object['name']}",
-                    "step": f"Action {i+1}/{count}"
-                })
+            frames.append({
+                "frame": i,
+                "x": np.random.uniform(1, 9),
+                "y": np.random.uniform(1, 9),
+                "size": np.random.uniform(20, 40),
+                "label": f"Action {i+1}",
+                "color": f"hsl({i*40}, 100%, 50%)"
+            })
         
         df = pd.DataFrame(frames)
         
-        # Create the target object (static)
-        target_df = pd.DataFrame([{
-            "x": 7,
-            "y": 5,
-            "character": story_object["emoji"],
-            "color": story_object["color"],
-            "size": 30,
-            "info": f"Target: {story_object['name']}",
-            "step": "Target"
-        }])
-        
         # Create animated scatter plot
         fig = px.scatter(
             df, 
             x="x", 
             y="y", 
             animation_frame="frame",
-            text="character",
+            text="label",
             size="size",
             size_max=45,
-            color_discrete_sequence=[character["color"]],
+            color="color",
+            color_discrete_sequence=px.colors.qualitative.Alphabet,
             range_x=[0, 10],
             range_y=[0, 10],
-            hover_name="info",
-            hover_data={"step": True}
-        )
-        
-        # Add the target object
-        fig.add_trace(go.Scatter(
-            x=target_df["x"],
-            y=target_df["y"],
-            text=target_df["character"],
-            mode="text",
-            textfont=dict(size=40, color=story_object["color"]),
-            hoverinfo="text",
-            hovertext=target_df["info"],
-            showlegend=False
-        ))
-        
-        # Add concept indicator
-        fig.add_annotation(
-            x=0.5,
-            y=1.05,
-            xref="paper",
-            yref="paper",
-            text=f"{CONCEPTS[concept]['emoji']} {CONCEPTS[concept]['name']}",
-            showarrow=False,
-            font=dict(size=24, color=concept_color)
+            title=f"Interactive Animation: {story[:40]}{'...' if len(story) > 40 else ''}"
         )
         
         # Customize layout
         fig.update_layout(
-            title=f'Animation: {story[:60]}{"..." if len(story) > 60 else ""}',
             showlegend=False,
             plot_bgcolor='rgba(240,248,255,1)',
             paper_bgcolor='rgba(240,248,255,1)',
@@ -541,20 +438,12 @@ def create_interactive_animation(story, concepts):
             height=600,
             xaxis=dict(showgrid=False, zeroline=False, visible=False),
             yaxis=dict(showgrid=False, zeroline=False, visible=False),
-            updatemenus=[dict(
-                type="buttons",
-                buttons=[dict(
-                    label="Play",
-                    method="animate",
-                    args=[None, {"frame": {"duration": 100, "redraw": True}}]
-                )]
-            )]
         )
         
         fig.update_traces(
-            textfont_size=40,
+            textfont_size=20,
             textposition='middle center',
-            hoverlabel=dict(bgcolor="white", font_size=16)
+            marker=dict(sizemode='diameter')
         )
         
         return fig
@@ -563,125 +452,21 @@ def create_interactive_animation(story, concepts):
         st.error(f"Interactive animation error: {str(e)}")
         return None
 
-def text_to_speech(text, filename="story_audio.mp3"):
-    """Convert text to speech using gTTS"""
-    try:
-        tts = gTTS(text=text, lang='en')
-        tts.save(filename)
-        return filename
-    except Exception as e:
-        st.error(f"Audio creation error: {str(e)}")
-        return None
-
-def autoplay_audio(file_path):
-    """Autoplay audio in Streamlit"""
-    with open(file_path, "rb") as f:
-        data = f.read()
-        b64 = base64.b64encode(data).decode()
-        md = f"""
-            <audio autoplay>
-            <source src="data:audio/mp3;base64,{b64}" type="audio/mp3">
-            </audio>
-            """
-        st.markdown(md, unsafe_allow_html=True)
-
-def generate_animation_code(story, concepts):
-    """Generate Python animation code based on story"""
-    try:
-        # Get story elements
-        character, story_object = extract_story_elements(story)
-        count = extract_count_from_story(story)
-        concept = concepts[0] if concepts else "loop"
-        
-        # Get the appropriate template
-        template = ANIMATION_TEMPLATES.get(concept, ANIMATION_TEMPLATES["loop"])
-        
-        # Fill in the template
-        condition_desc = random.choice(["sunny", "raining", "dark"])
-        code = template["code"].format(
-            emoji=character["emoji"],
-            color=character["color"],
-            story=story[:80] + ('...' if len(story) > 80 else ''),
-            count=count,
-            condition_desc=condition_desc,
-            target_emoji=story_object["emoji"],
-            concept_color=CONCEPTS[concept]["color"]
-        )
-        
-        return code, template["description"], template["visual_cue"]
-    
-    except Exception as e:
-        return f"# Error generating code\nprint('Could not generate code: {str(e)}')", "", ""
-
-def create_story_image(story):
-    """Create a story image with Matplotlib"""
-    try:
-        # Create figure
-        fig, ax = plt.subplots(figsize=(10, 6))
-        ax.set_facecolor('#f0f8ff')
-        ax.set_xlim(0, 10)
-        ax.set_ylim(0, 10)
-        ax.axis('off')
-        
-        # Add title
-        ax.text(5, 8, '✨ Your Story ✨', fontsize=24, 
-                ha='center', color='purple', fontweight='bold')
-        
-        # Add story text (wrapped)
-        words = story.split()
-        lines = []
-        current_line = []
-        
-        for word in words:
-            if len(' '.join(current_line + [word])) < 60:
-                current_line.append(word)
-            else:
-                lines.append(' '.join(current_line))
-                current_line = [word]
-        
-        if current_line:
-            lines.append(' '.join(current_line))
-        
-        for i, line in enumerate(lines):
-            ax.text(5, 6.5 - i*0.7, line, fontsize=14, ha='center', color='#333')
-        
-        # Add decoration
-        ax.text(2, 2, '🐰', fontsize=40, ha='center')
-        ax.text(8, 2, '🐉', fontsize=40, ha='center')
-        ax.text(5, 1, 'Created with StoryCoder', fontsize=12, 
-                ha='center', style='italic', color='gray')
-        
-        # Save to buffer
-        buf = io.BytesIO()
-        plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0.5, dpi=150)
-        buf.seek(0)
-        plt.close()
-        return buf
-    
-    except Exception as e:
-        st.error(f"Image generation error: {str(e)}")
-        return None
-
-def generate_animation_preview(concept):
-    """Get a preview GIF for the animation concept"""
-    try:
-        # Return a sample GIF for the concept
-        return ANIMATION_EXAMPLES.get(concept, "https://i.imgur.com/7zQY1eE.gif")
-    except:
-        return "https://i.imgur.com/7zQY1eE.gif"
-
 def main():
     """Main application function"""
     st.title("🧙‍♂️ StoryCoder - Learn Python Through Stories!")
     st.subheader("Turn your story into an animation and discover coding secrets!")
     
+    # Load AI models
+    ai_models = load_models()
+    
     # Initialize session state
     if 'story' not in st.session_state:
         st.session_state.story = ""
     if 'concepts' not in st.session_state:
         st.session_state.concepts = []
-    if 'visualization' not in st.session_state:
-        st.session_state.visualization = None
+    if 'story_scene' not in st.session_state:
+        st.session_state.story_scene = None
     if 'interactive_animation' not in st.session_state:
         st.session_state.interactive_animation = None
     if 'animation_code' not in st.session_state:
@@ -714,7 +499,7 @@ def main():
         if st.button("🔄 Reset"):
             st.session_state.story = ""
             st.session_state.concepts = []
-            st.session_state.visualization = None
+            st.session_state.story_scene = None
             st.session_state.interactive_animation = None
             st.session_state.animation_code = ""
             st.session_state.code_description = ""
@@ -744,24 +529,29 @@ def main():
                     with st.spinner("🧠 Analyzing your story for coding concepts..."):
                         st.session_state.concepts = analyze_story(story)
                     
-                    with st.spinner("🎨 Creating your visualization..."):
-                        st.session_state.visualization = create_visualization(
-                            story, st.session_state.concepts
+                    # Get the main concept
+                    main_concept = st.session_state.concepts[0] if st.session_state.concepts else "variable"
+                    
+                    with st.spinner("🎨 Generating AI story scene..."):
+                        st.session_state.story_scene = create_story_scene(
+                            story, main_concept, ai_models
                         )
-                        
+                    
                     with st.spinner("🚀 Creating interactive animation..."):
                         st.session_state.interactive_animation = create_interactive_animation(
-                            story, st.session_state.concepts
+                            story, main_concept
                         )
-                        
+                    
                     with st.spinner("💻 Generating animation code..."):
                         st.session_state.animation_code, st.session_state.code_description, st.session_state.visual_cue = generate_animation_code(
-                            story, st.session_state.concepts
+                            story, main_concept
                         )
                     
                     with st.spinner("🔊 Creating audio narration..."):
                         audio_text = f"Your story: {story}. This demonstrates the programming concept: {st.session_state.code_description}"
-                        st.session_state.audio_file = text_to_speech(audio_text)
+                        st.session_state.audio_file = text_to_speech_custom(
+                            audio_text, ai_models
+                        )
                     
                     st.session_state.active_tab = "animation"
                     st.rerun()
@@ -772,15 +562,21 @@ def main():
             with col1:
                 st.caption("Loop Example")
                 st.code('"A dragon breathes fire 5 times at the castle"', language="text")
-                st.image(generate_animation_preview("loop"), use_column_width=True, caption="Loop Animation Preview")
+                st.image(create_animation_preview("", "loop"), 
+                         use_column_width=True, 
+                         caption="Loop Animation Preview")
             with col2:
                 st.caption("Conditional Example")
                 st.code('"If it rains, the cat stays inside, else it goes out"', language="text")
-                st.image(generate_animation_preview("conditional"), use_column_width=True, caption="Conditional Animation Preview")
+                st.image(create_animation_preview("", "conditional"), 
+                         use_column_width=True, 
+                         caption="Conditional Animation Preview")
             with col3:
                 st.caption("Function Example")
                 st.code('"A wizard casts a spell to make flowers grow"', language="text")
-                st.image(generate_animation_preview("function"), use_column_width=True, caption="Function Animation Preview")
+                st.image(create_animation_preview("", "function"), 
+                         use_column_width=True, 
+                         caption="Function Animation Preview")
     
     # Animation tab
     elif st.session_state.active_tab == "animation":
@@ -791,30 +587,41 @@ def main():
             st.session_state.active_tab = "story"
             st.rerun()
         
-        # Display visualization
-        st.subheader("🎨 Story Visualization")
-        if st.session_state.visualization:
-            st.image(st.session_state.visualization, use_container_width=True)
+        # Display AI-generated scene
+        st.subheader("🖼️ AI-Generated Story Scene")
+        if st.session_state.story_scene:
+            st.image(st.session_state.story_scene, use_container_width=True)
         else:
-            st.warning("Visualization couldn't be generated. Showing story image instead.")
-            story_img = create_story_image(st.session_state.story)
-            if story_img:
-                st.image(story_img, use_container_width=True)
+            st.warning("Scene couldn't be generated. Showing example instead.")
+            concept = st.session_state.concepts[0] if st.session_state.concepts else "loop"
+            st.image(create_animation_preview("", concept), 
+                     use_container_width=True)
         
         # Display interactive animation
-        st.subheader("🚀 Interactive Animation")
+        st.subheader("🎮 Interactive Animation")
         if st.session_state.interactive_animation:
             st.plotly_chart(st.session_state.interactive_animation, use_container_width=True)
         else:
             st.info("Interactive animation preview. The real animation would run on your computer!")
-            st.image(generate_animation_preview(st.session_state.concepts[0] if st.session_state.concepts else "loop"), 
+            concept = st.session_state.concepts[0] if st.session_state.concepts else "loop"
+            st.image(create_animation_preview("", concept), 
                     use_container_width=True)
         
+        # Animation concept preview
+        st.subheader("📽️ Animation Concept Preview")
+        if st.session_state.concepts:
+            concept = st.session_state.concepts[0]
+            st.image(create_animation_preview("", concept), 
+                     caption=f"{CONCEPTS[concept]['name']} Animation Example", 
+                     use_container_width=True)
+        else:
+            st.info("Animation preview would appear here")
+        
         # Play audio narration
         st.subheader("🔊 Story Narration")
         if st.session_state.audio_file:
             audio_bytes = open(st.session_state.audio_file, 'rb').read()
-            st.audio(audio_bytes, format='audio/mp3')
+            st.audio(audio_bytes, format='audio/wav')
             
             if st.button("▶️ Play Automatically"):
                 autoplay_audio(st.session_state.audio_file)
@@ -851,7 +658,7 @@ def main():
                     """, unsafe_allow_html=True)
                     
                     # Show animation preview for the concept
-                    st.image(generate_animation_preview(concept), 
+                    st.image(create_animation_preview("", concept), 
                              caption=f"{details['name']} Animation Example", 
                              use_column_width=True)
         
@@ -891,7 +698,7 @@ def main():
             # Show what the animation would look like
             st.subheader("🎬 What Your Animation Would Look Like")
             concept = st.session_state.concepts[0] if st.session_state.concepts else "loop"
-            st.image(generate_animation_preview(concept), 
+            st.image(create_animation_preview("", concept), 
                     caption="This is similar to what your animation would look like", 
                     use_column_width=True)
         else: