unet workflow

lib_comfyui/default_workflow_types.py

@@ -20,6 +20,16 @@
     tabs='img2img',
     default_workflow=default_workflows_dir / 'process_latent.json',
 )
+before_unet_workflow_type = external_code.WorkflowType(
+    base_id='before_unet',
+    display_name='Before unet',
+    default_workflow=default_workflows_dir / 'process_latent.json',
+)
+after_unet_workflow_type = external_code.WorkflowType(
+    base_id='after_unet',
+    display_name='After unet',
+    default_workflow=default_workflows_dir / 'process_latent.json',
+)
 postprocess_workflow_type = external_code.WorkflowType(
     base_id='postprocess',
     display_name='Postprocess',
@@ -37,6 +47,8 @@ def add_default_workflow_types():
         sandbox_tab_workflow_type,
         postprocess_workflow_type,
         postprocess_latent_workflow_type,
+        before_unet_workflow_type,
+        after_unet_workflow_type,
         preprocess_workflow_type,
         preprocess_latent_workflow_type,
     ]

lib_comfyui/webui/callbacks.py

@@ -9,6 +9,8 @@ def register_callbacks():
     script_callbacks.on_ui_settings(on_ui_settings)
     script_callbacks.on_app_started(on_app_started)
     script_callbacks.on_script_unloaded(on_script_unloaded)
+    script_callbacks.on_cfg_denoiser(on_cfg_denoiser)
+    script_callbacks.on_cfg_denoised(on_cfg_denoised)
 
 
 @ipc.restrict_to_process('webui')
@@ -32,3 +34,13 @@ def on_script_unloaded():
     workflow_patcher.clear_patches()
     global_state.is_ui_instantiated = False
     external_code.clear_workflow_types()
+
+
+@ipc.restrict_to_process('webui')
+def on_cfg_denoiser(params):
+    workflow_patcher.before_unet(params)
+
+
+@ipc.restrict_to_process('webui')
+def on_cfg_denoised(params):
+    workflow_patcher.after_unet(params)

lib_comfyui/webui/workflow_patcher.py

@@ -2,7 +2,6 @@
 import torch
 import torchvision.transforms.functional as F
 from lib_comfyui import ipc, global_state, default_workflow_types, external_code
-from modules import shared
 
 
 __original_create_sampler = None
@@ -84,3 +83,21 @@ def p_img2img_init(*args, original_function, p_ref, **kwargs):
         p_ref.init_images = [F.to_pil_image(image_tensor * 255) for image_tensor in preprocessed_images]
 
     return original_function(*args, **kwargs)
+
+
+def before_unet(params):
+    if getattr(global_state, 'enabled', True):
+        params.x = torch.stack(external_code.run_workflow(
+            workflow_type=default_workflow_types.before_unet_workflow_type,
+            tab=global_state.current_tab,
+            batch_input=params.x.to(device='cpu'),
+        )).to(device=params.x.device)
+
+
+def after_unet(params):
+    if getattr(global_state, 'enabled', True):
+        params.x = torch.stack(external_code.run_workflow(
+            workflow_type=default_workflow_types.after_unet_workflow_type,
+            tab=global_state.current_tab,
+            batch_input=params.x.to(device='cpu'),
+        )).to(device=params.x.device)

scripts/comfyui.py

@@ -1,13 +1,13 @@
 import gradio as gr
-from modules import scripts, ui
+from modules import scripts, ui, processing
 from lib_comfyui import comfyui_context, global_state, platform_utils, external_code, default_workflow_types, comfyui_process
 from lib_comfyui.webui import callbacks, settings, workflow_patcher
 from lib_comfyui.comfyui import iframe_requests
 import functools
 
 
 class ComfyUIScript(scripts.Script):
-    def get_xxx2img_str(self, is_img2img: bool = None):
+    def get_tab_str(self, is_img2img: bool = None):
         if is_img2img is None:
             is_img2img = self.is_img2img
         return "img2img" if is_img2img else "txt2img"
@@ -24,18 +24,18 @@ def ui(self, is_img2img):
             return self.get_alwayson_ui(is_img2img)
 
     def get_alwayson_ui(self, is_img2img: bool):
-        xxx2img = self.get_xxx2img_str(is_img2img)
+        tab = self.get_tab_str(is_img2img)
 
         with gr.Row():
             queue_front = gr.Checkbox(
                 label='Queue front',
                 elem_id=self.elem_id('queue_front'),
                 value=True,
             )
-            workflow_types = external_code.get_workflow_types(xxx2img)
+            workflow_types = external_code.get_workflow_types(tab)
             first_workflow_type = workflow_types[0]
             workflow_type_ids = {
-                workflow_type.display_name: workflow_type.get_ids(xxx2img)[0]
+                workflow_type.display_name: workflow_type.get_ids(tab)[0]
                 for workflow_type in workflow_types
             }
             workflow_display_name = gr.Dropdown(
@@ -73,7 +73,7 @@ def get_alwayson_ui(self, is_img2img: bool):
                 _js='reloadComfyuiIFrames'
             )
 
-        self.setup_infotext_updates(workflow_types, xxx2img)
+        self.setup_infotext_updates(workflow_types, tab)
 
         return queue_front,
 
@@ -103,7 +103,7 @@ def get_iframes_html(self, is_img2img: bool, first_workflow_type_id: str) -> str
         comfyui_client_url = settings.get_comfyui_client_url()
 
         iframes = []
-        for workflow_type_id in external_code.get_workflow_type_ids(self.get_xxx2img_str(is_img2img)):
+        for workflow_type_id in external_code.get_workflow_type_ids(self.get_tab_str(is_img2img)):
             html_classes = ['comfyui-embedded-widget']
             if workflow_type_id == first_workflow_type_id:
                 html_classes.append('comfyui-embedded-widget-display')
@@ -128,6 +128,7 @@ def process(self, p, queue_front, **kwargs):
             return
 
         global_state.queue_front = queue_front
+        global_state.current_tab = self.get_tab_str(isinstance(p, processing.StableDiffusionProcessingImg2Img))
         workflow_patcher.patch_processing(p)
 
     def postprocess_batch_list(self, p, pp, *args, **kwargs):
@@ -138,7 +139,7 @@ def postprocess_batch_list(self, p, pp, *args, **kwargs):
 
         batch_results = external_code.run_workflow(
             workflow_type=default_workflow_types.postprocess_workflow_type,
-            tab=self.get_xxx2img_str(),
+            tab=self.get_tab_str(),
             batch_input=pp.images,
         )
 
@@ -150,7 +151,7 @@ def postprocess_batch_list(self, p, pp, *args, **kwargs):
         pp.images.clear()
         pp.images.extend(batch_results)
 
-        iframe_requests.extend_infotext_with_comfyui_workflows(p, self.get_xxx2img_str())
+        iframe_requests.extend_infotext_with_comfyui_workflows(p, self.get_tab_str())
 
 
 callbacks.register_callbacks()