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 = """
AnimeGANv2 Cartoonizer
"""
# =====================================================
# HOME PAGE
# =====================================================
@app.get("/", response_class=HTMLResponse)
async def home():
return HTML
# =====================================================
# CARTOON API
# =====================================================
@app.post("/cartoon")
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
)