From 5c15a247882c6a6b3831558da107b2b6f85f2dcd Mon Sep 17 00:00:00 2001
From: btian <ben.tian@graviti.cn>
Date: Wed, 31 Jul 2024 14:37:34 +0800
Subject: [PATCH] fix: main model is missed in first generation after page
 refreshed

---
 iz_helpers/run.py           | 7 ++++++-
 iz_helpers/ui.py            | 2 +-
 javascript/infinite-zoom.js | 4 ++++
 3 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/iz_helpers/run.py b/iz_helpers/run.py
index b05b6d7..7a48e79 100644
--- a/iz_helpers/run.py
+++ b/iz_helpers/run.py
@@ -176,6 +176,7 @@ def outpaint_steps(
 def create_zoom(
     request: gr.Request,
     id_task,
+    model_title: str,
     common_prompt_pre,
     prompts_array,
     common_prompt_suf,
@@ -200,12 +201,16 @@ def create_zoom(
     upscale_do,
     upscaler_name,
     upscale_by,
-    main_sd_model,
     inpainting_denoising_strength=1,
     inpainting_full_res=0,
     inpainting_padding=0,
     progress=None,
 ):
+    model_title_prefix = 'model_title('
+    if model_title.startswith(model_title_prefix) and model_title.endswith(')'):
+        main_sd_model = model_title[len(model_title_prefix):-1]
+    else:
+        main_sd_model = model_title
     result = [None, None, "", "", ""]
     for i in range(batchcount):
         print(f"{id_task}: Batch {i+1}/{batchcount}")
diff --git a/iz_helpers/ui.py b/iz_helpers/ui.py
index 97637c2..b97428b 100644
--- a/iz_helpers/ui.py
+++ b/iz_helpers/ui.py
@@ -309,6 +309,7 @@ def on_ui_tabs():
             _js="iz_submit",
             inputs=[
                 id_task,
+                main_sd_model,
                 main_common_prompt_pre,
                 main_prompts,
                 main_common_prompt_suf,
@@ -333,7 +334,6 @@ def on_ui_tabs():
                 upscale_do,
                 upscaler_name,
                 upscale_by,
-                main_sd_model,
             ],
             outputs=[output_video, out_image, generation_info, html_info, html_log, upgrade_info],
         )
diff --git a/javascript/infinite-zoom.js b/javascript/infinite-zoom.js
index a68c484..e6eb96f 100644
--- a/javascript/infinite-zoom.js
+++ b/javascript/infinite-zoom.js
@@ -27,8 +27,12 @@ async function iz_submit() {
     var id = randomId();
     localStorage.setItem("iz_task_id", id);
 
+    const mainModel = gradioApp().querySelector("#sd_model_checkpoint_dropdown input");
+
     var res = Array.from(arguments);
     res[0] = id;
+    res[1] = `model_title(${mainModel.value})`;
+
     return res;
 }