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(""" """, unsafe_allow_html=True) st.markdown('

🎨 Anime Style Converter

', unsafe_allow_html=True) st.markdown('

Convert your images and videos to anime style using AnimeGANv2

', 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).")