cont

* conditionally send 100 Continue instead

__init__.py

@@ -32,6 +32,7 @@ class ExternalToolingNodes(ComfyExtension):
             krita.KritaMaskLayer,
             krita.Parameter,
             krita.KritaStyle,
+            krita.KritaStyleAndPrompt,
         ]
 
 

api.py

@@ -56,7 +56,9 @@ model_names = {
     "ACEStep": "ace-step",
     "Omnigen2": "omnigen2",
     "QwenImage": "qwen-image",
+    "ErnieImage": "ernie-image",
     "Flux2": "flux2",
+    "Anima": "anima",
 }
 
 gguf_architectures = {
@@ -121,6 +123,9 @@ def inspect_safetensors(filename: str, model_type: str, is_checkpoint: bool):
                     raw_name = base_model.__class__.__name__
                     if raw_name == "SDXL":
                         model_type = base_model.model_type(cfg).name.lower().replace("_", "-")
+                    if raw_name == "Flux2":
+                        hidden_size = unet_config.get("hidden_size", 0)
+                        model_type = {3072: "klein-4b", 4096: "klein-9b"}.get(hidden_size, "dev")
 
             if not raw_name:
                 return {"base_model": "unknown"}
@@ -190,7 +195,6 @@ def inspect_gguf(filename: str, model_type: str):
         else:  # stable-diffusion.cpp, requires conversion. not handled for now
             return {"base_model": "flux", "is_inpaint": False}
 
-        # Detect Chroma (modified Flux)
         if arch_str == "flux" and any(
             t.name.startswith("distilled_guidance_layer")
             for t in itertools.islice(reader.tensors, 5)
@@ -204,10 +208,27 @@ def inspect_gguf(filename: str, model_type: str):
                     arch_str = "z-image"
                     break
 
+        # Detect Flux variants
+        result_type = None
+        if arch_str == "flux":
+            for t in reader.tensors:
+                if t.name.startswith("distilled_guidance_layer"):
+                    arch_str = "chroma"
+                    break
+                elif t.name == "double_stream_modulation_img.lin.weight":
+                    arch_str = "flux2"
+                    if t.shape[0] == 3072:
+                        result_type = "klein-4b"
+                    elif t.shape[0] == 4096:
+                        result_type = "klein-9b"
+                    break
+
         result = {
             "base_model": gguf_architectures.get(arch_str, arch_str),
             "is_inpaint": False,
         }
+        if result_type is not None:
+            result["type"] = result_type
         try:
             if file_type := reader.get_field("general.file_type"):
                 result["quant"] = file_type.contents().lower()
@@ -327,11 +348,11 @@ if _server is not None:
         except Exception as e:
             return web.json_response(dict(error=str(e)), status=500)
 
-    @_server.routes.put("/api/etn/image/{id}")
     async def put_image(request: web.Request):
         try:
             id = request.match_info.get("id", "")
             if id in image_cache:
+                await request.release()  # Consume and discard the data to avoid connection abort
                 return web.json_response(dict(status="cached"), status=200)
 
             content_type = request.headers.get("Content-Type", "application/octet-stream")
@@ -344,6 +365,17 @@ if _server is not None:
         except Exception as e:
             return web.json_response(dict(error=str(e)), status=500)
 
+    async def _put_image_expect_handler(request: web.Request):
+        if request.match_info.get("id", "") in image_cache:
+            # Skip "100 Continue" since we don't need the data, return 200 immediately.
+            return web.json_response(dict(status="cached"), status=200)
+        # otherwise run default aiohttp handler
+        return None
+
+    _server.app.router.add_route(
+        "PUT", "/api/etn/image/{id}", put_image, expect_handler=_put_image_expect_handler
+    )
+
     @_server.routes.put("/api/etn/upload/{folder_name}/{filename}")
     async def upload(request: web.Request):
         folder_name = request.match_info.get("folder_name", "")

krita.py

@@ -1,15 +1,16 @@
 import sys
-import torch
-import numpy as np
 from enum import Enum
 from pathlib import Path
 from typing import Any, NamedTuple
-from PIL import Image
 
-import server
 import comfy.samplers
+import numpy as np
+import server
+import torch
 from comfy.comfy_types.node_typing import IO
 from comfy_api.latest import io
+from PIL import Image
+
 from .nodes import SendImageWebSocket
 
 
@@ -102,7 +103,7 @@ class KritaOutput(io.ComfyNode):
         )
 
     @classmethod
-    def execute(
+    def execute(  # type: ignore
         cls,
         images: torch.Tensor,
         x: int = 0,
@@ -141,7 +142,7 @@ class KritaSendText(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, value: Any, name: str, type: str):
+    def execute(cls, value: Any, name: str, type: str):  # type: ignore
         mime = {
             "text": "text/plain",
             "markdown": "text/markdown",
@@ -169,12 +170,13 @@ class KritaCanvas(io.ComfyNode):
                 io.Int.Output(display_name="width"),
                 io.Int.Output(display_name="height"),
                 io.Int.Output(display_name="seed"),
+                io.Mask.Output(display_name="mask"),
             ],
         )
 
     @classmethod
-    def execute(cls):
-        return io.NodeOutput(_placeholder_image(), 512, 512, 0)
+    def execute(cls, **kwargs):
+        return io.NodeOutput(_placeholder_image(), 512, 512, 0, torch.ones(1, 512, 512))
 
 
 class SelectionContext(Enum):
@@ -236,7 +238,7 @@ class KritaImageLayer(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, name: str):
+    def execute(cls, name: str):  # type: ignore
         return io.NodeOutput(_placeholder_image(), torch.ones(1, 512, 512))
 
 
@@ -254,7 +256,7 @@ class KritaMaskLayer(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, name: str):
+    def execute(cls, name: str):  # type: ignore
         return io.NodeOutput(torch.ones(1, 512, 512))
 
 
@@ -289,7 +291,7 @@ class Parameter(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, name: str, type: str, default, min=0.0, max=1.0):
+    def execute(cls, name: str, type: str, default, min=0.0, max=1.0):  # type: ignore
         if type == "number":
             return io.NodeOutput(float(default))
         elif type == "number (integer)":
@@ -326,5 +328,37 @@ class KritaStyle(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, name: str, sampler_preset: str):
+    def execute(cls, name: str, sampler_preset: str):  # type: ignore
+        raise NotImplementedError("This workflow must be started from Krita!")
+
+
+class KritaStyleAndPrompt(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="ETN_KritaStyleAndPrompt",
+            display_name="Krita Style & Prompt",
+            category="krita",
+            inputs=[
+                io.Combo.Input("sampler_preset", options=["auto", "regular", "live"]),
+            ],
+            outputs=[
+                io.Model.Output(display_name="model (with loras)"),
+                io.Clip.Output(display_name="clip"),
+                io.Vae.Output(display_name="vae"),
+                io.String.Output(display_name="positive prompt (evaluated)"),
+                io.String.Output(display_name="negative prompt (evaluated)"),
+                io.Combo.Output(
+                    display_name="sampler name", options=comfy.samplers.KSampler.SAMPLERS
+                ),
+                io.Combo.Output(
+                    display_name="scheduler", options=comfy.samplers.KSampler.SCHEDULERS
+                ),
+                io.Int.Output(display_name="steps"),
+                io.Float.Output(display_name="guidance"),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, name: str, sampler_preset: str):  # type: ignore
         raise NotImplementedError("This workflow must be started from Krita!")

nsfw.py

@@ -1,5 +1,4 @@
 from __future__ import annotations
-from weakref import ref as WeakRef
 from pathlib import Path
 from tqdm import tqdm
 import torch
@@ -39,6 +38,10 @@ class CLIPSafetyChecker(PreTrainedModel):
         self.concept_embeds_weights = nn.Parameter(torch.ones(17), requires_grad=False)
         self.special_care_embeds_weights = nn.Parameter(torch.ones(3), requires_grad=False)
 
+        # Model requires post_init after transformers v4.57.3
+        if hasattr(self, "post_init"):
+            self.post_init()
+
     def forward(self, clip_input, images: Tensor, sensitivity: float):
         with torch.no_grad():
             image_batch = self.vision_model(clip_input)[1]

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-tooling-nodes"
 description = "Provides nodes and server API extensions geared towards using ComfyUI as a backend for external tools."
-version = "3.1.0"
+version = "3.1.4"
 license = { file = "LICENSE" }
 
 [project.urls]
@@ -13,7 +13,7 @@ line-length = 100
 preview = true
 
 [tool.ruff.lint]
-ignore = ["E741"]
+ignore = ["E741", "BLE001"]
 
 [tool.black]
 line-length = 100

tile.py

@@ -9,9 +9,13 @@ IntArray = npt.NDArray[np.int_]
 
 
 class TileLayout:
-    def __init__(self, image: Tensor, min_tile_size: int, padding: int, blending: int):
-        assert all([x % 8 == 0 for x in image.shape[-3:-1]]), "Image size must be divisible by 8"
-        assert min_tile_size % 8 == 0, "Tile size must be divisible by 8"
+    def __init__(
+        self, image: Tensor, min_tile_size: int, padding: int, blending: int, multiple: int
+    ):
+        assert all([x % multiple == 0 for x in image.shape[-3:-1]]), (
+            "Image size must be divisible by multiple"
+        )
+        assert min_tile_size % multiple == 0, "Tile size must be divisible by multiple"
         assert blending <= padding, "Blending must be smaller than padding"
 
         self.image_size: IntArray = np.array(image.shape[-3:-1])
@@ -21,7 +25,7 @@ class TileLayout:
 
         image_size_with_overlap = self.image_size + (self.tile_count - 1) * 2 * padding
         tile_size = np.ceil(image_size_with_overlap / self.tile_count)
-        self.tile_size: IntArray = (np.ceil(tile_size / 8) * 8).astype(int)
+        self.tile_size: IntArray = (np.ceil(tile_size / multiple) * multiple).astype(int)
 
     def size(self, coord: IntArray):
         return self.end(coord) - self.start(coord)
@@ -84,13 +88,14 @@ class CreateTileLayout(io.ComfyNode):
                 io.Int.Input("min_tile_size", default=512, min=64, max=8192, step=8),
                 io.Int.Input("padding", default=32, min=0, max=8192, step=8),
                 io.Int.Input("blending", default=8, min=0, max=256, step=8),
+                io.Int.Input("multiple", default=8, min=1, max=1024, step=1),
             ],
             outputs=[io.Custom("TileLayout").Output(display_name="layout")],
         )
 
     @classmethod
-    def execute(cls, image: Tensor, min_tile_size: int, padding: int, blending: int):
-        return io.NodeOutput(TileLayout(image, min_tile_size, padding, blending))
+    def execute(cls, image: Tensor, min_tile_size: int, padding: int, blending: int, multiple: int):
+        return io.NodeOutput(TileLayout(image, min_tile_size, padding, blending, multiple))
 
 
 class ExtractImageTile(io.ComfyNode):