app.py

#!/usr/bin/env python3
"""
Ultimate Anime Animation Generator with Pose Control - CPU Optimized
"""

import os
import time
import numpy as np
import gradio as gr
from PIL import Image
import torch
import cv2
import imageio
import mediapipe as mp
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from diffusers.utils import export_to_video
from transformers import pipeline
from tqdm import tqdm
import warnings
import logging

# Suppress warnings for cleaner output
warnings.filterwarnings("ignore")

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Best anime models optimized for CPU
MODEL_PRIORITY_LIST = [
    ("AnimeDiffusion", "Ojimi/anime-kawai-diffusion"),
    ("AnythingV5", "andite/anything-v5.0"),
    ("Waifu Diffusion v1.4", "hakurei/waifu-diffusion-v1-4"),
    ("Anime Stable Diffusion", "digiplay/animeStableDiffusion_v122")
]

# Animation constants optimized for CPU
DEFAULT_PROMPT = "masterpiece, best quality, 1girl, school uniform, classroom, sunlight"
MAX_DURATION = 8  # seconds (CPU limit)
MAX_FPS = 10  # Optimized for CPU
MAX_KEYFRAMES = 6  # To prevent OOM
POSE_SEQUENCES = {
    "Walking Cycle": ["standing", "walking_start", "walking_mid", "walking_end"],
    "Running Cycle": ["standing", "running_start", "running_mid", "running_end"],
    "Sitting": ["standing", "sitting_down", "sitting"],
    "Jumping": ["standing", "jumping_start", "jumping_mid", "jumping_end", "landing"],
    "Custom": []
}
POSE_MODIFIERS = list(POSE_SEQUENCES.keys()) + ["fighting stance", "reading", "writing", "thinking", "surprised"]
ANGLE_MODIFIERS = ["front", "side", "3/4", "back", "high angle", "low angle"]
EMOTION_MODIFIERS = ["smiling", "neutral", "serious", "angry", "surprised", "sad", "happy"]

class AnimeAnimationStudio:
    def __init__(self):
        self.device = "cpu"
        self.dtype = torch.float32
        self.pipe = None
        self.current_model = None
        self.pose_detector = mp.solutions.pose.Pose(
            static_image_mode=True,
            model_complexity=1,
            enable_segmentation=False,
            min_detection_confidence=0.5
        )
        self.load_model()
        
    def load_model(self, model_name: str = None) -> bool:
        """Load the best available anime model with error handling"""
        models_to_try = MODEL_PRIORITY_LIST.copy()
        if model_name:
            # Prioritize the requested model
            models_to_try = [m for m in models_to_try if m[0] == model_name] + models_to_try
            
        for name, repo in models_to_try:
            try:
                logger.info(f"Loading model: {name} ({repo})")
                
                # Clear previous model
                if self.pipe is not None:
                    del self.pipe
                    torch.cuda.empty_cache() if torch.cuda.is_available() else None
                
                # Load new model with optimizations
                self.pipe = StableDiffusionPipeline.from_pretrained(
                    repo,
                    torch_dtype=self.dtype,
                    safety_checker=None,
                    requires_safety_checker=False
                )
                self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
                    self.pipe.scheduler.config
                )
                self.pipe = self.pipe.to(self.device)
                self.current_model = name
                logger.info(f"Successfully loaded {name}")
                return True
            except Exception as e:
                logger.warning(f"Model load failed for {name}: {str(e)}")
                continue
        
        logger.error("All model loading attempts failed")
        return False

    def detect_pose(self, image: Image.Image):
        """Detect human pose from reference image"""
        if image is None:
            return None
            
        try:
            # Convert to OpenCV format
            image_np = np.array(image)
            results = self.pose_detector.process(image_np)
            
            if not results.pose_landmarks:
                return None
                
            # Extract key points
            key_points = []
            for landmark in results.pose_landmarks.landmark:
                key_points.append({
                    'x': landmark.x,
                    'y': landmark.y,
                    'visibility': landmark.visibility
                })
                
            return key_points
        except Exception as e:
            logger.error(f"Pose detection failed: {str(e)}")
            return None

    def generate_frame(self, prompt: str, negative_prompt: str = "", 
                      seed: int = -1, pose_description: str = "") -> Image.Image:
        """Generate a single anime frame with pose control"""
        if self.pipe is None:
            logger.error("No model loaded!")
            return None
            
        try:
            # Enhance prompt with pose information
            full_prompt = f"{prompt}, {pose_description}, masterpiece, best quality, anime style"
            
            generator = torch.Generator(device=self.device)
            if seed > 0:
                generator.manual_seed(seed)
                
            return self.pipe(
                prompt=full_prompt,
                negative_prompt=negative_prompt,
                generator=generator,
                num_inference_steps=25,  # More steps for better quality
                guidance_scale=8.0,      # Higher guidance for anime
                width=512,
                height=512
            ).images[0]
        except Exception as e:
            logger.error(f"Frame generation failed: {str(e)}")
            return None

    def generate_keyframes(self, base_prompt: str, negative_prompt: str, 
                          seed: int, pose_sequence: list, emotion: str,
                          progress) -> list:
        """Generate key animation frames based on pose sequence"""
        keyframes = []
        status_messages = []
        
        for i, pose in progress.tqdm(enumerate(pose_sequence), desc="Generating keyframes", total=len(pose_sequence)):
            # Create detailed pose description
            pose_desc = f"{pose} pose" if pose != "Custom" else base_prompt
            
            # Add emotion to the pose
            if emotion and emotion != "none":
                pose_desc += f", {emotion} expression"
            
            frame = None
            for attempt in range(2):  # Retry once if fails
                frame = self.generate_frame(
                    prompt=base_prompt,
                    negative_prompt=negative_prompt,
                    seed=seed + i if seed > 0 else -1,
                    pose_description=pose_desc
                )
                if frame is not None:
                    break
                self.load_model()  # Try reloading model if failed
                
            if frame is None:
                return None, f"Keyframe {i+1} generation failed"
                
            keyframes.append(np.array(frame))
            status_messages.append(f"Generated {pose} pose")
        
        return keyframes, "\n".join(status_messages)

    def create_animation(self, keyframes: list, fps: int, 
                         transition: str, progress) -> str:
        """Create animation from keyframes with smooth transitions"""
        try:
            # Create transitions between keyframes
            frames = []
            total_segments = len(keyframes) - 1
            transition_frames = max(1, int(fps * 0.5))  # Half-second transitions
            
            for i in progress.tqdm(range(total_segments), desc="Creating animation"):
                # Add the current keyframe
                frames.append(keyframes[i])
                
                # Create transition between current and next keyframe
                start_frame = keyframes[i]
                end_frame = keyframes[i+1]
                
                for j in range(transition_frames):
                    alpha = j / transition_frames
                    
                    if transition == "Crossfade":
                        # Simple crossfade
                        blended = cv2.addWeighted(
                            start_frame, 1 - alpha,
                            end_frame, alpha,
                            0
                        )
                    else:  # Slide effect
                        # Calculate slide offset
                        offset = int(alpha * 100)
                        
                        # Create a new frame with sliding content
                        blended = np.zeros_like(start_frame)
                        w, h, _ = start_frame.shape
                        
                        # Start frame slides out
                        blended[:, :w-offset] = start_frame[:, offset:]
                        
                        # End frame slides in
                        blended[:, w-offset:] = end_frame[:, :offset]
                        
                    frames.append(blended)
            
            # Add the last keyframe
            frames.append(keyframes[-1])
            
            # Save as video
            output_path = f"anime_animation_{int(time.time())}.mp4"
            with imageio.get_writer(output_path, fps=fps, macro_block_size=1) as writer:
                for frame in frames:
                    writer.append_data(frame)
            
            return output_path, None
        except Exception as e:
            return None, f"Animation creation failed: {str(e)}"

    def generate_animation(
        self,
        prompt: str,
        duration: float,
        fps: int,
        negative_prompt: str,
        seed: int,
        pose_sequence: list,
        emotion: str,
        transition: str,
        reference_image: Image.Image,
        progress=gr.Progress()
    ) -> tuple:
        """Main animation generation workflow"""
        start_time = time.time()
        
        # Validate inputs
        if duration <= 0 or fps <= 0:
            return None, "Duration and FPS must be positive"
            
        duration = min(duration, MAX_DURATION)
        fps = min(fps, MAX_FPS)
        
        # Use pose detection if reference image is provided
        pose_description = ""
        if reference_image is not None:
            pose_data = self.detect_pose(reference_image)
            if pose_data:
                pose_description = "specific pose from reference"
                logger.info("Using pose from reference image")
        
        # Generate keyframes
        keyframes, status = self.generate_keyframes(
            prompt, 
            negative_prompt, 
            seed, 
            pose_sequence,
            emotion,
            progress
        )
        if keyframes is None:
            return None, status
        
        # Create animation from keyframes
        video_path, error = self.create_animation(
            keyframes, 
            fps, 
            transition,
            progress
        )
        
        if video_path is None:
            return None, error
        
        # Final status
        time_taken = time.time() - start_time
        status = f"✨ Animation created in {time_taken:.1f}s\nModel: {self.current_model}\n{status}"
        
        return video_path, status

def create_ui():
    studio = AnimeAnimationStudio()
    
    # Custom CSS for anime-style UI
    custom_css = """
    .gradio-container {
        background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);
        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
    }
    .gr-button {
        background: linear-gradient(45deg, #ff7aa2, #ff9ec4) !important;
        border: none !important;
        color: white !important;
        font-weight: bold !important;
        border-radius: 20px !important;
        padding: 12px 25px !important;
        transition: all 0.3s ease !important;
    }
    .gr-button:hover {
        transform: translateY(-2px);
        box-shadow: 0 7px 14px rgba(255, 122, 162, 0.3) !important;
    }
    .gr-box {
        background-color: rgba(255, 255, 255, 0.08) !important;
        border: 1px solid rgba(255, 255, 255, 0.1) !important;
        border-radius: 10px !important;
        color: white !important;
    }
    .gr-input, .gr-textbox, .gr-number, .gr-slider, .gr-dropdown {
        background-color: rgba(26, 26, 46, 0.7) !important;
        border: 1px solid #4e4376 !important;
        color: white !important;
        border-radius: 8px !important;
    }
    h1, h2, h3, h4, .gr-label {
        color: #ff9ec4 !important;
    }
    .gr-progress-bar {
        background: linear-gradient(45deg, #ff7aa2, #ff9ec4) !important;
    }
    .tabs {
        background: rgba(26, 26, 46, 0.7) !important;
        border-radius: 10px;
        padding: 15px;
    }
    .tab-nav {
        background: transparent !important;
    }
    """
    
    with gr.Blocks(title="Anime Animation Studio", css=custom_css, theme=gr.themes.Default()) as demo:
        # Header with logo and title
        with gr.Row():
            gr.HTML("""
            <div style="text-align: center; width: 100%;">
                <h1 style="font-size: 2.5rem; margin-bottom: 10px; background: linear-gradient(45deg, #ff7aa2, #ff9ec4, #b8e1ff);
                    -webkit-background-clip: text; -webkit-text-fill-color: transparent;">
                    🎬 Ultimate Anime Animation Studio
                </h1>
                <p style="color: #b8e1ff; font-size: 1.1rem;">
                    Create professional anime animations from text prompts - Powered by AI
                </p>
            </div>
            """)
        
        with gr.Row():
            # Left column - Inputs
            with gr.Column(scale=4):
                with gr.Tab("Animation Settings"):
                    # Model selection
                    model_dropdown = gr.Dropdown(
                        label="🎨 Anime Model",
                        choices=[m[0] for m in MODEL_PRIORITY_LIST],
                        value=studio.current_model or MODEL_PRIORITY_LIST[0][0],
                        interactive=True
                    )
                    
                    # Prompt inputs
                    with gr.Row():
                        with gr.Column(scale=3):
                            prompt_textbox = gr.Textbox(
                                label="📝 Scene Description",
                                value=DEFAULT_PROMPT,
                                lines=3,
                                placeholder="Describe your anime scene in detail..."
                            )
                        with gr.Column(scale=1):
                            emotion_dropdown = gr.Dropdown(
                                label="😄 Character Emotion",
                                choices=EMOTION_MODIFIERS,
                                value="smiling"
                            )
                    
                    negative_prompt_textbox = gr.Textbox(
                        label="🚫 Negative Prompts",
                        value="blurry, low quality, distorted, text, watermark, signature, bad anatomy",
                        lines=2,
                        placeholder="What to avoid in the animation..."
                    )
                    
                    # Pose sequence
                    pose_sequence = gr.CheckboxGroup(
                        label="💃 Pose Sequence (drag to reorder)",
                        choices=POSE_MODIFIERS,
                        value=["Walking Cycle"],
                        interactive=True
                    )
                    
                    # Reference image
                    reference_image = gr.Image(
                        label="🖼️ Reference Pose (optional)",
                        type="pil",
                        interactive=True
                    )
                
                with gr.Tab("Advanced Settings"):
                    with gr.Row():
                        with gr.Column():
                            duration_slider = gr.Slider(
                                label=f"⏱️ Duration (seconds, max {MAX_DURATION})",
                                minimum=1,
                                maximum=MAX_DURATION,
                                value=4,
                                step=0.5
                            )
                            fps_slider = gr.Slider(
                                label=f"🎞️ FPS (max {MAX_FPS})",
                                minimum=3,
                                maximum=MAX_FPS,
                                value=6,
                                step=1
                            )
                        with gr.Column():
                            transition_radio = gr.Radio(
                                label="🔄 Transition Effect",
                                choices=["Crossfade", "Slide"],
                                value="Crossfade"
                            )
                            seed_number = gr.Number(
                                label="🌱 Seed (-1 for random)",
                                value=-1,
                                precision=0
                            )
                    
                    gr.Markdown("### 💡 Tips for Professional Results")
                    gr.Markdown("- Use detailed descriptions of characters, outfits, and settings")
                    gr.Markdown("- Start with 3-4 poses for smooth animations")
                    gr.Markdown("- Add emotions to bring characters to life")
                    gr.Markdown("- Use reference images for specific poses")
                
                generate_button = gr.Button(
                    "✨ Generate Anime Animation", 
                    variant="primary",
                    size="lg"
                )
            
            # Right column - Output
            with gr.Column(scale=3):
                output_video = gr.Video(
                    label="🎥 Generated Animation",
                    format="mp4",
                    interactive=False
                )
                
                status_text = gr.Textbox(
                    label="📊 Status",
                    value=f"Loaded: {studio.current_model}" if studio.current_model else "Ready to create anime!",
                    interactive=False
                )
                
                with gr.Accordion("Example Prompts", open=False):
                    gr.Examples(
                        examples=[
                            ["masterpiece, best quality, 1girl, school uniform, classroom, sunlight, looking at viewer"],
                            ["cyberpunk city, anime girl with neon katana, rain, glowing eyes, dynamic pose"],
                            ["fantasy forest, elf archer, glowing bow, magical particles, aiming at target"],
                            ["sci-fi spaceship bridge, captain's chair, confident female captain, stars through window"]
                        ],
                        inputs=prompt_textbox,
                        label="Click to apply example prompts"
                    )
        
        # Model change handler
        def on_model_change(model_name):
            if studio.load_model(model_name):
                return f"Model changed to: {studio.current_model}"
            return "Model change failed!"
        
        model_dropdown.change(
            on_model_change,
            inputs=model_dropdown,
            outputs=status_text
        )
        
        # Generation handler
        generate_button.click(
            studio.generate_animation,
            inputs=[
                prompt_textbox,
                duration_slider,
                fps_slider,
                negative_prompt_textbox,
                seed_number,
                pose_sequence,
                emotion_dropdown,
                transition_radio,
                reference_image
            ],
            outputs=[output_video, status_text]
        )
    
    return demo

if __name__ == "__main__":
    # Set seed for reproducibility
    torch.manual_seed(42)
    
    demo = create_ui()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        enable_queue=True,
        share=False,
        show_error=True,
        favicon_path="anime_icon.png"
    )