442 lines
17 KiB
Python
442 lines
17 KiB
Python
"""Nodes for fooocus using in ComfyUI"""
|
|
import asyncio
|
|
import time
|
|
import uuid
|
|
|
|
import folder_paths
|
|
import os
|
|
from modules.patch import PatchSettings, patch_all, patch_settings
|
|
patch_all()
|
|
# 会重写 modules.model_sampling.ModelSamplingDiscrete._register_schedule
|
|
# 不执行 patch_all() 会导致报错
|
|
|
|
from modules.flags import scheduler_list, sampler_list
|
|
from node_utils import params_to_params, AlwaysEqualProxy
|
|
from apis.utils.call_worker import binary_output
|
|
from ldm_patched.modules.model_management import cleanup_models, unload_all_models
|
|
|
|
|
|
any_type = AlwaysEqualProxy("*")
|
|
|
|
root_dir = os.path.dirname(os.path.realpath(__file__))
|
|
preset_list = [p.split('.')[0] for p in os.listdir(f'{root_dir}/presets') if p.endswith('json')]
|
|
preset_list.append("initial")
|
|
|
|
preset_list = list(set(preset_list))
|
|
|
|
upscale_method = [
|
|
"Disabled",
|
|
"Vary (Subtle)",
|
|
"Vary (Strong)",
|
|
"Upscale (1.5x)",
|
|
"Upscale (2x)",
|
|
"Upscale (Fast 2x)",
|
|
"Upscale (Custom)"
|
|
]
|
|
mask_model = [
|
|
"u2net",
|
|
"u2netp",
|
|
"u2net_human_seg",
|
|
"u2net_cloth_seg",
|
|
"silueta",
|
|
"isnet-general-use",
|
|
"isnet-anime",
|
|
"sam"
|
|
]
|
|
|
|
|
|
class ClearVram:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"anything": (any_type, {})},
|
|
"optional": {},
|
|
"hidden": {"unique_id": "UNIQUE_ID", "extra_pnginfo": "EXTRA_PNGINFO",}}
|
|
|
|
RETURN_TYPES = (any_type,)
|
|
RETURN_NAMES = ("output",)
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "empty_cache"
|
|
CATEGORY = "Fooocus"
|
|
|
|
def empty_cache(self, anything, unique_id=None, extra_pnginfo=None):
|
|
time.sleep(1)
|
|
cleanup_models()
|
|
unload_all_models()
|
|
return (anything,)
|
|
|
|
@staticmethod
|
|
def IS_CHANGED(**kwargs):
|
|
return uuid.uuid4().hex
|
|
|
|
|
|
class FooocusSettings:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"optional": {
|
|
'adaptive_cfg': ("FLOAT", {"default": 7.0, "min": 1.0, "max": 30.0, 'step': 0.1}),
|
|
'adm_scaler_positive': ("FLOAT", {"default": 1.5, "min": 0.0, "max": 3.0, 'step': 0.1}),
|
|
'adm_scaler_negative': ("FLOAT", {"default": 0.8, "min": 0.0, "max": 3.0, 'step': 0.1}),
|
|
'adm_scaler_end': ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, 'step': 0.1}),
|
|
'black_out_nsfw': (["enable", "disable"], {"default": "disable"}),
|
|
'canny_low_threshold': ("INT", {"default": 64, "min": 1, "max": 255}),
|
|
'canny_high_threshold': ("INT", {"default": 128, "min": 1, "max": 255}),
|
|
'controlnet_softness': ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, 'step': 0.01}),
|
|
'clip_skip': ("INT", {"default": 2, "min": 1, "max": 12}),
|
|
'debugging_cn_preprocessor': (["enable", "disable"], {"default": "disable"}),
|
|
'debugging_dino': (["enable", "disable"], {"default": "disable"}),
|
|
'dino_erode_or_dilate': ("INT", {"default": 0, "min": -64, "max": 64}),
|
|
'debugging_enhance_masks_checkbox': (["enable", "disable"], {"default": "disable"}),
|
|
'debugging_inpaint_preprocessor': (["enable", "disable"], {"default": "disable"}),
|
|
'disable_preview': (["enable", "disable"], {"default": "disable"}),
|
|
'disable_intermediate_results': (["enable", "disable"], {"default": "disable"}),
|
|
'disable_seed_increment': (["enable", "disable"], {"default": "disable"}),
|
|
'freeu_enabled': (["enable", "disable"], {"default": "disable"}),
|
|
'freeu_b1': ("FLOAT", {"default": 1.01}),
|
|
'freeu_b2': ("FLOAT", {"default": 1.02}),
|
|
'freeu_s1': ("FLOAT", {"default": 0.99}),
|
|
'freeu_s2': ("FLOAT", {"default": 0.95}),
|
|
'generate_image_grid': (["enable", "disable"], {"default": "disable"}),
|
|
'mixing_image_prompt_and_vary_upscale': (["enable", "disable"], {"default": "disable"}),
|
|
'mixing_image_prompt_and_inpaint': (["enable", "disable"], {"default": "disable"}),
|
|
'output_format': (["png", "jpeg", "webp"], {"default": "png", "multiline": False}),
|
|
'overwrite_step': ("INT", {"default": -1, "min": -1, "max": 999999}),
|
|
'overwrite_switch': ("INT", {"default": -1, "min": -1, "max": 999999}),
|
|
'overwrite_width': ("INT", {"default": -1, "min": -1, "max": 2048}),
|
|
'overwrite_height': ("INT", {"default": -1, "min": -1, "max": 2048}),
|
|
'overwrite_vary_strength': ("FLOAT", {"default": -1, "min": -1, "max": 1.0}),
|
|
'overwrite_upscale_strength': ("FLOAT", {"default": -1, "min": -1, "max": 1.0}),
|
|
'read_wildcards_in_order': (["enable", "disable"], {"default": "disable"}),
|
|
'skipping_cn_preprocessor': (["enable", "disable"], {"default": "disable"})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("FOOOCUS_SETTINGS",)
|
|
RETURN_NAMES = ("settings",)
|
|
|
|
FUNCTION = "process"
|
|
CATEGORY = "Fooocus"
|
|
|
|
@staticmethod
|
|
def process(**kwargs):
|
|
return (kwargs,)
|
|
|
|
@staticmethod
|
|
def IS_CHANGED(**kwargs):
|
|
return uuid.uuid4().hex
|
|
|
|
|
|
class LoraStacks:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
max_lora_num = 5
|
|
inputs = {
|
|
"required": {
|
|
"num_loras": ("INT", {"default": 5, "min": 1, "max": max_lora_num}),
|
|
},
|
|
"optional": {
|
|
"optional_lora_stack": ("LORA_STACKS",),
|
|
},
|
|
}
|
|
|
|
for i in range(1, max_lora_num+1):
|
|
inputs["optional"][f"lora_{i}_enable"] = (
|
|
["enable", "disable"], {"default": "disable"})
|
|
inputs["optional"][f"lora_{i}_name"] = (
|
|
["None"] + folder_paths.get_filename_list("loras"), {"default": "None"})
|
|
inputs["optional"][f"lora_{i}_strength"] = (
|
|
"FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01})
|
|
return inputs
|
|
|
|
RETURN_TYPES = ("LORA_STACKS",)
|
|
RETURN_NAMES = ("lora_stacks",)
|
|
FUNCTION = "stack"
|
|
|
|
CATEGORY = "Fooocus"
|
|
|
|
def stack(self, num_loras, optional_lora_stack=None, **kwargs):
|
|
|
|
loras = []
|
|
|
|
if optional_lora_stack is not None:
|
|
loras.extend([l for l in optional_lora_stack if l[0] != "None"])
|
|
|
|
# Import Lora values
|
|
for i in range(1, num_loras + 1):
|
|
lora_name = kwargs.get(f"lora_{i}_name")
|
|
|
|
if not lora_name or lora_name == "None":
|
|
continue
|
|
|
|
lora_enable = kwargs.get(f"lora_{i}_enable")
|
|
lora_strength = kwargs.get(f"lora_{i}_strength")
|
|
loras.append((lora_enable, lora_name, lora_strength))
|
|
|
|
return (loras,)
|
|
|
|
|
|
class ImagePrompts:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
image_prompts = {
|
|
"optional": {
|
|
"image_prompts": ("IMAGE_PROMPTS",)
|
|
}
|
|
}
|
|
|
|
for i in range(1, 5):
|
|
image_prompts["optional"][f"cn_img_{i}"] = ("STRING", {"default": "None", "multiline": False})
|
|
image_prompts["optional"][f"cn_stop_{i}"] = ("FLOAT", {"default": 0.6, "min": 0.0, "max": 1.0, "step": 0.01})
|
|
image_prompts["optional"][f"cn_weight_{i}"] = ("FLOAT", {"default": 0.6, "min": 0.0, "max": 2.0, "step": 0.01})
|
|
image_prompts["optional"][f"cn_type_{i}"] = (
|
|
["ImagePrompt", "FaceSwap", "PyraCanny", "CPDS"], {"default": "ImagePrompt"})
|
|
return image_prompts
|
|
|
|
RETURN_TYPES = ("IMAGE_PROMPTS",)
|
|
RETURN_NAMES = ("image_prompts",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Fooocus"
|
|
|
|
@staticmethod
|
|
def process(**kwargs):
|
|
return (kwargs,)
|
|
|
|
@staticmethod
|
|
def IS_CHANGED(**kwargs):
|
|
return uuid.uuid4().hex
|
|
|
|
|
|
class EnhanceControl:
|
|
def __init__(self):
|
|
pass
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
control = {
|
|
"required": {
|
|
"enhance_enabled": (["enable", "disable"], {"default": "disable"}),
|
|
"enhance_mask_dino_prompt": ("STRING", {"default": "", "multiline": False}),
|
|
"enhance_prompt": ("STRING", {"default": "", "multiline": False}),
|
|
"enhance_negative_prompt": ("STRING", {"default": "", "multiline": False}),
|
|
"enhance_mask_model": (mask_model, {"default": "sam"}),
|
|
"enhance_mask_cloth_category": ("STRING", {"default": "full", "multiline": False}),
|
|
"enhance_mask_sam_model": (["vit_b", "vit_h", "vit_l"], {"default": "vit_b"}),
|
|
"enhance_mask_text_threshold": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0}),
|
|
"enhance_mask_box_threshold": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0}),
|
|
"enhance_mask_sam_max_detections": ("INT", {"default": 0, "min": 0, "max": 10}),
|
|
"enhance_inpaint_disable_initial_latent": (["enable", "disable"], {"default": "disable"}),
|
|
"enhance_inpaint_engine": (["v2.6",], {"default": "v2.6"}),
|
|
"enhance_inpaint_strength": ("FLOAT", {"default": 1, "min": 0.0, "max": 1.0}),
|
|
"enhance_inpaint_respective_field": ("FLOAT", {"default": 0.618, "min": 0.0, "max": 1.0}),
|
|
"enhance_inpaint_erode_or_dilate": ("INT", {"default": 0, "min": -64, "max": 64}),
|
|
"enhance_mask_invert": (["enable", "disable"], {"default": "disable"})
|
|
},
|
|
}
|
|
return control
|
|
RETURN_TYPES = ("ENHANCE_CONTROL",)
|
|
RETURN_NAMES = ("enhance_control",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Fooocus"
|
|
@staticmethod
|
|
def process(**kwargs):
|
|
return (kwargs,)
|
|
|
|
@staticmethod
|
|
def IS_CHANGED(**kwargs):
|
|
return uuid.uuid4().hex
|
|
|
|
|
|
class EnhanceControls:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
enhance_control = {
|
|
"required": {
|
|
"enhance_image": ("STRING", {"default": "", "multiline": False}),
|
|
"enhance_enable": (["enable", "disable"], {"default": "disable"}),
|
|
"enhance_uov_method": (upscale_method, {"default": "Disabled"}),
|
|
"enhance_uov_processing_order": (["Before First Enhancement", "After Last Enhancement"], {"default": "Before First Enhancement"}),
|
|
"enhance_uov_prompt_type": (["Last Filled Enhancement Prompts", "Original Prompts"], {"default": "Original Prompts"})
|
|
},
|
|
"optional": {
|
|
"enhance_ctrl_1": ("ENHANCE_CONTROL",),
|
|
"enhance_ctrl_2": ("ENHANCE_CONTROL",),
|
|
"enhance_ctrl_3": ("ENHANCE_CONTROL",)
|
|
}
|
|
}
|
|
return enhance_control
|
|
RETURN_TYPES = ("ENHANCE_CONTROLS",)
|
|
RETURN_NAMES = ("enhance_controls",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Fooocus"
|
|
|
|
@staticmethod
|
|
def process(**kwargs):
|
|
return (kwargs,)
|
|
|
|
@staticmethod
|
|
def IS_CHANGED(**kwargs):
|
|
return uuid.uuid4().hex
|
|
|
|
|
|
class FooocusSampler:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
'prompt': ("STRING", {"default": "A beautiful girl", "multiline": True}),
|
|
'negative_prompt': ("STRING", {"default": "", "multiline": True}),
|
|
'style_selections': ("STRING", {"default": "Fooocus V2, Fooocus Enhance, Fooocus Sharp", "multiline": False}),
|
|
'aspect_ratios_selection': ("STRING", {"default": "1152*896", "multiline": False}),
|
|
'performance_selection': (["Quality", "Speed", "Extreme Speed", "Lightning", "Hyper-SD"], {"default": "Speed"}),
|
|
'image_number': ("INT", {"default": 1, "min": 1, "max": 999}),
|
|
'image_seed': ("INT", {"default": -1, "min": -1, "max": 2**63-1}),
|
|
'sharpness': ("FLOAT", {"default": 2.0, "min": 0.0, "max": 30.0}),
|
|
'guidance_scale': ("FLOAT", {"default": 7.0, "min": 1.0, "max": 30.0}),
|
|
'base_model_name': ("STRING", {"default": "juggernautXL_v8Rundiffusion.safetensors", "multiline": False}),
|
|
'refiner_model_name': ("STRING", {"default": "None", "multiline": False}),
|
|
'refiner_switch': ("FLOAT", {"default": 0.8, "min": 0.1, "max": 1}),
|
|
'refiner_swap_method': (["joint", "separate", "vae"], {"default": "joint"}),
|
|
'preset': (preset_list, {"default": "initial"}),
|
|
'sampler_name': (sampler_list, {"default": "dpmpp_2m_sde_gpu"}),
|
|
'scheduler_name': (scheduler_list, {"default": "karras"}),
|
|
'vae_name': ("STRING", {"default": "Default (model)", "multiline": False}),
|
|
'loras': ("LORA_STACKS",),
|
|
'inpaint_outpaint': ("INPAINT_OUTPAINT",),
|
|
'upscale_vary': ("UPSCALE_VARY",),
|
|
'image_prompts': ("IMAGE_PROMPTS",),
|
|
'enhance_ctrls': ("ENHANCE_CONTROLS",),
|
|
'settings': ("FOOOCUS_SETTINGS",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Fooocus"
|
|
|
|
@staticmethod
|
|
def process(**kwargs):
|
|
request = params_to_params(kwargs)
|
|
image = asyncio.run(binary_output(request))
|
|
return (image,)
|
|
|
|
@staticmethod
|
|
def IS_CHANGED(**kwargs):
|
|
return uuid.uuid4().hex
|
|
|
|
|
|
class InpaintOutpaint:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
inpaint_outpaint = {
|
|
"required": {
|
|
"input_image": ("STRING", {"default": "", "multiline": False}),
|
|
"input_mask": ("STRING", {"default": "", "multiline": False}),
|
|
"inpaint_prompt": ("STRING", {"default": "", "multiline": False}),
|
|
"inpaint_disable_initial_latent": (["enable", "disable"], {"default": "disable"}),
|
|
"inpaint_engine": (["v2.6",], {"default": "v2.6"}),
|
|
"inpaint_strength": ("FLOAT", {"default": 1, "min": 0.0, "max": 1.0}),
|
|
"inpaint_respective_field": ("FLOAT", {"default": 0.618, "min": 0.0, "max": 1.0}),
|
|
"inpaint_erode_or_dilate": ("INT", {"default": 0, "min": -64, "max": 64}),
|
|
"inpaint_mask_invert": (["enable", "disable"], {"default": "disable"}),
|
|
"outpaint_left": (["enable", "disable"], {"default": "disable"}),
|
|
"outpaint_left_distance": ("INT", {"default": 0, "min": 0}),
|
|
"outpaint_right": (["enable", "disable"], {"default": "disable"}),
|
|
"outpaint_right_distance": ("INT", {"default": 0, "min": 0}),
|
|
"outpaint_top": (["enable", "disable"], {"default": "disable"}),
|
|
"outpaint_top_distance": ("INT", {"default": 0, "min": 0}),
|
|
"outpaint_bottom": (["enable", "disable"], {"default": "disable"}),
|
|
"outpaint_bottom_distance": ("INT", {"default": 0, "min": 0}),
|
|
}
|
|
}
|
|
return inpaint_outpaint
|
|
|
|
RETURN_TYPES = ("INPAINT_OUTPAINT",)
|
|
RETURN_NAMES = ("inpaint_outpaint",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Fooocus"
|
|
|
|
@staticmethod
|
|
def process(**kwargs):
|
|
return (kwargs,)
|
|
|
|
@staticmethod
|
|
def IS_CHANGED(**kwargs):
|
|
return uuid.uuid4().hex
|
|
|
|
|
|
class UpscaleVary:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"uov_input_image": ("STRING", {"default": "", "multiline": False}),
|
|
"uov_method": (upscale_method, {"default": "Disabled"}),
|
|
"upscale_multiple": ("FLOAT", {"default": 2.0, "min": 1.0, "max": 4.0})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("UPSCALE_VARY",)
|
|
RETURN_NAMES = ("upscale_vary",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Fooocus"
|
|
|
|
@staticmethod
|
|
def process(**kwargs):
|
|
return (kwargs,)
|
|
|
|
@staticmethod
|
|
def IS_CHANGED(**kwargs):
|
|
return uuid.uuid4().hex
|
|
|
|
|
|
# A dictionary that contains all nodes you want to export with their names
|
|
# NOTE: names should be globally unique
|
|
NODE_CLASS_MAPPINGS = {
|
|
"FooocusSettings": FooocusSettings,
|
|
"FooocusSampler": FooocusSampler,
|
|
"LoraStacks": LoraStacks,
|
|
"ImagePrompts": ImagePrompts,
|
|
"EnhanceControl": EnhanceControl,
|
|
"EnhanceControls": EnhanceControls,
|
|
"InpaintOutpaint": InpaintOutpaint,
|
|
"UpscaleVary": UpscaleVary,
|
|
"ClearVram": ClearVram,
|
|
}
|
|
|
|
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"FooocusSettings": "Fooocus Settings",
|
|
"FooocusSampler": "Fooocus Sampler",
|
|
"LoraStacks": "LoraStacks",
|
|
"ImagePrompts": "Image Prompts",
|
|
"EnhanceControl": "Enhance Control",
|
|
"EnhanceControls": "Enhance Controls",
|
|
"InpaintOutpaint": "Inpaint Outpaint",
|
|
"UpscaleVary": "Upscale Vary",
|
|
"ClearVram": "Clear Vram",
|
|
}
|