From 53bb458601e1036879ca1553bdfdb0563f1a3fe2 Mon Sep 17 00:00:00 2001 From: InformEthics Date: Wed, 27 Nov 2024 23:32:07 +0000 Subject: [PATCH] Update Mflux_Pro.py Add new variables to ControlNetPipeline which are important for ControlNet only Reworked ControlNetLoader node to communicate with and set the new ControlNetPipeline correctly Changing the MfluxImg2Img Node name for clarity and also changing the strength variable to image_strength --- Mflux_Comfy/Mflux_Pro.py | 85 +++++++++++++++++++++++----------------- 1 file changed, 50 insertions(+), 35 deletions(-) diff --git a/Mflux_Comfy/Mflux_Pro.py b/Mflux_Comfy/Mflux_Pro.py index 170b36f..1544f5a 100644 --- a/Mflux_Comfy/Mflux_Pro.py +++ b/Mflux_Comfy/Mflux_Pro.py @@ -6,16 +6,16 @@ import torch from mflux.controlnet.controlnet_util import ControlnetUtil -class MfluxImagePipeline: - def __init__(self, image_path, strength): +class MfluxImg2ImgPipeline: + def __init__(self, image_path, image_strength): self.image_path = image_path - self.strength = strength + self.image_strength = image_strength def clear_cache(self): self.image_path = None - self.strength = None + self.image_strength = None -class MfluxLoadImage: +class MfluxImg2Img: @classmethod def INPUT_TYPES(cls): input_dir = folder_paths.get_input_directory() @@ -24,35 +24,35 @@ class MfluxLoadImage: return { "required": { "image": (sorted(files), {"image_upload": True}), - "strength": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01}), + "image_strength": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01}), } } CATEGORY = "MFlux/Pro" - RETURN_TYPES = ("MfluxImagePipeline", "INT", "INT") - RETURN_NAMES = ("image", "width", "height") + RETURN_TYPES = ("MfluxImg2ImgPipeline", "INT", "INT") + RETURN_NAMES = ("img2img", "width", "height") FUNCTION = "load_and_process" - def load_and_process(self, image, strength): + def load_and_process(self, image, image_strength): image_path = folder_paths.get_annotated_filepath(image) with Image.open(image_path) as img: width, height = img.size - return MfluxImagePipeline(image_path, strength), width, height + return MfluxImagePipeline(image_path, image_strength), width, height @classmethod - def IS_CHANGED(cls, image, strength): + def IS_CHANGED(cls, image, image_strength): image_hash = hash(image) - strength_rounded = round(float(strength), 2) + strength_rounded = round(float(image_strength), 2) return (image_hash, strength_rounded) @classmethod - def VALIDATE_INPUTS(cls, image, strength): + def VALIDATE_INPUTS(cls, image, image_strength): if not folder_paths.exists_annotated_filepath(image): return "Invalid image file: {}".format(image) - if not isinstance(strength, (int, float)): + if not isinstance(image_strength, (int, float)): return "Strength must be a number" return True @@ -115,14 +115,20 @@ class MfluxLorasLoader: return (MfluxLorasPipeline(list(lora_paths), list(lora_scales)),) class MfluxControlNetPipeline: - def __init__(self, model_selection, save_canny=False): + def __init__(self, model_selection, control_image_path, control_strength, save_canny=False): self.model_selection = model_selection + self.control_image_path = control_image_path + self.control_strength = control_strength self.save_canny = save_canny + def clear_cache(self): self.model_selection = None + self.control_image_path = None + self.control_strength = None self.save_canny = False + class MfluxControlNetLoader: @classmethod def INPUT_TYPES(cls): @@ -130,30 +136,32 @@ class MfluxControlNetLoader: "InstantX/FLUX.1-dev-Controlnet-Canny", ] + input_dir = folder_paths.get_input_directory() + files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))] + return { "required": { "model_selection": (controlnet_models, {"default": "InstantX/FLUX.1-dev-Controlnet-Canny"}), + "image": (sorted(files), {"image_upload": True}), + "control_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), "save_canny": ("BOOLEAN", {"default": False, "label_on": "True", "label_off": "False"}), - }, - "optional": { - "image": ("MfluxImagePipeline", ), } + } - CATEGORY = "MFlux/Pro" - RETURN_TYPES = ("MfluxControlNetPipeline", "IMAGE") - RETURN_NAMES = ("ControlNet", "preprocessed_image") + CATEGORY = "MFlux" + RETURN_TYPES = ("MfluxControlNetPipeline", "INT", "INT", "IMAGE",) + RETURN_NAMES = ("ControlNet", "width", "height", "preprocessed_image") FUNCTION = "load_and_select" - def load_and_select(self, model_selection, save_canny, image=None): + def load_and_select(self, model_selection, image, control_strength, save_canny): - image_path = None - strength = None - if image: - image_path = image.image_path - strength = image.strength + control_image_path = folder_paths.get_annotated_filepath(image) - with Image.open(image_path) as img: + with Image.open(control_image_path) as img: + width, height = img.size + + with Image.open(control_image_path) as img: canny_image = ControlnetUtil.preprocess_canny(img) canny_image_np = np.array(canny_image).astype(np.float32) / 255.0 canny_tensor = torch.from_numpy(canny_image_np) @@ -161,16 +169,17 @@ class MfluxControlNetLoader: if canny_tensor.dim() == 3: canny_tensor = canny_tensor.unsqueeze(0) - return MfluxControlNetPipeline(model_selection, save_canny), canny_tensor + return MfluxControlNetPipeline(model_selection, control_image_path, control_strength, save_canny), width, height, canny_tensor @classmethod - def IS_CHANGED(cls, image, model_selection, save_canny): - return str(image.image_path) + str(model_selection) + str(save_canny) + def IS_CHANGED(cls, model_selection, image, control_strength, save_canny): + control_image_path = folder_paths.get_annotated_filepath(image) + control_strength = round(float(control_strength), 2) + control_image_hash = hash(image) + return str(control_image_hash) + str(model_selection) + str(control_strength) + str(save_canny) @classmethod - def VALIDATE_INPUTS(cls, image, model_selection, save_canny): - if image and not folder_paths.exists_annotated_filepath(image.image_path): - return "Invalid image file: {}".format(image.image_path) + def VALIDATE_INPUTS(cls, image, model_selection, control_strength, save_canny): available_models = [ "InstantX/FLUX.1-dev-Controlnet-Canny", @@ -178,7 +187,13 @@ class MfluxControlNetLoader: if model_selection not in available_models: return f"Invalid model selection: {model_selection}" + if not folder_paths.exists_annotated_filepath(image): + return "Invalid control image file: {}".format(image) + + if not isinstance(control_strength, (int, float)): + return "Strength must be a number" + if not isinstance(save_canny, bool): return "save_canny must be a boolean value" - return True \ No newline at end of file + return True