from fastapi import FastAPI, File, UploadFile
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
import shutil
from io import BytesIO
import torch
from PIL import Image

from vtoonify_model import cacartoon1

app = FastAPI()



model = cacartoon1(device='cuda' if torch.cuda.is_available() else 'cpu')

def generate_cartoon(image_bytes: bytes) -> bytes:
    image = Image.open(BytesIO(image_bytes))
    cartoon_image = model.generate_cartoon(image)
    with BytesIO() as output:
        cartoon_image.save(output, format="PNG")
        return output.getvalue()

@app.post("/upload/")
async def upload_image(file: UploadFile = File(...)):
    contents = await file.read()
    result_bytes = generate_cartoon(contents)
    return {"result": result_bytes}

app.mount("/", StaticFiles(directory="AB", html=True), name="static")

@app.get("/")
def index() -> FileResponse:
    return FileResponse(path="/app/AB/index.html", media_type="text/html")