JobManager: Re-merge Sygil-Dev#611

PR Sygil-Dev#611 seems to have got lost in the shuffle after
the transition to 'dev'.

This commit re-merges the feature branch. This adds
support for viewing preview images as the image
generates, as well as cancelling in-progress images
and a couple fixes and clean-ups.

frontend/job_manager.py

@@ -1,7 +1,7 @@
 ''' 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
@@ -30,7 +30,17 @@ 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
  removed_output_idxs: List[int] = field(default_factory=list)
@@ -76,14 +86,21 @@ 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
 
 
@@ -102,11 +119,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 +147,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 +169,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 +211,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 +263,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 +276,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]:
@@ -233,7 +292,7 @@ def _call_func(self, func_key: FuncKey, session_key: str) -> List[Component]:
  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
 
  # Filter the function output for any removed outputs
  filtered_output = []
@@ -254,12 +313,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 +333,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 +361,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 +370,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 +430,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 +460,43 @@ 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!"}
+ 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)
  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]