feat: update vanillagan + dcgan

app.py

@@ -1,28 +1,48 @@
 from fastapi import FastAPI, Query
+from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse
-from utils import load_model_pt, generate_image_stylegan, load_model_pkl, generate_image_from_pkl, generate_image_vanillagan
-import random
+from utils import load_model_pt, generate_image_stylegan, load_model_pkl, generate_image_from_pkl, generate_image_from_onnx
 
 app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware, 
+    allow_origins=["*"],
+    allow_methods=["GET"],
+    allow_headers=["*"],
+    allow_credentials=False, # gapake creds
+)
 styleganv2 = load_model_pkl("styleganv2.pkl")
-gan = load_model_pt("generator_VanillaGAN.pt", model_type="vanillagan")
 
 @app.get("/")
 def root():
-    return {"message": "Welcome to the FastAPI StyleGAN API"}
+    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(-1)):
-    if seed == -1: 
-        seed = random.randint(0, 65535)
+def generate_styleganv2(seed: int = Query(0)):
     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_vanillagan(gan)
+    image_stream = generate_image_from_onnx("vanillagan.onnx", model='vanillagan')
     return StreamingResponse(image_stream, media_type="image/png")