feat: add vanilla gan

app.py

@@ -1,10 +1,16 @@
 from fastapi import FastAPI, Query
 from fastapi.responses import StreamingResponse
-from utils import load_model_pt, generate_image_stylegan, load_model_pkl, generate_image_from_pkl
+from utils import load_model_pt, generate_image_stylegan, load_model_pkl, generate_image_from_pkl, generate_image_vanillagan
+import random
 
 app = FastAPI()
-stylegan = load_model_pt("model_128.pt")
+stylegan = load_model_pt("model_128.pt", model_type="stylegan")
 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"}
 
 @app.get("/ping")
 def ping():
@@ -16,7 +22,14 @@ def generate_stylegan():
     return StreamingResponse(image_stream, media_type="image/png")
 
 @app.get("/generate/styleganv2")
-def generate_styleganv2(seed: int = Query(0)):
+def generate_styleganv2(seed: int = Query(-1)):
+    if seed == -1: 
+        seed = random.randint(0, 65535)
     image_stream = generate_image_from_pkl(styleganv2, seed=seed, trunc=1)
     return StreamingResponse(image_stream, media_type="image/png")
 
+@app.get("/generate/vanillagan")
+def generate_vanillagan():
+    image_stream = generate_image_vanillagan(gan)
+    return StreamingResponse(image_stream, media_type="image/png")
+

model.py → stylegan_model.py

@@ -1,7 +1,7 @@
-from torch import nn, optim
+from torch import nn
 import torch
 from torch.nn import functional as F
-from typing import Any, Callable, Optional
+from typing import Optional
 import math
 
 class WSLinear(nn.Module):

utils.py

@@ -1,4 +1,5 @@
-from model import StyleGAN
+from stylegan_model import StyleGAN
+from vanillagan_model import VanillaGAN
 import torch
 from io import BytesIO
 from torchvision.utils import save_image
@@ -11,10 +12,14 @@ LATENT_FEATURES = 512
 RESOLUTION = 128
 
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-def load_model_pt(path='model_128.pt'):
-    model = StyleGAN(LATENT_FEATURES, RESOLUTION).to(DEVICE)
-    last_checkpoint = torch.load(path, map_location=DEVICE)
-    model.load_state_dict(last_checkpoint['generator'], strict=False)
+def load_model_pt(path='model_128.pt',model_type='stylegan'):
+    if model_type == "stylegan":
+        model = StyleGAN(LATENT_FEATURES, RESOLUTION).to(DEVICE)
+        last_checkpoint = torch.load(path, map_location=DEVICE)
+        model.load_state_dict(last_checkpoint['generator'], strict=False)
+    elif model_type == "vanillagan":
+        model = VanillaGAN(RESOLUTION, LATENT_FEATURES).to(DEVICE)
+        model.load_state_dict(torch.load(path, map_location=DEVICE))
     model.eval()
     return model
 
@@ -28,6 +33,16 @@ def generate_image_stylegan(generator, steps=5, alpha=1.0):
         save_image(image, buffer, format='PNG')
         buffer.seek(0)
         return buffer
+
+def generate_image_vanillagan(generator):
+    with torch.no_grad():
+        image = generator(torch.randn(1, LATENT_FEATURES, device=DEVICE)).view(1, 3, RESOLUTION, RESOLUTION)
+        image = (image * 0.5 + 0.5).clamp(0, 1) 
+
+        buffer = BytesIO()
+        save_image(image, buffer, format='PNG')
+        buffer.seek(0)
+        return buffer
     
 def load_model_pkl(path='styleganv2.pkl'):
     with open(path, 'rb') as f:

vanillagan_model.py

@@ -0,0 +1,32 @@
+from torch import nn, optim
+import torch
+from torch.nn import functional as F
+from typing import Any, Callable, Optional
+import math
+
+class VanillaGAN(nn.Module):
+    def __init__(self, resolution, latent_dim, hidden_dim=512, channels=3):
+        super(VanillaGAN, self).__init__()
+        output_dim = resolution * resolution * channels
+
+        self.layers = nn.Sequential(
+            self.gen_block(latent_dim, hidden_dim),
+            self.gen_block(hidden_dim, hidden_dim*2),
+             self.gen_block(hidden_dim*2, hidden_dim*2),
+            self.gen_block(hidden_dim*2, hidden_dim),
+            self.gen_block(hidden_dim, hidden_dim),
+            self.gen_block(hidden_dim, hidden_dim//2),
+            
+            nn.Linear(hidden_dim//2, output_dim),
+            nn.Tanh()
+        )
+
+    def gen_block(self, input_dim, output_dim):
+        return nn.Sequential(
+            nn.Linear(input_dim, output_dim, bias=False),
+            nn.BatchNorm1d(output_dim, 0.8),
+            nn.LeakyReLU(0.2)
+        )
+
+    def forward(self, x):
+        return self.layers(x)