From c14f7c4e7cb4f9995b3f1a44c77488fd31981e8e Mon Sep 17 00:00:00 2001
From: Vasyanator <135327897+Vasyanator@users.noreply.github.com>
Date: Tue, 27 Aug 2024 21:17:26 +0400
Subject: [PATCH] Add ui_block(), ui_tab() and ui_params() support

---
 modules/extensions.py | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/modules/extensions.py b/modules/extensions.py
index 6729b996..c15eb562 100644
--- a/modules/extensions.py
+++ b/modules/extensions.py
@@ -188,24 +188,42 @@ def _apply_custom_js():
 def create_extensions_block():
     to_display = []
     for extension, name in iterator():
-        if hasattr(extension, "ui") and not (hasattr(extension, 'params') and extension.params.get('is_tab', False)):
+        # Use ui_block if it is defined, otherwise use the old ui
+        if hasattr(extension, "ui_block"):
+            to_display.append((extension, name))
+        elif hasattr(extension, "ui") and not (hasattr(extension, 'params') and extension.params.get('is_tab', False)):
             to_display.append((extension, name))
 
-    # Creating the extension ui elements
     if len(to_display) > 0:
         with gr.Column(elem_id="extensions"):
             for row in to_display:
                 extension, _ = row
-                extension.ui()
+                if hasattr(extension, "ui_block"):
+                    extension.ui_block()
+                else:
+                    extension.ui()
 
 
 def create_extensions_tabs():
     for extension, name in iterator():
-        if hasattr(extension, "ui") and (hasattr(extension, 'params') and extension.params.get('is_tab', False)):
+        # Use ui_tab if it is defined, otherwise use the old ui with the is_tab parameter
+        if hasattr(extension, "ui_tab"):
+            display_name = getattr(extension, 'params', {}).get('display_name', name)
+            with gr.Tab(display_name, elem_classes="extension-tab"):
+                extension.ui_tab()
+        elif hasattr(extension, "ui") and (hasattr(extension, 'params') and extension.params.get('is_tab', False)):
             display_name = getattr(extension, 'params', {}).get('display_name', name)
             with gr.Tab(display_name, elem_classes="extension-tab"):
                 extension.ui()
 
+# Creates a tab in Parameters to hold the extension settings
+def create_extensions_params():
+    for extension, name in iterator():
+        if hasattr(extension, "ui_params"):
+            display_name = getattr(extension, 'params', {}).get('display_name', name)
+            with gr.Tab(display_name):
+                extension.ui_params()
+
 
 EXTENSION_MAP = {
     "input": partial(_apply_string_extensions, "input_modifier"),