[refactor] Flatten Ming-flash-omni-2.0 imagegen args (#4587)

Signed-off-by: Yuanheng <jonathan.zhaoyh@gmail.com>
This commit is contained in:
Yuanheng Zhao authored and GitHub committed 2026-06-24 17:01:58 +08:00
1 parent eb81914c64
commit fb3abb38fb
4 files changed
+69 -71

No files matched your search

@@ -88,7 +88,7 @@ curl -s http://127.0.0.1:8091/v1/chat/completions \
### Optional knobs
Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_args.image_gen`:
Pass image-gen overrides as flat keys on the diffusion-stage `sampling_params_list[1].extra_args`:
| Key | Default | Description |
| --- | --- | --- |
@@ -97,25 +97,48 @@ Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_
| `cfg` | 2.0 | Classifier-free guidance scale. |
| `seed` | 42 | Per-request RNG seed (deterministic when ByT5 is also seed-stable). |
| `byte5_text` | (auto) | Override the glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "...". ` format. |
| `negative_prompt` | empty | Real CFG negative conditioning (must be set on **stage-0 thinker** `sampling_params` so `expand_cfg_prompts` spawns the companion). |
| `negative_prompt` | empty | Real CFG negative conditioning (set on **stage-0 thinker** `extra_args` so `expand_cfg_prompts` spawns the companion). |
Example with all the knobs:
```bash
curl http://127.0.0.1:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "Jonathan1909/Ming-flash-omni-2.0",
"modalities": ["image"],
"sampling_params_list": [
{"temperature":0.4,"top_p":0.9,"top_k":1,"max_tokens":1,"seed":42,
"extra_args":{"image_gen":{"negative_prompt":"ugly, blurry, distorted"}}},
{"seed":42,"extra_args":{"image_gen":{
"steps":6,"cfg":1.5,"height":512,"width":512,"seed":123,
"byte5_text":["理解与生成统一"]}}}
],
"messages": [{"role": "user", "content": "Draw a poster."}]
}' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen_knobs.png
-H "Content-Type: application/json" \
-d '{
"model": "Jonathan1909/Ming-flash-omni-2.0",
"modalities": ["image"],
"sampling_params_list": [
{
"temperature": 0.4,
"top_p": 0.9,
"top_k": 1,
"max_tokens": 1,
"seed": 42,
"extra_args": {
"negative_prompt": "ugly, blurry, distorted"
}
},
{
"seed": 42,
"extra_args": {
"steps": 6,
"cfg": 1.5,
"height": 512,
"width": 512,
"seed": 123,
"byte5_text": ["理解与生成统一"]
}
}
],
"messages": [
{
"role": "user",
"content": "Draw a poster."
}
]
}' \
| jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
| base64 -d > ming_imagegen_knobs.png
```
### img2img (reference image)
@@ -198,7 +198,7 @@ class MingImagePipeline(ZImagePipeline):
self.condition_encoder.load_from_checkpoint(model_path)
# Optional ByT5 glyph/text encoder. Only loaded when the checkpoint
# ``extra_args.image_gen.byte5_text`` requests will be ignored.
# ships a byt5/ subfolder; otherwise byte5_text requests are ignored
byte5_dir = Path(model_path) / "byt5"
if byte5_dir.exists():
self.byte5 = MingByT5Encoder.from_checkpoint(byte5_dir, device=self.device, dtype=dtype)
@@ -217,13 +217,12 @@ class MingImagePipeline(ZImagePipeline):
"""Resolve byte5 glyph texts.
Two sources, in order of priority:
1. ``extra["byte5_text"]`` — auto-extracted from the user prompt's
quoted spans by ``thinker2imagegen`` (already wrapped as
``'Text "<glyph>". '`` by Ming's ``get_text_from_prompt``).
2. ``sampling_params.extra_args["image_gen"]["byte5_text"]`` — an
explicit override for programmatic callers; raw strings without
the ``Text "..."`` wrapper are auto-wrapped here to match the
distribution ByT5 was trained on.
1. `extra["byte5_text"]`: auto-extracted from the user prompt's quoted spans
by thinker2imagegen (already wrapped as `'Text "<glyph>". '`
by Ming's get_text_from_prompt).
2. `sampling_params.extra_args["byte5_text"]`:
raw strings without the `Text "..."` wrapper are auto-wrapped here
to match the distribution ByT5 was trained on.
"""
# Source 1: auto-extracted, already wrapped. Return as-is if non-empty.
raw = extra.get("byte5_text")
@@ -236,7 +235,7 @@ class MingImagePipeline(ZImagePipeline):
# Source 2: explicit override — wrap raw strings so the byte5 encoder
# sees the same ``Text "<glyph>". `` format Ming used during training.
raw = ((getattr(sampling_params, "extra_args", None) or {}).get("image_gen") or {}).get("byte5_text")
raw = (getattr(sampling_params, "extra_args", None) or {}).get("byte5_text")
if isinstance(raw, str):
raw = [raw]
if isinstance(raw, list):
@@ -366,21 +365,23 @@ class MingImagePipeline(ZImagePipeline):
negative_cap_feats = torch.cat((negative_cap_feats, torch.zeros_like(byte5_feats)), dim=1)
logger.debug("[MingImagePipeline.forward] byte5 cat'd: cap_feats=%s", tuple(cap_feats.shape))
# Sampling knobs: extra_args.image_gen.* > sampling_params.* > MingImageGenConfig defaults.
# Sampling knobs, in priority order:
# top-level extra_args[key] > sampling_params.* attr >
# MingImageGenConfig default. Knobs live flat on extra_args
sp = req.sampling_params
cfg = self.image_gen_config
ig = (sp.extra_args or {}).get("image_gen") or {}
ea = sp.extra_args or {}
resolved: dict[str, Any] = {}
for ig_key, sp_attr, default in (
for ea_key, sp_attr, default in (
("height", "height", cfg.default_height),
("width", "width", cfg.default_width),
("steps", "num_inference_steps", cfg.num_inference_steps),
("cfg", "guidance_scale", cfg.guidance_scale),
("seed", "seed", None),
):
for v in (ig.get(ig_key), getattr(sp, sp_attr), default):
for v in (ea.get(ea_key), getattr(sp, sp_attr), default):
if v is not None:
resolved[ig_key] = v
resolved[ea_key] = v
break
height = int(resolved["height"])
@@ -393,7 +394,7 @@ class MingImagePipeline(ZImagePipeline):
# ``sp.generator`` causes two problems:
# (1) if the caller pre-seeded it with sp.seed (e.g. the top-level
# ``seed`` key on OmniDiffusionSamplingParams), any override via
# ``extra_args.image_gen.seed`` would be silently ignored; and
# `extra_args["seed"]` would be silently ignored; and
# (2) a persistent generator instance accumulates state across
# requests → same-seed replays produce different outputs.
if seed is not None:
@@ -435,7 +436,7 @@ class MingImagePipeline(ZImagePipeline):
num_inference_steps,
guidance_scale,
seed,
ig,
ea,
None if ref_latent is None else tuple(ref_latent.shape),
)
try:
@@ -486,21 +486,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
for k, v in engine_prompt_image.items()
}
# Forward image_gen extra_args from diffusion-stage sampling
# params into the prompt dict so that thinker2imagegen can
# access them (e.g. byte5_text for ByT5 glyph rendering).
if hasattr(request, "sampling_params_list") and request.sampling_params_list:
for sp_raw in request.sampling_params_list:
ea = (
(sp_raw.get("extra_args") or {})
if isinstance(sp_raw, dict)
else (getattr(sp_raw, "extra_args", None) or {})
)
ig = ea.get("image_gen") or {}
if ig:
tprompt["image_gen_extra_args"] = ig
break
engine_prompts = [tprompt]
# Store height/width for applying to diffusion stage sampling params later
_image_gen_height = height
@@ -104,10 +104,9 @@ def expand_cfg_prompts(
) -> list[_CfgExpandedPrompt]:
"""Expand a text-to-image request into one CFG-text companion (opt-in).
Triggers only when a non-empty
``sampling_params.extra_args["image_gen"]["negative_prompt"]`` is set on
the stage-0 params; otherwise returns ``[]`` and the pipeline falls back
to zero negative (Ming's default behavior).
Triggers only when a non-empty `negative_prompt` is set flat on the stage-0 params
(`sampling_params.extra_args`); otherwise returns an empty list
and the pipeline falls back to zero negative (Ming's default behavior).
"""
if not isinstance(prompt, dict):
return []
@@ -115,8 +114,7 @@ def expand_cfg_prompts(
return []
extra_args = getattr(sampling_params, "extra_args", None) or {}
image_gen_args = extra_args.get("image_gen") or {}
negative = image_gen_args.get("negative_prompt")
negative = extra_args.get("negative_prompt")
if not isinstance(negative, str) or not negative.strip():
return []
@@ -328,7 +326,7 @@ def _resolve_token_ids_from_stage_or_defaults(
def _extract_byte5_from_sampling_params(sampling_params: Any) -> list[str] | None:
"""Read ``byte5_text`` from the diffusion-stage sampling_params.
Looks up ``sampling_params.extra_args["image_gen"]["byte5_text"]`` (the
Looks up `sampling_params.extra_args["byte5_text"]` key (the
explicit API surface for ByT5 glyph text). Returns ``None`` if absent or
malformed, so callers can fall back to other sources.
"""
@@ -337,10 +335,9 @@ def _extract_byte5_from_sampling_params(sampling_params: Any) -> list[str] | Non
extra = getattr(sampling_params, "extra_args", None)
if not isinstance(extra, dict):
return None
image_gen = extra.get("image_gen")
if not isinstance(image_gen, dict):
return None
texts = image_gen.get("byte5_text")
texts = extra.get("byte5_text")
if isinstance(texts, str):
texts = [texts]
if isinstance(texts, list) and texts:
return [t for t in texts if isinstance(t, str)]
return None
@@ -361,10 +358,9 @@ def thinker2imagegen(
CFG negative conditioning. Unknown-suffix outputs are skipped.
``sampling_params`` is the diffusion stage's own SamplingParams, supplied
by the orchestrator. ByT5 explicit ``byte5_text`` is read from
``sampling_params.extra_args.image_gen.byte5_text`` (preferred); falls
back to ``prompt.image_gen_extra_args.byte5_text``, then to auto-extraction
from quoted prompt text.
by the orchestrator. ByT5 explicit ``byte5_text`` is read from the flat
`sampling_params.extra_args.byte5_text` key (preferred); otherwise it is
auto-extracted from quoted prompt text.
"""
thinker_outputs = source_outputs
image_patch_token_id, image_end_token_id, num_query_tokens = _resolve_token_ids_from_stage_or_defaults(stage=None)
@@ -417,17 +413,10 @@ def thinker2imagegen(
if ref_image is not None:
extra["reference_image"] = ref_image
# ByT5 glyph text: prefer sampling_params (stage-1 explicit API), then
# prompt.image_gen_extra_args (serving_chat path), then auto-extract
# from quoted prompt text.
# ByT5 glyph text: prefer the flat byte5_text on sampling_params
# (stage-1 explicit API), else auto-extract from quoted prompt text.
prompt_text = prompt.get("prompt", "")
sp_byte5 = _extract_byte5_from_sampling_params(sampling_params)
if sp_byte5:
byte5_texts: list[str] | None = sp_byte5
else:
ig_extra = prompt.get("image_gen_extra_args") or {}
cand = ig_extra.get("byte5_text")
byte5_texts = cand if isinstance(cand, list) and cand else None
byte5_texts: list[str] | None = _extract_byte5_from_sampling_params(sampling_params)
if byte5_texts:
extra["byte5_text"] = [