Update utils.py

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
@@ -12,10 +13,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
 
@@ -29,6 +34,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:
@@ -58,7 +73,7 @@ def generate_image_from_pkl(generator, seed=0, trunc=1):
 def generate_image_from_onnx(path='model_128.onnx', model=None):
     if model is None: 
         return ValueError("Model not provided.")
-    if model == 'progan':
+    if model == 'progan' or model== 'dcgan':
         z = np.random.randn(1, 512, 1, 1).astype(np.float32)
     else:
         z = np.random.randn(1, 512).astype(np.float32)
@@ -68,6 +83,9 @@ def generate_image_from_onnx(path='model_128.onnx', model=None):
     image = inference_session.run(None, {input_name: z})[0]
 
     image = image.squeeze(0)
+    
+    if model == "vanillagan":
+        image = image.reshape(3, 128, 128)
     image = (image * 0.5 + 0.5) * 255
     image = image.astype(np.uint8)
     image = np.transpose(image, (1, 2, 0))
@@ -77,4 +95,4 @@ def generate_image_from_onnx(path='model_128.onnx', model=None):
     image.save(buffer, format='PNG')
     buffer.seek(0)
 
-    return buffer
+    return buffer