Update utils.py

utils.py

@@ -1,5 +1,4 @@
-from stylegan_model import StyleGAN
-from vanillagan_model import VanillaGAN
+from model import StyleGAN
 import torch
 from io import BytesIO
 from torchvision.utils import save_image
@@ -7,19 +6,16 @@ import numpy as np
 import legacy
 from PIL import Image
 import time
+import onnxruntime as ort
 
 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_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))
+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)
     model.eval()
     return model
 
@@ -33,16 +29,6 @@ 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:
@@ -68,3 +54,27 @@ def generate_image_from_pkl(generator, seed=0, trunc=1):
     print(f"Image generation time: {end - start:.2f} seconds")
 
     return buffer
+
+def generate_image_from_onnx(path='model_128.onnx', model=None):
+    if model is None: 
+        return ValueError("Model not provided.")
+    if model == 'progan':
+        z = np.random.randn(1, 512, 1, 1).astype(np.float32)
+    else:
+        z = np.random.randn(1, 512).astype(np.float32)
+    inference_session = ort.InferenceSession(path)
+    input_name = inference_session.get_inputs()[0].name
+
+    image = inference_session.run(None, {input_name: z})[0]
+
+    image = image.squeeze(0)
+    image = (image * 0.5 + 0.5) * 255
+    image = image.astype(np.uint8)
+    image = np.transpose(image, (1, 2, 0))
+    image = Image.fromarray(image, 'RGB')
+
+    buffer = BytesIO()
+    image.save(buffer, format='PNG')
+    buffer.seek(0)
+
+    return buffer