Tile Captioned Sampling: Add grid full-size image for reference for VLM on current tile. Improves detection of copped ambiguous objects.

Author: WASasquatch

modules/interface/tile_activity.py

@@ -1,7 +1,7 @@
 """The tiles a tiled sampler is working on, announced as they change.
 
-:data:`EVENT` carries ``node_id``, ``prompt_id``, ``phase``, ``tiles``, ``step``, ``steps`` and
-``boxes``, each box ``[top, bottom, left, right]`` in fractions of the frame.
+:data:`EVENT` carries ``node_id``, ``prompt_id``, ``phase``, ``tiles``, ``step``, ``steps``,
+``visited`` and ``boxes``, each box ``[top, bottom, left, right]`` in fractions of the frame.
 """
 
 from __future__ import annotations
@@ -39,27 +39,51 @@ def __init__(self, boxes: list, node_id=None, prompt_id=None):
         self.node_id = node_id if node_id is not None else executing_node_id()
         self.prompt_id = prompt_id if prompt_id is not None else executing_prompt_id()
         self.last = None
+        self.step = 0
+        self.steps = 0
+        self.visited = set()
 
-    def __call__(self, phase: str, tiles=(), step: int = 0, steps: int = 0) -> None:
+    def begin(self, steps: int) -> None:
+        """Start counting a sampling run.
+
+        Args:
+            steps: Steps in the run.
+        """
+        self.steps, self.step, self.visited = int(steps), 0, set()
+
+    def finished(self, step: int) -> None:
+        """Note a step the sampler finished, which opens the next.
+
+        Args:
+            step: The finished step, from 0.
+        """
+        self.step, self.visited = min(int(step) + 1, max(self.steps - 1, 0)), set()
+
+    def __call__(self, phase: str, tiles=(), step=None, steps=None) -> None:
         """Announce the tiles being worked on.
 
         Args:
             phase: :data:`PREPARE`, :data:`OVERVIEW`, :data:`CAPTION`, :data:`SAMPLE` or
                 :data:`DONE`.
             tiles: Indices of the tiles, row by row.
-            step: The step being taken, from 0.
-            steps: Steps in the run.
+            step: The step being taken from 0, the counted one when left out.
+            steps: Steps in the run, the counted number when left out.
         """
         if self.node_id is None:
             return
         active = sorted({int(tile) for tile in tiles})
-        key = (phase, tuple(active), int(step))
+        step = self.step if step is None else int(step)
+        steps = self.steps if steps is None else int(steps)
+        if phase == SAMPLE:
+            self.visited.update(active)
+        key = (phase, tuple(active), step)
         if key == self.last:
             return
         self.last = key
         fields = {
             "node_id": str(self.node_id), "prompt_id": self.prompt_id or "", "phase": phase,
-            "tiles": active, "step": int(step), "steps": int(steps), "boxes": self.boxes,
+            "tiles": active, "step": step, "steps": steps, "boxes": self.boxes,
+            "visited": len(self.visited) if phase == SAMPLE else 0,
         }
         try:
             from server import PromptServer

modules/latent/h3_writer.py

@@ -314,8 +314,9 @@ def make_writer(clip, seed: int, sampling: Sampling, system: str = ""):
         system: The rules every question is asked under, empty for none.
 
     Returns:
-        ``write(prompt, image=None, max_length=1024, temperature=None, ruled=True)`` answering
-        text; ``temperature`` replaces the sampling's own, ``ruled`` asks under ``system``.
+        ``write(prompt, image=None, max_length=1024, temperature=None, ruled=True, images=None)``
+        answering text; ``temperature`` replaces the sampling's own, ``ruled`` asks under
+        ``system``, ``images`` hands several pictures in place of ``image``.
     """
     import comfy.model_management
     import comfy.utils
@@ -328,15 +329,16 @@ def make_writer(clip, seed: int, sampling: Sampling, system: str = ""):
     def quiet(*args, **kwargs):
         comfy.model_management.throw_exception_if_processing_interrupted()
 
-    def write(prompt, image=None, max_length=1024, temperature=None, ruled=True):
+    def write(prompt, image=None, max_length=1024, temperature=None, ruled=True, images=None):
         asked = str(prompt)
         extra = {}
         if ruled and rules:
             if honoured:
                 extra["system_prompt"] = rules
             else:
                 asked = f"Follow these rules.\n\n{rules}\n\n----------\n\n{asked}"
-        tokens = clip.tokenize(asked, image=image, min_length=1, thinking=bool(sampling.thinking), **extra)
+        pictures = {"images": list(images)} if images else {"image": image}
+        tokens = clip.tokenize(asked, min_length=1, thinking=bool(sampling.thinking), **pictures, **extra)
         budget = int(max_length) + (THINKING_TOKENS if sampling.thinking else 0)
         heat = sampling.temperature if temperature is None else min(float(temperature), sampling.temperature)
         # The node's own bar counts its steps; the model's per-token bar only checks for Stop.

modules/sampling/tile_captions.py

@@ -21,30 +21,38 @@
     "OVERVIEW_CONTEXT",
     "OVERVIEW_QUESTION",
     "SYSTEM_CONTEXT",
+    "TILE_CONTEXT",
     "TILE_QUESTION",
+    "TILE_QUESTION_PANELS",
     "caption_tiles",
     "encode",
     "require_captioner",
     "takes_pictures",
+    "takes_several_pictures",
 ]
 
 logger = log.get_logger("sampling.tile_captions")
 
 #: The rules every tile caption is written under.
 DEFAULT_SYSTEM_PROMPT = (
     "You write the prompt an image model uses to repaint a crop of a larger picture.\n"
-    "- Look only at the attached crop. Name each subject and object visible in it, largest "
-    "first, including any only partly in frame, and say which part of it shows. Then describe "
-    "their surfaces, textures, colours and lighting, and any legible text in double quotes.\n"
-    "- The overview describes the whole picture, and most of what it names lies outside this "
-    "crop. Use it only to name and style what the crop shows. Never name anything that is not "
-    "visible in the crop.\n"
+    "- You are shown the whole picture with its tiles ruled and numbered and the crop's tile "
+    "outlined in yellow, and the crop itself at full detail. Describe only what lies inside "
+    "the crop.\n"
+    "- Name each subject and object inside the crop, largest first, including any only partly "
+    "in frame. For each, say which part of it shows, where that part sits within the frame and "
+    "how much of the frame it fills. For a person or animal only partly in the crop, name "
+    "exactly which part of the body shows and whose it is, as the whole picture shows it. Then "
+    "describe their surfaces, textures, colours and lighting, and any legible text in double "
+    "quotes.\n"
+    "- Use the whole picture and the overview only to recognise and name what the crop shows. "
+    "Most of what they show lies outside the crop. Never name anything that is not inside it.\n"
     "- Keep the medium and style of the whole picture.\n"
     "- Where the crop holds only empty background, sky, water or a plain surface, name that "
     "surface and its colour in a few words.\n"
     "- Describe what is there directly. Write one plain paragraph of at most 60 words, with no "
-    "preamble and no list. Never write the words tile, crop, section, corner, overview, "
-    "attached or image, and never say where in the picture the crop sits."
+    "preamble and no list. Never write the words tile, crop, overview, outlined, numbered, "
+    "attached or image, and never say where the crop sits in the whole picture."
 )
 
 #: What the model is asked about the whole picture.
@@ -60,14 +68,35 @@
     "Call things what the prompt calls them.\n\n"
 )
 
+#: What precedes each tile question where the picture's own prompt is known; ``{context}`` is
+#: filled in.
+TILE_CONTEXT = "The whole picture was made from this prompt, given for names only: {context}\n\n"
+
 #: What is added to the system prompt where the picture's own prompt is known.
 SYSTEM_CONTEXT = "\n\nThe original prompt and context of the image was:\n\n{context}"
 
-#: What the model is asked about each tile; ``{overview}`` and ``{place}`` are filled in.
+#: What the model is asked about each tile, shown as two pictures; ``{overview}``,
+#: ``{number}`` and ``{place}`` are filled in.
 TILE_QUESTION = (
-    "Overview of the whole picture: {overview}\n\n"
-    "The attached picture is a crop of that picture ({place}). Write the prompt for everything "
-    "visible in the attached picture."
+    "Overview of the whole picture, given for names only: {overview}\n\n"
+    "The first attached picture is the whole picture with its tiles ruled and numbered; tile "
+    "{number} ({place}) is outlined in yellow. The second is tile {number} at full detail.\n\n"
+    "Write the prompt for the second picture only. Start with its largest subject. For "
+    "everything cut off by an edge of the second picture, say which edge it touches and which "
+    "part of it shows. Leave out anything outside tile {number}, even where the overview or "
+    "the prompt names it, and copy no sentence from either."
+)
+
+#: What the model is asked about each tile, shown as one picture of two panels.
+TILE_QUESTION_PANELS = (
+    "Overview of the whole picture, given for names only: {overview}\n\n"
+    "The attached picture has two panels. On the left is the whole picture with its tiles "
+    "ruled and numbered, tile {number} ({place}) outlined in yellow; on the right is tile "
+    "{number} at full detail.\n\n"
+    "Write the prompt for the right panel only. Start with its largest subject. For "
+    "everything cut off by an edge of the right panel, say which edge it touches and which "
+    "part of it shows. Leave out anything outside tile {number}, even where the overview or "
+    "the prompt names it, and copy no sentence from either."
 )
 
 #: Longest side in pixels of the picture the overview is written from.
@@ -76,6 +105,9 @@
 #: Longest side in pixels of the picture a tile caption is written from.
 TILE_SIDE = 768
 
+#: Longest side in pixels of the whole picture shown beside each tile.
+CONTEXT_SIDE = 768
+
 #: Most tokens an overview may run to.
 OVERVIEW_TOKENS = 200
 
@@ -117,6 +149,22 @@ def require_captioner(captioner, clip, node: str) -> None:
         )
 
 
+def _flat(tokens) -> list:
+    """Every entry of a tokenizer's answer, in order."""
+    batches = tokens.values() if isinstance(tokens, dict) else [tokens]
+    return [item for batch in batches for row in batch for item in row]
+
+
+def _pictures_in(tokens) -> int:
+    """How many entries of a tokenizer's answer carry a picture rather than a token id."""
+    count = 0
+    for item in _flat(tokens):
+        head = item[0] if isinstance(item, tuple) and item else item
+        if isinstance(head, dict):
+            count += 1
+    return count
+
+
 def takes_pictures(clip) -> bool:
     """Whether a CLIP's tokenizer turns a picture into tokens.
 
@@ -126,16 +174,38 @@ def takes_pictures(clip) -> bool:
     Returns:
         True where tokenizing with a picture answers different tokens than without one.
     """
-    def flat(tokens):
-        batches = tokens.values() if isinstance(tokens, dict) else [tokens]
-        return [item for batch in batches for row in batch for item in row]
+    try:
+        plain = clip.tokenize("probe", min_length=1)
+        seen = clip.tokenize("probe", image=torch.zeros(1, 64, 64, 3), min_length=1)
+    except Exception:
+        return False
+    return len(_flat(seen)) != len(_flat(plain)) or _pictures_in(seen) > 0
+
+
+def takes_several_pictures(clip) -> bool:
+    """Whether a CLIP's tokenizer takes two pictures into one question.
+
+    Args:
+        clip: A ComfyUI ``CLIP``.
 
+    Returns:
+        True where two pictures handed together both arrive in the tokens.
+    """
     try:
-        plain = flat(clip.tokenize("probe", min_length=1))
-        seen = flat(clip.tokenize("probe", image=torch.zeros(1, 64, 64, 3), min_length=1))
+        probe = torch.zeros(1, 64, 64, 3)
+        return _pictures_in(clip.tokenize("probe", images=[probe, probe], min_length=1)) == 2
     except Exception:
         return False
-    return len(seen) != len(plain) or any(not isinstance(item, (int, tuple)) for item in seen)
+
+
+def _panels(context: torch.Tensor, crop: torch.Tensor) -> torch.Tensor:
+    """The context picture and the crop side by side at one height, a dark gap between them."""
+    height = int(context.shape[1])
+    width = max(1, round(crop.shape[2] * height / max(1, crop.shape[1])))
+    resized = torch.nn.functional.interpolate(crop.movedim(-1, 1), size=(height, width), mode="bicubic",
+                                              align_corners=False, antialias=True).movedim(1, -1)
+    gap = torch.zeros(1, height, max(4, height // 64), 3)
+    return torch.cat([context, gap, resized.clamp(0.0, 1.0)], dim=2)
 
 
 def _fit(picture: torch.Tensor, side: int) -> torch.Tensor:
@@ -181,14 +251,19 @@ def _tidy(text: str) -> str:
     return reply
 
 
-def _ask(write, store: OrderedDict, picture: torch.Tensor, question: str, system: str,
+def _ask(write, store: OrderedDict, pictures: list, question: str, system: str,
          budget: int, ruled: bool) -> str:
-    """One answer about a picture, read from the cache where the same question was asked before."""
-    key = (_digest(picture), question, system if ruled else "", int(budget))
+    """One answer about pictures, read from the cache where the same question was asked before."""
+    key = (tuple(_digest(picture) for picture in pictures), question, system if ruled else "",
+           int(budget))
     if key in store:
         store.move_to_end(key)
         return store[key]
-    answer = _tidy(write(question, image=picture, max_length=budget, ruled=ruled))
+    if len(pictures) > 1:
+        reply = write(question, images=pictures, max_length=budget, ruled=ruled)
+    else:
+        reply = write(question, image=pictures[0], max_length=budget, ruled=ruled)
+    answer = _tidy(reply)
     store[key] = answer
     while len(store) > MOST_CACHED:
         store.popitem(last=False)
@@ -223,25 +298,30 @@ def caption_tiles(captioner, image: torch.Tensor, grid: tiles.Grid, system_promp
     write = h3_writer.make_writer(captioner, 0, SAMPLING, system)
     store = _store(captioner)
     lead = OVERVIEW_CONTEXT.format(context=context) if context else ""
+    tile_lead = TILE_CONTEXT.format(context=context) if context else ""
     if announce is not None:
         announce("overview")
-    overview = _ask(write, store, _fit(image, OVERVIEW_SIDE), lead + OVERVIEW_QUESTION, system,
+    overview = _ask(write, store, [_fit(image, OVERVIEW_SIDE)], lead + OVERVIEW_QUESTION, system,
                     OVERVIEW_TOKENS, ruled=False)
     if progress is not None:
         progress.update(1)
     height, width = int(image.shape[1]), int(image.shape[2])
+    several = takes_several_pictures(captioner)
+    whole = _fit(image, CONTEXT_SIDE)
     captions = []
     for index, (h0, h1, w0, w1) in enumerate(grid.boxes()):
         comfy.model_management.throw_exception_if_processing_interrupted()
         if announce is not None:
             announce("caption", [index])
         top, bottom = int(round(h0 * grid.scale)), min(height, int(round(h1 * grid.scale)))
         left, right = int(round(w0 * grid.scale)), min(width, int(round(w1 * grid.scale)))
-        crop = image[:, top:max(bottom, top + 1), left:max(right, left + 1)]
-        question = lead + TILE_QUESTION.format(overview=overview or "not available",
-                                               place=tiles.position(grid, index))
-        captions.append(_ask(write, store, _fit(crop, TILE_SIDE), question, system, TILE_TOKENS,
-                             ruled=True))
+        crop = _fit(image[:, top:max(bottom, top + 1), left:max(right, left + 1)], TILE_SIDE)
+        context_picture = tiles.numbered(whole, grid, index, CONTEXT_SIDE)
+        template = TILE_QUESTION if several else TILE_QUESTION_PANELS
+        pictures = [context_picture, crop] if several else [_panels(context_picture, crop)]
+        question = tile_lead + template.format(overview=overview or "not available",
+                                               number=index + 1, place=tiles.position(grid, index))
+        captions.append(_ask(write, store, pictures, question, system, TILE_TOKENS, ruled=True))
         if progress is not None:
             progress.update(1)
     return overview, captions

modules/sampling/tile_refine.py

@@ -25,6 +25,7 @@
     "prepare",
     "publish",
     "release",
+    "tracked",
     "tile_settings",
     "upscaled",
 ]
@@ -227,6 +228,26 @@ def release(*keep) -> None:
         logger.debug("The other models could not be released: %s", error)
 
 
+def tracked(callback, announce, steps: int):
+    """A sampler callback that also counts finished steps for the node's panel.
+
+    Args:
+        callback: The callback the sampler is handed, called as ``callback(step, x0, x, total)``.
+        announce: The frame's announcer.
+        steps: Steps in the run.
+
+    Returns:
+        The wrapped callback.
+    """
+    announce.begin(steps)
+
+    def counted(step, x0, x, total_steps):
+        announce.finished(step)
+        return callback(step, x0, x, total_steps)
+
+    return counted
+
+
 def _publish_plan(picture: torch.Tensor, grid: tiles.Grid) -> None:
     """Publish the tile plan over a picture, for the node's own panel."""
     from ..interface import preview