From fa58996f37b336a0cb7be9a64c5716422645763e Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Tue, 1 Aug 2023 23:53:01 +0530
Subject: [PATCH] sdkit 1.0.157 - tensorRT build configuration from the UI;
 clamp images to 8 instead of 64 pixels

---
 scripts/check_modules.py  |  2 +-
 ui/easydiffusion/types.py |  3 +++
 ui/media/js/main.js       | 17 +++++++++++++++++
 ui/media/js/parameters.js | 21 ++++++++++++++++-----
 4 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/scripts/check_modules.py b/scripts/check_modules.py
index 51ecf3e7..4513e87b 100644
--- a/scripts/check_modules.py
+++ b/scripts/check_modules.py
@@ -18,7 +18,7 @@ os_name = platform.system()
 modules_to_check = {
     "torch": ("1.11.0", "1.13.1", "2.0.0"),
     "torchvision": ("0.12.0", "0.14.1", "0.15.1"),
-    "sdkit": "1.0.156",
+    "sdkit": "1.0.157",
     "stable-diffusion-sdkit": "2.1.4",
     "rich": "12.6.0",
     "uvicorn": "0.19.0",
diff --git a/ui/easydiffusion/types.py b/ui/easydiffusion/types.py
index 181a9505..bf5772f0 100644
--- a/ui/easydiffusion/types.py
+++ b/ui/easydiffusion/types.py
@@ -226,6 +226,9 @@ def convert_legacy_render_req_to_new(old_req: dict):
         model_params["stable-diffusion"] = {
             "clip_skip": bool(old_req.get("clip_skip", False)),
             "convert_to_tensorrt": bool(old_req.get("convert_to_tensorrt", False)),
+            "trt_build_config": old_req.get(
+                "trt_build_config", {"batch_size_range": (1, 2), "dimensions_range": [(768, 1024)]}
+            ),
         }
 
     # move the filter params
diff --git a/ui/media/js/main.js b/ui/media/js/main.js
index 0db31abe..21e55923 100644
--- a/ui/media/js/main.js
+++ b/ui/media/js/main.js
@@ -1452,6 +1452,22 @@ function getCurrentUserRequest() {
     if (testDiffusers.checked && document.getElementById("toggle-tensorrt-install").innerHTML == "Uninstall") {
         // TRT is installed
         newTask.reqBody.convert_to_tensorrt = document.querySelector("#convert_to_tensorrt").checked
+        let trtBuildConfig = {
+            batch_size_range: [
+                parseInt(document.querySelector("#trt-build-min-batch").value),
+                parseInt(document.querySelector("#trt-build-max-batch").value),
+            ],
+            dimensions_range: [],
+        }
+
+        let sizes = [512, 768, 1024, 1280, 1536]
+        sizes.forEach((i) => {
+            let el = document.querySelector("#trt-build-res-" + i)
+            if (el.checked) {
+                trtBuildConfig["dimensions_range"].push([i, i + 256])
+            }
+        })
+        newTask.reqBody.trt_build_config = trtBuildConfig
     }
     if (controlnetModelField.value !== "" && IMAGE_REGEX.test(controlImagePreview.src)) {
         newTask.reqBody.use_controlnet_model = controlnetModelField.value
@@ -2383,6 +2399,7 @@ function packagesUpdate(event) {
 
     if (document.getElementById("toggle-tensorrt-install").innerHTML == "Uninstall") {
         document.querySelector("#enable_trt_config").classList.remove("displayNone")
+        document.querySelector("#trt-build-config").classList.remove("displayNone")
 
         if (!trtSettingsForced) {
             // settings for demo
diff --git a/ui/media/js/parameters.js b/ui/media/js/parameters.js
index 19892c94..8b66c461 100644
--- a/ui/media/js/parameters.js
+++ b/ui/media/js/parameters.js
@@ -248,7 +248,18 @@ var PARAMETERS = [
         label: "NVIDIA TensorRT",
         note: `Faster image generation by converting your Stable Diffusion models to the NVIDIA TensorRT format. You can choose the
                models to convert. Download size: approximately 2 GB.<br/><br/>
-               <b>Early access version:</b> support for LoRA is still under development.`,
+               <b>Early access version:</b> support for LoRA is still under development.
+               <div id="trt-build-config" class="displayNone">
+                    <h3>Build Config:</h3>
+                    Min batch size: <input id="trt-build-min-batch" type="number" min="1" value="1" /><br/>
+                    Max batch size: <input id="trt-build-max-batch" type="number" min="1" value="1" /><br/><br/>
+                    <b>Build for resolutions</b>:<br/>
+                    <input id="trt-build-res-512" type="checkbox" value="1" /> 512x512 to 768x768<br/>
+                    <input id="trt-build-res-768" type="checkbox" value="1" checked /> 768x768 to 1024x1024<br/>
+                    <input id="trt-build-res-1024" type="checkbox" value="1" /> 1024x1024 to 1280x1280<br/>
+                    <input id="trt-build-res-1280" type="checkbox" value="1" /> 1280x1280 to 1536x1536<br/>
+                    <input id="trt-build-res-1536" type="checkbox" value="1" /> 1536x1536 to 1792x1792<br/>
+               </div>`,
         icon: "fa-angles-up",
         render: () => '<button id="toggle-tensorrt-install" class="primaryButton">Install</button>',
         table: installExtrasTable,
@@ -453,8 +464,8 @@ async function getAppConfig() {
             document.querySelectorAll("#sampler_name option.diffusers-only").forEach((option) => {
                 option.style.display = "none"
             })
-            customWidthField.step=64
-            customHeightField.step=64
+            customWidthField.step = 64
+            customHeightField.step = 64
         } else {
             document.querySelector("#lora_model_container").style.display = ""
             document.querySelector("#tiling_container").style.display = ""
@@ -465,8 +476,8 @@ async function getAppConfig() {
             document.querySelector("#clip_skip_config").classList.remove("displayNone")
             document.querySelector("#embeddings-button").classList.remove("displayNone")
             document.querySelector("#negative-embeddings-button").classList.remove("displayNone")
-            customWidthField.step=8
-            customHeightField.step=8
+            customWidthField.step = 8
+            customHeightField.step = 8
         }
 
         console.log("get config status response", config)