github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jaret Burkett
2026-06-05 11:14:36 -06:00
parent 6fd2253932
commit c994398850
6 changed files with 220 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
extensions_built_in/captioner/Ideogram4Captioner.py
View File

@@ -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:

21
extensions_built_in/captioner/prompts/ideogram4_caption_prompt.py
View File

@@ -19,16 +19,16 @@ You are CAPTIONING a real image, not imagining one. Describe ONLY what is visibl
## OUTPUT CONTRACT — exactly three top-level keys, in this order:
```json
{"aspect_ratio":"W:H","high_level_description":"...","compositional_deconstruction":{"background":"...","elements":[ ... ]}}
{"high_level_description":"...","style_description":{"aesthetics":"...","lighting":"...","photo":"...","medium":"...","color_palette":["#RRGGBB"]},"compositional_deconstruction":{"background":"...","elements":[ ... ]}}
```
- Emit a SINGLE-LINE MINIFIED JSON object — no markdown fences, no commentary, no other top-level keys.
- Preserve non-ASCII characters as-is (CJK, Cyrillic, Devanagari, Arabic, accented Latin). Never escape with `\\uNNNN`, transliterate, or replace `café` with `cafe`.
- Use SINGLE quotes for embedded text references in prose fields (`'Joe's Diner'`, not `\\"Joe's Diner\\"`). The `text` field of text elements is the exception — that field holds the verbatim characters visible in the image, may use any characters, and follows QUOTED SPAN FIDELITY below.
### `aspect_ratio` (first field, always required)
### Target aspect ratio (input only — never emit it)
The exact target aspect ratio is GIVEN to you in the user message as `W:H`. Echo it VERBATIM. Never recompute it, never emit `auto`, never override it from your own reading of the frame. Every bbox you emit is normalized to THIS aspect ratio.
The user message gives the image's aspect ratio as `W:H`. Use it ONLY to size your bounding boxes correctly (a box is square only on a square frame). Do NOT emit an `aspect_ratio` key — it is not part of the output.
### `high_level_description` — observational summary (50-word hard cap)
@@ -42,19 +42,24 @@ The exact target aspect ratio is GIVEN to you in the user message as `W:H`. Echo
GOOD: `A full-action shot of a male soccer player in a red kit and black Adidas cleats kicking a soccer ball on a green turf field, with a blurred crowd in the stadium background.`
BAD (over-specifies): `A male soccer player captured mid-kick on a bright green grass pitch, right leg fully extended through the follow-through at the precise moment his black-and-white studded boot makes contact with a white-and-black size-5 ball...`
## IDENTIFY THE MEDIUM
## STYLE DESCRIPTION — the `style_description` block (always required)
State the medium accurately in HLD/background prose as natural framing: `photograph | illustration | 3D render | graphic design` (and the specific style when obvious — `35mm film photograph`, `flat vector illustration`, `Pixar-style 3D render`, `watercolor illustration`). Read it from the image; do not impose a default. Name a recognizable style ONCE, briefly — do not append invented technique detail.
A nested object capturing the image's overall look, OBSERVED from the image (never invented). Exactly these five keys:
- `aesthetics` — the overall mood/aesthetic in a short phrase (`Cinematic, minimal, serene.` / `Bright, playful, high-energy.`).
- `lighting` — the actual lighting: direction, quality, contrast, and the colour of the light. Describe a warm-coloured source concretely (`amber pool from a candle`) but never use the bare word `warm` as a grade.
- `photo` — the medium-specific capture/render spec. Photograph → camera/film look, framing, grain, focus (`35mm film still, 16:9 framing, subtle grain, shallow depth of field`). Other media → the rendering technique (`flat vector, clean edges` / `octane 3D render, soft global illumination` / `loose watercolor on textured paper`).
- `medium` — one short phrase: `Photograph.` / `Illustration.` / `3D render.` / `Graphic design.` Read it from the image; do not impose a default.
- `color_palette` — an array of the image's DOMINANT colours as hex strings (`"#1B3A5C"`), up to 16, ordered most → least dominant. Sample the colours actually present; do not invent colours that are not there.
## ELEMENTS — what they are, what they're not
Each element is one of:
```
{"type":"obj","bbox":[x1,y1,x2,y2],"desc":"..."}
{"type":"text","bbox":[x1,y1,x2,y2],"text":"LINE ONE\\nLINE TWO","desc":"..."}
{"type":"obj","bbox":[x1,y1,x2,y2],"color_palette":["#RRGGBB"],"desc":"..."}
{"type":"text","bbox":[x1,y1,x2,y2],"color_palette":["#RRGGBB"],"text":"LINE ONE\\nLINE TWO","desc":"..."}
```
`bbox` is optional per-element (see BBOX section below).
`bbox` and `color_palette` are both OPTIONAL per-element. `bbox`: see BBOX section below. `color_palette`: up to 5 hex strings of that element's own dominant colours — include it when the element has distinctive colours worth pinning (a red jacket, a brand logo, coloured text), omit it for colour-neutral elements.
### SINGLE SUBJECT = SINGLE ELEMENT

26
extensions_built_in/captioner/prompts/ideogram4_prompt.py
View File

@@ -10,19 +10,16 @@ You convert a natural-language user idea into a structured JSON caption an image
## OUTPUT CONTRACT — exactly three top-level keys, in this order:
```json
{"aspect_ratio":"W:H","high_level_description":"...","compositional_deconstruction":{"background":"...","elements":[ ... ]}}
{"high_level_description":"...","style_description":{"aesthetics":"...","lighting":"...","photo":"...","medium":"...","color_palette":["#RRGGBB"]},"compositional_deconstruction":{"background":"...","elements":[ ... ]}}
```
- Emit a SINGLE-LINE MINIFIED JSON object — no markdown fences, no commentary, no other top-level keys.
- Preserve non-ASCII characters as-is (CJK, Cyrillic, Devanagari, Arabic, accented Latin). Never escape with `\uNNNN`, transliterate, or replace `café` with `cafe`.
- Use SINGLE quotes for embedded text references in prose fields (`'Joe's Diner'`, not `\"Joe's Diner\"`). The `text` field of text elements is the exception — that field holds the user's verbatim characters, may use any characters, and follows QUOTED SPAN FIDELITY below.
### `aspect_ratio` (first field, always required)
### Target aspect ratio (input only — never emit it)
A string in `W:H` form with positive integers (`1:1`, `16:9`, `9:16`, `4:5`, `3:1`, `2:3`, etc.).
- If the user message gives a concrete `W:H`, echo it verbatim.
- If the user message says `auto`, pick a concrete ratio that matches the medium and composition (panoramic subjects → wide ratios like `16:9` or `3:1`; portrait subjects → tall like `9:16` or `4:5`; designed artifacts → format conventions like `2:3` book cover, `3:4` poster; ambiguous → `1:1`). NEVER emit the literal string `auto`.
- The aspect ratio you commit to drives every bbox decision. Pick it first.
The user message gives a target aspect ratio as `W:H` (or `auto`). Use it ONLY to drive your bounding-box decisions — a box is square only on a square frame, so the ratio shapes every bbox. Do NOT emit an `aspect_ratio` key; it is not part of the output.
### `high_level_description` — observational summary (50-word hard cap)
@@ -36,15 +33,26 @@ A string in `W:H` form with positive integers (`1:1`, `16:9`, `9:16`, `4:5`, `3:
GOOD: `A full-action shot of a male soccer player in a red kit and black Adidas cleats kicking a soccer ball on a green turf field, with a blurred crowd in the stadium background.`
BAD (over-specifies): `A male soccer player captured mid-kick on a bright green grass pitch, right leg fully extended through the follow-through at the precise moment his black-and-white studded boot makes contact with a white-and-black size-5 ball...`
### `style_description` — the global look block (always required)
A nested object with exactly these five keys:
- `aesthetics` — overall mood/aesthetic in a short phrase (`Cinematic, minimal, serene.`).
- `lighting` — direction, quality, contrast, and colour of the light. Describe a warm-coloured source concretely (`amber sun low at the horizon`); never use the bare word `warm` as a grade.
- `photo` — the medium-specific capture/render spec. Photograph → camera/film look, framing, grain, focus (`35mm motion-picture film still, 16:9 framing, subtle grain`). Other media → the rendering technique (`flat vector, clean edges`; `octane 3D render`; `loose watercolor on textured paper`).
- `medium` — one short phrase: `Photograph.` / `Illustration.` / `3D render.` / `Graphic design.`
- `color_palette` — an array of the dominant colours as hex strings (`"#1B3A5C"`), up to 16, ordered most → least dominant. This conditions the image's colours directly, so commit to the actual hexes you intend.
Name a recognized style ONCE here (see PLANNING → Style commitment); do not append invented technique detail on top of a well-known style name.
## ELEMENTS — what they are, what they're not
Each element is one of:
```
{"type":"obj","bbox":[y1,x1,y2,x2],"desc":"..."}
{"type":"text","bbox":[y1,x1,y2,x2],"text":"LINE ONE\nLINE TWO","desc":"..."}
{"type":"obj","bbox":[y1,x1,y2,x2],"color_palette":["#RRGGBB"],"desc":"..."}
{"type":"text","bbox":[y1,x1,y2,x2],"color_palette":["#RRGGBB"],"text":"LINE ONE\nLINE TWO","desc":"..."}
```
`bbox` is optional per-element (see BBOX section below).
`bbox` and `color_palette` are both OPTIONAL per-element. `bbox`: see BBOX section below. `color_palette`: up to 5 hex strings steering that element's own dominant colours — include it when the element has a distinctive colour (a red jacket, a brand logo, coloured text), omit it otherwise.
### SINGLE SUBJECT = SINGLE ELEMENT

23
extensions_built_in/captioner/prompts/ideogram4_upsample_prompt.py
View File

@@ -17,29 +17,40 @@ You convert a user prompt into a structured JSON caption an image renderer can c
## OUTPUT CONTRACT — exactly three top-level keys, in this order:
```json
{"aspect_ratio":"W:H","high_level_description":"...","compositional_deconstruction":{"background":"...","elements":[ ... ]}}
{"high_level_description":"...","style_description":{"aesthetics":"...","lighting":"...","photo":"...","medium":"...","color_palette":["#RRGGBB"]},"compositional_deconstruction":{"background":"...","elements":[ ... ]}}
```
- Emit a SINGLE-LINE MINIFIED JSON object — no markdown fences, no commentary, no other top-level keys.
- Preserve non-ASCII characters as-is (CJK, Cyrillic, Arabic, accented Latin). Never escape them as unicode code-point sequences or transliterate.
- Use SINGLE quotes for embedded text references in prose fields (`'Joe's Diner'`). The `text` field is the exception — it holds verbatim characters.
### `aspect_ratio` (first field)
### Target aspect ratio (input only — never emit it)
The target ratio is given. Echo it VERBATIM. If it is `auto`, pick a concrete `W:H` that fits the composition (portrait subject → tall, panoramic → wide, ambiguous → `1:1`). Never emit `auto`.
The user message gives a target aspect ratio as `W:H` (or `auto`). Use it ONLY to size your bounding boxes correctly (a box is square only on a square frame). Do NOT emit an `aspect_ratio` key — it is not part of the output.
### `high_level_description` (50-word cap)
One short sentence, reads like a natural prompt, starts with the subject — no "this image shows". Names the subject(s), any trigger/name verbatim, and the overall composition. Don't enumerate fine detail.
## STYLE DESCRIPTION — the `style_description` block (always required)
A nested object with exactly these five keys, filled FROM the prompt:
- `aesthetics` — the overall mood/aesthetic in a short phrase.
- `lighting` — the lighting (direction, quality, colour). Describe a warm-coloured source concretely; never use the bare word `warm` as a grade.
- `photo` — the medium-specific capture/render spec (photograph → camera/film look, framing, grain, focus; other media → the rendering technique).
- `medium` — one short phrase: `Photograph.` / `Illustration.` / `3D render.` / `Graphic design.`
- `color_palette` — an array of dominant colours as hex strings (`"#1B3A5C"`), up to 16, ordered most → least dominant.
Respect FIDELITY: if the prompt NAMES a style, medium, artist, or look, put it in these fields BY NAME (e.g. `medium`/`photo`/`aesthetics`) and do NOT invent its characteristics. Pull lighting and colours from what the prompt states. In faithful mode, only commit to a value the prompt implies, keeping the rest minimal; in creative mode you may infer fitting style values — but never elaborate a named style and never override what the user gave.
## ELEMENTS
Each element is one of:
```
{"type":"obj","bbox":[y1,x1,y2,x2],"desc":"..."}
{"type":"text","bbox":[y1,x1,y2,x2],"text":"LINE ONE\nLINE TWO","desc":"..."}
{"type":"obj","bbox":[y1,x1,y2,x2],"color_palette":["#RRGGBB"],"desc":"..."}
{"type":"text","bbox":[y1,x1,y2,x2],"color_palette":["#RRGGBB"],"text":"LINE ONE\nLINE TWO","desc":"..."}
```
`bbox` is optional per element (see BBOX).
`bbox` and `color_palette` are both OPTIONAL per element. `bbox`: see BBOX. `color_palette`: up to 5 hex strings of that element's dominant colours — include it when the prompt gives the element a distinctive colour (a red jacket, coloured text), otherwise omit.
- **One coherent subject = ONE element.** A person, animal, vehicle, building, or plant is a single element; its parts are attributes of that element's `desc`, never separate elements. Multiple distinct subjects = multiple elements (one each).
- **`desc`:** identity first, then only the attributes the user gave (or that the structure plainly needs). For a named person/trigger: name + action/pose/placement ONLY, no appearance. For a generic un-named subject, you may state the concrete attributes the prompt implies, but do not invent an identity or backstory.