diff --git a/.gitignore b/.gitignore index 6010b30..dea414e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ model_zoo/ outputs/ +*benchmark_tmp.csv # Byte-compiled / optimized / DLL files __pycache__/ @@ -130,6 +131,7 @@ venv/ ENV/ env.bak/ venv.bak/ +.venv*/ # Spyder project settings .spyderproject @@ -160,4 +162,4 @@ cython_debug/ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ \ No newline at end of file +#.idea/ diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..de288e1 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "python.formatting.provider": "black" +} \ No newline at end of file diff --git a/README.md b/README.md index eea53d1..eff5b26 100644 --- a/README.md +++ b/README.md @@ -98,6 +98,32 @@ for idx, im in enumerate(images): im.save(f"{idx:06}.png") ``` +## Benchmarking inference + +Detailed benchmark documentation can be found [here](./docs/benchmark.md). + +### Setup + +Before running the benchmark, make sure you have completed the repository [installation steps](#installation). + +You will then need to set the huggingface access token: +1. Create a user account on HuggingFace and generate an access token. +2. Set your huggingface access token as the `ACCESS_TOKEN` environment variable: +``` +export ACCESS_TOKEN= +``` + +### Usage + +Launch the benchmark script to append benchmark results to the existing [benchmark.csv](./benchmark.csv) results file: +``` +python ./scripts/benchmark.py +``` + +### Results + +Stable Diffusion Text2Image Latency (seconds) + ## Links - [Captioned Pokémon dataset](https://huggingface.co/datasets/lambdalabs/pokemon-blip-captions) diff --git a/benchmark.csv b/benchmark.csv new file mode 100644 index 0000000..5a5e51d --- /dev/null +++ b/benchmark.csv @@ -0,0 +1,58 @@ +Intel(R) Core(TM) i7-6850K CPU @ 3.60GHz,single,pytorch,1,458.97,0.0 +Intel(R) Core(TM) i7-6850K CPU @ 3.60GHz,single,onnx,1,286.13,0.0 +NVIDIA GeForce RTX 3090,single,pytorch,1,7.96,7.72 +NVIDIA GeForce RTX 3090,half,pytorch,1,4.83,4.54 +NVIDIA GeForce RTX 3090,single,pytorch,2,14.49,11 +NVIDIA GeForce RTX 3090,half,pytorch,2,8.42,8.75 +NVIDIA GeForce RTX 3090,single,pytorch,4,27.94,17.69 +NVIDIA GeForce RTX 3090,half,pytorch,4,15.87,15.36 +NVIDIA GeForce RTX 3090,single,pytorch,8,-1.0,-1.0 +NVIDIA GeForce RTX 3090,half,pytorch,8,-1.0,-1.0 +NVIDIA RTX A5500,single,pytorch,1,8.55,7.69 +NVIDIA RTX A5500,half,pytorch,1,5.05,4.58 +NVIDIA RTX A5500,single,pytorch,2,15.71,11 +NVIDIA RTX A5500,half,pytorch,2,9.37,8.8 +NVIDIA RTX A5500,single,pytorch,4,30.51,17.69 +NVIDIA RTX A5500,half,pytorch,4,16.97,15.33 +NVIDIA RTX A5500,single,pytorch,8,-1.0,-1.0 +NVIDIA RTX A5500,half,pytorch,8,-1.0,-1.0 +AMD EPYC 7352 24-Core Processor,single,pytorch,1,529.93,0.0 +AMD EPYC 7352 24-Core Processor,single,onnx,1,223.19,0.0 +NVIDIA GeForce RTX 3080,single,pytorch,4,-1.0,-1.0 +NVIDIA GeForce RTX 3080,half,pytorch,4,-1.0,-1.0 +NVIDIA GeForce RTX 3080,single,pytorch,1,-1.0,-1.0 +NVIDIA GeForce RTX 3080,half,pytorch,1,5.59,4.52 +NVIDIA GeForce RTX 3080,single,pytorch,2,-1.0,-1.0 +NVIDIA GeForce RTX 3080,half,pytorch,2,-1.0,-1.0 +NVIDIA A100 80GB PCIe,single,pytorch,1,6.39,7.75 +NVIDIA A100 80GB PCIe,half,pytorch,1,3.74,4.55 +NVIDIA A100 80GB PCIe,single,pytorch,2,11.12,11.05 +NVIDIA A100 80GB PCIe,half,pytorch,2,5.72,8.77 +NVIDIA A100 80GB PCIe,single,pytorch,4,20.18,17.63 +NVIDIA A100 80GB PCIe,half,pytorch,4,10.04,15.34 +NVIDIA A100 80GB PCIe,single,pytorch,8,38.88,30.88 +NVIDIA A100 80GB PCIe,half,pytorch,8,18.68,28.47 +NVIDIA A100 80GB PCIe,single,pytorch,16,76.92,57.46 +NVIDIA A100 80GB PCIe,half,pytorch,16,36.67,54.73 +NVIDIA A100 80GB PCIe,half,pytorch,28,63.88,78.78 +NVIDIA RTX A6000,single,pytorch,1,8.09,7.75 +NVIDIA RTX A6000,half,pytorch,1,5.03,4.53 +NVIDIA RTX A6000,single,pytorch,2,14.86,10.98 +NVIDIA RTX A6000,half,pytorch,2,9.03,8.79 +NVIDIA RTX A6000,single,pytorch,4,27.92,17.62 +NVIDIA RTX A6000,half,pytorch,4,17.0,15.34 +NVIDIA RTX A6000,single,pytorch,8,53.95,30.88 +NVIDIA RTX A6000,half,pytorch,8,32.57,28.51 +NVIDIA RTX A6000,half,pytorch,16,63.16,46.11 +Quadro RTX 8000,single,pytorch,1,12.3,7.71 +Quadro RTX 8000,half,pytorch,1,5.93,4.52 +Quadro RTX 8000,single,pytorch,2,24.42,9.16 +Quadro RTX 8000,half,pytorch,2,10.92,7.02 +Quadro RTX 8000,single,pytorch,4,42.56,15.58 +Quadro RTX 8000,half,pytorch,4,21.24,12.39 +Quadro RTX 8000,single,pytorch,8,76.96,23.11 +Quadro RTX 8000,half,pytorch,8,40.52,20.98 +Quadro RTX 8000,single,pytorch,16,152.55,42.47 +Quadro RTX 8000,half,pytorch,16,80.31,38.18 +Quadro RTX 8000,single,pytorch,32,-1.0,-1.0 +Quadro RTX 8000,half,pytorch,32,-1.0,-1.0 diff --git a/docs/benchmark.md b/docs/benchmark.md new file mode 100644 index 0000000..9d8e92c --- /dev/null +++ b/docs/benchmark.md @@ -0,0 +1,112 @@ +# Benchmarking Diffuser Models + +We present a benchmark of [Stable Diffusion](https://huggingface.co/CompVis/stable-diffusion) model inference. This text2image model uses a text prompt as input and outputs an image of resolution `512x512`. + +Our experiments analyze inference performance in terms of speed, memory consumption, throughput, and quality of the output images. We look at how different choices in hardware (GPU model, GPU vs CPU) and software (single vs half precision, pytorch vs onnxruntime) affect inference performance. + +For reference, we will be providing benchmark results for the following GPU devices: A100 80GB PCIe, RTX3090, RTXA5500, RTXA6000, RTX3080, RTX8000. Please refer to the ["Reproducing the experiments"](#reproducing-the-experiments) section for details on running these experiments in your own environment. + + +## Inference speed + +The figure below shows the latency at inference when using different hardware and precision for generating a single image using the (arbitrary) text prompt: *"a photo of an astronaut riding a horse on mars"*. + +Stable Diffusion Text2Image Latency (seconds) + + +We find that: +* The inference latencies range between `3.74` to `5.56` seconds across our tested Ampere GPUs, including the consumer 3080 card to the flagship A100 80GB card. +* Half-precision reduces the latency by about `40%` for Ampere GPUs, and by `52%` for the previous generation `RTX8000` GPU. + +We believe Ampere GPUs enjoy a relatively "smaller" speedup from half-precision due to their use of `TF32`. For readers who are not familiar with `TF32`, it is a [`19-bit` format](https://blogs.nvidia.com/blog/2020/05/14/tensorfloat-32-precision-format/) that has been used as the default single-precision data type on Ampere GPUs for major deep learning frameworks such as PyTorch and TensorFlow. One can expect half-precision's speedup over `FP32` to be bigger since it is a true `32-bit` format. + + +We run these same inference jobs CPU devices to put in perspective the inference speed performance observed on GPU. + +Stable Diffusion Text2Image GPU v CPU + + +We note that: +* GPUs are significantly faster -- by one or two orders of magnitudes depending on the precisions. +* `onnxruntime` can reduce the latency for CPU by about `40%` to `50%`, depending on the type of CPUs. + +ONNX currently does not have [stable support](https://github.com/huggingface/diffusers/issues/489) for Huggingface diffusers. +We will investigate `onnxruntime-gpu` in future benchmarks. + + + + +## Memory + +We also measure the memory consumption of running stable diffusion inference. + +Stable Diffusion Text2Image Memory (GB) + +Memory usage is observed to be consistent across all tested GPUs: +* It takes about `7.7 GB` GPU memory to run single-precision inference with batch size one. +* It takes about `4.5 GB` GPU memory to run half-precision inference with batch size one. + + + + +## Throughput + +Latency measures how quickly a _single_ input can be processed, which is critical to online applications that don't tolerate even the slightest delay. However, some (offline) applications may focus on "throughput", which measures the total volume of data processed in a fixed amount of time. + + +Our throughput benchmark pushes the batch size to the maximum for each GPU, and measures the number of images they can process per minute. The reason for maximizing the batch size is to keep tensor cores busy so that computation can dominate the workload, avoiding any non-computational bottlenecks. + +We run a series of throughput experiment in pytorch with half-precision and using the maximum batch size that can be used for each GPU: + +Stable Diffusion Text2Image Throughput (images/minute) + +We note: +* Once again, A100 80GB is the top performer and has the highest throughput. +* The gap between A100 80GB and other cards in terms of throughput can be explained by the larger maximum batch size that can be used on this card. + + +As a concrete example, the chart below shows how A100 80GB's throughput increases by `64%` when we changed the batch size from 1 to 28 (the largest without causing an out of memory error). It is also interesting to see that the increase is not linear and flattens out when batch size reaches a certain value, at which point the tensor cores on the GPU are saturated and any new data in the GPU memory will have to be queued up before getting their own computing resources. + +Stable Diffusion Text2Image Batch size vs Throughput (images/minute) + + +## Precision + +We are curious about whether half-precision introduces degradations to the quality of the output images. To test this out, we fixed the text prompt as well as the "latent" input vector and fed them to the single-precision model and the half-precision model. We ran the inference for 100 steps and saved both models' outputs at each step, as well as the difference map: + +![Evolution of precision v degradation across 100 steps](./pictures/benchmark_sd_precision_history.gif) + +Our observation is that there are indeed visible differences between the single-precision output and the half-precision output, especially in the early steps. The differences often decrease with the number of steps, but might not always vanish. + +Interestingly, such a difference may not imply artifacts in half-precision's outputs. For example, in step 70, the picture below shows half-precision didn't produce the artifact in the single-precision output (an extra front leg): + +![Precision v Degradation at step 70](./pictures/benchmark_sd_precision_step_70.png) + +--- + +## Reproducing the experiments + +You can use this [Lambda Diffusers](https://github.com/LambdaLabsML/lambda-diffusers) repository to reproduce the results presented in this article. + +## Setup + +Before running the benchmark, make sure you have completed the repository [installation steps](../README.md#installation). + +You will then need to set the huggingface access token: +1. Create a user account on HuggingFace and generate an access token. +2. Set your huggingface access token as the `ACCESS_TOKEN` environment variable: +``` +export ACCESS_TOKEN= +``` + +## Usage + +Launch the `benchmark.py` script to append benchmark results to the existing [benchmark.csv](../benchmark.csv) results file: +``` +python ./scripts/benchmark.py +``` + +Lauch the `benchmark_quality.py` script to compare the output of single-precision and half-precision models: +``` +python ./scripts/benchmark_quality.py +``` \ No newline at end of file diff --git a/docs/pictures/FreeMono.ttf b/docs/pictures/FreeMono.ttf new file mode 100644 index 0000000..7485f9e Binary files /dev/null and b/docs/pictures/FreeMono.ttf differ diff --git a/docs/pictures/benchmark_sd_precision_history.gif b/docs/pictures/benchmark_sd_precision_history.gif new file mode 100644 index 0000000..7cc1fde Binary files /dev/null and b/docs/pictures/benchmark_sd_precision_history.gif differ diff --git a/docs/pictures/benchmark_sd_precision_step_70.png b/docs/pictures/benchmark_sd_precision_step_70.png new file mode 100644 index 0000000..b7e5359 Binary files /dev/null and b/docs/pictures/benchmark_sd_precision_step_70.png differ diff --git a/docs/pictures/pretty_benchmark_sd_txt2img_batchsize_vs_throughput.png b/docs/pictures/pretty_benchmark_sd_txt2img_batchsize_vs_throughput.png new file mode 100644 index 0000000..82f2b99 Binary files /dev/null and b/docs/pictures/pretty_benchmark_sd_txt2img_batchsize_vs_throughput.png differ diff --git a/docs/pictures/pretty_benchmark_sd_txt2img_gpu_vs_cpu.png b/docs/pictures/pretty_benchmark_sd_txt2img_gpu_vs_cpu.png new file mode 100644 index 0000000..9e633de Binary files /dev/null and b/docs/pictures/pretty_benchmark_sd_txt2img_gpu_vs_cpu.png differ diff --git a/docs/pictures/pretty_benchmark_sd_txt2img_latency.png b/docs/pictures/pretty_benchmark_sd_txt2img_latency.png new file mode 100644 index 0000000..124bb40 Binary files /dev/null and b/docs/pictures/pretty_benchmark_sd_txt2img_latency.png differ diff --git a/docs/pictures/pretty_benchmark_sd_txt2img_mem.png b/docs/pictures/pretty_benchmark_sd_txt2img_mem.png new file mode 100644 index 0000000..ea6a2f8 Binary files /dev/null and b/docs/pictures/pretty_benchmark_sd_txt2img_mem.png differ diff --git a/docs/pictures/pretty_benchmark_sd_txt2img_throughput.png b/docs/pictures/pretty_benchmark_sd_txt2img_throughput.png new file mode 100644 index 0000000..978c06b Binary files /dev/null and b/docs/pictures/pretty_benchmark_sd_txt2img_throughput.png differ diff --git a/requirements.txt b/requirements.txt index b41456b..9089b38 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,9 @@ -torch -torchvision -transformers -ftfy -Pillow -diffusers --e . \ No newline at end of file +--extra-index-url https://download.pytorch.org/whl/cu116 torch + +transformers==4.22.1 +ftfy==6.1.1 +Pillow==9.2.0 +diffusers==0.3.0 +onnxruntime==1.12.1 +scikit-image==0.19.3 +-e . diff --git a/scripts/benchmark.py b/scripts/benchmark.py new file mode 100644 index 0000000..9158177 --- /dev/null +++ b/scripts/benchmark.py @@ -0,0 +1,215 @@ +import os +import subprocess +import multiprocessing +import argparse +import pathlib +import csv +from contextlib import nullcontext +import torch +from torch import autocast +from diffusers import StableDiffusionPipeline, StableDiffusionOnnxPipeline + +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + +prompt = "a photo of an astronaut riding a horse on mars" + +def get_inference_pipeline(precision, backend): + """ + returns HuggingFace diffuser pipeline + cf https://github.com/huggingface/diffusers#text-to-image-generation-with-stable-diffusion + note: could not download from CompVis/stable-diffusion-v1-4 (access restricted) + """ + + assert precision in ("half", "single"), "precision in ['half', 'single']" + assert backend in ("pytorch", "onnx"), "backend in ['pytorch', 'onnx']" + + if backend == "pytorch": + pipe = StableDiffusionPipeline.from_pretrained( + "CompVis/stable-diffusion-v1-4", + revision="main" if precision == "single" else "fp16", + use_auth_token=os.environ["ACCESS_TOKEN"], + torch_dtype=torch.float32 if precision == "single" else torch.float16, + ) + pipe = pipe.to(device) + else: + pipe = StableDiffusionOnnxPipeline.from_pretrained( + "CompVis/stable-diffusion-v1-4", + use_auth_token=os.environ["ACCESS_TOKEN"], + revision="onnx", + provider="CPUExecutionProvider" if device.type=="cpu" else "CUDAExecutionProvider", + torch_dtype=torch.float32 if precision == "single" else torch.float16, + ) + + # Disable safety + disable_safety = True + if disable_safety: + + def null_safety(images, **kwargs): + return images, False + + pipe.safety_checker = null_safety + return pipe + + +def do_inference(pipe, n_samples, precision, num_inference_steps): + torch.cuda.empty_cache() + context = autocast if (device.type == "cuda" and precision == 'half') else nullcontext + with context("cuda"): + images = pipe(prompt=[prompt] * n_samples, num_inference_steps=num_inference_steps).images + + return images + + +def get_inference_time(pipe, n_samples, n_repeats, precision, num_inference_steps): + from torch.utils.benchmark import Timer + timer = Timer( + stmt="do_inference(pipe, n_samples, precision, num_inference_steps)", + setup="from __main__ import do_inference", + globals={"pipe": pipe, "n_samples": n_samples, "precision": precision, "num_inference_steps": num_inference_steps}, + num_threads=multiprocessing.cpu_count() + ) + profile_result = timer.timeit( + n_repeats + ) # benchmark.Timer performs 2 iterations for warmup + return round(profile_result.mean, 2) + + +def get_inference_memory(pipe, n_samples, precision, num_inference_steps): + if not torch.cuda.is_available(): + return 0 + + torch.cuda.empty_cache() + context = autocast if (device.type == "cuda" and precision == 'half') else nullcontext + with context("cuda"): + images = pipe(prompt=[prompt] * n_samples, num_inference_steps=num_inference_steps).images + + mem = torch.cuda.memory_reserved() + return round(mem / 1e9, 2) + + +def run_benchmark(n_repeats, n_samples, precision, backend, num_inference_steps): + """ + * n_repeats: nb datapoints for inference latency benchmark + * n_samples: number of samples to generate (~ batch size) + * precision: 'half' or 'single' (use fp16 or fp32 tensors) + + returns: + dict like {'memory usage': 17.70, 'latency': 86.71'} + """ + + pipe = get_inference_pipeline(precision, backend) + + logs = { + "memory": 0.00 if device.type=="cpu" else get_inference_memory(pipe, n_samples, precision, num_inference_steps), + "latency": get_inference_time(pipe, n_samples, n_repeats, precision, num_inference_steps), + } + print(f"n_samples: {n_samples}\tprecision: {precision}\tbackend: {backend}") + print(logs, "\n") + return logs + + +def get_device_description(): + """ + returns descriptor of cuda device such as + 'NVIDIA RTX A6000' + """ + if device.type == "cpu": + name = subprocess.check_output( + "grep -m 1 'model name' /proc/cpuinfo", + shell=True + ).decode("utf-8") + name = " ".join(name.split(" ")[2:]).strip() + return name + else: + return torch.cuda.get_device_name() + + +def run_benchmark_grid(grid, n_repeats, num_inference_steps): + """ + * grid : dict like + { + "n_samples": (1, 2), + "precision": ("single", "half"), + } + * n_repeats: nb datapoints for inference latency benchmark + """ + + csv_fpath = pathlib.Path(__file__).parent.parent / "benchmark_tmp.csv" + # create benchmark.csv if not exists + if not os.path.isfile(csv_fpath): + header = ["device", "precision", "runtime", "n_samples", "latency", "memory"] + with open(csv_fpath, "w") as f: + writer = csv.writer(f) + writer.writerow(header) + + # append new benchmark results to it if benchmark_tmp.csv already exists + with open(csv_fpath, "a") as f: + writer = csv.writer(f) + device_desc = get_device_description() + for n_samples in grid["n_samples"]: + for precision in grid["precision"]: + for backend in grid["backend"]: + try: + new_log = run_benchmark( + n_repeats=n_repeats, + n_samples=n_samples, + precision=precision, + backend=backend, + num_inference_steps=num_inference_steps + ) + except Exception as e: + if "CUDA out of memory" in str(e) or "Failed to allocate memory" in str(e): + print(str(e)) + torch.cuda.empty_cache() + new_log = { + "latency": -1.00, + "memory": -1.00 + } + else: + raise e + + latency = new_log["latency"] + memory = new_log["memory"] + new_row = [device_desc, precision, backend, n_samples, latency, memory] + writer.writerow(new_row) + + +if __name__ == "__main__": + + parser = argparse.ArgumentParser() + + parser.add_argument( + "--samples", + default="1", + type=str, + help="Comma sepearated list of batch sizes (number of samples)" + ) + + parser.add_argument( + "--steps", + default=50, + type=int, + help="Number of diffusion steps." + ) + + parser.add_argument( + "--repeats", + default=3, + type=int, + help="Number of repeats.", + ) + + args = parser.parse_args() + + grid = { + "n_samples": tuple(map(int, args.samples.split(","))), + # Only use single-precision for cpu because "LayerNormKernelImpl" not implemented for 'Half' on cpu, + # Remove autocast won't help. Ref: + # https://github.com/CompVis/stable-diffusion/issues/307 + "precision": ("single",) if device.type == "cpu" else ("single", "half"), + # Only use onnx for cpu, until issues are fixed by upstreams. Ref: + # https://github.com/huggingface/diffusers/issues/489#issuecomment-1261577250 + # https://github.com/huggingface/diffusers/pull/440 + "backend": ("pytorch", "onnx") if device.type == "cpu" else ("pytorch",) + } + run_benchmark_grid(grid, n_repeats=args.repeats, num_inference_steps=args.steps) diff --git a/scripts/benchmark_quality.py b/scripts/benchmark_quality.py new file mode 100644 index 0000000..03151d4 --- /dev/null +++ b/scripts/benchmark_quality.py @@ -0,0 +1,119 @@ +import os +from platform import mac_ver +import numpy as np +from PIL import Image, ImageDraw, ImageFont +from skimage.metrics import structural_similarity as ssim +from skimage.metrics import normalized_root_mse as nmse +from skimage.metrics import peak_signal_noise_ratio as psnr + +import torch +from torch import autocast +from diffusers import StableDiffusionPipeline + +model_id = "CompVis/stable-diffusion-v1-4" +device = "cuda" +prompt = "a photo of an astronaut riding a horse on mars" +output_folder = "_".join(prompt.split(" ")) +os.makedirs(output_folder, exist_ok=True) + +num_images = 1 +width = 512 +height = 512 +min_inference_steps = 10 +max_inference_steps = 100 + +list_ssim = [] +list_nmse = [] +list_psnr = [] + +# Create piplines for single and half-precision +pipe = StableDiffusionPipeline.from_pretrained( + model_id, + use_auth_token=True, + torch_dtype=torch.float32) +pipe = pipe.to(device) + +pipe_half = StableDiffusionPipeline.from_pretrained( + model_id, + revision="fp16", + torch_dtype=torch.float16, + use_auth_token=True) +pipe_half = pipe_half.to(device) + +# Generate latent vectors +generator = torch.Generator(device=device) +latents = None +seeds = [] +for _ in range(num_images): + # Get a new random seed, store it and use it as the generator state + seed = generator.seed() + seeds.append(seed) + generator = generator.manual_seed(seed) + + image_latents = torch.randn( + (1, pipe.unet.in_channels, height // 8, width // 8), + generator = generator, + device = device + ) + latents = image_latents if latents is None else torch.cat((latents, image_latents)) + +for num_inference_steps in range(min_inference_steps, max_inference_steps, 5): + # Inference with single and half-precision + + torch.cuda.empty_cache() + images = pipe( + [prompt] * num_images, + guidance_scale=7.5, + latents = latents, + num_inference_steps = num_inference_steps + )["sample"] + + torch.cuda.empty_cache() + with torch.autocast(device): + images_half = pipe_half( + [prompt] * num_images, + guidance_scale=7.5, + latents = latents, + num_inference_steps = num_inference_steps + )["sample"] + + m_ssim = [] + m_nmse = [] + m_psnr = [] + + for idx, (image, image_half) in enumerate(zip(images, images_half)): + # Need to convert to float because uint8 can't store negative value + np_image = np.float32(np.asarray(image)) / 255.0 + np_image_half = np.float32(np.asarray(image_half)) / 255.0 + np_image_diff = np.absolute(np.float32(np.asarray(image)) - np.float32(np.asarray(image_half))) + + # Compute quantitative metrics + m_ssim.append(ssim(np_image, np_image_half, channel_axis=2)) + m_nmse.append(nmse(np_image, np_image_half)) + m_psnr.append(psnr(np_image, np_image_half)) + im_diff = Image.fromarray(np.uint8(np_image_diff)) + + # Compose results in a single output image + dst = Image.new('RGB', (image.width + image_half.width + im_diff.width, image.height)) + dst.paste(image, (0, 0)) + dst.paste(image_half, (image.width, 0)) + dst.paste(im_diff, (image.width + image_half.width, 0)) + I1 = ImageDraw.Draw(dst) + font = ImageFont.truetype('../docs/pictures/FreeMono.ttf', 16) + I1.text((32, 32), "Single Precision", font=font, fill=(255, 255, 255)) + I1.text((image.width + 32, 32), "Half Precision", font=font, fill=(255, 255, 255)) + I1.text((image.width + image_half.width + 32, 32), "Step " + str(num_inference_steps), font=font, fill=(255, 255, 255)) + dst.save(output_folder + "/" + str(idx) + "_" + str(num_inference_steps) + ".png") + + list_ssim.append(sum(m_ssim) / len(m_ssim)) + list_nmse.append(sum(m_nmse) / len(m_nmse)) + list_psnr.append(sum(m_psnr) / len(m_psnr)) + +print("ssim: ") +print(list_ssim) + +print("nmse: ") +print(list_nmse) + +print("psnr: ") +print(list_psnr)