import numpy as np
import base64
import io
# import torch
from PIL import Image
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.mount('/static', StaticFiles(directory='static'), name='static')

# load model
# send to device

@app.get('/generate') # rate limit this
async def generate():
	# create random latent vector
	# forward latent vector to generator
	# convert result tensor into image with ToPILImage()
	# return image
	arr = np.random.randint(0, 256, (256, 256, 3)).astype(np.uint8)
	pil_img = Image.fromarray(arr)
	buffer = io.BytesIO()
	pil_img.save(buffer, format='PNG')
	buffer.seek(0)
	img_str = base64.b64encode(buffer.read()).decode('utf-8')
	return JSONResponse({ 'image': img_str })