Re-merge #611 - View/Cancel in-progress diffusions

frontend/job_manager.py

@@ -1,14 +1,15 @@
 ''' Provides simple job management for gradio, allowing viewing and stopping in-progress multi-batch generations '''
 from __future__ import annotations
 import gradio as gr
-from gradio.components import Component, Gallery
+from gradio.components import Component, Gallery, Slider
 from threading import Event, Timer
 from typing import Callable, List, Dict, Tuple, Optional, Any
 from dataclasses import dataclass, field
 from functools import partial
 from PIL.Image import Image
 import uuid
 import traceback
+import time
 
 
 @dataclass(eq=True, frozen=True)
@@ -30,9 +31,21 @@ class JobInfo:
     session_key: str
     job_token: Optional[int] = None
     images: List[Image] = field(default_factory=list)
+    active_image: Image = None
+    rec_steps_enabled: bool = False
+    rec_steps_imgs: List[Image] = field(default_factory=list)
+    rec_steps_intrvl: int = None
+    rec_steps_to_gallery: bool = False
+    rec_steps_to_file: bool = False
     should_stop: Event = field(default_factory=Event)
+    refresh_active_image_requested: Event = field(default_factory=Event)
+    refresh_active_image_done: Event = field(default_factory=Event)
+    stop_cur_iter: Event = field(default_factory=Event)
+    active_iteration_cnt: int = field(default_factory=int)
     job_status: str = field(default_factory=str)
     finished: bool = False
+    started: bool = False
+    timestamp: float = None
     removed_output_idxs: List[int] = field(default_factory=list)
 
 
@@ -76,18 +89,27 @@ def wrap_func(
         '''
         return self._job_manager._wrap_func(
             func=func, inputs=inputs, outputs=outputs,
-            refresh_btn=self._refresh_btn, stop_btn=self._stop_btn, status_text=self._status_text
+            job_ui=self
         )
 
     _refresh_btn: gr.Button
     _stop_btn: gr.Button
     _status_text: gr.Textbox
     _stop_all_session_btn: gr.Button
     _free_done_sessions_btn: gr.Button
+    _active_image: gr.Image
+    _active_image_stop_btn: gr.Button
+    _active_image_refresh_btn: gr.Button
+    _rec_steps_intrvl_sldr: gr.Slider
+    _rec_steps_checkbox: gr.Checkbox
+    _save_rec_steps_to_gallery_chkbx: gr.Checkbox
+    _save_rec_steps_to_file_chkbx: gr.Checkbox
     _job_manager: JobManager
 
 
 class JobManager:
+    JOB_MAX_START_TIME = 5.0  # How long can a job be stuck 'starting' before assuming it isn't running
+
     def __init__(self, max_jobs: int):
         self._max_jobs: int = max_jobs
         self._avail_job_tokens: List[Any] = list(range(max_jobs))
@@ -102,11 +124,23 @@ def draw_gradio_ui(self) -> JobManagerUi:
         '''
         assert gr.context.Context.block is not None, "draw_gradio_ui must be called within a 'gr.Blocks' 'with' context"
         with gr.Tabs():
-            with gr.TabItem("Current Session"):
+            with gr.TabItem("Job Controls"):
                 with gr.Row():
-                    stop_btn = gr.Button("Stop", elem_id="stop", variant="secondary")
-                    refresh_btn = gr.Button("Refresh", elem_id="refresh", variant="secondary")
+                    stop_btn = gr.Button("Stop All Batches", elem_id="stop", variant="secondary")
+                    refresh_btn = gr.Button("Refresh Finished Batches", elem_id="refresh", variant="secondary")
                 status_text = gr.Textbox(placeholder="Job Status", interactive=False, show_label=False)
+                with gr.Row():
+                    active_image_stop_btn = gr.Button("Skip Active Batch", variant="secondary")
+                    active_image_refresh_btn = gr.Button("View Batch Progress", variant="secondary")
+                active_image = gr.Image(type="pil", interactive=False, visible=False, elem_id="active_iteration_image")
+            with gr.TabItem("Batch Progress Settings"):
+                with gr.Row():
+                    record_steps_checkbox = gr.Checkbox(value=False, label="Enable Batch Progress Grid")
+                    record_steps_interval_slider = gr.Slider(
+                        value=3, label="Record Interval (steps)", minimum=1, maximum=25, step=1)
+                with gr.Row() as record_steps_box:
+                    steps_to_gallery_checkbox = gr.Checkbox(value=False, label="Save Progress Grid to Gallery")
+                    steps_to_file_checkbox = gr.Checkbox(value=False, label="Save Progress Grid to File")
             with gr.TabItem("Maintenance"):
                 with gr.Row():
                     gr.Markdown(
@@ -118,9 +152,15 @@ def draw_gradio_ui(self) -> JobManagerUi:
                     free_done_sessions_btn = gr.Button(
                         "Clear Finished Jobs", elem_id="clear_finished", variant="secondary"
                     )
+
         return JobManagerUi(_refresh_btn=refresh_btn, _stop_btn=stop_btn, _status_text=status_text,
                             _stop_all_session_btn=stop_all_sessions_btn, _free_done_sessions_btn=free_done_sessions_btn,
-                            _job_manager=self)
+                            _active_image=active_image, _active_image_stop_btn=active_image_stop_btn,
+                            _active_image_refresh_btn=active_image_refresh_btn,
+                            _rec_steps_checkbox=record_steps_checkbox,
+                            _save_rec_steps_to_gallery_chkbx=steps_to_gallery_checkbox,
+                            _save_rec_steps_to_file_chkbx=steps_to_file_checkbox,
+                            _rec_steps_intrvl_sldr=record_steps_interval_slider, _job_manager=self)
 
     def clear_all_finished_jobs(self):
         ''' Removes all currently finished jobs, across all sessions.
@@ -134,6 +174,7 @@ def stop_all_jobs(self):
         for session in self._sessions.values():
             for job in session.jobs.values():
                 job.should_stop.set()
+                job.stop_cur_iter.set()
 
     def _get_job_token(self, block: bool = False) -> Optional[int]:
         ''' Attempts to acquire a job token, optionally blocking until available '''
@@ -175,6 +216,26 @@ def _stop_wrapped_func(self, func_key: FuncKey, session_key: str) -> List[Compon
         job_info.should_stop.set()
         return "Stopping after current batch finishes"
 
+    def _refresh_cur_iter_func(self, func_key: FuncKey, session_key: str) -> List[Component]:
+        ''' Updates information from the active iteration '''
+        session_info, job_info = self._get_call_info(func_key, session_key)
+        if job_info is None:
+            return [None, f"Session {session_key} was not running function {func_key}"]
+
+        job_info.refresh_active_image_requested.set()
+        if job_info.refresh_active_image_done.wait(timeout=20.0):
+            job_info.refresh_active_image_done.clear()
+            return [gr.Image.update(value=job_info.active_image, visible=True), f"Sample iteration {job_info.active_iteration_cnt}"]
+        return [gr.Image.update(visible=False), "Timed out getting image"]
+
+    def _stop_cur_iter_func(self, func_key: FuncKey, session_key: str) -> List[Component]:
+        ''' Marks that the active iteration should be stopped'''
+        session_info, job_info = self._get_call_info(func_key, session_key)
+        if job_info is None:
+            return [None, f"Session {session_key} was not running function {func_key}"]
+        job_info.stop_cur_iter.set()
+        return [gr.Image.update(visible=False), "Stopping current iteration"]
+
     def _get_call_info(self, func_key: FuncKey, session_key: str) -> Tuple[SessionInfo, JobInfo]:
         ''' Helper to get the SessionInfo and JobInfo. '''
         session_info = self._sessions.get(session_key, None)
@@ -207,7 +268,8 @@ def _run_queued_jobs(self) -> None:
 
     def _pre_call_func(
             self, func_key: FuncKey, output_dummy_obj: Component, refresh_btn: gr.Button, stop_btn: gr.Button,
-            status_text: gr.Textbox, session_key: str) -> List[Component]:
+            status_text: gr.Textbox, active_image: gr.Image, active_refresh_btn: gr.Button, active_stop_btn: gr.Button,
+            session_key: str) -> List[Component]:
         ''' Called when a job is about to start '''
         session_info, job_info = self._get_call_info(func_key, session_key)
 
@@ -219,7 +281,9 @@ def _pre_call_func(
         return {output_dummy_obj: triggerChangeEvent(),
                 refresh_btn: gr.Button.update(variant="primary", value=refresh_btn.value),
                 stop_btn: gr.Button.update(variant="primary", value=stop_btn.value),
-                status_text: gr.Textbox.update(value="Generation has started. Click 'Refresh' for updates")
+                status_text: gr.Textbox.update(value="Generation has started. Click 'Refresh' to see finished images, 'View Batch Progress' for active images"),
+                active_refresh_btn: gr.Button.update(variant="primary", value=active_refresh_btn.value),
+                active_stop_btn: gr.Button.update(variant="primary", value=active_stop_btn.value),
                 }
 
     def _call_func(self, func_key: FuncKey, session_key: str) -> List[Component]:
@@ -228,24 +292,26 @@ def _call_func(self, func_key: FuncKey, session_key: str) -> List[Component]:
         if session_info is None or job_info is None:
             return []
 
+        job_info.started = True
         try:
+            if job_info.should_stop.is_set():
+                raise Exception(f"Job {job_info} requested a stop before execution began")
             outputs = job_info.func(*job_info.inputs, job_info=job_info)
         except Exception as e:
             job_info.job_status = f"Error: {e}"
             print(f"Exception processing job {job_info}: {e}\n{traceback.format_exc()}")
-            outputs = []
+            raise
+        finally:
+            job_info.finished = True
+            session_info.finished_jobs[func_key] = session_info.jobs.pop(func_key)
+            self._release_job_token(job_info.job_token)
 
         # Filter the function output for any removed outputs
         filtered_output = []
         for idx, output in enumerate(outputs):
             if idx not in job_info.removed_output_idxs:
                 filtered_output.append(output)
 
-        job_info.finished = True
-        session_info.finished_jobs[func_key] = session_info.jobs.pop(func_key)
-
-        self._release_job_token(job_info.job_token)
-
         # The wrapper added a dummy JSON output. Append a random text string
         # to fire the dummy objects 'change' event to notify that the job is done
         filtered_output.append(triggerChangeEvent())
@@ -254,12 +320,16 @@ def _call_func(self, func_key: FuncKey, session_key: str) -> List[Component]:
 
     def _post_call_func(
             self, func_key: FuncKey, output_dummy_obj: Component, refresh_btn: gr.Button, stop_btn: gr.Button,
-            status_text: gr.Textbox, session_key: str) -> List[Component]:
+            status_text: gr.Textbox, active_image: gr.Image, active_refresh_btn: gr.Button, active_stop_btn: gr.Button,
+            session_key: str) -> List[Component]:
         ''' Called when a job completes '''
         return {output_dummy_obj: triggerChangeEvent(),
                 refresh_btn: gr.Button.update(variant="secondary", value=refresh_btn.value),
                 stop_btn: gr.Button.update(variant="secondary", value=stop_btn.value),
-                status_text: gr.Textbox.update(value="Generation has finished!")
+                status_text: gr.Textbox.update(value="Generation has finished!"),
+                active_refresh_btn: gr.Button.update(variant="secondary", value=active_refresh_btn.value),
+                active_stop_btn: gr.Button.update(variant="secondary", value=active_stop_btn.value),
+                active_image: gr.Image.update(visible=False)
                 }
 
     def _update_gallery_event(self, func_key: FuncKey, session_key: str) -> List[Component]:
@@ -270,21 +340,17 @@ def _update_gallery_event(self, func_key: FuncKey, session_key: str) -> List[Com
         if session_info is None or job_info is None:
             return []
 
-        if job_info.finished:
-            session_info.finished_jobs.pop(func_key)
-
         return job_info.images
 
-    def _wrap_func(
-            self, func: Callable, inputs: List[Component], outputs: List[Component],
-            refresh_btn: gr.Button = None, stop_btn: gr.Button = None,
-            status_text: Optional[gr.Textbox] = None) -> Tuple[Callable, List[Component]]:
+    def _wrap_func(self, func: Callable, inputs: List[Component],
+                   outputs: List[Component],
+                   job_ui: JobManagerUi) -> Tuple[Callable, List[Component]]:
         ''' handles JobManageUI's wrap_func'''
 
         assert gr.context.Context.block is not None, "wrap_func must be called within a 'gr.Blocks' 'with' context"
 
         # Create a unique key for this job
-        func_key = FuncKey(job_id=uuid.uuid4(), func=func)
+        func_key = FuncKey(job_id=uuid.uuid4().hex, func=func)
 
         # Create a unique session key (next gradio release can use gr.State, see https://gradio.app/state_in_blocks/)
         if self._session_key is None:
@@ -302,9 +368,6 @@ def _wrap_func(
                 del outputs[idx]
                 break
 
-        # Add the session key to the inputs
-        inputs += [self._session_key]
-
         # Create dummy objects
         update_gallery_obj = gr.JSON(visible=False, elem_id="JobManagerDummyObject")
         update_gallery_obj.change(
@@ -314,21 +377,49 @@ def _wrap_func(
             queue=False
         )
 
-        if refresh_btn:
-            refresh_btn.variant = 'secondary'
-            refresh_btn.click(
+        if job_ui._refresh_btn:
+            job_ui._refresh_btn.variant = 'secondary'
+            job_ui._refresh_btn.click(
                 partial(self._refresh_func, func_key),
                 [self._session_key],
-                [update_gallery_obj, status_text],
+                [update_gallery_obj, job_ui._status_text],
                 queue=False
             )
 
-        if stop_btn:
-            stop_btn.variant = 'secondary'
-            stop_btn.click(
+        if job_ui._stop_btn:
+            job_ui._stop_btn.variant = 'secondary'
+            job_ui._stop_btn.click(
                 partial(self._stop_wrapped_func, func_key),
                 [self._session_key],
-                [status_text],
+                [job_ui._status_text],
+                queue=False
+            )
+
+        if job_ui._active_image and job_ui._active_image_refresh_btn:
+            job_ui._active_image_refresh_btn.click(
+                partial(self._refresh_cur_iter_func, func_key),
+                [self._session_key],
+                [job_ui._active_image, job_ui._status_text],
+                queue=False
+            )
+
+        if job_ui._active_image_stop_btn:
+            job_ui._active_image_stop_btn.click(
+                partial(self._stop_cur_iter_func, func_key),
+                [self._session_key],
+                [job_ui._active_image, job_ui._status_text],
+                queue=False
+            )
+
+        if job_ui._stop_all_session_btn:
+            job_ui._stop_all_session_btn.click(
+                self.stop_all_jobs, [], [],
+                queue=False
+            )
+
+        if job_ui._free_done_sessions_btn:
+            job_ui._free_done_sessions_btn.click(
+                self.clear_all_finished_jobs, [], [],
                 queue=False
             )
 
@@ -346,7 +437,8 @@ def _wrap_func(
         # Since some parameters are optional it makes sense to use the 'dict' return value type, which requires
         # the Component as a key... so group together the UI components that the event listeners are going to update
         # to make it easy to append to function calls and outputs
-        job_ui_params = [refresh_btn, stop_btn, status_text]
+        job_ui_params = [job_ui._refresh_btn, job_ui._stop_btn, job_ui._status_text,
+                         job_ui._active_image, job_ui._active_image_refresh_btn, job_ui._active_image_stop_btn]
         job_ui_outputs = [comp for comp in job_ui_params if comp is not None]
 
         # Here a chain is constructed that will make a 'pre' call, a 'run' call, and a 'post' call,
@@ -375,27 +467,50 @@ def _wrap_func(
             queue=False
         )
 
+        # Add any components that we want the runtime values for
+        added_inputs = [self._session_key, job_ui._rec_steps_checkbox, job_ui._save_rec_steps_to_gallery_chkbx,
+                        job_ui._save_rec_steps_to_file_chkbx, job_ui._rec_steps_intrvl_sldr]
+
         # Now replace the original function with one that creates a JobInfo and triggers the dummy obj
+        def wrapped_func(*wrapped_inputs):
+            # Remove the added_inputs (pop opposite order of list)
 
-        def wrapped_func(*inputs):
-            session_key = inputs[-1]
-            inputs = inputs[:-1]
+            wrapped_inputs = list(wrapped_inputs)
+            rec_steps_interval: int = wrapped_inputs.pop()
+            save_rec_steps_file: bool = wrapped_inputs.pop()
+            save_rec_steps_grid: bool = wrapped_inputs.pop()
+            record_steps_enabled: bool = wrapped_inputs.pop()
+            session_key: str = wrapped_inputs.pop()
+            job_inputs = tuple(wrapped_inputs)
 
             # Get or create a session for this key
             session_info = self._sessions.setdefault(session_key, SessionInfo())
 
             # Is this session already running this job?
             if func_key in session_info.jobs:
-                return {status_text: "This session is already running that function!"}
+                job_info = session_info.jobs[func_key]
+                # If the job seems stuck in 'starting' then go ahead and toss it
+                if not job_info.started and time.time() > job_info.timestamp + JobManager.JOB_MAX_START_TIME:
+                    job_info.should_stop.set()
+                    job_info.stop_cur_iter.set()
+                    session_info.jobs.pop(func_key)
+                    return {job_ui._status_text: "Canceled possibly hung job. Try again"}
+                return {job_ui._status_text: "This session is already running that function!"}
+
+            # Is this a new run of a previously finished job? Clear old info
+            if func_key in session_info.finished_jobs:
+                session_info.finished_jobs.pop(func_key)
 
             job_token = self._get_job_token(block=False)
-            job = JobInfo(inputs=inputs, func=func, removed_output_idxs=removed_idxs, session_key=session_key,
-                          job_token=job_token)
+            job = JobInfo(
+                inputs=job_inputs, func=func, removed_output_idxs=removed_idxs, session_key=session_key,
+                job_token=job_token, rec_steps_enabled=record_steps_enabled, rec_steps_intrvl=rec_steps_interval,
+                rec_steps_to_gallery=save_rec_steps_grid, rec_steps_to_file=save_rec_steps_file, timestamp=time.time())
             session_info.jobs[func_key] = job
 
             ret = {pre_call_dummyobj: triggerChangeEvent()}
             if job_token is None:
-                ret[status_text] = "Job is queued"
+                ret[job_ui._status_text] = "Job is queued"
             return ret
 
-        return wrapped_func, inputs, [pre_call_dummyobj, status_text]
+        return wrapped_func, inputs + added_inputs, [pre_call_dummyobj, job_ui._status_text]