102 lines
3.4 KiB
Python
102 lines
3.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
# AR image-grid patch size (pixels per visual token). Mirrors the constant the
|
|
# bespoke MammothModa2 example used to convert image dimensions -> AR grid dims.
|
|
_PATCH_SIZE = 16
|
|
|
|
# Shared by Preview and Dev (see ``t2i_generation_config.json``).
|
|
_EOL_TOKEN_ID = 152064
|
|
_VISUAL_TOKEN_START_ID = 152072
|
|
_VISUAL_TOKEN_END_ID = 168456
|
|
|
|
_AR_SYSTEM_PROMPT = "You are a helpful image generator."
|
|
|
|
|
|
def build_text_to_image_prompt(
|
|
prompt: str,
|
|
negative_prompt: str | None,
|
|
height: int | None = None,
|
|
width: int | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Build the MammothModa2 AR-stage prompt for text-to-image generation.
|
|
|
|
Reproduces the prompt string and the structural ``additional_information``
|
|
that the former bespoke example (``run_mammothmoda2_t2i.py``) constructed,
|
|
deriving the AR image grid from ``height`` / ``width``.
|
|
|
|
Model-specific sampling knobs (``text_guidance_scale``, ``cfg_range``,
|
|
``num_inference_steps``) flow separately via ``extra_body`` -> ``extra_args``.
|
|
The structural token ids are included in ``additional_information`` so the
|
|
AR sampler can apply the same constraints for both Preview and Dev.
|
|
|
|
MammothModa2 t2i uses classifier-free guidance via ``text_guidance_scale``
|
|
and has no explicit negative-prompt path, so ``negative_prompt`` is accepted
|
|
for signature compatibility but not injected.
|
|
"""
|
|
h = height or 1024
|
|
w = width or 1024
|
|
ar_height = h // _PATCH_SIZE
|
|
ar_width = w // _PATCH_SIZE
|
|
|
|
return {
|
|
"prompt": (
|
|
f"<|im_start|>system\n{_AR_SYSTEM_PROMPT}<|im_end|>\n"
|
|
f"<|im_start|>user\n{prompt}<|im_end|>\n"
|
|
"<|im_start|>assistant\n"
|
|
f"<|image start|>{ar_width}*{ar_height}<|image token|>"
|
|
),
|
|
"additional_information": {
|
|
"omni_task": ["t2i"],
|
|
"ar_width": [ar_width],
|
|
"ar_height": [ar_height],
|
|
"image_height": [h],
|
|
"image_width": [w],
|
|
"eol_token_id": [_EOL_TOKEN_ID],
|
|
"visual_token_start_id": [_VISUAL_TOKEN_START_ID],
|
|
"visual_token_end_id": [_VISUAL_TOKEN_END_ID],
|
|
},
|
|
}
|
|
|
|
|
|
MAMMOTHMODA2_PREVIEW_EXTRA_BODY_PARAMS = frozenset(
|
|
{
|
|
"text_guidance_scale",
|
|
"cfg_range",
|
|
# MammothModa2's DiT stage consumes inputs via the kwargs interface rather
|
|
# than OmniDiffusionRequest, so the standard --num-inference-steps flag does
|
|
# not reach it; it is routed through extra_body like the CFG knobs.
|
|
"num_inference_steps",
|
|
}
|
|
)
|
|
|
|
|
|
def build_x_to_text_prompt(
|
|
model: str,
|
|
prompt: str,
|
|
has_image: bool,
|
|
) -> tuple[dict[str, Any], list[int] | None]:
|
|
"""Build MammothModa2's chat prompt for text-output tasks."""
|
|
del model
|
|
vision = "<|vision_start|><|image_pad|><|vision_end|>" if has_image else ""
|
|
return (
|
|
{
|
|
"prompt": (
|
|
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
|
|
f"<|im_start|>user\n{vision}{prompt}<|im_end|>\n"
|
|
"<|im_start|>assistant\n"
|
|
),
|
|
"modalities": ["text"],
|
|
"additional_information": {"omni_task": ["chat"]},
|
|
},
|
|
None,
|
|
)
|
|
|
|
|
|
MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS = frozenset()
|
|
MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True
|