Update the Ideogram 4 prompt generation/parsing/ui to handle the updated format notes better.
This commit is contained in:
Jaret Burkett committed
1 parent
fcccc0fbd2
commit
86b19589a0
10 files changed
+497
-155
No files matched your search
@@ -0,0 +1,282 @@
|
||||
"""Shared helpers for Ideogram-4 structured JSON captions.
|
||||
|
||||
This is the single source of truth for the caption schema so the captioner, the
|
||||
prompt upsampler, the dataloader, and the model encoder all agree. It encodes the
|
||||
official Ideogram-4 rules and, crucially, MIGRATES the old caption format we used
|
||||
before those rules were published into the new one ("digest" old, emit new).
|
||||
|
||||
Official schema (summary):
|
||||
- three top-level keys: high_level_description (optional), style_description
|
||||
(optional), compositional_deconstruction (required).
|
||||
- style_description holds EXACTLY ONE of `photo` (photographs) or `art_style`
|
||||
(illustration/painting/3D/graphic design), never both. Key order is strict and
|
||||
branch-dependent:
|
||||
photo branch: aesthetics, lighting, photo, medium, color_palette
|
||||
non-photo branch: aesthetics, lighting, medium, art_style, color_palette
|
||||
- medium is one of: photograph, illustration, 3d_render, painting, graphic_design
|
||||
- color_palette: UPPERCASE #RRGGBB only, up to 16 per image / 5 per element.
|
||||
- elements, strict key order:
|
||||
obj: type, bbox, desc, color_palette
|
||||
text: type, bbox, text, desc, color_palette
|
||||
bbox is optional, normalized 0-1000, [y_min, x_min, y_max, x_max], top-left.
|
||||
- serialize compact: separators=(",", ":"), ensure_ascii=False (no \\uXXXX).
|
||||
|
||||
The OLD format we previously emitted differed by: always using `photo` (even for
|
||||
non-photo media), putting `color_palette` before `desc`/`text`, title-cased medium
|
||||
with a trailing period ("Illustration."), and lowercase / 3-digit hex. Every
|
||||
function here accepts the old shape and returns the new one.
|
||||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
from collections import OrderedDict
|
||||
|
||||
MAX_IMAGE_PALETTE = 16 # style_description.color_palette
|
||||
MAX_ELEMENT_PALETTE = 5 # per-element color_palette
|
||||
|
||||
# Canonical medium tokens (official set).
|
||||
MEDIUM_OPTIONS = [
|
||||
"photograph",
|
||||
"illustration",
|
||||
"3d_render",
|
||||
"painting",
|
||||
"graphic_design",
|
||||
]
|
||||
|
||||
# Map common variants (including our old "Title." style) to the canonical token.
|
||||
# Anything not listed is treated as a custom medium and preserved verbatim.
|
||||
_MEDIUM_ALIASES = {
|
||||
"photograph": "photograph",
|
||||
"photo": "photograph",
|
||||
"illustration": "illustration",
|
||||
"3d render": "3d_render",
|
||||
"3d_render": "3d_render",
|
||||
"3d-render": "3d_render",
|
||||
"3drender": "3d_render",
|
||||
"render": "3d_render",
|
||||
"3d": "3d_render",
|
||||
"painting": "painting",
|
||||
"graphic design": "graphic_design",
|
||||
"graphic_design": "graphic_design",
|
||||
"graphic-design": "graphic_design",
|
||||
"graphic": "graphic_design",
|
||||
}
|
||||
|
||||
_HEX6_RE = re.compile(r"^#[0-9a-fA-F]{6}$")
|
||||
_HEX3_RE = re.compile(r"^#[0-9a-fA-F]{3}$")
|
||||
|
||||
|
||||
def canon_medium(medium):
|
||||
"""Canonicalize a medium string to an official token when recognized,
|
||||
otherwise return it stripped (custom mediums are allowed, preserved as-is)."""
|
||||
if not isinstance(medium, str):
|
||||
return medium
|
||||
key = medium.strip().rstrip(".").strip().lower()
|
||||
if key in _MEDIUM_ALIASES:
|
||||
return _MEDIUM_ALIASES[key]
|
||||
return medium.strip()
|
||||
|
||||
|
||||
def is_photo_medium(medium):
|
||||
"""True for the photograph branch (uses `photo`), False for the art_style branch."""
|
||||
return canon_medium(medium) == "photograph"
|
||||
|
||||
|
||||
def normalize_hex(color):
|
||||
"""Return an UPPERCASE #RRGGBB string, expanding #RGB -> #RRGGBB. None if invalid."""
|
||||
if not isinstance(color, str):
|
||||
return None
|
||||
s = color.strip()
|
||||
if _HEX6_RE.match(s):
|
||||
return "#" + s[1:].upper()
|
||||
if _HEX3_RE.match(s):
|
||||
return "#" + "".join(ch * 2 for ch in s[1:]).upper()
|
||||
return None
|
||||
|
||||
|
||||
def sanitize_palette(palette, max_len):
|
||||
"""Keep unique, valid, UPPERCASE 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:
|
||||
h = normalize_hex(c)
|
||||
if h is None or h in seen:
|
||||
continue
|
||||
seen.add(h)
|
||||
out.append(h)
|
||||
if len(out) >= max_len:
|
||||
break
|
||||
return out or None
|
||||
|
||||
|
||||
def normalize_style(style):
|
||||
"""Reorder/clean style_description into the correct branch (photo vs art_style)
|
||||
with the strict key order, canonical medium, and uppercase palette. Accepts the
|
||||
old shape (always `photo`) and migrates it based on the medium."""
|
||||
if not isinstance(style, dict):
|
||||
return style
|
||||
|
||||
raw_medium = style.get("medium")
|
||||
medium = canon_medium(raw_medium) if raw_medium is not None else None
|
||||
has_photo = bool(style.get("photo"))
|
||||
has_art = bool(style.get("art_style"))
|
||||
|
||||
# Decide the branch. A recognized medium is authoritative; otherwise infer from
|
||||
# whichever render key the (old) data already had, defaulting to photo.
|
||||
if medium in MEDIUM_OPTIONS:
|
||||
photo_branch = medium == "photograph"
|
||||
elif has_art and not has_photo:
|
||||
photo_branch = False
|
||||
else:
|
||||
photo_branch = True
|
||||
|
||||
photo_val = style.get("photo") if has_photo else None
|
||||
art_val = style.get("art_style") if has_art else None
|
||||
|
||||
out = OrderedDict()
|
||||
if "aesthetics" in style:
|
||||
out["aesthetics"] = style["aesthetics"]
|
||||
if "lighting" in style:
|
||||
out["lighting"] = style["lighting"]
|
||||
|
||||
if photo_branch:
|
||||
# aesthetics, lighting, photo, medium, color_palette
|
||||
val = photo_val if photo_val is not None else art_val
|
||||
if val is not None:
|
||||
out["photo"] = val
|
||||
if medium is not None:
|
||||
out["medium"] = medium
|
||||
else:
|
||||
# aesthetics, lighting, medium, art_style, color_palette
|
||||
if medium is not None:
|
||||
out["medium"] = medium
|
||||
val = art_val if art_val is not None else photo_val
|
||||
if val is not None:
|
||||
out["art_style"] = val
|
||||
|
||||
pal = sanitize_palette(style.get("color_palette"), MAX_IMAGE_PALETTE)
|
||||
if pal is not None:
|
||||
out["color_palette"] = pal
|
||||
|
||||
# Preserve any unexpected extra keys at the end rather than dropping them.
|
||||
for k, v in style.items():
|
||||
if k not in (
|
||||
"aesthetics",
|
||||
"lighting",
|
||||
"photo",
|
||||
"art_style",
|
||||
"medium",
|
||||
"color_palette",
|
||||
):
|
||||
out[k] = v
|
||||
return out
|
||||
|
||||
|
||||
def normalize_element(el):
|
||||
"""Reorder an element's keys to the strict schema order and uppercase its
|
||||
palette. obj: type, bbox, desc, color_palette. text: type, bbox, text, desc,
|
||||
color_palette. bbox is kept verbatim (already [y1,x1,y2,x2] in stored form)."""
|
||||
if not isinstance(el, dict):
|
||||
return el
|
||||
etype = el.get("type", "obj")
|
||||
out = OrderedDict()
|
||||
out["type"] = etype
|
||||
if el.get("bbox") is not None:
|
||||
out["bbox"] = el["bbox"]
|
||||
if etype == "text":
|
||||
if "text" in el:
|
||||
out["text"] = el["text"]
|
||||
if "desc" in el:
|
||||
out["desc"] = el["desc"]
|
||||
else:
|
||||
if "desc" in el:
|
||||
out["desc"] = el["desc"]
|
||||
pal = sanitize_palette(el.get("color_palette"), MAX_ELEMENT_PALETTE)
|
||||
if pal is not None:
|
||||
out["color_palette"] = pal
|
||||
# Preserve any extras (e.g. future keys) at the end.
|
||||
for k, v in el.items():
|
||||
if k not in out and k != "color_palette":
|
||||
out[k] = v
|
||||
return out
|
||||
|
||||
|
||||
def normalize_caption_dict(data):
|
||||
"""Normalize a parsed caption dict in place-ish: drop input-only aspect_ratio,
|
||||
enforce top-level key order, normalize style (photo/art_style branch) and every
|
||||
element. Returns a new OrderedDict. Accepts old-format captions and emits new."""
|
||||
if not isinstance(data, dict):
|
||||
return data
|
||||
data.pop("aspect_ratio", None) # input-only context, never part of output
|
||||
|
||||
out = OrderedDict()
|
||||
if "high_level_description" in data:
|
||||
out["high_level_description"] = data["high_level_description"]
|
||||
if "style_description" in data:
|
||||
out["style_description"] = normalize_style(data["style_description"])
|
||||
|
||||
decon = data.get("compositional_deconstruction")
|
||||
if isinstance(decon, dict):
|
||||
nd = OrderedDict()
|
||||
if "background" in decon:
|
||||
nd["background"] = decon["background"]
|
||||
els = decon.get("elements")
|
||||
if isinstance(els, list):
|
||||
nd["elements"] = [normalize_element(e) for e in els]
|
||||
for k, v in decon.items():
|
||||
if k not in ("background", "elements"):
|
||||
nd[k] = v
|
||||
out["compositional_deconstruction"] = nd
|
||||
elif decon is not None:
|
||||
out["compositional_deconstruction"] = decon
|
||||
|
||||
for k, v in data.items():
|
||||
if k not in (
|
||||
"high_level_description",
|
||||
"style_description",
|
||||
"compositional_deconstruction",
|
||||
"aspect_ratio",
|
||||
):
|
||||
out[k] = v
|
||||
return out
|
||||
|
||||
|
||||
def is_ideogram_caption_str(text):
|
||||
"""True if text parses as a JSON object with a compositional_deconstruction block."""
|
||||
t = (text or "").strip()
|
||||
if not t.startswith("{"):
|
||||
return False
|
||||
try:
|
||||
d = json.loads(t)
|
||||
except Exception:
|
||||
return False
|
||||
return isinstance(d, dict) and isinstance(
|
||||
d.get("compositional_deconstruction"), dict
|
||||
)
|
||||
|
||||
|
||||
def to_model_string(data):
|
||||
"""Serialize a caption dict to the compact, model-ready string the renderer wants."""
|
||||
return json.dumps(data, ensure_ascii=False, separators=(",", ":"))
|
||||
|
||||
|
||||
def digest_caption_string(text):
|
||||
"""Parse, normalize (migrating old format), and return the compact model-ready
|
||||
string. Returns the input unchanged if it is not an Ideogram structured caption
|
||||
(plain-text captions pass straight through)."""
|
||||
t = (text or "").strip()
|
||||
if not t.startswith("{"):
|
||||
return text
|
||||
try:
|
||||
data = json.loads(t, object_pairs_hook=OrderedDict)
|
||||
except Exception:
|
||||
return text
|
||||
if not (
|
||||
isinstance(data, dict)
|
||||
and isinstance(data.get("compositional_deconstruction"), dict)
|
||||
):
|
||||
return text
|
||||
return to_model_string(normalize_caption_dict(data))
|
||||
Reference in new issue
Block a user