work on startup profile display

4 files changed, 50 insertions, 6 deletions

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index 40f388a5..ecffc206 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -7,6 +7,8 @@ from typing import Optional, Dict, Any
 from fastapi import FastAPI

 from gradio import Blocks

 

+from modules import timer

+

 

 def report_exception(c, job):

     print(f"Error executing callback {job} for {c.script}", file=sys.stderr)

@@ -123,6 +125,7 @@ def app_started_callback(demo: Optional[Blocks], app: FastAPI):
     for c in callback_map['callbacks_app_started']:

         try:

             c.callback(demo, app)

+            timer.startup_timer.record(c.script)

         except Exception:

             report_exception(c, 'app_started_callback')

 

diff --git a/modules/scripts.py b/modules/scripts.py
index c902804b..7ef1a8f8 100644
--- a/modules/scripts.py
+++ b/modules/scripts.py
@@ -6,7 +6,7 @@ from collections import namedtuple
 

 import gradio as gr

 

-from modules import shared, paths, script_callbacks, extensions, script_loading, scripts_postprocessing

+from modules import shared, paths, script_callbacks, extensions, script_loading, scripts_postprocessing, timer

 

 AlwaysVisible = object()

 

@@ -270,6 +270,7 @@ def load_scripts():
         finally:

             sys.path = syspath

             current_basedir = paths.script_path

+            timer.startup_timer.record(scriptfile.filename)

 

     global scripts_txt2img, scripts_img2img, scripts_postproc

 

diff --git a/modules/timer.py b/modules/timer.py
index ba92be33..da99e49f 100644
--- a/modules/timer.py
+++ b/modules/timer.py
@@ -1,11 +1,30 @@
 import time

 

 

+class TimerSubcategory:

+    def __init__(self, timer, category):

+        self.timer = timer

+        self.category = category

+        self.start = None

+        self.original_base_category = timer.base_category

+

+    def __enter__(self):

+        self.start = time.time()

+        self.timer.base_category = self.original_base_category + self.category + "/"

+

+    def __exit__(self, exc_type, exc_val, exc_tb):

+        elapsed_for_subcategroy = time.time() - self.start

+        self.timer.base_category = self.original_base_category

+        self.timer.add_time_to_record(self.original_base_category + self.category, elapsed_for_subcategroy)

+        self.timer.record(self.category)

+

+

 class Timer:

     def __init__(self):

         self.start = time.time()

         self.records = {}

         self.total = 0

+        self.base_category = ''

 

     def elapsed(self):

         end = time.time()

@@ -13,18 +32,29 @@ class Timer:
         self.start = end

         return res

 

-    def record(self, category, extra_time=0):

-        e = self.elapsed()

+    def add_time_to_record(self, category, amount):

         if category not in self.records:

             self.records[category] = 0

 

-        self.records[category] += e + extra_time

+        self.records[category] += amount

+

+    def record(self, category, extra_time=0):

+        e = self.elapsed()

+

+        self.add_time_to_record(self.base_category + category, e + extra_time)

+

         self.total += e + extra_time

 

+    def subcategory(self, name):

+        self.elapsed()

+

+        subcat = TimerSubcategory(self, name)

+        return subcat

+

     def summary(self):

         res = f"{self.total:.1f}s"

 

-        additions = [x for x in self.records.items() if x[1] >= 0.1]

+        additions = [(category, time_taken) for category, time_taken in self.records.items() if time_taken >= 0.1 and '/' not in category]

         if not additions:

             return res

 

@@ -34,5 +64,13 @@ class Timer:
 

         return res

 

+    def dump(self):

+        return {'total': self.total, 'records': self.records}

+

     def reset(self):

         self.__init__()

+

+

+startup_timer = Timer()

+

+startup_record = None

diff --git a/modules/ui.py b/modules/ui.py
index 82820ab5..5174da63 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -13,7 +13,7 @@ import numpy as np
 from PIL import Image, PngImagePlugin  # noqa: F401

 from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call, wrap_gradio_call

 

-from modules import sd_hijack, sd_models, localization, script_callbacks, ui_extensions, deepbooru, sd_vae, extra_networks, ui_common, ui_postprocessing, progress, ui_loadsave

+from modules import sd_hijack, sd_models, localization, script_callbacks, ui_extensions, deepbooru, sd_vae, extra_networks, ui_common, ui_postprocessing, progress, ui_loadsave, timer

 from modules.ui_components import FormRow, FormGroup, ToolButton, FormHTML

 from modules.paths import script_path, data_path

 

@@ -1901,3 +1901,5 @@ def setup_ui_api(app):
     app.add_api_route("/internal/quicksettings-hint", quicksettings_hint, methods=["GET"], response_model=List[QuicksettingsHint])

 

     app.add_api_route("/internal/ping", lambda: {}, methods=["GET"])

+

+    app.add_api_route("/internal/profile-startup", lambda: timer.startup_record, methods=["GET"])