Merge f9bc8cc781 into 6880614319

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

AGENTS.md

@@ -127,6 +127,8 @@
 - Do not add unnecessary `try`/`except` blocks. Use them for optional dependency,
   platform, or backend capability detection only when the program has a useful
   fallback. Prefer specific exception types when changing new code.
+- If a library version is pinned in `requirements.txt`, do not add code to
+  ComfyUI to handle older versions of that library.
 - Remove any workarounds for PyTorch versions that ComfyUI no longer officially
   supports. Deprecated workarounds include catching an exception and rerunning
   the same op with the input cast to float. If a workaround does not have a

comfy/ldm/chroma/model.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 
 import torch
 from torch import Tensor, nn
-from einops import rearrange, repeat
+from einops import rearrange
 import comfy.patcher_extension
 import comfy.ldm.common_dit
 
@@ -267,29 +267,54 @@ class Chroma(nn.Module):
             img = self.final_layer(img, vec=final_mod)  # (N, T, patch_size ** 2 * out_channels)
         return img
 
-    def forward(self, x, timestep, context, guidance, control=None, transformer_options={}, **kwargs):
+    def forward(self, x, timestep, context, guidance=None, ref_latents=None, control=None, transformer_options={}, **kwargs):
         return comfy.patcher_extension.WrapperExecutor.new_class_executor(
             self._forward,
             self,
             comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
-        ).execute(x, timestep, context, guidance, control, transformer_options, **kwargs)
+        ).execute(x, timestep, context, guidance, ref_latents, control, transformer_options, **kwargs)
 
-    def _forward(self, x, timestep, context, guidance, control=None, transformer_options={}, **kwargs):
-        bs, c, h, w = x.shape
-        x = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size))
-
-        img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=self.patch_size, pw=self.patch_size)
+    def _forward(self, x, timestep, context, guidance=None, ref_latents=None, control=None, transformer_options={}, **kwargs):
+        bs, c, h_orig, w_orig = x.shape
 
+        h_len = ((h_orig + (self.patch_size // 2)) // self.patch_size)
+        w_len = ((w_orig + (self.patch_size // 2)) // self.patch_size)
+        img, img_ids = comfy.ldm.common_dit.process_img(x, patch_size=self.patch_size, transformer_options=transformer_options)
         if img.ndim != 3 or context.ndim != 3:
             raise ValueError("Input img and txt tensors must have 3 dimensions.")
+        img_tokens = img.shape[1]
+        if ref_latents is not None:
+            h = 0
+            w = 0
+            index = 0
+            ref_latents_method = kwargs.get("ref_latents_method", "offset")
+            for ref in ref_latents:
+                if ref_latents_method == "index":
+                    index += 1
+                    h_offset = 0
+                    w_offset = 0
+                elif ref_latents_method == "uxo":
+                    index = 0
+                    h_offset = h_len * self.patch_size + h
+                    w_offset = w_len * self.patch_size + w
+                    h += ref.shape[-2]
+                    w += ref.shape[-1]
+                else:
+                    index = 1
+                    h_offset = 0
+                    w_offset = 0
+                    if ref.shape[-2] + h > ref.shape[-1] + w:
+                        w_offset = w
+                    else:
+                        h_offset = h
+                    h = max(h, ref.shape[-2] + h_offset)
+                    w = max(w, ref.shape[-1] + w_offset)
 
-        h_len = ((h + (self.patch_size // 2)) // self.patch_size)
-        w_len = ((w + (self.patch_size // 2)) // self.patch_size)
-        img_ids = torch.zeros((h_len, w_len, 3), device=x.device, dtype=x.dtype)
-        img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype).unsqueeze(1)
-        img_ids[:, :, 2] = img_ids[:, :, 2] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype).unsqueeze(0)
-        img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+                kontext, kontext_ids = comfy.ldm.common_dit.process_img(ref, index=index, h_offset=h_offset, w_offset=w_offset, patch_size=self.patch_size, transformer_options=transformer_options)
+                img = torch.cat([img, kontext], dim=1)
+                img_ids = torch.cat([img_ids, kontext_ids], dim=1)
 
         txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
         out = self.forward_orig(img, img_ids, context, txt_ids, timestep, guidance, control, transformer_options, attn_mask=kwargs.get("attention_mask", None))
-        return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=self.patch_size, pw=self.patch_size)[:,:,:h,:w]
+        out = out[:, :img_tokens]
+        return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=self.patch_size, pw=self.patch_size)[:,:,:h_orig,:w_orig]

comfy/ldm/common_dit.py

@@ -1,4 +1,5 @@
 import torch
+from einops import rearrange, repeat
 import comfy.rmsnorm
 
 
@@ -14,3 +15,32 @@ def pad_to_patch_size(img, patch_size=(2, 2), padding_mode="circular"):
 
 
 rms_norm = comfy.rmsnorm.rms_norm
+
+def process_img(x, index=0, h_offset=0, w_offset=0, patch_size=2, transformer_options={}, num_axes=3):
+    bs, c, h, w = x.shape
+    x = pad_to_patch_size(x, (patch_size, patch_size))
+
+    img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size)
+    h_len = ((h + (patch_size // 2)) // patch_size)
+    w_len = ((w + (patch_size // 2)) // patch_size)
+
+    h_offset = ((h_offset + (patch_size // 2)) // patch_size)
+    w_offset = ((w_offset + (patch_size // 2)) // patch_size)
+
+    steps_h = h_len
+    steps_w = w_len
+
+    rope_options = transformer_options.get("rope_options", None)
+    if rope_options is not None:
+        h_len = (h_len - 1.0) * rope_options.get("scale_y", 1.0) + 1.0
+        w_len = (w_len - 1.0) * rope_options.get("scale_x", 1.0) + 1.0
+
+        index += rope_options.get("shift_t", 0.0)
+        h_offset += rope_options.get("shift_y", 0.0)
+        w_offset += rope_options.get("shift_x", 0.0)
+
+    img_ids = torch.zeros((steps_h, steps_w, num_axes), device=x.device, dtype=x.dtype)
+    img_ids[:, :, 0] = img_ids[:, :, 1] + index
+    img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(h_offset, h_len - 1 + h_offset, steps=steps_h, device=x.device, dtype=x.dtype).unsqueeze(1)
+    img_ids[:, :, 2] = img_ids[:, :, 2] + torch.linspace(w_offset, w_len - 1 + w_offset, steps=steps_w, device=x.device, dtype=x.dtype).unsqueeze(0)
+    return img, repeat(img_ids, "h w c -> b (h w) c", b=bs)

comfy/ldm/flux/model.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 
 import torch
 from torch import Tensor, nn
-from einops import rearrange, repeat
+from einops import rearrange
 import comfy.ldm.common_dit
 import comfy.patcher_extension
 
@@ -311,36 +311,6 @@ class Flux(nn.Module):
         img = self.final_layer(img, vec_orig, **extra_kwargs)  # (N, T, patch_size ** 2 * out_channels)
         return img
 
-    def process_img(self, x, index=0, h_offset=0, w_offset=0, transformer_options={}):
-        bs, c, h, w = x.shape
-        patch_size = self.patch_size
-        x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch_size, patch_size))
-
-        img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size)
-        h_len = ((h + (patch_size // 2)) // patch_size)
-        w_len = ((w + (patch_size // 2)) // patch_size)
-
-        h_offset = ((h_offset + (patch_size // 2)) // patch_size)
-        w_offset = ((w_offset + (patch_size // 2)) // patch_size)
-
-        steps_h = h_len
-        steps_w = w_len
-
-        rope_options = transformer_options.get("rope_options", None)
-        if rope_options is not None:
-            h_len = (h_len - 1.0) * rope_options.get("scale_y", 1.0) + 1.0
-            w_len = (w_len - 1.0) * rope_options.get("scale_x", 1.0) + 1.0
-
-            index += rope_options.get("shift_t", 0.0)
-            h_offset += rope_options.get("shift_y", 0.0)
-            w_offset += rope_options.get("shift_x", 0.0)
-
-        img_ids = torch.zeros((steps_h, steps_w, len(self.params.axes_dim)), device=x.device, dtype=torch.float32)
-        img_ids[:, :, 0] = img_ids[:, :, 1] + index
-        img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(h_offset, h_len - 1 + h_offset, steps=steps_h, device=x.device, dtype=torch.float32).unsqueeze(1)
-        img_ids[:, :, 2] = img_ids[:, :, 2] + torch.linspace(w_offset, w_len - 1 + w_offset, steps=steps_w, device=x.device, dtype=torch.float32).unsqueeze(0)
-        return img, repeat(img_ids, "h w c -> b (h w) c", b=bs)
-
     def forward(self, x, timestep, context, y=None, guidance=None, ref_latents=None, control=None, transformer_options={}, **kwargs):
         return comfy.patcher_extension.WrapperExecutor.new_class_executor(
             self._forward,
@@ -354,7 +324,8 @@ class Flux(nn.Module):
 
         h_len = ((h_orig + (patch_size // 2)) // patch_size)
         w_len = ((w_orig + (patch_size // 2)) // patch_size)
-        img, img_ids = self.process_img(x, transformer_options=transformer_options)
+        num_axes = len(self.params.axes_dim)
+        img, img_ids = comfy.ldm.common_dit.process_img(x, patch_size=patch_size, transformer_options=transformer_options, num_axes=num_axes)
         img_tokens = img.shape[1]
         timestep_zero_index = None
         if ref_latents is not None:
@@ -386,7 +357,7 @@ class Flux(nn.Module):
                     h = max(h, ref.shape[-2] + h_offset)
                     w = max(w, ref.shape[-1] + w_offset)
 
-                kontext, kontext_ids = self.process_img(ref, index=index, h_offset=h_offset, w_offset=w_offset, transformer_options=transformer_options)
+                kontext, kontext_ids = comfy.ldm.common_dit.process_img(ref, index=index, h_offset=h_offset, w_offset=w_offset, patch_size=patch_size, transformer_options=transformer_options, num_axes=num_axes)
                 img = torch.cat([img, kontext], dim=1)
                 img_ids = torch.cat([img_ids, kontext_ids], dim=1)
                 ref_num_tokens.append(kontext.shape[1])

comfy_extras/nodes_text_overlay.py

@@ -0,0 +1,150 @@
+import numpy as np
+import torch
+from PIL import Image as PILImage, ImageColor, ImageDraw, ImageFont
+from typing_extensions import override
+
+from comfy_api.latest import ComfyExtension, IO
+
+
+class TextOverlay(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="TextOverlay",
+            display_name="Draw Text Overlay",
+            category="text",
+            description="Draw text overlay on an image or batch of images.",
+            search_aliases=["text", "label", "caption", "subtitle", "watermark", "title", "addlabel", "overlay"],
+            inputs=[
+                IO.Image.Input("images"),
+                IO.String.Input("text", multiline=True, default=""),
+                IO.Float.Input("font_size", default=5.0, min=0.5, max=50.0, step=0.5, tooltip="Font size as a percentage of the image height."),
+                IO.Color.Input("color", default="#ffffff", tooltip="Color of the text."),
+                IO.Combo.Input("position", options=["top", "bottom"], default="top"),
+                IO.Combo.Input("align", options=["left", "center", "right"], default="left"),
+                IO.Boolean.Input("outline", default=True, tooltip="Draw a black outline around the text."),
+            ],
+            outputs=[IO.Image.Output(display_name="images")],
+        )
+
+    @classmethod
+    def execute(cls, images, text, font_size, color, position, align, outline) -> IO.NodeOutput:
+        if text.strip() == "":
+            return IO.NodeOutput(images)
+
+        text = text.replace("\\n", "\n").replace("\\t", "\t")
+
+        text_rgba = cls.parse_color_to_rgba(color)
+        outline_rgba = (0, 0, 0, 255) if outline else (0, 0, 0, 0)
+
+        # Render the overlay once and composite it across all frames in the batch
+        height = images.shape[1]
+        width = images.shape[2]
+        overlay_rgb, overlay_alpha = cls.render_overlay_text(width, height, text, position, align, font_size, text_rgba, outline_rgba)
+        overlay_rgb = overlay_rgb.to(device=images.device, dtype=images.dtype)
+        overlay_alpha = overlay_alpha.to(device=images.device, dtype=images.dtype)
+
+        result = images * (1.0 - overlay_alpha) + overlay_rgb * overlay_alpha
+        return IO.NodeOutput(result)
+
+    @staticmethod
+    def parse_color_to_rgba(color_string):
+        parsed = ImageColor.getrgb(color_string)
+
+        if len(parsed) == 3:
+            return (*parsed, 255)
+
+        return parsed
+
+    @classmethod
+    def render_overlay_text(cls, width, height, text, position, align, font_size, text_rgba, outline_rgba):
+        line_spacing = 1.2
+        margin_percent = 1.0
+        min_font_percent = 2.0
+        min_font_pixels = 10
+        outline_thickness_factor = 0.04
+
+        # Draw onto a transparent layer so the result can be alpha-composited over any frame.
+        layer = PILImage.new("RGBA", (width, height), (0, 0, 0, 0))
+        draw = ImageDraw.Draw(layer)
+
+        margin = int(round(margin_percent / 100.0 * min(width, height)))
+        max_width = max(1, width - 2 * margin)
+        max_height = max(1, height - 2 * margin)
+
+        # Font scales with resolution, then shrinks to fit the height.
+        size = max(1, int(round(font_size / 100.0 * height)))
+        floor = min(size, max(min_font_pixels, int(round(min_font_percent / 100.0 * height))))
+
+        while True:
+            font = ImageFont.load_default(size=size)
+            stroke = max(1, int(round(size * outline_thickness_factor))) if outline_rgba[3] > 0 else 0
+            block = "\n".join(cls.wrap_text(text, font, max_width))
+            # convert line spacing to pixel spacing
+            single = draw.textbbox((0, 0), "Ay", font=font, stroke_width=stroke)
+            double = draw.multiline_textbbox((0, 0), "Ay\nAy", font=font, spacing=0, stroke_width=stroke)
+            natural_advance = (double[3] - double[1]) - (single[3] - single[1])
+            pixel_spacing = int(round(size * line_spacing - natural_advance))
+            box = draw.multiline_textbbox((0, 0), block, font=font, spacing=pixel_spacing, stroke_width=stroke)
+            block_height = box[3] - box[1]
+
+            if block_height <= max_height or size <= floor:
+                break
+
+            size = max(floor, int(size * 0.9))
+
+        anchor_h, x = {"left": ("l", margin), "center": ("m", width / 2), "right": ("r", width - margin)}[align]
+
+        # Offset y so the rendered text sits flush against the margin
+        if position == "bottom":
+            y = height - margin - box[3]
+        else:
+            y = margin - box[1]
+
+        draw.multiline_text((x, y), block, font=font, fill=text_rgba, anchor=anchor_h + "a",
+                            align=align, spacing=pixel_spacing, stroke_width=stroke, stroke_fill=outline_rgba)
+
+        overlay = np.array(layer).astype(np.float32) / 255.0
+        overlay_rgb = torch.from_numpy(overlay[:, :, :3])
+        overlay_alpha = torch.from_numpy(overlay[:, :, 3:4])
+        return overlay_rgb, overlay_alpha
+
+    @staticmethod
+    def wrap_text(text, font, max_width):
+        lines = []
+        for raw_line in text.split("\n"):
+            words = raw_line.split()
+            if not words:
+                lines.append("")
+                continue
+            current = ""
+            # Break the line into words and split words that are too long
+            for word in words:
+                while font.getlength(word) > max_width and len(word) > 1:
+                    cut = 1
+                    while cut < len(word) and font.getlength(word[:cut + 1]) <= max_width:
+                        cut += 1
+                    if current:
+                        lines.append(current)
+                        current = ""
+                    lines.append(word[:cut])
+                    word = word[cut:]
+                candidate = word if not current else current + " " + word
+                if not current or font.getlength(candidate) <= max_width:
+                    current = candidate
+                else:
+                    lines.append(current)
+                    current = word
+            if current:
+                lines.append(current)
+        return lines
+
+
+class TextOverlayExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
+        return [TextOverlay]
+
+
+async def comfy_entrypoint() -> TextOverlayExtension:
+    return TextOverlayExtension()

nodes.py

@@ -2478,6 +2478,7 @@ async def init_builtin_extra_nodes():
         "nodes_glsl.py",
         "nodes_lora_debug.py",
         "nodes_textgen.py",
+        "nodes_text_overlay.py",
         "nodes_color.py",
         "nodes_toolkit.py",
         "nodes_replacements.py",