Merge e1c91db633 into ea9bd9fc74

javascript/progressbar.js

@@ -124,8 +124,8 @@ function setTitle(progress){
 }
 
 
-function randomId(){
+function randomId(prefix=null){
-    return "task(" + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7)+")"
+    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

javascript/ui.js

@@ -137,7 +137,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)
 
@@ -154,7 +154,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)
     })
@@ -336,3 +336,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 progress info 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
+
+}

modules/call_queue.py

@@ -7,7 +7,7 @@ import time
 from modules import shared, progress
 
 queue_lock = threading.Lock()
-
+queue_lock_condition = threading.Condition(lock=queue_lock)
 
 def wrap_queued_call(func):
     def f(*args, **kwargs):
@@ -37,6 +37,7 @@ def wrap_gradio_gpu_call(func, extra_outputs=None):
                 res = func(*args, **kwargs)
             finally:
                 progress.finish_task(id_task)
+                progress.set_last_task_result(id_task, res)
 
             shared.state.end()
 

modules/progress.py

@@ -4,7 +4,9 @@ import time
 
 import gradio as gr
 from pydantic import BaseModel, Field
+from typing import List
 
+from modules import call_queue
 from modules.shared import opts
 
 import modules.shared as shared
@@ -36,6 +38,34 @@ def finish_task(id_task):
 def add_task_to_queue(id_job):
     pending_tasks[id_job] = time.time()
 
+last_task_id = None
+last_task_result = None
+
+def set_last_task_result(id_job, result):
+  global last_task_id
+  global last_task_result
+
+  last_task_id = id_job
+  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
+      with call_queue.queue_lock_condition:
+        call_queue.queue_lock_condition.wait_for(lambda: t_task == last_task_id)
+
+      return last_task_result
+
+
+class CurrentTaskResponse(BaseModel):
+  current_task: str = Field(default=None, title="Task ID", description="id of the current progress task")
 
 class ProgressRequest(BaseModel):
     id_task: str = Field(default=None, title="Task ID", description="id of the task to get progress for")
@@ -56,6 +86,8 @@ class ProgressResponse(BaseModel):
 def setup_progress_api(app):
     return app.add_api_route("/internal/progress", progressapi, methods=["POST"], response_model=ProgressResponse)
 
+def setup_current_task_api(app):
+    return app.add_api_route("/internal/current_task", current_task_api, methods=["GET"], response_model=CurrentTaskResponse)
 
 def progressapi(req: ProgressRequest):
     active = req.id_task == current_task
@@ -97,3 +129,5 @@ def progressapi(req: ProgressRequest):
 
     return ProgressResponse(active=active, queued=queued, completed=completed, progress=progress, eta=eta, live_preview=live_preview, id_live_preview=id_live_preview, textinfo=shared.state.textinfo)
 
+def current_task_api():
+  return CurrentTaskResponse(current_task=current_task)

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)
 
@@ -305,6 +306,7 @@ def create_toprow(is_img2img):
                 interrupt = gr.Button('Interrupt', elem_id=f"{id_part}_interrupt")
                 skip = gr.Button('Skip', elem_id=f"{id_part}_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(),
@@ -341,7 +343,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):
@@ -458,7 +460,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)
@@ -586,6 +588,19 @@ def create_ui():
             txt2img_prompt.submit(**txt2img_args)
             submit.click(**txt2img_args)
 
+
+            restore_progress.click(
+              fn=lambda: restore_progress_call('txt2img'),
+              _js="() => restoreProgress('txt2img')",
+              inputs=[],
+              outputs=[
+                  txt2img_gallery,
+                  generation_info,
+                  html_info,
+                  html_log,
+              ]
+            )
+
             res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height])
 
             txt_prompt_img.change(
@@ -656,7 +671,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)
 
@@ -904,6 +919,19 @@ def create_ui():
 
             img2img_prompt.submit(**img2img_args)
             submit.click(**img2img_args)
+
+            restore_progress.click(
+              fn=lambda: restore_progress_call('img2img'),
+              _js="() => restoreProgress('img2img')",
+              inputs=[],
+              outputs=[
+                  img2img_gallery,
+                  generation_info,
+                  html_info,
+                  html_log,
+              ]
+            )
+
             res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height])
 
             img2img_interrogate.click(

webui.py

@@ -232,6 +232,7 @@ def webui():
         app.add_middleware(GZipMiddleware, minimum_size=1000)
 
         modules.progress.setup_progress_api(app)
+        modules.progress.setup_current_task_api(app)
 
         if launch_api:
             create_api(app)