app.py

import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from transformers import pipeline
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import io
import base64
import json
import re
import time
import os
from typing import List, Dict, Tuple, Optional
import warnings
warnings.filterwarnings("ignore")

# Global variables for models
sd_pipe = None
tts_pipe = None

def initialize_models():
    """Initialize AI models on first use"""
    global sd_pipe, tts_pipe
    
    if sd_pipe is None:
        print("Loading Stable Diffusion model...")
        sd_pipe = StableDiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            torch_dtype=torch.float32,  # Use float32 for CPU
            safety_checker=None,
            requires_safety_checker=False
        )
        sd_pipe.scheduler = DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config)
        sd_pipe = sd_pipe.to("cpu")
        sd_pipe.enable_attention_slicing()
    
    if tts_pipe is None:
        try:
            print("Loading TTS model...")
            tts_pipe = pipeline("text-to-speech", model="microsoft/speecht5_tts", device=-1)
        except:
            print("TTS model not available, continuing without audio...")
            tts_pipe = None

class StorySegmenter:
    """Handles story segmentation and prompt enhancement"""
    
    @staticmethod
    def split_story(story: str, max_segments: int = 15) -> List[str]:
        """Split story into segments suitable for 10-second videos"""
        # Split by sentences first
        sentences = re.split(r'[.!?]+', story)
        sentences = [s.strip() for s in sentences if s.strip()]
        
        segments = []
        current_segment = ""
        
        for sentence in sentences:
            # If adding this sentence would make segment too long, start new segment
            if len(current_segment + " " + sentence) > 200 or len(segments) >= max_segments:
                if current_segment:
                    segments.append(current_segment.strip())
                    current_segment = sentence
                else:
                    segments.append(sentence)
            else:
                current_segment += (" " + sentence if current_segment else sentence)
        
        # Add final segment
        if current_segment:
            segments.append(current_segment.strip())
        
        return segments[:max_segments]
    
    @staticmethod
    def enhance_prompt(segment: str, character_name: str, character_traits: str, 
                      style: str, scene_context: str = "") -> str:
        """Enhance segment prompt with character and style information"""
        enhanced = f"{segment}. "
        enhanced += f"Character: {character_name} ({character_traits}). "
        enhanced += f"Style: {style}, high quality, detailed. "
        if scene_context:
            enhanced += f"Scene context: {scene_context}. "
        
        # Add negative prompts
        enhanced += "NEGATIVE: blurry, low quality, distorted, bad anatomy"
        
        return enhanced

class ConsistencyManager:
    """Manages visual consistency across segments"""
    
    def __init__(self, base_seed: int = 42):
        self.base_seed = base_seed
        self.character_prompt = ""
        self.scene_context = ""
        self.last_scene_elements = []
    
    def get_segment_seed(self, segment_index: int) -> int:
        """Get consistent seed for segment"""
        return self.base_seed + segment_index
    
    def update_context(self, segment: str):
        """Update scene context based on current segment"""
        # Simple context extraction - in production would use NLP
        if any(word in segment.lower() for word in ['house', 'home', 'room', 'kitchen']):
            self.scene_context = "indoor domestic setting"
        elif any(word in segment.lower() for word in ['forest', 'tree', 'nature', 'outdoor']):
            self.scene_context = "outdoor natural setting"
        elif any(word in segment.lower() for word in ['city', 'street', 'building']):
            self.scene_context = "urban setting"

class VideoGenerator:
    """Handles video/image generation for each segment"""
    
    def __init__(self):
        self.consistency_manager = ConsistencyManager()
    
    def generate_segment_image(self, enhanced_prompt: str, seed: int, 
                             width: int = 512, height: int = 512) -> Image.Image:
        """Generate image for a story segment"""
        initialize_models()
        
        if sd_pipe is None:
            # Fallback: create a placeholder image
            return self.create_placeholder_image(enhanced_prompt, width, height)
        
        try:
            generator = torch.Generator().manual_seed(seed)
            
            # Generate image
            with torch.no_grad():
                result = sd_pipe(
                    prompt=enhanced_prompt,
                    negative_prompt="blurry, low quality, distorted, bad anatomy, ugly",
                    num_inference_steps=20,  # Reduced for faster generation
                    guidance_scale=7.5,
                    width=width,
                    height=height,
                    generator=generator
                )
            
            return result.images[0]
        
        except Exception as e:
            print(f"Error generating image: {e}")
            return self.create_placeholder_image(enhanced_prompt, width, height)
    
    def create_placeholder_image(self, prompt: str, width: int, height: int) -> Image.Image:
        """Create a placeholder image when generation fails"""
        img = Image.new('RGB', (width, height), color='lightblue')
        draw = ImageDraw.Draw(img)
        
        # Try to load a font, fallback to default if not available
        try:
            font = ImageFont.truetype("arial.ttf", 20)
        except:
            font = ImageFont.load_default()
        
        # Wrap text
        words = prompt[:100].split()
        lines = []
        current_line = ""
        
        for word in words:
            if len(current_line + word) < 40:
                current_line += word + " "
            else:
                lines.append(current_line)
                current_line = word + " "
        if current_line:
            lines.append(current_line)
        
        # Draw text
        y_offset = height // 2 - (len(lines) * 25) // 2
        for line in lines:
            bbox = draw.textbbox((0, 0), line, font=font)
            text_width = bbox[2] - bbox[0]
            x_offset = (width - text_width) // 2
            draw.text((x_offset, y_offset), line, fill='black', font=font)
            y_offset += 25
        
        return img

class AudioGenerator:
    """Handles audio generation for segments"""
    
    @staticmethod
    def generate_segment_audio(text: str, speaker_id: int = 0) -> Optional[bytes]:
        """Generate audio for a text segment"""
        initialize_models()
        
        if tts_pipe is None:
            return None
        
        try:
            # Generate audio
            audio_data = tts_pipe(text)
            
            # Convert to bytes (simplified - in production would handle proper audio format)
            if 'audio' in audio_data:
                # Convert audio array to bytes representation
                audio_array = np.array(audio_data['audio'])
                return audio_array.tobytes()
            
        except Exception as e:
            print(f"Error generating audio: {e}")
        
        return None

def process_story(story_text: str, character_name: str, character_traits: str, 
                 style: str, enable_voiceover: bool, reference_image: Optional[Image.Image] = None,
                 progress=gr.Progress()) -> Tuple[List[Image.Image], List[str], str]:
    """Main processing function"""
    
    if not story_text.strip():
        return [], [], "Please enter a story to generate."
    
    if not character_name.strip():
        character_name = "Main Character"
    
    if not character_traits.strip():
        character_traits = "detailed character design"
    
    # Initialize components
    segmenter = StorySegmenter()
    video_gen = VideoGenerator()
    audio_gen = AudioGenerator()
    
    # Step 1: Split story into segments
    progress(0.1, "Splitting story into segments...")
    segments = segmenter.split_story(story_text)
    
    if not segments:
        return [], [], "Could not create segments from the story."
    
    progress(0.2, f"Created {len(segments)} segments")
    
    # Step 2: Generate content for each segment
    generated_images = []
    generated_audio_info = []
    
    for i, segment in enumerate(segments):
        progress((0.2 + 0.7 * (i / len(segments))), f"Generating segment {i+1}/{len(segments)}")
        
        # Update consistency context
        video_gen.consistency_manager.update_context(segment)
        
        # Enhance prompt
        enhanced_prompt = segmenter.enhance_prompt(
            segment, character_name, character_traits, style,
            video_gen.consistency_manager.scene_context
        )
        
        # Generate image
        seed = video_gen.consistency_manager.get_segment_seed(i)
        image = video_gen.generate_segment_image(enhanced_prompt, seed)
        generated_images.append(image)
        
        # Generate audio if enabled
        if enable_voiceover:
            audio_bytes = audio_gen.generate_segment_audio(segment)
            audio_info = f"Audio generated for segment {i+1}" if audio_bytes else f"Audio generation failed for segment {i+1}"
        else:
            audio_info = f"No audio (voiceover disabled)"
        
        generated_audio_info.append(audio_info)
        
        # Small delay to prevent overwhelming the system
        time.sleep(0.1)
    
    progress(1.0, "Generation complete!")
    
    # Create summary
    summary = f"""
## Generation Summary

**Story Segments**: {len(segments)}
**Character**: {character_name} ({character_traits})
**Style**: {style}
**Voiceover**: {'Enabled' if enable_voiceover else 'Disabled'}

### Segments Generated:
"""
    
    for i, segment in enumerate(segments):
        summary += f"\n**Segment {i+1}**: {segment[:100]}{'...' if len(segment) > 100 else ''}"
    
    return generated_images, generated_audio_info, summary

def create_interface():
    """Create the Gradio interface"""
    
    with gr.Blocks(title="AI Video Story Generator", theme=gr.themes.Soft()) as interface:
        
        gr.Markdown("""
        # 🎬 AI Video Story Generator
        
        Generate animated story segments with consistent characters and scenes. 
        Running on free CPU tier - generates 10-15 segments of ~10 seconds each.
        """)
        
        with gr.Row():
            with gr.Column(scale=2):
                story_input = gr.Textbox(
                    label="📖 Story Text",
                    placeholder="Enter your story here (up to 50,000 characters)...",
                    lines=10,
                    max_lines=20
                )
                
                with gr.Row():
                    character_name = gr.Textbox(
                        label="👤 Character Name",
                        placeholder="e.g., Alice, Hero, The Detective",
                        value="Main Character"
                    )
                    character_traits = gr.Textbox(
                        label="✨ Character Traits",
                        placeholder="e.g., young woman, brown hair, blue dress",
                        value="detailed character design"
                    )
                
                with gr.Row():
                    style = gr.Dropdown(
                        label="🎨 Visual Style",
                        choices=[
                            "anime style",
                            "realistic",
                            "cartoon style",
                            "fantasy art",
                            "sci-fi concept art",
                            "children's book illustration",
                            "comic book style"
                        ],
                        value="anime style"
                    )
                    enable_voiceover = gr.Checkbox(
                        label="🔊 Enable Voiceover",
                        value=False
                    )
                
                reference_image = gr.Image(
                    label="🖼️ Reference Image (Optional)",
                    type="pil"
                )
            
            with gr.Column(scale=1):
                generate_btn = gr.Button(
                    "🎬 Generate Story Video",
                    variant="primary",
                    size="lg"
                )
                
                gr.Markdown("""
                ### 📋 Instructions:
                1. Enter your story text
                2. Define your main character
                3. Choose visual style
                4. Enable voiceover if desired
                5. Click Generate!
                
                ### ⚡ Free Tier Limits:
                - Max 15 segments (~3 minutes)
                - CPU processing (slower)
                - Sequential generation
                """)
        
        # Results section
        with gr.Row():
            summary_output = gr.Markdown(label="📊 Generation Summary")
        
        # Generated content gallery
        with gr.Row():
            generated_gallery = gr.Gallery(
                label="🖼️ Generated Segments",
                show_label=True,
                elem_id="gallery",
                columns=3,
                rows=2,
                height="auto"
            )
        
        with gr.Row():
            audio_info = gr.Textbox(
                label="🔊 Audio Information",
                lines=5,
                interactive=False
            )
        
        # Download section
        gr.Markdown("""
        ### 📥 Download Instructions:
        1. Right-click on any image to save individual segments
        2. Use external tools to combine segments into final video
        3. Audio files (if generated) can be downloaded separately
        
        ### 🔧 Recommended Tools for Combining:
        - **Free**: OpenShot, DaVinci Resolve, Blender
        - **Online**: Kapwing, Canva Video Editor
        - **Command Line**: FFmpeg
        """)
        
        # Event handlers
        generate_btn.click(
            fn=process_story,
            inputs=[
                story_input,
                character_name,
                character_traits,
                style,
                enable_voiceover,
                reference_image
            ],
            outputs=[
                generated_gallery,
                audio_info,
                summary_output
            ],
            show_progress=True
        )
        
        # Example stories
        gr.Markdown("""
        ### 📚 Example Stories to Try:
        
        **Fantasy Adventure**: "A brave knight discovers a magical forest where the trees whisper ancient secrets. She meets a wise dragon who offers to teach her the old magic. Together they must stop an evil sorcerer from destroying the realm."
        
        **Sci-Fi Mystery**: "On a space station orbiting Mars, Detective Chen investigates strange disappearances. The security cameras show nothing, but she notices the artificial gravity fluctuating. Her investigation leads to a discovery that changes everything."
        
        **Children's Tale**: "Little Bear couldn't sleep because of the thunder. He decided to visit his friend Owl, who lived in the big oak tree. Owl taught him that storms bring rain for flowers, and showed him how lightning dances across the sky."
        """)
    
    return interface

# Create and launch the interface
if __name__ == "__main__":
    interface = create_interface()
    interface.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )