Version bump

Make up/downsample block targeting more resiliant in RAUNet
Change publish workflow to trigger on release
2024-08-16 13:06:57 -06:00 · 2024-08-15 19:06:56 -06:00 · 2024-08-15 03:12:19 -06:00 · 2024-08-15 02:57:11 -06:00 · 2024-08-15 02:54:51 -06:00 · 2024-08-15 02:39:49 -06:00
4 changed files with 189 additions and 25 deletions
@@ -0,0 +1,16 @@
+name: Publish to Comfy registry
+on:
+  workflow_dispatch:
+  release: { types: ["published"] }
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+      - name: Publish Custom Node
+        uses: Comfy-Org/publish-node-action@main
+        with:
+          personal_access_token: ${{ secrets.COMFYORG_REGISTRY_API_KEY }}
@@ -29,22 +29,45 @@ class ShiftSize(WindowSize):

 class ApplyMSWMSAAttention:
    RETURN_TYPES = ("MODEL",)
+    OUTPUT_TOOLTIPS = ("Model patched with the MSW-MSA attention effect.",)
    FUNCTION = "patch"
    CATEGORY = "model_patches/unet"
+    DESCRIPTION = "This node applies an attention patch which _may_ slightly improve quality especially when generating at high resolutions. It is a large performance increase on SD1.x, may improve performance on SDXL. This is the advanced version of the node with more parameters, use ApplyMSWMSAAttentionSimple if this seems too complex. NOTE: Only supports SD1.x, SD2.x and SDXL."

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
-                "input_blocks": ("STRING", {"default": "1,2"}),
-                "middle_blocks": ("STRING", {"default": ""}),
-                "output_blocks": ("STRING", {"default": "9,10,11"}),
+                "input_blocks": (
+                    "STRING",
+                    {
+                        "default": "1,2",
+                        "tooltip": "Comma-separated list of input blocks to patch. Default is for SD1.x, you can try 4,5 for SDXL",
+                    },
+                ),
+                "middle_blocks": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "tooltip": "Comma-separated list of middle blocks to patch. Generally not recommended.",
+                    },
+                ),
+                "output_blocks": (
+                    "STRING",
+                    {
+                        "default": "9,10,11",
+                        "tooltip": "Comma-separated list of output blocks to patch. Default is for SD1.x, you can try 5,4 for SDXL",
+                    },
+                ),
                "time_mode": (
                    (
                        "percent",
                        "timestep",
                        "sigma",
                    ),
+                    {
+                        "tooltip": "Time mode controls how to interpret the values in start_time and end_time.",
+                    },
                ),
                "start_time": (
                    "FLOAT",
@@ -54,6 +77,7 @@ class ApplyMSWMSAAttention:
                        "max": 999.0,
                        "round": False,
                        "step": 0.01,
+                        "tooltip": "Time the MSW-MSA attention effect starts applying - value is inclusive.",
                    },
                ),
                "end_time": (
@@ -64,9 +88,15 @@ class ApplyMSWMSAAttention:
                        "max": 999.0,
                        "round": False,
                        "step": 0.01,
+                        "tooltip": "Time the MSW-MSA attention effect ends - value is inclusive.",
+                    },
+                ),
+                "model": (
+                    "MODEL",
+                    {
+                        "tooltip": "Model to patch with the MSW-MSA attention effect.",
                    },
                ),
-                "model": ("MODEL",),
            },
        }

@@ -235,15 +265,27 @@ class ApplyMSWMSAAttention:

 class ApplyMSWMSAAttentionSimple:
    RETURN_TYPES = ("MODEL",)
+    OUTPUT_TOOLTIPS = ("Model patched with the MSW-MSA attention effect.",)
    FUNCTION = "go"
    CATEGORY = "model_patches/unet"
+    DESCRIPTION = "This node applies an attention patch which _may_ slightly improve quality especially when generating at high resolutions. It is a large performance increase on SD1.x, may improve performance on SDXL. This is the simplified version of the node with less parameters. Use ApplyMSWMSAAttention if you require more control. NOTE: Only supports SD1.x, SD2.x and SDXL."

    @classmethod
    def INPUT_TYPES(cls) -> dict:
        return {
            "required": {
-                "model_type": (("SD15", "SDXL"),),
-                "model": ("MODEL",),
+                "model_type": (
+                    ("SD15", "SDXL"),
+                    {
+                        "tooltip": "Model type being patched. Choose SD15 for SD 1.4, SD 2.x.",
+                    },
+                ),
+                "model": (
+                    "MODEL",
+                    {
+                        "tooltip": "Model to patch with the MSW-MSA attention effect.",
+                    },
+                ),
            },
        }

@@ -54,7 +54,7 @@ GLOBAL_STATE: HDState
 class HDState:
    def __init__(self):
        self.no_controlnet_workaround = (
-            os.environ.get("JANKHIDIFFUSION_NO_CONTROLNET_WORKAROUND") is not None
+            "JANKHIDIFFUSION_NO_CONTROLNET_WORKAROUND" in os.environ
        )
        self.controlnet_scale_args = {"mode": "bilinear", "align_corners": False}
        self.patched_freeu_advanced = False
@@ -216,17 +216,41 @@ def forward_downsample(

 class ApplyRAUNet:
    RETURN_TYPES = ("MODEL",)
+    OUTPUT_TOOLTIPS = ("Model patched with the RAUNet effect.",)
    FUNCTION = "patch"
    CATEGORY = "model_patches/unet"
+    DESCRIPTION = "This node is used to enable generation at higher resolutions than a model was trained for with less artifacts or other negative effects. This is the advanced version with more tuneable parameters, use ApplyRAUNetSimple if this seems too complex. NOTE: Only supports SD1.x, SD2.x and SDXL."

    @classmethod
    def INPUT_TYPES(cls) -> dict:
        return {
            "required": {
-                "model": ("MODEL",),
-                "input_blocks": ("STRING", {"default": "3"}),
-                "output_blocks": ("STRING", {"default": "8"}),
-                "time_mode": (("percent", "timestep", "sigma"),),
+                "model": (
+                    "MODEL",
+                    {
+                        "tooltip": "Model to be patched with the RAUNet effect.",
+                    },
+                ),
+                "input_blocks": (
+                    "STRING",
+                    {
+                        "default": "3",
+                        "tooltip": "Comma-separated list of input Downsample blocks. The default of 3 will work with SD1.x and SDXL.",
+                    },
+                ),
+                "output_blocks": (
+                    "STRING",
+                    {
+                        "default": "8",
+                        "tooltip": "Comma-separated list of output Upsample blocks. The default is for SD1.x, for SDXL use 5.",
+                    },
+                ),
+                "time_mode": (
+                    ("percent", "timestep", "sigma"),
+                    {
+                        "tooltip": "Time mode controls how to interpret the values in start_time and end_time.",
+                    },
+                ),
                "start_time": (
                    "FLOAT",
                    {
@@ -235,6 +259,7 @@ class ApplyRAUNet:
                        "max": 999.0,
                        "round": False,
                        "step": 0.01,
+                        "tooltip": "Time normal RAUNet effects start applying - value is inclusive.",
                    },
                ),
                "end_time": (
@@ -245,9 +270,15 @@ class ApplyRAUNet:
                        "max": 999.0,
                        "round": False,
                        "step": 0.01,
+                        "tooltip": "Time normal RAUNet effects end - value is inclusive.",
+                    },
+                ),
+                "upscale_mode": (
+                    UPSCALE_METHODS,
+                    {
+                        "tooltip": "Method used when upscaling latents in output Upscale blocks.",
                    },
                ),
-                "upscale_mode": (UPSCALE_METHODS,),
                "ca_start_time": (
                    "FLOAT",
                    {
@@ -256,6 +287,7 @@ class ApplyRAUNet:
                        "max": 999.0,
                        "round": False,
                        "step": 0.01,
+                        "tooltip": "Time normal cross-attention effects start applying - value is inclusive..",
                    },
                ),
                "ca_end_time": (
@@ -266,22 +298,52 @@ class ApplyRAUNet:
                        "max": 999.0,
                        "round": False,
                        "step": 0.01,
+                        "tooltip": "Time normal cross-attention effects end - value is inclusive.",
+                    },
+                ),
+                "ca_input_blocks": (
+                    "STRING",
+                    {
+                        "default": "4",
+                        "tooltip": "Comma separated list of input cross-attention blocks. Default is for SD1.x, for SDXL you can try using 2 (or just disable it).",
+                    },
+                ),
+                "ca_output_blocks": (
+                    "STRING",
+                    {
+                        "default": "8",
+                        "tooltip": "Comma-separated list of output cross-attention blocks. Default is for SD1.x, for SDXL you can try using 7 (or just disable it).",
+                    },
+                ),
+                "ca_upscale_mode": (
+                    UPSCALE_METHODS,
+                    {
+                        "tooltip": "Mode used when upscaling latents in output cross-attention blocks.",
                    },
                ),
-                "ca_input_blocks": ("STRING", {"default": "4"}),
-                "ca_output_blocks": ("STRING", {"default": "8"}),
-                "ca_upscale_mode": (UPSCALE_METHODS,),
                "ca_downscale_mode": (
                    ("avg_pool2d", *UPSCALE_METHODS),
-                    {"default": "avg_pool2d"},
+                    {
+                        "default": "avg_pool2d",
+                        "tooltip": "Mode used when downscaling latents in output cross-attention blocks (use avg_pool2d for normal Hidiffusion behavior).",
+                    },
                ),
                "ca_downscale_factor": (
                    "FLOAT",
-                    {"default": 2.0, "min": 0.01, "step": 0.1, "round": False},
+                    {
+                        "default": 2.0,
+                        "min": 0.01,
+                        "step": 0.1,
+                        "round": False,
+                        "tooltip": "Factor to downscale with in cross-attention, 2.0 means downscale to half size. Must be an integer when using ca_downscale_mode avg_pool2d.",
+                    },
                ),
                "two_stage_upscale_mode": (
                    ("disabled", *UPSCALE_METHODS),
-                    {"default": "disabled"},
+                    {
+                        "default": "disabled",
+                        "tooltip": "When upscaling in output Upscale blocks (non-NA), do half the upscale with this mode and half with the normal upscale mode. May produce a different effect, isn't necessarily better.",
+                    },
                ),
            },
        }
@@ -391,13 +453,22 @@ class ApplyRAUNet:
            model.set_model_output_block_patch(output_block_patch)

        for block_type, block_index in use_blocks:
-            subidx, block_fun = (
-                (0, forward_downsample)
-                if block_type == "input"
-                else (2, forward_upsample)
+            main_block = model.get_model_object(
+                f"diffusion_model.{block_type}_blocks.{block_index}",
            )
-            block_name = f"diffusion_model.{block_type}_blocks.{block_index}.{subidx}"
+            block_fun, expected_class = (
+                (forward_downsample, openaimodel.Downsample)
+                if block_type == "input"
+                else (forward_upsample, openaimodel.Upsample)
+            )
+            block_name = f"diffusion_model.{block_type}_blocks.{block_index}.{len(main_block) - 1}"
            block = model.get_model_object(block_name)
+            if not isinstance(block, expected_class):
+                block_type_name = getattr(type(block), "__name__", "unknown")
+                error_message = (
+                    f"User error: {block_type} {block_index} requires targeting an {expected_class.__name__} block but got block of type {block_type_name} instead.",
+                )
+                raise ValueError(error_message)  # noqa: TRY004
            model.add_object_patch(
                f"{block_name}.forward",
                partial(block_fun, block_index, block, block.forward, hdconfig),
@@ -410,33 +481,54 @@ class ApplyRAUNet:

 class ApplyRAUNetSimple:
    RETURN_TYPES = ("MODEL",)
+    OUTPUT_TOOLTIPS = ("Model patched with the RAUNet effect.",)
    FUNCTION = "patch"
    CATEGORY = "model_patches/unet"
+    DESCRIPTION = "This node is used to enable generation at higher resolutions than a model was trained for with less artifacts or other negative effects. This is the simplified version with less parameters, use ApplyRAUNet if you require more control. NOTE: Only supports SD1.x, SD2.x and SDXL."

    @classmethod
    def INPUT_TYPES(cls) -> dict:
        return {
            "required": {
-                "model": ("MODEL",),
-                "model_type": (("SD15", "SDXL"),),
+                "model": (
+                    "MODEL",
+                    {
+                        "tooltip": "Model to be patched with the RAUNet effect.",
+                    },
+                ),
+                "model_type": (
+                    ("SD15", "SDXL"),
+                    {
+                        "tooltip": "Model type being patched. Choose SD15 for SD 1.4 or SD 2.x.",
+                    },
+                ),
                "res_mode": (
                    (
                        "high (1536-2048)",
                        "low (1024 or lower)",
                        "ultra (over 2048)",
                    ),
+                    {
+                        "tooltip": "Resolution mode hint, does not have to correspond to the actual size.",
+                    },
                ),
                "upscale_mode": (
                    (
                        "default",
                        *UPSCALE_METHODS,
                    ),
+                    {
+                        "tooltip": "Method used when upscaling latents in output Upsample blocks.",
+                    },
                ),
                "ca_upscale_mode": (
                    (
                        "default",
                        *UPSCALE_METHODS,
                    ),
+                    {
+                        "tooltip": "Method used when upscaling latents in cross attention blocks.",
+                    },
                ),
            },
        }
@@ -0,0 +1,14 @@
+[project]
+name = "comfyui_jankhidiffusion"
+description = "Janky implementation of HiDiffusion for ComfyUI. Enables generating at resolutions higher than what the model was trained for. Only supports SD 1.x (maybe 2.x) and SDXL."
+version = "0.8.1"
+license = { file = "LICENSE" }
+
+[project.urls]
+Repository = "https://github.com/blepping/comfyui_jankhidiffusion"
+#  Used by Comfy Registry https://comfyregistry.org
+
+[tool.comfy]
+PublisherId = "blepping"
+DisplayName = "comfyui_jankhidiffusion"
+Icon = ""
Author	SHA1	Message	Date
blepping	64090c80b7	Version bump	2024-08-16 13:06:57 -06:00
blepping	4f48873f98	Make up/downsample block targeting more resiliant in RAUNet	2024-08-15 19:06:56 -06:00
blepping	922a400f6f	Change publish workflow to trigger on release	2024-08-15 03:12:19 -06:00
blepping	7548ad6d07	Merge pull request #19 from blepping/comfyorg_publish Set up Comfy Registry publishing	2024-08-15 02:57:11 -06:00
blepping	028a831031	Set up Comfy Registry publishing	2024-08-15 02:54:51 -06:00
blepping	4e8ef65a7e	Slight consistency tweak for tooltip text.	2024-08-15 02:39:49 -06:00
blepping	3d33f3f7e7	Merge pull request #11 from haohaocreates/publish Add Github Action for Publishing to Comfy Registry	2024-08-14 11:07:16 -06:00
blepping	be31421715	Merge pull request #12 from haohaocreates/pyproject Add pyproject.toml for Custom Node Registry	2024-08-14 11:06:41 -06:00
blepping	c1f7e80f78	Fix output blocks tooltip in ApplyRAUNet node	2024-08-14 10:56:44 -06:00
blepping	00e41dc6b1	Add tooltips metadata to nodes	2024-08-14 10:54:24 -06:00
blepping	4925c89a31	Merge pull request #18 from blepping/refactor * Refactor RAUNet code to avoid monkeypatching Upsample/Downsample blocks (by pamparamm) * Move two_stage_upscale toggle into two_stage_upscale_mode (by pamparamm) * Refactor RAUNet code to avoid monkeypatching forward_timestep_embed * Allow setting a downscale factor and mode for CA downsampling in advanced RAUNet node * Make it so MSW-MSA attention failing due to size mismatches is a warning rather than hard error	2024-08-13 03:57:49 -06:00
haohaocreates	156ae752e1	chore(pyproject): Add pyproject.toml for Custom Node Registry	2024-05-22 19:18:35 -04:00
haohaocreates	89e2ce4c44	chore(publish): Add Github Action for Publishing to Comfy Registry	2024-05-22 19:18:31 -04:00