Spaces:
Sleeping
Sleeping
File size: 1,711 Bytes
70a1838 bfc3e04 70a1838 6948715 4c63c9a 70a1838 bfc3e04 70a1838 a3d2818 bfc3e04 70a1838 bfc3e04 70a1838 4c63c9a 70a1838 bfc3e04 a3d2818 bfc3e04 a3d2818 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | from fastapi import FastAPI, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from utils import load_model_pkl, generate_image_from_pkl, generate_image_from_onnx
from random import randint
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["GET"],
allow_headers=["*"],
allow_credentials=False, # gapake creds
)
styleganv2 = load_model_pkl("styleganv2.pkl")
@app.get("/")
def root():
return {"message": "Welcome to the StyleGAN API!"}
@app.get("/ping")
def ping():
return {"status": "pong"}
@app.get("/generate/stylegan")
def generate_stylegan_onnx():
image_stream = generate_image_from_onnx("stylegan.onnx", model='stylegan')
return StreamingResponse(image_stream, media_type="image/png")
@app.get("/generate/styleganv2")
def generate_styleganv2(seed: int = Query(None)):
if seed is None:
seed = randint(0, 2**32 - 1)
image_stream = generate_image_from_pkl(styleganv2, seed=seed, trunc=1)
return StreamingResponse(image_stream, media_type="image/png")
@app.get("/generate/progan")
def generate_progan():
image_stream = generate_image_from_onnx("progan.onnx", model='progan')
return StreamingResponse(image_stream, media_type="image/png")
@app.get("/generate/dcgan")
def generate_dcgan():
image_stream = generate_image_from_onnx("batik_dcgan.onnx", model='dcgan')
return StreamingResponse(image_stream, media_type="image/png")
@app.get("/generate/vanillagan")
def generate_vanillagan():
image_stream = generate_image_from_onnx("vanillagan.onnx", model='vanillagan')
return StreamingResponse(image_stream, media_type="image/png")
|