app.py

import gradio as gr
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from transformers import AutoProcessor, AutoModelForVision2Seq, AutoModelForCausalLM, AutoTokenizer
import torch
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import textwrap
import os
import gc
import re
import psutil
from datetime import datetime
import spaces
from kokoro import KModel, KPipeline
import soundfile as sf
import math

def clear_memory():
    """Helper function to clear both CUDA and system memory, safe for Spaces environment"""
    gc.collect()
    
    # Only perform CUDA operations if we're in a GPU task context
    if hasattr(spaces, "current_task") and spaces.current_task and torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()
    
    process = psutil.Process(os.getpid())
    if hasattr(process, 'memory_info'):
        process.memory_info().rss
    
    gc.collect(generation=0)
    gc.collect(generation=1)
    gc.collect(generation=2)
    
    # Only log GPU stats if we're in a GPU task context
    if hasattr(spaces, "current_task") and spaces.current_task and torch.cuda.is_available():
        print(f"GPU Memory allocated: {torch.cuda.memory_allocated()/1024**2:.2f} MB")
        print(f"GPU Memory cached: {torch.cuda.memory_reserved()/1024**2:.2f} MB")
    print(f"CPU RAM used: {process.memory_info().rss/1024**2:.2f} MB")


# Initialize models at startup - only the lightweight ones
print("Loading models...")

# Load SmolVLM for image analysis
processor_vlm = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM-500M-Instruct")
model_vlm = AutoModelForVision2Seq.from_pretrained(
    "HuggingFaceTB/SmolVLM-500M-Instruct", 
    torch_dtype=torch.bfloat16
).to("cuda")

# Load SmolLM2 for story and prompt generation
checkpoint = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
tokenizer_lm = AutoTokenizer.from_pretrained(checkpoint)
model_lm = AutoModelForCausalLM.from_pretrained(checkpoint).to("cuda")

# Initialize Kokoro TTS pipeline
pipeline = KPipeline(lang_code='a')  # 'a' for American English

# Load Stable Diffusion model at startup
pipe_sd = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
).to("cuda")
pipe_sd.scheduler = DPMSolverMultistepScheduler.from_config(pipe_sd.scheduler.config)
pipe_sd.enable_attention_slicing()

@torch.inference_mode()
#@spaces.GPU(duration=30)
def generate_image():
    """Generate a random landscape image."""
    clear_memory()
        
    default_prompt = "a beautiful, professional landscape photograph"
    default_negative_prompt = "blurry, bad quality, distorted, deformed"
    default_steps = 30
    default_guidance = 7.5
    default_seed = torch.randint(0, 2**32 - 1, (1,)).item()
    
    generator = torch.Generator("cuda").manual_seed(default_seed)
    
    try:
        image = pipe_sd(
            prompt=default_prompt,
            negative_prompt=default_negative_prompt,
            num_inference_steps=default_steps,
            guidance_scale=default_guidance,
            generator=generator,
        ).images[0]
        

        clear_memory()
        return image
        
    except Exception as e:
        print(f"Error generating image: {e}")
        clear_memory()
        return None

@torch.inference_mode()
#@spaces.GPU(duration=30)
def analyze_image(image):
    if image is None:
        return "Please generate an image first."
    
    clear_memory()
    
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": "Describe this image and Be brief but descriptive."}
            ]
        }
    ]
    
    try:
        prompt = processor_vlm.apply_chat_template(messages, add_generation_prompt=True)
        
        inputs = processor_vlm(
            text=prompt,
            images=[image],
            return_tensors="pt"
        ).to('cuda')
        
        outputs = model_vlm.generate(
            input_ids=inputs.input_ids,
            pixel_values=inputs.pixel_values,
            attention_mask=inputs.attention_mask,
            num_return_sequences=1,
            no_repeat_ngram_size=2,
            max_new_tokens=500,  
            min_new_tokens=10 
        )
        
        description = processor_vlm.decode(outputs[0], skip_special_tokens=True)
        description = re.sub(r".*?Assistant:\s*", "", description, flags=re.DOTALL).strip()
        
        # Split into sentences and take only the first three
        sentences = re.split(r'(?<=[.!?])\s+', description)
        description = ' '.join(sentences[:9])
        
        clear_memory()
        return description
        
    except Exception as e:
        print(f"Error analyzing image: {e}")
        clear_memory()
        return "Error analyzing image. Please try again."


def format_story_paragraphs(story_text, max_paragraphs=9):
    """Formats the story by evenly distributing sentences over a max of 9 paragraphs."""
    
    # Remove unwanted tokens and artifacts
    story_text = story_text.replace("<|im_end|>", "").strip()
    
    # Split into sentences
    sentences = re.split(r'(?<=[.!?])\s+', story_text)
    sentences = [s.strip() for s in sentences if s.strip()]  # Remove empty sentences
    
    # Determine optimal sentence distribution
    total_sentences = len(sentences)
    num_paragraphs = min(max_paragraphs, total_sentences)  # Ensure we do not exceed max paragraphs
    sentences_per_paragraph = math.ceil(total_sentences / num_paragraphs)  # Distribute evenly
    
    # Group sentences into paragraphs
    paragraphs = []
    current_paragraph = []
    
    for sentence in sentences:
        current_paragraph.append(sentence)
        
        # Ensure even distribution across paragraphs
        if len(current_paragraph) == sentences_per_paragraph:
            paragraphs.append(" ".join(current_paragraph))
            current_paragraph = []
        
        # Stop if we reach max paragraphs
        if len(paragraphs) == max_paragraphs:
            break

    # If there are leftover sentences, add them as a final paragraph
    if current_paragraph:
        paragraphs.append(" ".join(current_paragraph))

    # Join paragraphs with a double newline
    return "\n\n".join(paragraphs)

#@spaces.GPU(duration=60)
def generate_story(image_description):
    """Runs story generation safely on Hugging Face Spaces with Stateless GPU"""
    clear_memory()
    
    # Define story prompt
    story_prompt = f"""Write a short children's story (about 200 words) based on this scene: {image_description}

    Requirements:
    1. Main character: An English bulldog named Champ
    2. Include these values: confidence, teamwork, caring, and hope
    3. Theme: "Doing the right thing is important"
    4. Keep it simple and engaging for young children
    5. End with a simple moral lesson"""

    try:
        # Move CUDA initialization inside the function
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        messages = [{"role": "user", "content": story_prompt}]
        input_text = tokenizer_lm.apply_chat_template(messages, tokenize=False)
        
        inputs = tokenizer_lm.encode(input_text, return_tensors="pt").to(device)  # Move to CUDA inside function
        
        outputs = model_lm.generate(
            inputs,
            max_new_tokens=750,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            repetition_penalty=1.2
        )
        
        # Decode and format the story
        story = tokenizer_lm.decode(outputs[0])
        story = clean_story_output(story)
        story = format_story_paragraphs(story)  # Evenly distribute sentences

        clear_memory()
        return story
        
    except Exception as e:
        print(f"Error generating story: {e}")
        clear_memory()
        return "Error generating story. Please try again."



@torch.inference_mode()
#@spaces.GPU(duration=30)
def generate_image_prompts(story_text):
    clear_memory()
    
    paragraphs = split_into_paragraphs(story_text)
    all_prompts = []
    prompt_instruction = '''Here is a story paragraph: {paragraph}

    Start your response with "Watercolor bulldog" and describe what Champ is doing in this scene. Include any friends. Add where it takes place and one mood detail. Keep it short.'''
    
    try:
        for i, paragraph in enumerate(paragraphs, 1):
            messages = [{"role": "user", "content": prompt_instruction.format(paragraph=paragraph)}]
            input_text = tokenizer_lm.apply_chat_template(messages, tokenize=False)
            
            inputs = tokenizer_lm.encode(input_text, return_tensors="pt").to("cuda")
            
            outputs = model_lm.generate(
                inputs,
                max_new_tokens=30,
                temperature=0.5,
                top_p=0.9,
                do_sample=True,
                repetition_penalty=1.2
            )
            
            prompt = process_generated_prompt(tokenizer_lm.decode(outputs[0]), paragraph)
            section = f"Paragraph {i}:\n{paragraph}\n\nScenery Prompt {i}:\n{prompt}\n\n{'='*50}"
            all_prompts.append(section)
            
            clear_memory()
        
        return '\n'.join(all_prompts)
        
    except Exception as e:
        print(f"Error generating prompts: {e}")
        clear_memory()
        return "Error generating prompts. Please try again."

@torch.inference_mode()
#@spaces.GPU(duration=60)
def generate_story_image(prompt, seed=-1):
    clear_memory()
    
    try:
        pipe_sd.load_lora_weights("Prof-Hunt/lora-bulldog")
        
        generator = torch.Generator("cuda")
        if seed != -1:
            generator.manual_seed(seed)
        else:
            generator.manual_seed(torch.randint(0, 2**32 - 1, (1,)).item())

        enhanced_prompt = f"{prompt}, watercolor style, children's book illustration, soft colors"

        image = pipe_sd(
            prompt=enhanced_prompt,
            negative_prompt="deformed, ugly, blurry, bad art, poor quality, distorted",
            num_inference_steps=50,
            guidance_scale=15,
            generator=generator
        ).images[0]

        pipe_sd.unload_lora_weights()
        clear_memory()
        return image

    except Exception as e:
        print(f"Error generating image: {e}")
        pipe_sd.unload_lora_weights()
        clear_memory()
        return None

@torch.inference_mode()
#@spaces.GPU(duration=180)
def generate_all_scenes(prompts_text):
    clear_memory()
    
    generated_images = []
    formatted_prompts = []
    progress_messages = []
    total_scenes = len([s for s in prompts_text.split('='*50) if s.strip()])
    
    def update_progress():
        """Create a progress message showing completed/total scenes"""
        completed = len(generated_images)
        message = f"Generated {completed}/{total_scenes} scenes\n\n"
        if progress_messages:
            message += "\n".join(progress_messages[-3:])  # Show last 3 status messages
        return message
    
    sections = prompts_text.split('='*50)
    
    for section_num, section in enumerate(sections, 1):
        if not section.strip():
            continue
        
        scene_prompt = None
        for line in section.split('\n'):
            if 'Scenery Prompt' in line:
                scene_num = line.split('Scenery Prompt')[1].split(':')[0].strip()
                next_line_index = section.split('\n').index(line) + 1
                if next_line_index < len(section.split('\n')):
                    scene_prompt = section.split('\n')[next_line_index].strip()
                    formatted_prompts.append(f"Scene {scene_num}: {scene_prompt}")
                break
        
        if scene_prompt:
            try:
                clear_memory()
                status_msg = f"🎨 Creating scene {section_num}: '{scene_prompt[:50]}...'"
                progress_messages.append(status_msg)
                
                # Yield progress update
                yield generated_images, "\n\n".join(formatted_prompts), update_progress()
                
                image = generate_story_image(scene_prompt)
                
                if image is not None:
                    # Convert PIL Image to numpy array with explicit mode conversion
                    pil_image = image if isinstance(image, Image.Image) else Image.fromarray(image)
                    pil_image = pil_image.convert('RGB')  # Ensure RGB mode
                    img_array = np.array(pil_image)
                    
                    # Verify array shape and type
                    if len(img_array.shape) == 3 and img_array.shape[2] == 3:
                        generated_images.append(img_array)
                        progress_messages.append(f"✅ Successfully completed scene {section_num}")
                    else:
                        progress_messages.append(f"❌ Error: Invalid image format for scene {section_num}")
                else:
                    progress_messages.append(f"❌ Failed to generate scene {section_num}")
                
                clear_memory()
                
            except Exception as e:
                error_msg = f"❌ Error generating scene {section_num}: {str(e)}"
                progress_messages.append(error_msg)
                clear_memory()
                continue
            
            # Yield progress update after each scene
            yield generated_images, "\n\n".join(formatted_prompts), update_progress()
    
    # Final status update
    if not generated_images:
        progress_messages.append("❌ No images were successfully generated")
    else:
        progress_messages.append(f"✅ Successfully completed all {len(generated_images)} scenes!")
    
    # Final yield
    yield generated_images, "\n\n".join(formatted_prompts), update_progress()
    
#@spaces.GPU(duration=60)
def add_text_to_scenes(gallery_images, prompts_text):
    """Add text overlays to all scenes"""
    print(f"Received gallery_images type: {type(gallery_images)}")
    print(f"Number of images in gallery: {len(gallery_images) if isinstance(gallery_images, list) else 0}")
    
    if not isinstance(gallery_images, list):
        print("Gallery images must be a list")
        return [], []

    clear_memory()
    
    # Process text sections
    sections = prompts_text.split('='*50)
    overlaid_images = []
    output_files = []
    
    # Create temporary directory for saving files
    temp_dir = "temp_book_pages"
    os.makedirs(temp_dir, exist_ok=True)
    
    for i, (img_data, section) in enumerate(zip(gallery_images, sections)):
        if not section.strip():
            continue
            
        print(f"\nProcessing image {i+1}:")
        print(f"Image data type: {type(img_data)}")
        
        try:
            # Handle tuple from Gradio gallery
            if isinstance(img_data, tuple):
                filepath = img_data[0] if isinstance(img_data[0], str) else None
                print(f"Found filepath: {filepath}")
                
                if filepath and os.path.exists(filepath):
                    print(f"Loading image from: {filepath}")
                    image = Image.open(filepath).convert('RGB')
                else:
                    print(f"Invalid filepath: {filepath}")
                    continue
            else:
                print(f"Unexpected image data type: {type(img_data)}")
                continue
            
            # Extract paragraph text
            lines = [line.strip() for line in section.split('\n') if line.strip()]
            paragraph = None
            for j, line in enumerate(lines):
                if line.startswith('Paragraph'):
                    if j + 1 < len(lines):
                        paragraph = lines[j + 1]
                        print(f"Found paragraph text for image {i+1}")
                        break
            
            if paragraph and image:
                # Add text overlay
                overlaid_img = overlay_text_on_image(image, paragraph)
                if overlaid_img is not None:
                    # Convert to numpy array for gallery display
                    overlaid_array = np.array(overlaid_img)
                    overlaid_images.append(overlaid_array)
                    
                    # Save file for download
                    output_path = os.path.join(temp_dir, f"panel_{i+1}.png")
                    overlaid_img.save(output_path)
                    output_files.append(output_path)
                    print(f"Successfully processed image {i+1}")
                else:
                    print(f"Failed to overlay text on image {i+1}")
            
        except Exception as e:
            print(f"Error processing image {i+1}: {str(e)}")
            import traceback
            print(traceback.format_exc())
            continue
    
    if not overlaid_images:
        print("No images were successfully processed")
    else:
        print(f"Successfully processed {len(overlaid_images)} images")
    
    clear_memory()
    return overlaid_images, output_files

def overlay_text_on_image(image, text):
    """Add black text with white outline for better visibility."""
    if image is None:
        return None
    
    try:
        # Ensure we're working with RGB mode
        img = image.convert('RGB')
        draw = ImageDraw.Draw(img)
        
        # Calculate font size based on image dimensions
        font_size = int(img.width * 0.025)
        try:
            font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", font_size)
        except:
            print("Using default font as DejaVuSans-Bold.ttf not found")
            font = ImageFont.load_default()

        # Calculate text positioning
        y_position = int(img.height * 0.005)
        x_margin = int(img.width * 0.005)
        available_width = img.width - (2 * x_margin)
        
        # Wrap text to fit image width
        wrapped_text = textwrap.fill(text, width=int(available_width / (font_size * 0.6)))
        
        # Add white outline to text for better readability
        outline_color = (255, 255, 255)
        text_color = (0, 0, 0)
        offsets = [-2, -1, 1, 2]

        # Draw text outline
        for dx in offsets:
            for dy in offsets:
                draw.multiline_text(
                    (x_margin + dx, y_position + dy),
                    wrapped_text,
                    font=font,
                    fill=outline_color
                )

        # Draw main text
        draw.multiline_text(
            (x_margin, y_position),
            wrapped_text,
            font=font,
            fill=text_color
        )

        return img
        
    except Exception as e:
        print(f"Error in overlay_text_on_image: {e}")
        return None
        
#@spaces.GPU(duration=60)
def generate_combined_audio_from_story(story_text, voice='af_heart', speed=1):
    """Generate a single audio file for all paragraphs in the story."""
    
    # Use existing paragraph logic (from your original function)
    paragraphs = []
    current_paragraph = []
    
    for line in story_text.split('\n'):
        line = line.strip()
        if not line:  # Empty line indicates paragraph break
            if current_paragraph:
                paragraphs.append(' '.join(current_paragraph))
                current_paragraph = []
        else:
            current_paragraph.append(line)
    
    if current_paragraph:
        paragraphs.append(' '.join(current_paragraph))

    # Combine audio for all paragraphs
    combined_audio = []
    
    for paragraph in paragraphs:
        if not paragraph.strip():
            continue  # Skip empty paragraphs

        generator = pipeline(
            paragraph,
            voice=voice,
            speed=speed,
            split_pattern=r'\n+'  # Split on newlines
        )

        for _, _, audio in generator:
            combined_audio.extend(audio)  # Append audio data

    # Convert combined audio to NumPy array
    combined_audio = np.array(combined_audio)

    # Save audio as a WAV file in Hugging Face's temporary directory
    filename = "/tmp/combined_story.wav"
    sf.write(filename, combined_audio, 24000)

    clear_memory()
    return filename  # Returning filename for Gradio File component 
    
# Helper functions
def clean_story_output(story):
    """Clean up the generated story text."""
    story = story.replace("<|im_end|>", "")
    
    story_start = story.find("Once upon")
    if story_start == -1:
        possible_starts = ["One day", "In a", "There was", "Champ"]
        for marker in possible_starts:
            story_start = story.find(marker)
            if story_start != -1:
                break
    
    if story_start != -1:
        story = story[story_start:]
    
    lines = story.split('\n')
    cleaned_lines = []
    for line in lines:
        line = line.strip()
        if line and not any(skip in line.lower() for skip in ['requirement', 'include these values', 'theme:', 'keep it simple', 'end with', 'write a']):
            if not line.startswith(('1.', '2.', '3.', '4.', '5.')):
                cleaned_lines.append(line)
    
    return '\n\n'.join(cleaned_lines).strip()

def split_into_paragraphs(text):
    """Split text into paragraphs."""
    paragraphs = []
    current_paragraph = []
    
    for line in text.split('\n'):
        line = line.strip()
        if not line:
            if current_paragraph:
                paragraphs.append(' '.join(current_paragraph))
                current_paragraph = []
        else:
            current_paragraph.append(line)
    
    if current_paragraph:
        paragraphs.append(' '.join(current_paragraph))
    
    return [p for p in paragraphs if not any(skip in p.lower() 
            for skip in ['requirement', 'include these values', 'theme:', 
                        'keep it simple', 'end with', 'write a'])]

def process_generated_prompt(prompt, paragraph):
    """Process and clean up generated image prompts."""
    prompt = prompt.replace("<|im_start|>", "").replace("<|im_end|>", "")
    prompt = prompt.replace("assistant", "").replace("system", "").replace("user", "")
    
    cleaned_lines = [line.strip() for line in prompt.split('\n')
                    if line.strip().lower().startswith("watercolor bulldog")]
    
    if cleaned_lines:
        prompt = cleaned_lines[0]
    else:
        setting = "quiet town" if "quiet town" in paragraph.lower() else "park"
        mood = "hopeful" if "wished" in paragraph.lower() else "peaceful"
        prompt = f"Watercolor bulldog watching friends play in {setting}, {mood} atmosphere."
    
    if not prompt.endswith('.'):
        prompt = prompt + '.'
    
    return prompt

def create_interface():
    # Define CSS for custom styling
    css = """
        /* Global styles */
        .gradio-container {
            background-color: #EBF8FF !important;
        }
        
        /* Custom button styling */
        .custom-button {
            background-color: #3B82F6 !important;
            color: white !important;
            border: none !important;
            border-radius: 8px !important;
            padding: 10px 20px !important;
            margin: 10px 0 !important;
            min-width: 200px !important;
        }
        
        .custom-button:hover {
            background-color: #2563EB !important;
        }
        
        /* Section styling */
        .section-content {
            background-color: white !important;
            border-radius: 12px !important;
            padding: 20px !important;
            margin: 10px 0 !important;
            box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05) !important;
        }
        
        /* AI Lesson box styling */
        .ai-lesson {
            background-color: #FEE2E2 !important;
            border-radius: 8px !important;
            padding: 15px !important;
            margin: 10px 0 !important;
            border: 1px solid #FCA5A5 !important;
        }
    """
    
    with gr.Blocks(css=css) as demo:
        gr.Markdown("""
        # 🎨 Tech Tales: AI Children's Story Creator
        
        Welcome to this educational AI story creation tool! This app demonstrates how multiple AI models 
        work together to create an illustrated children's story. Each step includes a brief AI lesson 
        to help you understand the technology being used.
        
        Let's create something magical! ✨

        ## 📌 [Open in Google Colab](https://colab.research.google.com/drive/18-4frRwaMxxqJrRjiXh1Q069R77INMeA?usp=sharing) 🚀

        
        """)
        
        # Step 1: Generate Landscape
        with gr.Row(elem_classes="section-content"):
            with gr.Column(elem_classes="ai-lesson"):
                gr.Markdown("""
                ### Step 1: Setting the Scene with AI 🖼️
                
                🤖 **AI Lesson: Text-to-Image Generation**
                We're using Stable Diffusion, a powerful AI model that turns text into images.
                How it works:
                - Starts with random noise and gradually refines it into an image
                - Uses millions of image-text pairs from its training
                - Combines understanding of both language and visual elements
                - Takes about 50 steps to create each image
                
                Real-world applications: Book illustrations, concept art, product visualization
                """)
            with gr.Column():
                generate_btn = gr.Button("1. Generate Random Landscape", elem_classes="custom-button")
                image_output = gr.Image(label="Your AI-Generated Landscape", type="pil", interactive=False)
        
        # Step 2: Analyze Scene
        with gr.Row(elem_classes="section-content"):
            with gr.Column(elem_classes="ai-lesson"):
                gr.Markdown("""
                ### Step 2: Teaching AI to See 👁️
                
                🤖 **AI Lesson: Vision-Language Models (VLM)**
                Our VLM acts like an AI art critic, understanding and describing images.
                How it works:
                - Processes images through neural networks
                - Identifies objects, scenes, colors, and relationships
                - Translates visual features into natural language
                - Uses attention mechanisms to focus on important details
                
                Real-world applications: Image search, accessibility tools, medical imaging
                """)
            with gr.Column():
                analyze_btn = gr.Button("2. Get Brief Description", elem_classes="custom-button")
                analysis_output = gr.Textbox(label="What the AI Sees", lines=3)
        
        # Step 3: Create Story
        with gr.Row(elem_classes="section-content"):
            with gr.Column(elem_classes="ai-lesson"):
                gr.Markdown("""
                ### Step 3: Crafting the Narrative 📖
                
                🤖 **AI Lesson: Large Language Models**
                Meet our AI storyteller! It uses a Large Language Model (LLM) to write creative stories.
                How it works:
                - Processes the scene description as context
                - Uses pattern recognition from millions of stories
                - Maintains narrative consistency and character development
                - Adapts its writing style for children
                
                Real-world applications: Content creation, creative writing, education
                """)
            with gr.Column():
                story_btn = gr.Button("3. Create Children's Story", elem_classes="custom-button")
                story_output = gr.Textbox(label="Your AI-Generated Story", lines=10)
        
        # Step 4: Generate Prompts
        with gr.Row(elem_classes="section-content"):
            with gr.Column(elem_classes="ai-lesson"):
                gr.Markdown("""
                ### Step 4: Planning the Illustrations 🎯
                
                🤖 **AI Lesson: Natural Language Processing**
                The AI breaks down the story into key scenes and creates optimal image prompts.
                How it works:
                - Analyzes story structure and pacing
                - Identifies key narrative moments
                - Generates specialized prompts for each scene
                - Ensures visual consistency across illustrations
                
                Real-world applications: Content planning, storyboarding, scene composition
                """)
            with gr.Column():
                prompts_btn = gr.Button("4. Generate Scene Prompts", elem_classes="custom-button")
                prompts_output = gr.Textbox(label="Scene Descriptions", lines=20)
        
        # Step 5: Generate Scenes
        with gr.Row(elem_classes="section-content"):
            with gr.Column(elem_classes="ai-lesson"):
                gr.Markdown("""
                ### Step 5: Bringing Scenes to Life 🎨
                
                🤖 **AI Lesson: Specialized Image Generation**
                Using a fine-tuned model to create consistent character illustrations.
                How it works:
                - Uses LoRA (Low-Rank Adaptation) for specialized training
                - Maintains consistent character appearance
                - Processes multiple scenes in parallel
                - Balances creativity with prompt adherence
                
                Real-world applications: Character design, animation, book illustration
                """)
            with gr.Column():
                generate_scenes_btn = gr.Button("5. Generate Story Scenes", elem_classes="custom-button")
                scene_progress = gr.Textbox(label="Generation Progress", lines=6, interactive=False)
                gallery = gr.Gallery(label="Story Scenes", columns=2, height="auto", interactive=False)
                scene_prompts_display = gr.Textbox(label="Scene Details", lines=8, interactive=False)
        
        # Step 6: Add Text
        with gr.Row(elem_classes="section-content"):
            with gr.Column(elem_classes="ai-lesson"):
                gr.Markdown("""
                ### Step 6: Creating Book Pages 📚
                
                🤖 **AI Lesson: Computer Vision & Layout**
                Combining images and text requires sophisticated layout algorithms.
                How it works:
                - Analyzes image composition for text placement
                - Adjusts font size and style for readability
                - Creates visual hierarchy between elements
                - Ensures consistent formatting across pages
                
                Real-world applications: Desktop publishing, web design, digital books
                """)
            with gr.Column():
                add_text_btn = gr.Button("6. Add Text to Scenes", elem_classes="custom-button")
                final_gallery = gr.Gallery(label="Final Book Pages", columns=2, height="auto", interactive=False)
                download_btn = gr.File(label="Download Your Story Book", file_count="multiple", interactive=False)
        
        # Step 7: Audio Generation
        with gr.Row(elem_classes="section-content"):
            with gr.Column(elem_classes="ai-lesson"):
                gr.Markdown("""
                ### Step 7: Adding Narration 🎧
                
                🤖 **AI Lesson: Text-to-Speech Synthesis**
                Converting our story into natural-sounding speech.
                How it works:
                - Uses neural networks for voice synthesis
                - Adds appropriate emotion and emphasis
                - Controls pacing and pronunciation
                - Maintains consistent voice throughout
                
                Real-world applications: Audiobooks, accessibility tools, virtual assistants
                """)
            with gr.Column():
                tts_btn = gr.Button("7. Read Story Aloud", elem_classes="custom-button")
                audio_output = gr.Audio(label="Story Narration")

        # Event handlers
        generate_btn.click(fn=generate_image, outputs=image_output)
        analyze_btn.click(fn=analyze_image, inputs=[image_output], outputs=analysis_output)
        story_btn.click(fn=generate_story, inputs=[analysis_output], outputs=story_output)
        prompts_btn.click(fn=generate_image_prompts, inputs=[story_output], outputs=prompts_output)
        generate_scenes_btn.click(
            fn=generate_all_scenes,
            inputs=[prompts_output],
            outputs=[gallery, scene_prompts_display, scene_progress]
        )
        add_text_btn.click(
            fn=add_text_to_scenes,
            inputs=[gallery, prompts_output],
            outputs=[final_gallery, download_btn]
        )
        tts_btn.click(fn=generate_combined_audio_from_story, inputs=[story_output], outputs=audio_output)

    return demo
if __name__ == "__main__":
    demo = create_interface()
    demo.launch()