feat(api): split up test scripts for diffusers and real esrgan

api/test-diffusers.py

@@ -0,0 +1,46 @@
+from diffusers import OnnxStableDiffusionPipeline
+from os import path
+
+import cv2
+import numpy as np
+import onnxruntime as ort
+import torch
+import time
+
+cfg = 8
+steps = 22
+height = 512
+width = 512
+
+model = path.join('..', 'models', 'stable-diffusion-onnx-v1-5')
+prompt = 'an astronaut eating a hamburger'
+output = path.join('..', 'outputs', 'test.png')
+
+print('generating test image...')
+pipe = OnnxStableDiffusionPipeline.from_pretrained(model, provider='DmlExecutionProvider', safety_checker=None)
+image = pipe(prompt, height, width, num_inference_steps=steps, guidance_scale=cfg).images[0]
+image.save(output)
+print('saved test image to %s' % output)
+
+
+upscale = path.join('..', 'outputs', 'test-large.png')
+esrgan = path.join('..', 'models', 'RealESRGAN_x4plus.onnx')
+
+print('upscaling test image...')
+sess = ort.InferenceSession(esrgan, providers=['DmlExecutionProvider'])
+
+in_image = cv2.imread(output, cv2.IMREAD_UNCHANGED)
+
+in_mat = cv2.cvtColor(in_image, cv2.COLOR_BGR2RGB)
+in_mat = np.transpose(in_mat, (2, 1, 0))[np.newaxis]
+in_mat = in_mat.astype(np.float32)
+in_mat = in_mat/255
+
+start_time = time.time()
+input_name = sess.get_inputs()[0].name
+output_name = sess.get_outputs()[0].name
+in_mat = torch.tensor(in_mat)
+out_mat = sess.run([output_name], {input_name: in_mat.cpu().numpy()})[0]
+elapsed_time = time.time() - start_time
+print(elapsed_time)
+print('upscaled test image to %s')

api/test-resrgan.py

@@ -0,0 +1,52 @@
+from os import path
+
+import cv2
+import numpy as np
+import onnxruntime as ort
+import torch
+import time
+
+cfg = 8
+steps = 22
+height = 512
+width = 512
+
+esrgan = path.join('..', 'models', 'RealESRGAN_x4plus.onnx')
+output = path.join('..', 'outputs', 'test.png')
+upscale = path.join('..', 'outputs', 'test-large.png')
+
+print('upscaling test image...')
+session = ort.InferenceSession(esrgan, providers=['DmlExecutionProvider'])
+
+in_image = cv2.imread(output, cv2.IMREAD_UNCHANGED)
+
+in_mat = cv2.cvtColor(in_image, cv2.COLOR_BGR2RGB)
+print('shape before', np.shape(in_mat))
+in_mat = np.transpose(in_mat, (2, 1, 0))[np.newaxis]
+print('shape after', np.shape(in_mat))
+in_mat = in_mat.astype(np.float32)
+in_mat = in_mat/255
+
+start_time = time.time()
+input_name = session.get_inputs()[0].name
+output_name = session.get_outputs()[0].name
+in_mat = torch.tensor(in_mat)
+out_mat = session.run([output_name], {
+    input_name: in_mat.cpu().numpy()
+})[0]
+elapsed_time = time.time() - start_time
+print(elapsed_time)
+
+print('output shape', np.shape(out_mat))
+out_mat = np.squeeze(out_mat, (0))
+print(np.shape(out_mat))
+out_mat = np.transpose(out_mat, (2, 1, 0))
+print(out_mat, np.shape(out_mat))
+out_mat = np.clip(out_mat, 0.0, 1.0)
+out_mat = out_mat * 255
+out_mat = out_mat.astype(np.uint8)
+out_image = cv2.cvtColor(out_mat, cv2.COLOR_RGB2BGR)
+
+cv2.imwrite(upscale, out_image)
+
+print('upscaled test image to %s' % upscale)

api/test-setup.py

@@ -1,15 +0,0 @@
-from diffusers import OnnxStableDiffusionPipeline
-from os import path
-
-cfg = 8
-steps = 22
-height = 512
-width = 512
-
-model = path.join('..', 'models', 'stable-diffusion-onnx-v1-5')
-prompt = 'an astronaut eating a hamburger'
-output = path.join('..', 'outputs', 'test.png')
-
-pipe = OnnxStableDiffusionPipeline.from_pretrained(model, provider='DmlExecutionProvider', safety_checker=None)
-image = pipe(prompt, height, width, num_inference_steps=steps, guidance_scale=cfg).images[0]
-image.save(output)