From b5adf240fb9e0df7f2a2cb0297bf16f1e08c4046 Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Thu, 2 Feb 2023 12:02:39 +0800 Subject: [PATCH 01/14] check in demo --- .gitignore | 1 + examples/README.md | 1 + examples/graph-cache/demo.sh | 4 ++ examples/graph-cache/infer.py | 79 +++++++++++++++++++++++++++++++++++ 4 files changed, 85 insertions(+) create mode 100644 examples/README.md create mode 100644 examples/graph-cache/demo.sh create mode 100644 examples/graph-cache/infer.py diff --git a/.gitignore b/.gitignore index dba001b..65387f8 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ huggingface_cache +log diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 0000000..4dedab2 --- /dev/null +++ b/examples/README.md @@ -0,0 +1 @@ +# Examples to run oneflow SD diff --git a/examples/graph-cache/demo.sh b/examples/graph-cache/demo.sh new file mode 100644 index 0000000..9e4ad6a --- /dev/null +++ b/examples/graph-cache/demo.sh @@ -0,0 +1,4 @@ +set -e +export ONEFLOW_NNGRAPH_ENABLE_PROGRESS_BAR=1 +python3 examples/graph-cache/infer.py --save +python3 examples/graph-cache/infer.py --laod diff --git a/examples/graph-cache/infer.py b/examples/graph-cache/infer.py new file mode 100644 index 0000000..f6380cd --- /dev/null +++ b/examples/graph-cache/infer.py @@ -0,0 +1,79 @@ +import os + +import argparse +import oneflow as torch +from diffusers import OneFlowStableDiffusionPipeline + + +def parse_args(): + parser = argparse.ArgumentParser(description="Simple demo of image generation.") + parser.add_argument( + "--prompt", type=str, default="a photo of an astronaut riding a horse on mars" + ) + parser.add_argument("--cached", type=str, default="./oneflow_saved_pipe") + parser.add_argument( + "--load", + default=False, + action="store_true", + help="If specified, load from cached", + ) + parser.add_argument( + "--save", + default=False, + action="store_true", + help="If specified, save to cached", + ) + args = parser.parse_args() + return args + + +args = parse_args() + +model = "CompVis/stable-diffusion-v1-4" +if args.load: + # Note: restore the cache by setting the pretrain path to a cache path + model = args.cached + print(f"will load pipe from: {args.cached}") +pipe = OneFlowStableDiffusionPipeline.from_pretrained( + model, + use_auth_token=True, + revision="fp16", + torch_dtype=torch.float16, + safety_checker=None, +) + +pipe = pipe.to("cuda") + +output_dir = "oneflow-sd-output" +os.makedirs(output_dir, exist_ok=True) +if args.save: + os.makedirs(args.cached, exist_ok=True) + +# Note: enable graph-related cache +pipe.set_unet_graphs_cache_size(10) +pipe.enable_graph_share_mem() +pipe.enable_save_graph() + + +def do_infer(n): + with torch.autocast("cuda"): + for i in [2, 1, 0]: + for j in [2, 1, 0]: + width = 768 + 128 * i + height = 768 + 128 * j + prompt = args.prompt + images = pipe(prompt, width=width, height=height).images + for i, image in enumerate(images): + prompt = prompt.strip().replace("\n", " ") + dst = os.path.join( + output_dir, f"{prompt[:100]}-{n}-{width}-{height}.png" + ) + image.save(dst) + + +for n in range(2): + do_infer(n) +if args.save: + # Note: graph cache will be saved with the weight + print(f"start saving pipe to: {args.cached}") + pipe.save_pretrained(args.cached) From ffeefd189f05f30cd5482f1d29264fcadebdef82 Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Thu, 2 Feb 2023 12:05:08 +0800 Subject: [PATCH 02/14] refine readme --- examples/README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/README.md b/examples/README.md index 4dedab2..5077115 100644 --- a/examples/README.md +++ b/examples/README.md @@ -1 +1,2 @@ # Examples to run oneflow SD +This directory contains scripts for non-throughput performance benchmarks/tests, for instance the performance of cache or compilation. From 01a595d0c525f8a6db31f8eaca024f10a47a5363 Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Thu, 2 Feb 2023 12:06:30 +0800 Subject: [PATCH 03/14] revert --- .gitignore | 1 - 1 file changed, 1 deletion(-) diff --git a/.gitignore b/.gitignore index 65387f8..dba001b 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1 @@ huggingface_cache -log From ecd536f6602240a4b197c390ada4571be4a3d21f Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Thu, 2 Feb 2023 12:07:15 +0800 Subject: [PATCH 04/14] update demo --- examples/graph-cache/demo.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/graph-cache/demo.sh b/examples/graph-cache/demo.sh index 9e4ad6a..8a9afd2 100644 --- a/examples/graph-cache/demo.sh +++ b/examples/graph-cache/demo.sh @@ -1,4 +1,4 @@ set -e export ONEFLOW_NNGRAPH_ENABLE_PROGRESS_BAR=1 -python3 examples/graph-cache/infer.py --save -python3 examples/graph-cache/infer.py --laod +time python3 examples/graph-cache/infer.py --save +time python3 examples/graph-cache/infer.py --laod From a806f358aa90b40e519fed12ad78c76daed32b3c Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Thu, 2 Feb 2023 12:09:03 +0800 Subject: [PATCH 05/14] refine --- examples/graph-cache/infer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/graph-cache/infer.py b/examples/graph-cache/infer.py index f6380cd..b7f3368 100644 --- a/examples/graph-cache/infer.py +++ b/examples/graph-cache/infer.py @@ -46,8 +46,6 @@ def parse_args(): output_dir = "oneflow-sd-output" os.makedirs(output_dir, exist_ok=True) -if args.save: - os.makedirs(args.cached, exist_ok=True) # Note: enable graph-related cache pipe.set_unet_graphs_cache_size(10) @@ -76,4 +74,6 @@ def do_infer(n): if args.save: # Note: graph cache will be saved with the weight print(f"start saving pipe to: {args.cached}") + os.makedirs(args.cached, exist_ok=True) pipe.save_pretrained(args.cached) + pipe.save_graph(args.cached) From 6b433992d8c953220c23d5ba0dcbe7dedf714bae Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Thu, 2 Feb 2023 12:11:41 +0800 Subject: [PATCH 06/14] update --- examples/graph-cache/infer.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/examples/graph-cache/infer.py b/examples/graph-cache/infer.py index b7f3368..ed82ba0 100644 --- a/examples/graph-cache/infer.py +++ b/examples/graph-cache/infer.py @@ -10,18 +10,19 @@ def parse_args(): parser.add_argument( "--prompt", type=str, default="a photo of an astronaut riding a horse on mars" ) - parser.add_argument("--cached", type=str, default="./oneflow_saved_pipe") + parser.add_argument("--cache", type=str, default="./oneflow-sd/graph_cache") + parser.add_argument("--model", type=str, default=".//oneflow-sd/model") parser.add_argument( "--load", default=False, action="store_true", - help="If specified, load from cached", + help="If specified, load from cache", ) parser.add_argument( "--save", default=False, action="store_true", - help="If specified, save to cached", + help="If specified, save to cache", ) args = parser.parse_args() return args @@ -32,8 +33,8 @@ def parse_args(): model = "CompVis/stable-diffusion-v1-4" if args.load: # Note: restore the cache by setting the pretrain path to a cache path - model = args.cached - print(f"will load pipe from: {args.cached}") + model = args.model + print(f"will load pipe from: {args.cache}") pipe = OneFlowStableDiffusionPipeline.from_pretrained( model, use_auth_token=True, @@ -43,6 +44,8 @@ def parse_args(): ) pipe = pipe.to("cuda") +if args.load: + pipe.load_graph(args.cache) output_dir = "oneflow-sd-output" os.makedirs(output_dir, exist_ok=True) @@ -73,7 +76,7 @@ def do_infer(n): do_infer(n) if args.save: # Note: graph cache will be saved with the weight - print(f"start saving pipe to: {args.cached}") - os.makedirs(args.cached, exist_ok=True) - pipe.save_pretrained(args.cached) - pipe.save_graph(args.cached) + print(f"start saving pipe to: {args.cache}") + os.makedirs(args.cache, exist_ok=True) + pipe.save_pretrained(args.cache) + pipe.save_graph(args.cache) From fd8f5092536e5f7f4c1c7cd50c26e0914364855b Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Thu, 2 Feb 2023 12:13:05 +0800 Subject: [PATCH 07/14] add --- examples/graph-cache/infer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/graph-cache/infer.py b/examples/graph-cache/infer.py index ed82ba0..b5543ce 100644 --- a/examples/graph-cache/infer.py +++ b/examples/graph-cache/infer.py @@ -75,8 +75,8 @@ def do_infer(n): for n in range(2): do_infer(n) if args.save: - # Note: graph cache will be saved with the weight print(f"start saving pipe to: {args.cache}") os.makedirs(args.cache, exist_ok=True) pipe.save_pretrained(args.cache) + # Note: save graph cache pipe.save_graph(args.cache) From bc2aac058c9633b403a81652ad489d3db1131dc4 Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Thu, 2 Feb 2023 12:16:25 +0800 Subject: [PATCH 08/14] add --- examples/graph-cache/infer.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/examples/graph-cache/infer.py b/examples/graph-cache/infer.py index b5543ce..f1f5bab 100644 --- a/examples/graph-cache/infer.py +++ b/examples/graph-cache/infer.py @@ -50,10 +50,13 @@ def parse_args(): output_dir = "oneflow-sd-output" os.makedirs(output_dir, exist_ok=True) -# Note: enable graph-related cache pipe.set_unet_graphs_cache_size(10) pipe.enable_graph_share_mem() -pipe.enable_save_graph() +# Note: enable saving/loading graph-related cache, these APIs are tricky and might be removed +if args.save: + pipe.enable_save_graph() +if args.load: + pipe.enable_load_graph() def do_infer(n): From 79939fadc2898618fe995f415f21430102a8c810 Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Thu, 2 Feb 2023 12:17:49 +0800 Subject: [PATCH 09/14] add --- examples/graph-cache/demo.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/graph-cache/demo.sh b/examples/graph-cache/demo.sh index 8a9afd2..a1814e6 100644 --- a/examples/graph-cache/demo.sh +++ b/examples/graph-cache/demo.sh @@ -1,4 +1,4 @@ set -e export ONEFLOW_NNGRAPH_ENABLE_PROGRESS_BAR=1 time python3 examples/graph-cache/infer.py --save -time python3 examples/graph-cache/infer.py --laod +time python3 examples/graph-cache/infer.py --load From c862fa3f403d03f54987f95af590cddeb2c8ccfc Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Thu, 2 Feb 2023 12:19:36 +0800 Subject: [PATCH 10/14] refine --- examples/graph-cache/infer.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/examples/graph-cache/infer.py b/examples/graph-cache/infer.py index f1f5bab..58a8c5f 100644 --- a/examples/graph-cache/infer.py +++ b/examples/graph-cache/infer.py @@ -11,7 +11,7 @@ def parse_args(): "--prompt", type=str, default="a photo of an astronaut riding a horse on mars" ) parser.add_argument("--cache", type=str, default="./oneflow-sd/graph_cache") - parser.add_argument("--model", type=str, default=".//oneflow-sd/model") + parser.add_argument("--model", type=str, default="./oneflow-sd/model") parser.add_argument( "--load", default=False, @@ -44,8 +44,6 @@ def parse_args(): ) pipe = pipe.to("cuda") -if args.load: - pipe.load_graph(args.cache) output_dir = "oneflow-sd-output" os.makedirs(output_dir, exist_ok=True) @@ -57,6 +55,7 @@ def parse_args(): pipe.enable_save_graph() if args.load: pipe.enable_load_graph() + pipe.load_graph(args.cache) def do_infer(n): From f380c1f9c323c9fc54769fd3e1d438e7ae22c4ff Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Thu, 2 Feb 2023 12:19:54 +0800 Subject: [PATCH 11/14] refine --- examples/graph-cache/infer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/graph-cache/infer.py b/examples/graph-cache/infer.py index 58a8c5f..d5f25f8 100644 --- a/examples/graph-cache/infer.py +++ b/examples/graph-cache/infer.py @@ -50,7 +50,7 @@ def parse_args(): pipe.set_unet_graphs_cache_size(10) pipe.enable_graph_share_mem() -# Note: enable saving/loading graph-related cache, these APIs are tricky and might be removed +# Note: enable saving/loading graph-related cache, these APIs are tricky and might be changed if args.save: pipe.enable_save_graph() if args.load: From 4371bd1e76f9f6bdd0b71c8a3c451a68ea1b2b72 Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Thu, 2 Feb 2023 12:53:06 +0800 Subject: [PATCH 12/14] fix --- examples/graph-cache/infer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/graph-cache/infer.py b/examples/graph-cache/infer.py index d5f25f8..2d8ec1c 100644 --- a/examples/graph-cache/infer.py +++ b/examples/graph-cache/infer.py @@ -79,6 +79,6 @@ def do_infer(n): if args.save: print(f"start saving pipe to: {args.cache}") os.makedirs(args.cache, exist_ok=True) - pipe.save_pretrained(args.cache) + pipe.save_pretrained(args.model) # Note: save graph cache pipe.save_graph(args.cache) From 561eea65558b726e0d5fef87b30c03e88250dcdc Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Thu, 2 Feb 2023 12:54:04 +0800 Subject: [PATCH 13/14] rm --- examples/graph-cache/infer.py | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/graph-cache/infer.py b/examples/graph-cache/infer.py index 2d8ec1c..630973a 100644 --- a/examples/graph-cache/infer.py +++ b/examples/graph-cache/infer.py @@ -54,7 +54,6 @@ def parse_args(): if args.save: pipe.enable_save_graph() if args.load: - pipe.enable_load_graph() pipe.load_graph(args.cache) From 5ea08c50f8a7616e3d55664825db9879e91321cd Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Thu, 2 Feb 2023 13:09:41 +0800 Subject: [PATCH 14/14] refine --- examples/graph-cache/infer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/graph-cache/infer.py b/examples/graph-cache/infer.py index 630973a..39aec60 100644 --- a/examples/graph-cache/infer.py +++ b/examples/graph-cache/infer.py @@ -76,8 +76,8 @@ def do_infer(n): for n in range(2): do_infer(n) if args.save: - print(f"start saving pipe to: {args.cache}") - os.makedirs(args.cache, exist_ok=True) pipe.save_pretrained(args.model) + print(f"saving cache to: {args.cache}") + os.makedirs(args.cache, exist_ok=True) # Note: save graph cache pipe.save_graph(args.cache)