app.py

import streamlit as st
import torch
import cv2
import numpy as np
import tempfile
import os
import subprocess
from PIL import Image
from io import BytesIO
from flask import Flask, request, jsonify, send_file
import threading
import base64

# ============================================
# FLASK APP SETUP
# ============================================
flask_app = Flask(__name__)

# ============================================
# LOAD MODEL (Shared between Flask & Streamlit)
# ============================================
@st.cache_resource
def load_model():
    """Load AnimeGANv2 model from local cache"""
    with st.spinner("🔄 Loading AnimeGANv2 model..."):
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        model = torch.hub.load(
            "bryandlee/animegan2-pytorch:main",
            "generator",
            pretrained="face_paint_512_v2",
            trust_repo=True,
            force_reload=False
        )
        model = model.to(device).eval()
        return model, device

# Global model variable (for Flask endpoints)
_model = None
_device = None

def get_model():
    global _model, _device
    if _model is None:
        _model, _device = load_model()
    return _model, _device

# ============================================
# IMAGE CONVERSION FUNCTION
# ============================================
def image_to_anime(input_image, size=512):
    """Convert single image to anime style"""
    model, device = get_model()
    
    img = input_image.resize((size, size))
    img_np = np.array(img).astype(np.float32) / 127.5 - 1
    img_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0).to(device)
    
    with torch.no_grad():
        out = model(img_tensor)
    
    out_np = out.squeeze(0).permute(1, 2, 0).cpu().numpy()
    out_np = (out_np + 1) * 127.5
    out_np = np.clip(out_np, 0, 255).astype(np.uint8)
    result = Image.fromarray(out_np).resize(input_image.size, Image.LANCZOS)
    return result

# ============================================
# VIDEO CONVERSION FUNCTION
# ============================================
def video_to_anime(video_bytes, progress_callback=None):
    """Convert video to anime style frame by frame"""
    
    with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp_input:
        tmp_input.write(video_bytes)
        input_path = tmp_input.name
    
    cap = cv2.VideoCapture(input_path)
    
    if not cap.isOpened():
        os.unlink(input_path)
        return None, 0
    
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    if total_frames == 0:
        cap.release()
        os.unlink(input_path)
        return None, 0
    
    # Limit frames for performance
    max_frames = 300
    if total_frames > max_frames:
        total_frames = max_frames
    
    output_path = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    
    frame_count = 0
    for i in range(total_frames):
        ret, frame = cap.read()
        if not ret:
            break
        
        try:
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame_pil = Image.fromarray(frame_rgb)
            anime_pil = image_to_anime(frame_pil)
            anime_frame = cv2.cvtColor(np.array(anime_pil), cv2.COLOR_RGB2BGR)
            out.write(anime_frame)
            frame_count += 1
        except Exception as e:
            continue
        
        if progress_callback:
            progress_callback((i + 1) / total_frames)
    
    cap.release()
    out.release()
    
    if frame_count == 0:
        os.unlink(input_path)
        os.unlink(output_path)
        return None, 0
    
    # Fix video with ffmpeg
    fixed_path = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
    ffmpeg_cmd = f'ffmpeg -i {output_path} -c:v libx264 -pix_fmt yuv420p -y {fixed_path} -loglevel error'
    os.system(ffmpeg_cmd)
    
    if os.path.exists(fixed_path) and os.path.getsize(fixed_path) > 0:
        os.unlink(output_path)
        output_path = fixed_path
    
    os.unlink(input_path)
    
    return output_path, frame_count

# ============================================
# FLASK API ENDPOINTS
# ============================================

@flask_app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint"""
    return jsonify({
        "status": "healthy",
        "service": "Anime Style Converter API",
        "model_loaded": _model is not None
    })

@flask_app.route('/convert/image', methods=['POST'])
def convert_image_api():
    """Convert image to anime style via API"""
    try:
        if 'image' not in request.files:
            return jsonify({"error": "No image file provided"}), 400
        
        image_file = request.files['image']
        original_img = Image.open(image_file).convert('RGB')
        anime_img = image_to_anime(original_img)
        
        buf = BytesIO()
        anime_img.save(buf, format='PNG')
        buf.seek(0)
        
        if request.args.get('format') == 'base64':
            encoded = base64.b64encode(buf.getvalue()).decode('utf-8')
            return jsonify({
                "status": "success",
                "image_base64": encoded,
                "format": "png"
            })
        else:
            return send_file(
                buf,
                mimetype='image/png',
                as_attachment=True,
                download_name='anime_output.png'
            )
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@flask_app.route('/convert/video', methods=['POST'])
def convert_video_api():
    """Convert video to anime style via API"""
    try:
        if 'video' not in request.files:
            return jsonify({"error": "No video file provided"}), 400
        
        video_file = request.files['video']
        video_bytes = video_file.read()
        
        output_path, frame_count = video_to_anime(video_bytes)
        
        if output_path is None:
            return jsonify({"error": "Video conversion failed"}), 500
        
        return send_file(
            output_path,
            mimetype='video/mp4',
            as_attachment=True,
            download_name='anime_video.mp4'
        )
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@flask_app.route('/convert/batch', methods=['POST'])
def batch_convert_api():
    """Batch convert multiple images via API"""
    try:
        if 'images' not in request.files:
            return jsonify({"error": "No images provided"}), 400
        
        images = request.files.getlist('images')
        results = []
        
        for idx, img_file in enumerate(images):
            original_img = Image.open(img_file).convert('RGB')
            anime_img = image_to_anime(original_img)
            
            buf = BytesIO()
            anime_img.save(buf, format='PNG')
            encoded = base64.b64encode(buf.getvalue()).decode('utf-8')
            
            results.append({
                "index": idx,
                "original_filename": img_file.filename,
                "image_base64": encoded
            })
        
        return jsonify({
            "status": "success",
            "count": len(results),
            "results": results
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# ============================================
# RUN FLASK IN BACKGROUND THREAD
# ============================================
def run_flask():
    flask_app.run(host='0.0.0.0', port=5000, debug=False, use_reloader=False)

# Start Flask in background
flask_thread = threading.Thread(target=run_flask, daemon=True)
flask_thread.start()

# ============================================
# STREAMLIT UI
# ============================================
st.set_page_config(
    page_title="Anime Style Converter",
    page_icon="🎨",
    layout="wide"
)

# Custom CSS
st.markdown("""
<style>
    .stApp {
        background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);
    }
    .main-title {
        text-align: center;
        color: #4ecdc4;
        font-size: 3em;
        margin-bottom: 10px;
    }
    .subtitle {
        text-align: center;
        color: #aaa;
        margin-bottom: 30px;
    }
    .stButton > button {
        background: #4ecdc4;
        color: white;
        border: none;
        padding: 10px 24px;
        border-radius: 8px;
        font-weight: bold;
        transition: transform 0.2s;
    }
    .stButton > button:hover {
        transform: scale(1.02);
        background: #45b7b1;
    }
    .api-info {
        background: #2d2d44;
        padding: 15px;
        border-radius: 10px;
        margin-top: 20px;
    }
</style>
""", unsafe_allow_html=True)

st.markdown('<h1 class="main-title">🎨 Anime Style Converter</h1>', unsafe_allow_html=True)
st.markdown('<p class="subtitle">Convert your images and videos to anime style using AnimeGANv2</p>', unsafe_allow_html=True)

# API Info in Sidebar
with st.sidebar:
    st.markdown("### 🔌 API Endpoints")
    st.markdown("""
    **Base URL:** `https://your-space.hf.space`
    
    **Endpoints:**
    - `GET /health` - Health check
    - `POST /convert/image` - Convert image
    - `POST /convert/video` - Convert video
    - `POST /convert/batch` - Batch convert images
    
    **Usage Examples:**
    
    ```bash
    # Convert Image
    curl -X POST https://your-space.hf.space/convert/image \\
      -F "image=@photo.jpg" -o anime.png
    
    # Convert Video
    curl -X POST https://your-space.hf.space/convert/video \\
      -F "video=@clip.mp4" -o anime.mp4
    ```
    """)
    
    # Show API status
    import requests
    try:
        health = requests.get("http://localhost:5000/health", timeout=2)
        if health.status_code == 200:
            st.success("✅ Flask API is running!")
        else:
            st.warning("⚠️ Flask API status unknown")
    except:
        st.warning("⚠️ Flask API starting...")

# Load model
try:
    model, device = get_model()
    st.success(f"✅ Model loaded on: {device}")
except Exception as e:
    st.error(f"❌ Failed to load model: {e}")
    st.stop()

# ============================================
# UI TABS
# ============================================
tab1, tab2 = st.tabs(["🖼️ Image Converter", "🎬 Video Converter"])

# TAB 1: IMAGE CONVERTER
with tab1:
    st.markdown("### 🖼️ Convert Image to Anime Style")
    
    col1, col2 = st.columns(2)
    
    with col1:
        uploaded_image = st.file_uploader(
            "📤 Upload an image",
            type=['jpg', 'jpeg', 'png', 'webp'],
            key="image_upload"
        )
        
        if uploaded_image is not None:
            original_img = Image.open(uploaded_image).convert('RGB')
            st.image(original_img, caption="📷 Original Image")
    
    with col2:
        if uploaded_image is not None:
            if st.button("✨ Convert to Anime", key="img_convert"):
                with st.spinner("🎨 Converting to anime style..."):
                    try:
                        anime_img = image_to_anime(original_img)
                        st.image(anime_img, caption="🎨 Anime Style")
                        
                        buf = BytesIO()
                        anime_img.save(buf, format="PNG")
                        byte_im = buf.getvalue()
                        
                        st.download_button(
                            label="📥 Download Anime Image",
                            data=byte_im,
                            file_name="anime_output.png",
                            mime="image/png"
                        )
                    except Exception as e:
                        st.error(f"Conversion failed: {e}")

# TAB 2: VIDEO CONVERTER
with tab2:
    st.markdown("### 🎬 Convert Video to Anime Style")
    
    st.info("""
    ⚠️ **Tips for best results:**
    - Upload videos shorter than 10 seconds (max 300 frames)
    - MP4 format recommended
    - Keep resolution under 720p for faster processing
    """)
    
    uploaded_video = st.file_uploader(
        "📤 Upload a video",
        type=['mp4', 'avi', 'mov', 'mkv'],
        key="video_upload"
    )
    
    if uploaded_video is not None:
        st.video(uploaded_video)
        
        with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp:
            tmp.write(uploaded_video.read())
            tmp_path = tmp.name
        
        cap = cv2.VideoCapture(tmp_path)
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        duration = total_frames / fps if fps > 0 else 0
        cap.release()
        os.unlink(tmp_path)
        
        uploaded_video.seek(0)
        
        st.markdown(f"**📹 Video Info:** {duration:.1f} seconds, {fps} fps, {total_frames} frames")
        
        if total_frames > 500:
            st.warning(f"⚠️ Video has {total_frames} frames. Please upload a shorter video (max 500 frames).")
        elif total_frames == 0:
            st.error("❌ Could not read video file.")
        else:
            if st.button("🎬 Convert to Anime Video", key="vid_convert"):
                uploaded_video.seek(0)
                
                progress_bar = st.progress(0)
                status_text = st.empty()
                
                def update_progress(progress):
                    progress_bar.progress(progress)
                
                output_path, frame_count = video_to_anime(uploaded_video.read(), update_progress)
                
                if output_path and frame_count > 0 and os.path.exists(output_path) and os.path.getsize(output_path) > 0:
                    status_text.text(f"✅ Converted {frame_count} frames!")
                    st.success(f"✅ Successfully converted {frame_count} frames!")
                    st.video(output_path)
                    
                    with open(output_path, 'rb') as f:
                        video_bytes = f.read()
                    
                    st.download_button(
                        label="📥 Download Anime Video",
                        data=video_bytes,
                        file_name="anime_video.mp4",
                        mime="video/mp4"
                    )
                    
                    os.unlink(output_path)
                else:
                    status_text.text("❌ Conversion failed")
                    st.error("❌ Video conversion failed. Please try a shorter video (5-10 seconds).")