Spaces:
Sleeping
Sleeping
| import os | |
| import cv2 | |
| import torch | |
| import uvicorn | |
| import numpy as np | |
| from fastapi import FastAPI, UploadFile, File, Form | |
| from fastapi.responses import HTMLResponse, JSONResponse | |
| from fastapi.staticfiles import StaticFiles | |
| from inference import Predictor | |
| from utils.image_processing import resize_image | |
| # ===================================================== | |
| # PERFORMANCE OPTIMIZATION | |
| # ===================================================== | |
| torch.set_grad_enabled(False) | |
| cv2.setNumThreads(0) | |
| # ===================================================== | |
| # CREATE OUTPUT FOLDER | |
| # ===================================================== | |
| os.makedirs("output", exist_ok=True) | |
| # ===================================================== | |
| # FASTAPI APP | |
| # ===================================================== | |
| app = FastAPI( | |
| title="AnimeGANv2 API", | |
| description="Fast Anime Cartoon Generator" | |
| ) | |
| app.mount("/output", StaticFiles(directory="output"), name="output") | |
| # ===================================================== | |
| # DEVICE | |
| # ===================================================== | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"Using device: {DEVICE}") | |
| # ===================================================== | |
| # STYLE MAP | |
| # ===================================================== | |
| STYLE_MAP = { | |
| "AnimeGAN_Hayao": "hayao", | |
| "AnimeGAN_Shinkai": "shinkai", | |
| "AnimeGANv2_Hayao": "hayao:v2", | |
| "AnimeGANv2_Shinkai": "shinkai:v2", | |
| "AnimeGANv2_Arcane": "arcane:v2", | |
| } | |
| # ===================================================== | |
| # MODEL CACHE | |
| # ===================================================== | |
| predictors = {} | |
| def get_predictor(style, imgsz): | |
| key = f"{style}_{imgsz}" | |
| if key not in predictors: | |
| print(f"Loading model: {style}") | |
| predictors[key] = Predictor( | |
| weight=STYLE_MAP[style], | |
| device=DEVICE, | |
| retain_color=False, | |
| imgsz=imgsz, | |
| ) | |
| return predictors[key] | |
| # ===================================================== | |
| # RESPONSIVE UI | |
| # ===================================================== | |
| HTML = """ | |
| <!DOCTYPE html> | |
| <html lang="en"> | |
| <head> | |
| <meta charset="UTF-8"> | |
| <meta | |
| name="viewport" | |
| content="width=device-width, initial-scale=1.0" | |
| /> | |
| <title>AnimeGANv2 Cartoonizer</title> | |
| <style> | |
| *{ | |
| margin:0; | |
| padding:0; | |
| box-sizing:border-box; | |
| font-family:Arial,sans-serif; | |
| } | |
| body{ | |
| background:#0f172a; | |
| color:white; | |
| min-height:100vh; | |
| padding:15px; | |
| } | |
| .container{ | |
| max-width:1200px; | |
| margin:auto; | |
| } | |
| .header{ | |
| text-align:center; | |
| margin-bottom:25px; | |
| } | |
| .header h1{ | |
| font-size:40px; | |
| margin-bottom:10px; | |
| } | |
| .header p{ | |
| opacity:0.8; | |
| } | |
| .grid{ | |
| display:grid; | |
| grid-template-columns:1fr 1fr; | |
| gap:20px; | |
| } | |
| .upload{ | |
| border:2px dashed #374151; | |
| padding:40px; | |
| border-radius:20px; | |
| text-align:center; | |
| cursor:pointer; | |
| transition:0.3s; | |
| } | |
| .upload:hover{ | |
| border-color:#7c3aed; | |
| } | |
| .upload input{ | |
| display:none; | |
| } | |
| .preview, | |
| .result{ | |
| width:100%; | |
| border-radius:15px; | |
| margin-top:20px; | |
| display:none; | |
| } | |
| select, | |
| button{ | |
| width:100%; | |
| padding:15px; | |
| margin-top:15px; | |
| border:none; | |
| border-radius:12px; | |
| font-size:16px; | |
| } | |
| select{ | |
| background:#1f2937; | |
| color:white; | |
| } | |
| button{ | |
| background:linear-gradient( | |
| 135deg, | |
| #7c3aed, | |
| #2563eb | |
| ); | |
| color:white; | |
| cursor:pointer; | |
| font-weight:bold; | |
| } | |
| button:hover{ | |
| opacity:0.9; | |
| } | |
| .loader{ | |
| display:none; | |
| text-align:center; | |
| margin-top:15px; | |
| } | |
| .download-btn{ | |
| display:none; | |
| text-decoration:none; | |
| } | |
| .footer{ | |
| text-align:center; | |
| margin-top:30px; | |
| opacity:0.6; | |
| } | |
| @media(max-width:900px){ | |
| .grid{ | |
| grid-template-columns:1fr; | |
| } | |
| .header h1{ | |
| font-size:28px; | |
| } | |
| } | |
| </style> | |
| </head> | |
| <body> | |
| <div class="container"> | |
| <div class="header"> | |
| <h1>AnimeGANv2 Cartoonizer</h1> | |
| <p> | |
| Fast AI Anime Image Generator | |
| </p> | |
| </div> | |
| <div class="grid"> | |
| <div class="card"> | |
| <label class="upload"> | |
| <input | |
| type="file" | |
| id="imageInput" | |
| accept="image/*" | |
| > | |
| <h2>Upload Image</h2> | |
| <p> | |
| PNG JPG JPEG WEBP | |
| </p> | |
| </label> | |
| <img | |
| id="preview" | |
| class="preview" | |
| > | |
| <select id="style"> | |
| <option value="AnimeGANv2_Arcane"> | |
| AnimeGANv2 Arcane | |
| </option> | |
| <option value="AnimeGANv2_Hayao"> | |
| AnimeGANv2 Hayao | |
| </option> | |
| <option value="AnimeGANv2_Shinkai"> | |
| AnimeGANv2 Shinkai | |
| </option> | |
| <option value="AnimeGAN_Hayao"> | |
| AnimeGAN Hayao | |
| </option> | |
| <option value="AnimeGAN_Shinkai"> | |
| AnimeGAN Shinkai | |
| </option> | |
| </select> | |
| <select id="imgsz"> | |
| <option value="512" selected> | |
| Mobile Fast 512px | |
| </option> | |
| <option value="768"> | |
| HD 768px | |
| </option> | |
| <option value="1024"> | |
| Full HD 1024px | |
| </option> | |
| </select> | |
| <button onclick="convertImage()"> | |
| Convert To Anime | |
| </button> | |
| <div | |
| class="loader" | |
| id="loader" | |
| > | |
| Processing... | |
| </div> | |
| </div> | |
| <div class="card"> | |
| <h2>Output</h2> | |
| <img | |
| id="resultImage" | |
| class="result" | |
| > | |
| <a | |
| id="downloadBtn" | |
| class="download-btn" | |
| download | |
| > | |
| <button> | |
| Download PNG | |
| </button> | |
| </a> | |
| </div> | |
| </div> | |
| <div class="footer"> | |
| FastAPI AnimeGANv2 | |
| </div> | |
| </div> | |
| <script> | |
| const imageInput = | |
| document.getElementById("imageInput"); | |
| const preview = | |
| document.getElementById("preview"); | |
| imageInput.addEventListener("change",function(){ | |
| const file = this.files[0]; | |
| if(file){ | |
| const reader = new FileReader(); | |
| reader.onload = function(e){ | |
| preview.src = e.target.result; | |
| preview.style.display = "block"; | |
| } | |
| reader.readAsDataURL(file); | |
| } | |
| }); | |
| async function convertImage(){ | |
| const file = | |
| imageInput.files[0]; | |
| if(!file){ | |
| alert("Upload image first"); | |
| return; | |
| } | |
| const style = | |
| document.getElementById("style").value; | |
| const imgsz = | |
| document.getElementById("imgsz").value; | |
| const formData = new FormData(); | |
| formData.append("file",file); | |
| formData.append("style",style); | |
| formData.append("imgsz",imgsz); | |
| document.getElementById("loader") | |
| .style.display = "block"; | |
| try{ | |
| const response = await fetch("/cartoon",{ | |
| method:"POST", | |
| body:formData | |
| }); | |
| const data = await response.json(); | |
| document.getElementById("loader") | |
| .style.display = "none"; | |
| if(data.status === "success"){ | |
| const resultImage = | |
| document.getElementById("resultImage"); | |
| resultImage.src = data.image_url; | |
| resultImage.style.display = "block"; | |
| const downloadBtn = | |
| document.getElementById("downloadBtn"); | |
| downloadBtn.href = data.download_url; | |
| downloadBtn.style.display = "block"; | |
| }else{ | |
| alert(data.message); | |
| } | |
| }catch(error){ | |
| document.getElementById("loader") | |
| .style.display = "none"; | |
| alert("Server Error"); | |
| } | |
| } | |
| </script> | |
| </body> | |
| </html> | |
| """ | |
| # ===================================================== | |
| # HOME PAGE | |
| # ===================================================== | |
| async def home(): | |
| return HTML | |
| # ===================================================== | |
| # CARTOON API | |
| # ===================================================== | |
| async def cartoon( | |
| file: UploadFile = File(...), | |
| style: str = Form(...), | |
| imgsz: int = Form(512), | |
| ): | |
| try: | |
| contents = await file.read() | |
| npimg = np.frombuffer( | |
| contents, | |
| np.uint8 | |
| ) | |
| image = cv2.imdecode( | |
| npimg, | |
| cv2.IMREAD_COLOR | |
| ) | |
| image = cv2.cvtColor( | |
| image, | |
| cv2.COLOR_BGR2RGB | |
| ) | |
| # ========================================= | |
| # MOBILE SIZE FAST OUTPUT | |
| # ========================================= | |
| image = resize_image( | |
| image, | |
| width=imgsz | |
| ) | |
| predictor = get_predictor( | |
| style, | |
| imgsz | |
| ) | |
| with torch.no_grad(): | |
| anime_image = predictor.transform(image)[0] | |
| # ========================================= | |
| # SAVE PNG OUTPUT | |
| # ========================================= | |
| filename = f"anime_{os.path.splitext(file.filename)[0]}.png" | |
| save_path = os.path.join( | |
| "output", | |
| filename | |
| ) | |
| cv2.imwrite( | |
| save_path, | |
| cv2.cvtColor( | |
| anime_image, | |
| cv2.COLOR_RGB2BGR | |
| ), | |
| [cv2.IMWRITE_PNG_COMPRESSION, 1] | |
| ) | |
| return JSONResponse({ | |
| "status":"success", | |
| "image_url": | |
| f"/output/{filename}", | |
| "download_url": | |
| f"/output/{filename}" | |
| }) | |
| except Exception as e: | |
| print(e) | |
| return JSONResponse({ | |
| "status":"error", | |
| "message":str(e) | |
| }) | |
| # ===================================================== | |
| # RUN SERVER | |
| # ===================================================== | |
| if __name__ == "__main__": | |
| uvicorn.run( | |
| app, | |
| host="0.0.0.0", | |
| port=7860 | |
| ) |