From db55818fe1c8f3346d518fcd3d8ca41d4ff84b0a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tr=E1=BA=A7n=20Nam=20Tu=E1=BA=A5n?= Date: Tue, 20 Aug 2024 01:10:55 +0700 Subject: [PATCH] Add Civitai sampler name --- py/capture.py | 94 +++++++++++++++++++++++++++++++++++++++++++----- py/nodes/node.py | 8 +++-- 2 files changed, 91 insertions(+), 11 deletions(-) diff --git a/py/capture.py b/py/capture.py index c01a2fc..914f4b6 100644 --- a/py/capture.py +++ b/py/capture.py @@ -70,7 +70,7 @@ class Capture: return inputs @classmethod - def gen_pnginfo_dict(cls, inputs_before_sampler_node, inputs_before_this_node): + def gen_pnginfo_dict(cls, inputs_before_sampler_node, inputs_before_this_node, save_civitai_sampler = False): pnginfo_dict = {} def update_pnginfo_dict(inputs, metafield, key): @@ -88,14 +88,18 @@ class Capture: update_pnginfo_dict(inputs_before_sampler_node, MetaField.STEPS, "Steps") sampler_names = inputs_before_sampler_node.get(MetaField.SAMPLER_NAME, []) - if len(sampler_names) > 0: - pnginfo_dict["Sampler"] = sampler_names[0][1] + schedulers = inputs_before_sampler_node.get(MetaField.SCHEDULER, []) - schedulers = inputs_before_sampler_node.get(MetaField.SCHEDULER, []) - if len(schedulers) > 0: - scheduler = schedulers[0][1] - if scheduler != "normal": - pnginfo_dict["Sampler"] += "_" + scheduler + if (save_civitai_sampler): + pnginfo_dict["Sampler"] = cls.get_sampler_for_civitai(sampler_names, schedulers) + else: + if len(sampler_names) > 0: + pnginfo_dict["Sampler"] = sampler_names[0][1] + + if len(schedulers) > 0: + scheduler = schedulers[0][1] + if scheduler != "normal": + pnginfo_dict["Sampler"] += "_" + scheduler update_pnginfo_dict(inputs_before_sampler_node, MetaField.CFG, "CFG scale") update_pnginfo_dict(inputs_before_sampler_node, MetaField.SEED, "Seed") @@ -213,3 +217,77 @@ class Capture: index += 1 return pnginfo_dict + + @classmethod + def get_sampler_for_civitai(cls, sampler_names, schedulers): + """ + Get the pretty sampler name for Civitai in the form of ` `. + - `dpmpp_2m` and `karras` will return `DPM++ 2M Karras` + + If there is a matching sampler name but no matching scheduler name, return only the matching sampler name. + - `dpmpp_2m` and `exponential` will return only `DPM++ 2M` + + if there is no matching sampler and scheduler name, return `_` + - `ipndm` and `normal` will return `ipndm` + - `ipndm` and `karras` will return `ipndm_karras` + + Reference: https://github.com/civitai/civitai/blob/main/src/server/common/constants.ts + + Last update: https://github.com/civitai/civitai/blob/a2e6d267eefe6f44811a640c570739bcb078e4a5/src/server/common/constants.ts#L138-L165 + """ + + def sampler_with_karras_exponential(sampler, scheduler): + match scheduler: + case "karras": + sampler += " Karras" + case "exponential": + sampler += " Exponential" + return sampler + + def sampler_with_karras(sampler, scheduler): + if scheduler == "karras": + return sampler + " Karras" + return sampler + + if len(sampler_names) > 0: + sampler = sampler_names[0][1] + if len(schedulers) > 0: + scheduler = schedulers[0][1] + + match sampler: + case "euler" | "euler_cfg_pp": + return "Euler" + case "euler_ancestral" | "euler_ancestral_cfg_pp": + return "Euler a" + case "heun" | "heunpp2": + return "Huen" + case "dpm_2": + return sampler_with_karras("DPM2", scheduler) + case "dpm_2_ancestral": + return sampler_with_karras("DPM2 a", scheduler) + case "lms": + return sampler_with_karras("LMS", scheduler) + case "dpm_fast": + return "DPM fast" + case "dpm_adaptive": + return "DPM adaptive" + case "dpmpp_2s_ancestral": + return sampler_with_karras("DPM++ 2S a", scheduler) + case "dpmpp_sde" | "dpmpp_sde_gpu": + return sampler_with_karras("DPM++ SDE", scheduler) + case "dpmpp_2m": + return sampler_with_karras("DPM++ 2M", scheduler) + case "dpmpp_2m_sde" | "dpmpp_2m_sde_gpu": + return sampler_with_karras("DPM++ 2M SDE", scheduler) + case "dpmpp_3m_sde" | "dpmpp_3m_sde_gpu": + return sampler_with_karras_exponential("DPM++ 3M SDE", scheduler) + case "lcm": + return "LCM" + case "ddim": + return "DDIM" + case "uni_pc" | "uni_pc_bh2": + return "UniPC" + + if scheduler == "normal": + return sampler + return sampler + "_" + scheduler \ No newline at end of file diff --git a/py/nodes/node.py b/py/nodes/node.py index 3e70a9c..790b69d 100644 --- a/py/nodes/node.py +++ b/py/nodes/node.py @@ -51,6 +51,7 @@ class SaveImageWithMetaData(BaseNode): "quality": ("INT", {"default": 100, "min": 1, "max": 100}), "save_workflow_json": ("BOOLEAN", {"default": False}), "add_counter_to_filename": ("BOOLEAN", {"default": True}), + "civitai_sampler": ("BOOLEAN", {"default": False}) }, "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, } @@ -73,11 +74,12 @@ class SaveImageWithMetaData(BaseNode): quality=100, save_workflow_json=False, add_counter_to_filename=True, + civitai_sampler = False, prompt=None, extra_pnginfo=None, ): pnginfo_dict_src = self.gen_pnginfo( - sampler_selection_method, sampler_selection_node_id + sampler_selection_method, sampler_selection_node_id, civitai_sampler ) results = list() for index, image in enumerate(images): @@ -160,7 +162,7 @@ class SaveImageWithMetaData(BaseNode): return {"ui": {"images": results}} @classmethod - def gen_pnginfo(cls, sampler_selection_method, sampler_selection_node_id): + def gen_pnginfo(cls, sampler_selection_method, sampler_selection_node_id, save_civitai_sampler): # get all node inputs inputs = Capture.get_inputs() @@ -185,7 +187,7 @@ class SaveImageWithMetaData(BaseNode): # generate PNGInfo from inputs pnginfo_dict = Capture.gen_pnginfo_dict( - inputs_before_sampler_node, inputs_before_this_node + inputs_before_sampler_node, inputs_before_this_node, save_civitai_sampler ) return pnginfo_dict