74 lines
1.9 KiB
Python
74 lines
1.9 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
|
||
|
|
|
||
|
|
|
||
|
|
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
|