Add Civitai sampler name

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 `<Sampler Name> <Scheduler name>`.
+            - `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 `<sampler_name>_<scheduler_name>`
+            - `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

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