add interrupt to processing

javascript/black-orange.css

@@ -93,6 +93,8 @@ svg.feather.feather-image, .feather .feather-image { display: none }
 #txt2img_gallery, #img2img_gallery, #extras_gallery { background: black !important; padding: 0; margin: 0; object-fit: contain; box-shadow: none; min-height: 0; }
 #txt2img_generate, #img2img_generate { height: 36px; border: none; border-radius: 0; min-height: 36px; padding: 0; }
 #txt2img_interrupt, #img2img_interrupt, #txt2img_skip, #img2img_skip { height: 36px; min-width: 116px; max-width: 116px; border: none; border-radius: 0; background-color: var(--inactive-color); margin-top: 46px; display: block !important; padding: 0; }
+#extras_generate, #extras_interrupt, #extras_skip { border: none; border-radius: 0; background-color: var(--inactive-color); }
+#extras_upscale { margin-top: 10px }
 #txt2img_progress_row > div { min-width: var(--left-column); max-width: var(--left-column); }
 #txt2img_results, #img2img_results, #extras_results { background-color: black; padding: 0; }
 #txt2img_seed_row { padding: 0; margin-top: 8px; }
@@ -104,7 +106,6 @@ svg.feather.feather-image, .feather .feather-image { display: none }
 #txtimg_hr_finalres { max-width: 200px; }
 #pnginfo_html2_info { margin-top: -18px; background-color: var(--input-background-fill); padding: var(--input-padding) }
 #txt2img_extra_refresh, #txt2img_extra_close { height: 1.7em; }
-#extras_generate { margin-top: 8px; }
 
 /* custom elements overrides */
 #steps-animation, #controlnet { border-width: 0; }

javascript/style.css

@@ -134,7 +134,7 @@ button.custom-button{
     position: absolute;
     width: 50%;
     height: 100%;
-    display: none;
+    display: block;
     background: #b4c0cc;
 }
 .gradio-button.generate-box-skip:hover, .gradio-button.generate-box-interrupt:hover{
@@ -644,3 +644,6 @@ footer {
 .theme-preview { display: none; position: fixed; border: 4px solid var(--neutral-600); box-shadow: 2px 2px 2px 2px var(--neutral-700); top: 0; bottom: 0; left: 0; right: 0; margin: auto; max-width: 75vw; z-index: 999; }
 
 #scripts_alwayson_txt2img, scripts_alwayson_img2img { display: grid }
+
+#extras_generate, #extras_interrupt, #extras_skip { display: block !important; position: relative; height: 36px; }
+#extras_upscale { margin-top: 10px }

modules/hashes.py

@@ -25,57 +25,43 @@ def cache(subsection):
             else:
                 with open(cache_filename, "r", encoding="utf8") as file:
                     cache_data = json.load(file)
-
     s = cache_data.get(subsection, {})
     cache_data[subsection] = s
-
     return s
 
 
 def calculate_sha256(filename):
     hash_sha256 = hashlib.sha256()
     blksize = 1024 * 1024
-
     with progress.open(filename, 'rb', description=f'Calculating model hash: [cyan]{filename}', auto_refresh=True) as f:
         for chunk in iter(lambda: f.read(blksize), b""):
             hash_sha256.update(chunk)
-
     return hash_sha256.hexdigest()
 
 
 def sha256_from_cache(filename, title):
     hashes = cache("hashes")
     ondisk_mtime = os.path.getmtime(filename)
-
     if title not in hashes:
         return None
-
     cached_sha256 = hashes[title].get("sha256", None)
     cached_mtime = hashes[title].get("mtime", 0)
-
     if ondisk_mtime > cached_mtime or cached_sha256 is None:
         return None
-
     return cached_sha256
 
 
 def sha256(filename, title):
     hashes = cache("hashes")
-
     sha256_value = sha256_from_cache(filename, title)
     if sha256_value is not None:
         return sha256_value
-
     if shared.cmd_opts.no_hashing:
         return None
-
     sha256_value = calculate_sha256(filename)
-
     hashes[title] = {
         "mtime": os.path.getmtime(filename),
         "sha256": sha256_value,
     }
-
     dump_cache()
-
     return sha256_value

modules/postprocessing.py

@@ -50,6 +50,9 @@ def run_postprocessing(extras_mode, image, image_folder: List[tempfile.NamedTemp
         outpath = opts.outdir_samples or opts.outdir_extras_samples
     infotext = ''
     for image, name, ext in zip(image_data, image_names, image_ext):
+        if shared.state.interrupted:
+            shared.log.debug('Postprocess interrupted')
+            break
         if image is None:
             continue
         shared.state.textinfo = name

modules/processing.py

@@ -603,6 +603,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             if state.skipped:
                 state.skipped = False
             if state.interrupted:
+                shared.log.debug(f'Process interrupted: {n}/{p.n_iter}')
                 break
             prompts = p.all_prompts[n * p.batch_size:(n + 1) * p.batch_size]
             negative_prompts = p.all_negative_prompts[n * p.batch_size:(n + 1) * p.batch_size]

modules/sd_models.py

@@ -248,6 +248,7 @@ def get_checkpoint_state_dict(checkpoint_info: CheckpointInfo, timer):
 
 
 def load_model_weights(model: torch.nn.Module, checkpoint_info: CheckpointInfo, state_dict, timer):
+    shared.log.debug(f'Model weights loading: {memory_stats()}')
     sd_model_hash = checkpoint_info.calculate_shorthash()
     timer.record("hash")
     shared.opts.data["sd_model_checkpoint"] = checkpoint_info.title

modules/shared.py

@@ -98,9 +98,11 @@ class State:
     server_start = None
 
     def skip(self):
+        log.debug('Skip requested')
         self.skipped = True
 
     def interrupt(self):
+        log.debug('Interrupt requested')
         self.interrupted = True
 
     def nextjob(self):

modules/ui_postprocessing.py

@@ -33,11 +33,19 @@ def create_ui():
             with gr.Row():
                 buttons = parameters_copypaste.create_buttons(["txt2img", "img2img", "inpaint"])
 
-            submit = gr.Button('Generate', elem_id="extras_generate", variant='primary')
+            # submit = gr.Button('Generate', elem_id="extras_generate", variant='primary') # TODO: add all
 
             script_inputs = scripts.scripts_postproc.setup_ui()
 
         with gr.Column():
+            id_part = 'extras'
+            with gr.Row(elem_id=f"{id_part}_generate_box", elem_classes="generate-box"):
+                submit = gr.Button('Generate', elem_id=f"{id_part}_generate", variant='primary')
+                interrupt = gr.Button('Stop', elem_id=f"{id_part}_interrupt", variant='secondary')
+                skip = gr.Button('Skip', elem_id=f"{id_part}_skip", variant='secondary')
+                skip.click(fn=lambda: shared.state.skip(), inputs=[], outputs=[])
+                interrupt.click(fn=lambda: shared.state.interrupt(), inputs=[], outputs=[])
+
             result_images, html_info_x, html_info, _html_log = ui_common.create_output_panel("extras", shared.opts.outdir_extras_samples)
             html_info = gr.HTML(elem_id="pnginfo_html_info")
             generation_info = gr.Textbox(elem_id="pnginfo_generation_info", label="Parameters", visible=False)

scripts/postprocessing_upscale.py

@@ -16,7 +16,7 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
         selected_tab = gr.State(value=0) # pylint: disable=abstract-class-instantiated
 
         with gr.Column():
-            with FormRow():
+            with FormRow(elem_id="extras_upscale"):
                 with gr.Tabs(elem_id="extras_resize_mode"):
                     with gr.TabItem('Scale by', elem_id="extras_scale_by_tab") as tab_scale_by:
                         upscaling_resize = gr.Slider(minimum=1.0, maximum=8.0, step=0.05, label="Resize", value=4, elem_id="extras_upscaling_resize")