From fb3abb38fb1916d0a963fe736b22285a7a508b48 Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Wed, 24 Jun 2026 17:01:58 +0800 Subject: [PATCH] [refactor] Flatten Ming-flash-omni-2.0 imagegen args (#4587) Signed-off-by: Yuanheng --- .../online_serving/ming_flash_omni/README.md | 53 +++++++++++++------ .../ming_flash_omni/pipeline_ming_imagegen.py | 33 ++++++------ vllm_omni/entrypoints/openai/serving_chat.py | 15 ------ .../stage_input_processors/ming_flash_omni.py | 39 +++++--------- 4 files changed, 69 insertions(+), 71 deletions(-) diff --git a/examples/online_serving/ming_flash_omni/README.md b/examples/online_serving/ming_flash_omni/README.md index 76b0bef9d..d60ba8156 100644 --- a/examples/online_serving/ming_flash_omni/README.md +++ b/examples/online_serving/ming_flash_omni/README.md @@ -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) diff --git a/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py index 0da26aff8..2c77351dd 100644 --- a/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py +++ b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py @@ -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 "". '`` 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 "". '` + 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 "". `` 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: diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index f5449837a..9e1fba61c 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -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 diff --git a/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py b/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py index 8f0df7269..0c38bb807 100644 --- a/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py @@ -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"] = [