From 23af24c3df0ab63eab8bb33b979c8bc53b0a70c5 Mon Sep 17 00:00:00 2001 From: ljleb Date: Thu, 10 Aug 2023 00:16:28 -0400 Subject: [PATCH 1/2] progress --- comfyui_custom_nodes/webui_io.py | 20 ++++++++++++++++++++ lib_comfyui/default_workflow_types.py | 12 ++++++++++++ lib_comfyui/webui/callbacks.py | 12 ++++++++++++ lib_comfyui/webui/workflow_patcher.py | 19 ++++++++++++++++++- scripts/comfyui.py | 19 ++++++++++--------- 5 files changed, 72 insertions(+), 10 deletions(-) diff --git a/comfyui_custom_nodes/webui_io.py b/comfyui_custom_nodes/webui_io.py index 6a333392..e5e05fb6 100644 --- a/comfyui_custom_nodes/webui_io.py +++ b/comfyui_custom_nodes/webui_io.py @@ -82,9 +82,29 @@ def set_images(self, latents): return [] +class InvertLatent: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "latents": ("LATENT", ), + }, + } + RETURN_TYPES = ("LATENT", ) + FUNCTION = "invert" + + CATEGORY = "webui" + + OUTPUT_NODE = False + + def invert(self, latents): + return [{'samples': -latents['samples']}] + + NODE_CLASS_MAPPINGS = { "ImageFromWebui": WebuiImageInput, "ImageToWebui": WebuiImageOutput, "LatentFromWebui": WebuiLatentInput, "LatentToWebui": WebuiLatentOutput, + "InvertLatent": InvertLatent, } diff --git a/lib_comfyui/default_workflow_types.py b/lib_comfyui/default_workflow_types.py index ac138db7..2edd0baa 100644 --- a/lib_comfyui/default_workflow_types.py +++ b/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, ] diff --git a/lib_comfyui/webui/callbacks.py b/lib_comfyui/webui/callbacks.py index 0ad25fda..e5b4b4e2 100644 --- a/lib_comfyui/webui/callbacks.py +++ b/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) diff --git a/lib_comfyui/webui/workflow_patcher.py b/lib_comfyui/webui/workflow_patcher.py index fc076fd1..9589edc0 100644 --- a/lib_comfyui/webui/workflow_patcher.py +++ b/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) diff --git a/scripts/comfyui.py b/scripts/comfyui.py index 7e361995..7f38fec2 100644 --- a/scripts/comfyui.py +++ b/scripts/comfyui.py @@ -1,5 +1,5 @@ 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 @@ -7,7 +7,7 @@ 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,7 +24,7 @@ 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( @@ -32,10 +32,10 @@ def get_alwayson_ui(self, is_img2img: bool): 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() From 21ea7b1a0eb2dfb9e599c8897212a3e5bbb2a84d Mon Sep 17 00:00:00 2001 From: ljleb Date: Thu, 10 Aug 2023 00:17:30 -0400 Subject: [PATCH 2/2] dont add my test stuff --- comfyui_custom_nodes/webui_io.py | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/comfyui_custom_nodes/webui_io.py b/comfyui_custom_nodes/webui_io.py index e5e05fb6..6a333392 100644 --- a/comfyui_custom_nodes/webui_io.py +++ b/comfyui_custom_nodes/webui_io.py @@ -82,29 +82,9 @@ def set_images(self, latents): return [] -class InvertLatent: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "latents": ("LATENT", ), - }, - } - RETURN_TYPES = ("LATENT", ) - FUNCTION = "invert" - - CATEGORY = "webui" - - OUTPUT_NODE = False - - def invert(self, latents): - return [{'samples': -latents['samples']}] - - NODE_CLASS_MAPPINGS = { "ImageFromWebui": WebuiImageInput, "ImageToWebui": WebuiImageOutput, "LatentFromWebui": WebuiLatentInput, "LatentToWebui": WebuiLatentOutput, - "InvertLatent": InvertLatent, }