add a button to restore the current progress

4 files changed, 80 insertions, 8 deletions

diff --git a/javascript/progressbar.js b/javascript/progressbar.js
index 4ac9b8db..7ba14192 100644
--- a/javascript/progressbar.js
+++ b/javascript/progressbar.js
@@ -59,8 +59,8 @@ function setTitle(progress){
 }
 
 
-function randomId(){
-    return "task(" + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7)+")"
+function randomId(prefix=null){
+    return "task(" + (prefix == null ? "" : prefix + "_") + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7)+")"
 }
 
 // starts sending progress requests to "/internal/progress" uri, creating progressbar above progressbarContainer element and
diff --git a/javascript/ui.js b/javascript/ui.js
index 4a440193..9fe884c0 100644
--- a/javascript/ui.js
+++ b/javascript/ui.js
@@ -163,7 +163,7 @@ function submit(){
     rememberGallerySelection('txt2img_gallery')
     showSubmitButtons('txt2img', false)
 
-    var id = randomId()
+    var id = randomId("txt2img")
     requestProgress(id, gradioApp().getElementById('txt2img_gallery_container'), gradioApp().getElementById('txt2img_gallery'), function(){
         showSubmitButtons('txt2img', true)
 
@@ -180,7 +180,7 @@ function submit_img2img(){
     rememberGallerySelection('img2img_gallery')
     showSubmitButtons('img2img', false)
 
-    var id = randomId()
+    var id = randomId("img2img")
     requestProgress(id, gradioApp().getElementById('img2img_gallery_container'), gradioApp().getElementById('img2img_gallery'), function(){
         showSubmitButtons('img2img', true)
     })
@@ -361,3 +361,35 @@ function selectCheckpoint(name){
     desiredCheckpointName = name;
     gradioApp().getElementById('change_checkpoint').click()
 }
+
+function restoreProgress (task_tag) {
+
+  if (task_tag) {
+    let successHandler = ({ current_task }) => {
+      if (current_task) {
+        let _task_tag = ["txt2img", "img2img"].find(t => current_task.startsWith(`task(${t}_`) && current_task.endsWith(")"))
+        if (!_task_tag) {
+          console.warn(`task tag ${current_task} not implemented yet`)
+          return
+        }
+        if (task_tag != _task_tag) return
+        showSubmitButtons(task_tag, false)
+        requestProgress(current_task, gradioApp().getElementById(`${task_tag}_gallery_container`), gradioApp().getElementById(`${task_tag}_gallery`), function(){
+          showSubmitButtons(task_tag, true)
+        })
+      }
+    }
+
+    let errorHandler = e => window.alert(`invalid internal api respsonse. message: ${e}`)
+
+    fetch("./internal/current_task")
+      .then(res => res.json())
+      .then(successHandler)
+      .catch(errorHandler)
+  }
+
+  var res = create_submit_args(arguments)
+  res[0] = 0
+  return res
+
+}
\ No newline at end of file
diff --git a/modules/progress.py b/modules/progress.py
index 27a336ad..36963c92 100644
--- a/modules/progress.py
+++ b/modules/progress.py
@@ -48,6 +48,20 @@ def set_last_task_result(id_job, result):
   last_task_result = result

 

 

+def restore_progress_call(task_tag):

+    if current_task is None or not current_task[5:-1].startswith(task_tag):

+

+      # image, generation_info, html_info, html_log

+      return tuple(list([None, None, None, None]))

+

+    else:

+

+      t_task = current_task

+      while t_task != last_task_id:

+        time.sleep(2.5)

+      return last_task_result

+

+

 class CurrentTaskResponse(BaseModel):

   current_task: str = Field(default=None, title="Task ID", description="id of the current progress task")

 

diff --git a/modules/ui.py b/modules/ui.py
index 627fbe0b..0133ee12 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -41,6 +41,7 @@ from modules.textual_inversion import textual_inversion
 import modules.hypernetworks.ui

 from modules.generation_parameters_copypaste import image_from_url_text

 import modules.extras

+from modules.progress import restore_progress_call

 

 warnings.filterwarnings("default" if opts.show_warnings else "ignore", category=UserWarning)

 

@@ -293,6 +294,7 @@ def create_toprow(is_img2img):
                 interrupt = gr.Button('Interrupt', elem_id=f"{id_part}_interrupt", elem_classes="generate-box-interrupt")

                 skip = gr.Button('Skip', elem_id=f"{id_part}_skip", elem_classes="generate-box-skip")

                 submit = gr.Button('Generate', elem_id=f"{id_part}_generate", variant='primary')

+                restore_progress = gr.Button('Restore Progress', elem_id=f"{id_part}_restore_progress")

 

                 skip.click(

                     fn=lambda: shared.state.skip(),

@@ -329,7 +331,7 @@ def create_toprow(is_img2img):
                 prompt_styles = gr.Dropdown(label="Styles", elem_id=f"{id_part}_styles", choices=[k for k, v in shared.prompt_styles.styles.items()], value=[], multiselect=True)

                 create_refresh_button(prompt_styles, shared.prompt_styles.reload, lambda: {"choices": [k for k, v in shared.prompt_styles.styles.items()]}, f"refresh_{id_part}_styles")

 

-    return prompt, prompt_styles, negative_prompt, submit, button_interrogate, button_deepbooru, prompt_style_apply, save_style, paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button

+    return prompt, prompt_styles, negative_prompt, submit, restore_progress, button_interrogate, button_deepbooru, prompt_style_apply, save_style, paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button

 

 

 def setup_progressbar(*args, **kwargs):

@@ -446,7 +448,7 @@ def create_ui():
     modules.scripts.scripts_txt2img.initialize_scripts(is_img2img=False)

 

     with gr.Blocks(analytics_enabled=False) as txt2img_interface:

-        txt2img_prompt, txt2img_prompt_styles, txt2img_negative_prompt, submit, _, _, txt2img_prompt_style_apply, txt2img_save_style, txt2img_paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button = create_toprow(is_img2img=False)

+        txt2img_prompt, txt2img_prompt_styles, txt2img_negative_prompt, submit, restore_progress, _, _, txt2img_prompt_style_apply, txt2img_save_style, txt2img_paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button = create_toprow(is_img2img=False)

 

         dummy_component = gr.Label(visible=False)

         txt_prompt_img = gr.File(label="", elem_id="txt2img_prompt_image", file_count="single", type="binary", visible=False)

@@ -578,6 +580,18 @@ def create_ui():
 

             res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height], show_progress=False)

 

+            restore_progress.click(

+              fn=lambda: restore_progress_call('txt2img'),

+              _js="() => restoreProgress('txt2img')",

+              inputs=[],

+              outputs=[

+                  txt2img_gallery,

+                  generation_info,

+                  html_info,

+                  html_log,

+              ]

+            )

+

             txt_prompt_img.change(

                 fn=modules.images.image_data,

                 inputs=[

@@ -646,7 +660,7 @@ def create_ui():
     modules.scripts.scripts_img2img.initialize_scripts(is_img2img=True)

 

     with gr.Blocks(analytics_enabled=False) as img2img_interface:

-        img2img_prompt, img2img_prompt_styles, img2img_negative_prompt, submit, img2img_interrogate, img2img_deepbooru, img2img_prompt_style_apply, img2img_save_style, img2img_paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button = create_toprow(is_img2img=True)

+        img2img_prompt, img2img_prompt_styles, img2img_negative_prompt, submit, restore_progress, img2img_interrogate, img2img_deepbooru, img2img_prompt_style_apply, img2img_save_style, img2img_paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button = create_toprow(is_img2img=True)

 

         img2img_prompt_img = gr.File(label="", elem_id="img2img_prompt_image", file_count="single", type="binary", visible=False)

 

@@ -898,6 +912,18 @@ def create_ui():
             submit.click(**img2img_args)

             res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height], show_progress=False)

 

+            restore_progress.click(

+              fn=lambda: restore_progress_call('img2img'),

+              _js="() => restoreProgress('img2img')",

+              inputs=[],

+              outputs=[

+                  img2img_gallery,

+                  generation_info,

+                  html_info,

+                  html_log,

+              ]

+            )

+

             img2img_interrogate.click(

                 fn=lambda *args: process_interrogate(interrogate, *args),

                 **interrogate_args,

@@ -1491,7 +1517,7 @@ def create_ui():
                 gr.HTML(shared.html("licenses.html"), elem_id="licenses")

 

             gr.Button(value="Show all pages", elem_id="settings_show_all_pages")

-            

+

 

         def unload_sd_weights():

             modules.sd_models.unload_model_weights()