From 2fd52ae230014906a7d6d5607f79b47f64e8f3e6 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 10 Jan 2026 02:35:04 +0000 Subject: [PATCH] fix(control): ensure state cleanup runs reliably Fix control tab getting stuck in running state after generation completes. --- modules/control/run.py | 2 -- modules/ui_control.py | 8 +++++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/modules/control/run.py b/modules/control/run.py index 726d3aa4f..393cb089c 100644 --- a/modules/control/run.py +++ b/modules/control/run.py @@ -679,7 +679,5 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg if len(info_txt) > 0: html_txt = html_txt + infotext_to_html(info_txt[0]) if is_generator: - jobid = shared.state.begin('UI') yield (output_images, blended_image, html_txt, output_filename) - shared.state.end(jobid) return (output_images, blended_image, html_txt, output_filename) diff --git a/modules/ui_control.py b/modules/ui_control.py index 5480c6dce..e22ec0957 100644 --- a/modules/ui_control.py +++ b/modules/ui_control.py @@ -103,12 +103,15 @@ def generate_click(job_id: str, state: str, active_tab: str, *args): for results in control_run(state, units, helpers.input_source, helpers.input_init, helpers.input_mask, active_tab, True, *args): progress.record_results(job_id, results) yield return_controls(results, t) + except GeneratorExit: + pass # Generator was closed by client - cleanup will run in finally except Exception as e: shared.log.error(f"Control exception: {e}") errors.display(e, 'Control') yield [None, None, None, None, f'Control: Exception: {e}', ''] - progress.finish_task(job_id) - shared.state.end(jobid) + finally: + progress.finish_task(job_id) + shared.state.end(jobid) def create_ui(_blocks: gr.Blocks=None): @@ -266,7 +269,6 @@ def create_ui(_blocks: gr.Blocks=None): prompt.submit(**select_dict) negative.submit(**select_dict) - btn_generate.click(**select_dict) for ctrl in [input_image, input_video, input_batch, input_folder, init_image, init_video, init_batch, init_folder, tab_image, tab_video, tab_batch, tab_folder, tab_image_init, tab_video_init, tab_batch_init, tab_folder_init]: if hasattr(ctrl, 'change'): ctrl.change(**select_dict)