From 679320a7640ffa2b6e4d1c5714cc53c3d3d7530d Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Tue, 28 Feb 2023 10:04:25 -0500
Subject: [PATCH] add profiler

---
 extensions-builtin/sd-extension-system-info |  2 +-
 launch.py                                   |  3 ++-
 modules/call_queue.py                       | 13 ++++++++++++-
 modules/processing.py                       |  8 ++++++++
 modules/shared.py                           |  1 +
 5 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/extensions-builtin/sd-extension-system-info b/extensions-builtin/sd-extension-system-info
index c78c59505..1c8ab36ef 160000
--- a/extensions-builtin/sd-extension-system-info
+++ b/extensions-builtin/sd-extension-system-info
@@ -1 +1 @@
-Subproject commit c78c5950517fc813b3701e9027c02ab156ccc2bd
+Subproject commit 1c8ab36ef2b41759295e3f6e79e24ad9da10debf
diff --git a/launch.py b/launch.py
index c09bc1211..6c165c0fa 100644
--- a/launch.py
+++ b/launch.py
@@ -240,7 +240,8 @@ def prepare_environment():
 
     stable_diffusion_commit_hash = os.environ.get('STABLE_DIFFUSION_COMMIT_HASH', "47b6b607fdd31875c9279cd2f4f16b92e4ea958e")
     taming_transformers_commit_hash = os.environ.get('TAMING_TRANSFORMERS_COMMIT_HASH', "24268930bf1dce879235a7fddd0b2355b84d7ea6")
-    k_diffusion_commit_hash = os.environ.get('K_DIFFUSION_COMMIT_HASH', "5b3af030dd83e0297272d861c19477735d0317ec")
+    # k_diffusion_commit_hash = os.environ.get('K_DIFFUSION_COMMIT_HASH', "5b3af030dd83e0297272d861c19477735d0317ec")
+    k_diffusion_commit_hash = os.environ.get('K_DIFFUSION_COMMIT_HASH', "b43db16749d51055f813255eea2fdf1def801919")
     codeformer_commit_hash = os.environ.get('CODEFORMER_COMMIT_HASH', "c5b4593074ba6214284d6acd5f1719b6c5d739af")
     blip_commit_hash = os.environ.get('BLIP_COMMIT_HASH', "48211a1594f1321b00f14c9f7a5b4813144b2fb9")
 
diff --git a/modules/call_queue.py b/modules/call_queue.py
index 92097c15e..f0c441a09 100644
--- a/modules/call_queue.py
+++ b/modules/call_queue.py
@@ -3,6 +3,7 @@ import sys
 import threading
 import traceback
 import time
+import cProfile, pstats, io
 
 from modules import shared, progress
 
@@ -44,7 +45,6 @@ def wrap_gradio_gpu_call(func, extra_outputs=None):
 
     return wrap_gradio_call(f, extra_outputs=extra_outputs, add_stats=True)
 
-
 def wrap_gradio_call(func, extra_outputs=None, add_stats=False):
     def f(*args, extra_outputs_array=extra_outputs, **kwargs):
         run_memmon = shared.opts.memmon_poll_rate > 0 and not shared.mem_mon.disabled and add_stats
@@ -53,7 +53,18 @@ def wrap_gradio_call(func, extra_outputs=None, add_stats=False):
         t = time.perf_counter()
 
         try:
+            if shared.cmd_opts.profile:
+                pr = cProfile.Profile()
+                pr.enable()
             res = list(func(*args, **kwargs))
+            if shared.cmd_opts.profile:
+                pr.disable()
+                s = io.StringIO()
+                ps = pstats.Stats(pr, stream=s)
+                ps.sort_stats(pstats.SortKey.CUMULATIVE)
+                # ps.strip_dirs()
+                ps.print_stats(15)
+                print('Profile:', s.getvalue())
         except Exception as e:
             # When printing out our debug argument list, do not print out more than a MB of text
             max_debug_str_len = 131072 # (1024*1024)/8
diff --git a/modules/processing.py b/modules/processing.py
index 2009d3bf8..e93505841 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -483,6 +483,14 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
             if k == 'sd_vae':
                 sd_vae.reload_vae_weights()
 
+        """
+        import torch.profiler
+        with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], record_shapes=True, with_modules=True) as prof:
+            with torch.profiler.record_function("process_images"):
+                res = process_images_inner(p)
+        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))
+        """
+
         res = process_images_inner(p)
 
     finally:
diff --git a/modules/shared.py b/modules/shared.py
index ec45fde3d..21102d4b6 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -109,6 +109,7 @@ parser.add_argument("--gradio-queue", action='store_true', help="Uses gradio que
 parser.add_argument("--skip-version-check", action='store_true', help="Do not check versions of torch and xformers")
 parser.add_argument("--no-hashing", action='store_true', help="disable sha256 hashing of checkpoints to help loading performance", default=False)
 parser.add_argument("--no-download-sd-model", action='store_true', help="don't download SD1.5 model even if no model is found in --ckpt-dir", default=False)
+parser.add_argument("--profile", action='store_true', help="run profiler", default=False)
 
 
 script_loading.preload_extensions(extensions.extensions_dir, parser)