Rework prompts and captioning systems for ideogram to more strictly match the format provided by ideogram.

extensions_built_in/captioner/Ideogram4Captioner.py

@@ -25,6 +25,11 @@ MIN_NEW_TOKENS = 3072
 # generator was trained on, instead of ugly fractions like 1023:768.
 MAX_AR_DENOMINATOR = 16
 
+# color_palette caps: the model often ignores these, so we enforce them.
+MAX_IMAGE_PALETTE = 16  # style_description.color_palette
+MAX_ELEMENT_PALETTE = 5  # per-element color_palette
+HEX_COLOR_RE = re.compile(r"^#(?:[0-9a-fA-F]{3}|[0-9a-fA-F]{6})$")
+
 
 class Ideogram4Captioner(Qwen3VLCaptioner):
     def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
@@ -104,22 +109,64 @@ class Ideogram4Captioner(Qwen3VLCaptioner):
         # stored order is [y1, x1, y2, x2]
         return [y1, x1, y2, x2]
 
-    def _normalize_caption(self, data: dict, aspect_ratio: str) -> dict:
-        """Validate/cleanup the parsed caption before storage, reordering each
-        bbox from model-native [x1,y1,x2,y2] to stored [y1,x1,y2,x2]."""
-        # Force the aspect ratio we computed; the model is told to echo it but
-        # we know the true value.
-        data["aspect_ratio"] = aspect_ratio
+    def _sanitize_palette(self, palette, max_len):
+        """Keep unique, valid hex colors in order, capped to max_len. Returns the
+        cleaned list, or None if nothing valid remains (drop the key)."""
+        if not isinstance(palette, (list, tuple)):
+            return None
+        seen = set()
+        out = []
+        for c in palette:
+            if not isinstance(c, str):
+                continue
+            c = c.strip()
+            if not HEX_COLOR_RE.match(c):
+                continue
+            key = c.lower()
+            if key in seen:
+                continue
+            seen.add(key)
+            out.append(c)
+            if len(out) >= max_len:
+                break
+        return out or None
+
+    def _normalize_caption(self, data: dict) -> dict:
+        """Validate/cleanup the parsed caption before storage: drop input-only
+        aspect_ratio, reorder bboxes to [y1,x1,y2,x2], and cap color palettes
+        (16 per image, 5 per element) since the model often exceeds them."""
+        # aspect_ratio is input-only context, not part of the output. Drop it if
+        # the model echoed it anyway.
+        data.pop("aspect_ratio", None)
+
+        style = data.get("style_description")
+        if isinstance(style, dict) and "color_palette" in style:
+            pal = self._sanitize_palette(style["color_palette"], MAX_IMAGE_PALETTE)
+            if pal is None:
+                style.pop("color_palette", None)
+            else:
+                style["color_palette"] = pal
+
         decon = data.get("compositional_deconstruction", {})
         elements = decon.get("elements", [])
         if isinstance(elements, list):
             for el in elements:
-                if isinstance(el, dict) and "bbox" in el:
+                if not isinstance(el, dict):
+                    continue
+                if "bbox" in el:
                     cleaned = self._convert_bbox(el["bbox"])
                     if cleaned is None:
                         el.pop("bbox", None)
                     else:
                         el["bbox"] = cleaned
+                if "color_palette" in el:
+                    pal = self._sanitize_palette(
+                        el["color_palette"], MAX_ELEMENT_PALETTE
+                    )
+                    if pal is None:
+                        el.pop("color_palette", None)
+                    else:
+                        el["color_palette"] = pal
         return data
 
     def get_caption_for_file(self, file_path: str) -> Optional[str]:
@@ -172,7 +219,7 @@ class Ideogram4Captioner(Qwen3VLCaptioner):
                 )
                 return output_text
 
-            data = self._normalize_caption(data, aspect_ratio)
+            data = self._normalize_caption(data)
             # Store pretty JSON for QC/editing; the dataloader minifies at load.
             return json.dumps(data, ensure_ascii=False, indent=2)
         except Exception as e: