add back postprocess_image cb

lib_comfyui/comfyui/iframe_requests.py

@@ -72,17 +72,17 @@ def start_workflow_sync(
     @ipc.restrict_to_process('webui')
     def validate_amount_of_nodes_or_throw(
         workflow_type_id: str,
-        max_amount_of_nodes: Sequence[Optional[int]]
+        max_amount_of_io_nodes: Sequence[Optional[int]]
     ) -> None:
-        if len(max_amount_of_nodes) != 2:
+        if len(max_amount_of_io_nodes) != 2:
             raise RuntimeError(f'Expected a sequence of length 2 for argument "max_amount_of_nodes", got {len(max_amount_of_nodes)} instead')
 
         workflow_graph = get_workflow_graph(workflow_type_id)
         node_types = [node['type'] for node in workflow_graph['nodes']]
         amount_of_from_webui_nodes = len([t for t in node_types if t == 'FromWebui'])
         amount_of_to_webui_nodes = len([t for t in node_types if t == 'ToWebui'])
-        max_from_webui_nodes = max_amount_of_nodes[0] if max_amount_of_nodes[0] is not None else amount_of_from_webui_nodes
-        max_to_webui_nodes = max_amount_of_nodes[1] if max_amount_of_nodes[1] is not None else amount_of_to_webui_nodes
+        max_from_webui_nodes = max_amount_of_io_nodes[0] if max_amount_of_io_nodes[0] is not None else amount_of_from_webui_nodes
+        max_to_webui_nodes = max_amount_of_io_nodes[1] if max_amount_of_io_nodes[1] is not None else amount_of_to_webui_nodes
 
         if amount_of_from_webui_nodes > max_from_webui_nodes:
             raise TooManyFromWebuiNodesError(f'Unable to run the workflow {workflow_type_id}. '

lib_comfyui/default_workflow_types.py

@@ -20,34 +20,41 @@
     default_workflow=external_code.AUTO_WORKFLOW,
     types='LATENT',
 )
+postprocess_latent_workflow_type = external_code.WorkflowType(
+    base_id='postprocess_latent',
+    display_name='Postprocess (latent)',
+    default_workflow=external_code.AUTO_WORKFLOW,
+    types='LATENT',
+)
 postprocess_workflow_type = external_code.WorkflowType(
     base_id='postprocess',
     display_name='Postprocess',
     default_workflow=external_code.AUTO_WORKFLOW,
     types='IMAGE',
 )
+postprocess_image_workflow_type = external_code.WorkflowType(
+    base_id='postprocess_image',
+    display_name='Postprocess image',
+    default_workflow=external_code.AUTO_WORKFLOW,
+    types='IMAGE',
+)
 before_save_image_workflow_type = external_code.WorkflowType(
     base_id='before_save_image',
     display_name='Before save image',
     default_workflow=external_code.AUTO_WORKFLOW,
     types='IMAGE',
 )
-postprocess_latent_workflow_type = external_code.WorkflowType(
-    base_id='postprocess_latent',
-    display_name='Postprocess (latent)',
-    default_workflow=external_code.AUTO_WORKFLOW,
-    types='LATENT',
-)
 
 
 def add_default_workflow_types():
     workflow_types = [
         sandbox_tab_workflow_type,
-        postprocess_workflow_type,
-        before_save_image_workflow_type,
-        postprocess_latent_workflow_type,
         preprocess_workflow_type,
         preprocess_latent_workflow_type,
+        postprocess_latent_workflow_type,
+        postprocess_workflow_type,
+        postprocess_image_workflow_type,
+        before_save_image_workflow_type,
     ]
 
     for workflow_type in workflow_types:

lib_comfyui/external_code/api.py

@@ -206,7 +206,7 @@ def run_workflow(
     batch_input: Any,
     queue_front: Optional[bool] = None,
     identity_on_error: Optional[bool] = False,
-    max_amount_of_nodes: Optional[Sequence[int]] = None
+    max_amount_of_io_nodes: Optional[Sequence[int]] = None
 ) -> List[Any]:
     """
     Run a comfyui workflow synchronously
@@ -221,7 +221,7 @@ def run_workflow(
             - if **input_types** is a str, **batch_input** should be a single value that should match the type expected by **input_types**.
         queue_front (Optional[bool]): Whether to queue the workflow before or after the currently queued workflows
         identity_on_error (Optional[bool]): Whether to return batch_input (converted to the type expected by workflow_type.types) instead of raising a RuntimeError when the workflow fails to run
-        max_amount_of_nodes: (Optional[Sequence[int]]): Maximum amount of From Webui and To Webui nodes present in the workflow, respectively. If either of the conditions are not met, a RuntimeError is raised
+        max_amount_of_io_nodes: (Optional[Sequence[int]]): Maximum amount of From Webui and To Webui nodes present in the workflow, respectively. If either of the conditions are not met, a RuntimeError is raised
     Returns:
         The outputs of the workflow, as a list.
         Each element of the list corresponds to one output node in the workflow.
@@ -248,18 +248,19 @@ def run_workflow(
     if queue_front is None:
         queue_front = getattr(global_state, 'queue_front', True)
 
-    if max_amount_of_nodes is None:
-        max_amount_of_nodes = [None, None]
+    if max_amount_of_io_nodes is None:
+        max_amount_of_io_nodes = [None, None]
 
     batch_input_args, input_types = _normalize_to_tuple(batch_input, workflow_type.input_types)
 
     if not candidate_ids:
-        raise ValueError(f'The workflow type {workflow_type.pretty_str()} does not exist on tab {tab}. Valid tabs for the given workflow type are: {workflow_type.tabs}')
+        raise ValueError(f'The workflow type {workflow_type.pretty_str()} does not exist on tab {tab}. '
+                         f'Valid tabs for the given workflow type are: {workflow_type.tabs}')
 
     workflow_type_id = candidate_ids[0]
 
     try:
-        ComfyuiIFrameRequests.validate_amount_of_nodes_or_throw(workflow_type_id, max_amount_of_nodes)
+        ComfyuiIFrameRequests.validate_amount_of_nodes_or_throw(workflow_type_id, max_amount_of_io_nodes)
     except RuntimeError as e:
         if not identity_on_error:
             raise e

lib_comfyui/webui/callbacks.py

@@ -60,7 +60,7 @@ def on_before_image_saved(params):
         tab=tab,
         batch_input=type_conversion.webui_image_to_comfyui([params.image]),
         identity_on_error=True,
-        max_amount_of_nodes=[None, 1]
+        max_amount_of_io_nodes=[None, 1]
     )
 
-    params.image.putdata(type_conversion.comfyui_image_to_webui(results[0], return_tensors=False)[0].getdata())
+    params.image = type_conversion.comfyui_image_to_webui(results[0], return_tensors=False)[0]

scripts/comfyui.py

@@ -54,6 +54,8 @@ def process(self, p, webui_client_id, queue_front, enabled_workflow_type_ids, *a
         patches.patch_processing(p)
 
     def postprocess_batch_list(self, p, pp, *args, **kwargs):
+        iframe_requests.extend_infotext_with_comfyui_workflows(p, self.get_tab())
+
         if not external_code.is_workflow_type_enabled(default_workflow_types.postprocess_workflow_type.get_ids(self.get_tab())[0]):
             return
 
@@ -79,7 +81,21 @@ def postprocess_batch_list(self, p, pp, *args, **kwargs):
 
         pp.images.clear()
         pp.images.extend(all_results)
-        iframe_requests.extend_infotext_with_comfyui_workflows(p, self.get_tab())
+
+    def postprocess_image(self, p, pp, *args):
+        if not external_code.is_workflow_type_enabled(
+                default_workflow_types.postprocess_image_workflow_type.get_ids(self.get_tab())[0]):
+            return
+
+        results = external_code.run_workflow(
+            workflow_type=default_workflow_types.postprocess_image_workflow_type,
+            tab=self.get_tab(),
+            batch_input=type_conversion.webui_image_to_comfyui([pp.image]),
+            identity_on_error=True,
+            max_amount_of_io_nodes=[None, 1]
+        )
+
+        pp.image = type_conversion.comfyui_image_to_webui(results[0], return_tensors=False)[0]
 
 
 def extract_contiguous_buckets(images, batch_size):