add generator

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
-from generator import generate_stylegan2
-from stylegan_generator import generate_stylegan
-from dcgan_generator import generate_dcgan
+from stylegan2_generator import generate_stylegan2
+from generator import generate_gan, generate_dcgan, generate_progan, generate_stylegan
 
 def generate_image(model_name):
     if model_name == "GAN / VanillaGAN":

dcgan_generator.py

@@ -4,7 +4,7 @@ import onnxruntime as ort
 from PIL import Image
 
 LATENT_FEATURES = 512
-MODEL_PATH = os.path.join("model", "dcgan.onnx")
+MODEL_PATH = os.path.join("model", "batik_dcgan.onnx")
 
 model = ort.InferenceSession(MODEL_PATH)
 input_name = model.get_inputs()[0].name

generator.py

@@ -1,33 +1,56 @@
-import os
-import sys
-
-BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-STYLEGAN2_DIR = os.path.join(BASE_DIR, "stylegan2")
-MODEL_PATH = os.path.join(BASE_DIR, "model", "network-snapshot-000560.pkl")
-
-sys.path.append(STYLEGAN2_DIR)
-
-import torch
-import legacy
-import dnnlib
 import numpy as np
+import onnxruntime as ort
 from PIL import Image
+import os
+import math
+
+def out(output):
+    image = output.squeeze()
+    image = (image * 0.5 + 0.5) * 255
+    image = image.astype(np.uint8)
+    image = np.transpose(image, (1, 2, 0))
+    return Image.fromarray(image, "RGB").resize((512, 512), Image.LANCZOS)
+
+def generate_gan():
+    return Image.new("RGB", (512, 512), color="gray")
+
+def generate_progan():
+    model_path = os.path.join("model", "batik_progan.onnx")
+    session = ort.InferenceSession(model_path)
+
+    noise = np.random.randn(1, 512, 1, 1).astype(np.float32)
+    alpha = np.array([1.0], dtype=np.float32)
+
+    output = session.run(None, {
+        'z': noise,
+        'alpha': alpha,
+    })
+    return out(output[0])
+
+def generate_dcgan():
+    model_path = os.path.join("model", "batik_dcgan.onnx")
+    session = ort.InferenceSession(model_path)
+    noise = np.random.randn(1, 512, 1, 1).astype(np.float32)
+    input_name = session.get_inputs()[0].name
+    output = session.run(None, {input_name: noise})
+    return out(output[0])
+
+def generate_stylegan():
+    model_path = os.path.join("model", "batik_stylegan.onnx")
+    session = ort.InferenceSession(model_path)
+
+    LATENT_FEATURES = 512
+    RESOLUTION = 256
+    LAST_INDEX = math.log2(RESOLUTION) - 2
+
+    z = np.random.randn(1, LATENT_FEATURES).astype(np.float32)
+    alpha = np.array([1.0], dtype=np.float32)
+    steps = np.array([LAST_INDEX], dtype=np.int64)
+
+    output = session.run(None, {
+        'z': z,
+        'alpha': alpha,
+        'steps': steps
+    })
+    return out(output[0])
 
-torch.autograd.set_grad_enabled(False)
-torch.backends.cudnn.benchmark = True
-
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-with dnnlib.util.open_url(MODEL_PATH) as f:
-    G = legacy.load_network_pkl(f)['G_ema'].to(device)
-
-def generate_stylegan2():
-    seed = np.random.randint(0, 2**32)
-    z = torch.from_numpy(np.random.RandomState(seed).randn(1, G.z_dim)).to(device)
-    label = torch.zeros([1, G.c_dim], device=device)
-    img = G(z, label, truncation_psi=1.0, noise_mode='const')
-    img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
-
-    pil_img = Image.fromarray(img[0].cpu().numpy(), 'RGB')
-    resized = pil_img.resize((512, 512), Image.LANCZOS)
-    return resized

progan_generator.py

@@ -0,0 +1,32 @@
+import os
+import math
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
+
+LATENT_FEATURES = 512
+RESOLUTION = 256
+LAST_INDEX = math.log2(RESOLUTION) - 2
+
+MODEL_PATH = os.path.join("model", "batik_stylegan.onnx")
+model = ort.InferenceSession(MODEL_PATH)
+
+alpha = np.array([1.0], dtype=np.float32)
+steps = np.array([LAST_INDEX], dtype=np.int64)
+
+def generate_stylegan():
+    z = np.random.randn(1, LATENT_FEATURES).astype(np.float32)
+
+    output = model.run(None, {
+        'z': z,
+        'alpha': alpha,
+        'steps': steps
+    })[0]
+
+    image = output.squeeze(0)
+    image = (image * 0.5 + 0.5) * 255
+    image = image.astype(np.uint8)
+    image = np.transpose(image, (1, 2, 0))
+    pil_img = Image.fromarray(image, 'RGB')
+    
+    return pil_img.resize((512, 512), Image.LANCZOS)

stylegan2_generator.py

@@ -0,0 +1,33 @@
+import os
+import sys
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+STYLEGAN2_DIR = os.path.join(BASE_DIR, "stylegan2")
+MODEL_PATH = os.path.join(BASE_DIR, "model", "network-snapshot-000560.pkl")
+
+sys.path.append(STYLEGAN2_DIR)
+
+import torch
+import legacy
+import dnnlib
+import numpy as np
+from PIL import Image
+
+torch.autograd.set_grad_enabled(False)
+torch.backends.cudnn.benchmark = True
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+with dnnlib.util.open_url(MODEL_PATH) as f:
+    G = legacy.load_network_pkl(f)['G_ema'].to(device)
+
+def generate_stylegan2():
+    seed = np.random.randint(0, 2**32)
+    z = torch.from_numpy(np.random.RandomState(seed).randn(1, G.z_dim)).to(device)
+    label = torch.zeros([1, G.c_dim], device=device)
+    img = G(z, label, truncation_psi=1.0, noise_mode='const')
+    img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+
+    pil_img = Image.fromarray(img[0].cpu().numpy(), 'RGB')
+    resized = pil_img.resize((512, 512), Image.LANCZOS)
+    return resized