Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
92 lines
2.4 KiB
Python
92 lines
2.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
|
|
|
|
from PIL import Image
|
|
|
|
|
|
def build_text_to_image_prompt(
|
|
prompt: str,
|
|
negative_prompt: str | None,
|
|
height: int | None = None,
|
|
width: int | None = None,
|
|
) -> dict[str, Any]:
|
|
text_prompt: dict[str, Any] = {
|
|
"prompt": f"<|im_start|>{prompt}<|im_end|>",
|
|
"modalities": ["image"],
|
|
"mm_processor_kwargs": {
|
|
"target_h": height,
|
|
"target_w": width,
|
|
"modalities": ["image"],
|
|
},
|
|
}
|
|
if negative_prompt is not None:
|
|
text_prompt["negative_prompt"] = negative_prompt
|
|
return text_prompt
|
|
|
|
|
|
def build_image_to_image_prompt(
|
|
prompt: str,
|
|
negative_prompt: str | None,
|
|
input_image: Image.Image | list[Image.Image],
|
|
height: int | None = None,
|
|
width: int | None = None,
|
|
) -> dict[str, Any]:
|
|
img_prompt: dict[str, Any] = {
|
|
"prompt": f"<|fim_middle|><|im_start|>{prompt}<|im_end|>",
|
|
"modalities": ["img2img"],
|
|
"multi_modal_data": {"img2img": input_image},
|
|
"mm_processor_kwargs": {"modalities": ["img2img"]},
|
|
}
|
|
if height is not None:
|
|
img_prompt["mm_processor_kwargs"]["target_h"] = height
|
|
if width is not None:
|
|
img_prompt["mm_processor_kwargs"]["target_w"] = width
|
|
if negative_prompt is not None:
|
|
img_prompt["negative_prompt"] = negative_prompt
|
|
return img_prompt
|
|
|
|
|
|
def build_x_to_text_prompt(
|
|
model: str,
|
|
prompt: str,
|
|
has_image: bool,
|
|
) -> tuple[dict[str, Any], list[int] | None]:
|
|
"""Build BAGEL's chat prompt for text-output tasks."""
|
|
del model
|
|
image_token = "<|image_pad|>\n" if has_image else ""
|
|
return (
|
|
{
|
|
"prompt": f"<|im_start|>user\n{image_token}{prompt}<|im_end|>\n<|im_start|>assistant\n",
|
|
"modalities": ["text"],
|
|
},
|
|
None,
|
|
)
|
|
|
|
|
|
BAGEL_EXTRA_BODY_PARAMS = frozenset(
|
|
{
|
|
"cfg_text_scale",
|
|
"cfg_img_scale",
|
|
"cfg_interval",
|
|
"cfg_renorm_type",
|
|
"cfg_renorm_min",
|
|
"negative_prompt",
|
|
"think",
|
|
"max_think_tokens",
|
|
"do_sample",
|
|
"text_temperature",
|
|
"timestep_shift",
|
|
}
|
|
)
|
|
BAGEL_EXTRA_OUTPUT_PARAMS = frozenset(
|
|
{
|
|
"text_output",
|
|
"think_text",
|
|
}
|
|
)
|
|
BAGEL_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True
|