474 lines
19 KiB
Python
474 lines
19 KiB
Python
"""
|
|
An high-level API client adapter that forwards ComfyUI inputs to vLLM-Omni's REST API,
|
|
and transforms the API responses back to ComfyUI formats.
|
|
|
|
The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@dougbtv).
|
|
Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
|
|
"""
|
|
|
|
from typing import Any
|
|
|
|
import aiohttp
|
|
import av.error
|
|
import torch
|
|
from comfy_api.input import AudioInput, VideoInput
|
|
|
|
from .format import (
|
|
audio_to_base64,
|
|
base64_to_audio,
|
|
base64_to_image_tensor,
|
|
bytes_to_audio,
|
|
image_tensor_to_base64,
|
|
image_tensor_to_png_bytes,
|
|
video_to_base64,
|
|
)
|
|
from .logger import get_logger, pretty_printer
|
|
from .models import lookup_model_spec
|
|
from .types import AudioFormat
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class VLLMOmniClient:
|
|
def __init__(self, base_url: str, timeout: float = 300.0):
|
|
self.base_url = base_url
|
|
self.timeout = aiohttp.ClientTimeout(total=timeout)
|
|
|
|
async def generate_image(
|
|
self,
|
|
*,
|
|
model: str,
|
|
prompt: str,
|
|
width: int,
|
|
height: int,
|
|
negative_prompt: str | None = None,
|
|
sampling_params: dict | None = None,
|
|
) -> torch.Tensor:
|
|
"""Run text-to-image generation via DALLE API"""
|
|
await self._check_model_exist(model)
|
|
|
|
size = f"{width}x{height}"
|
|
payload = {
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"size": size,
|
|
"response_format": "b64_json",
|
|
}
|
|
if negative_prompt:
|
|
payload["negative_prompt"] = negative_prompt
|
|
if sampling_params is not None:
|
|
# Only select specific sampling params
|
|
for k in (
|
|
"n",
|
|
"num_inference_steps",
|
|
"guidance_scale",
|
|
"true_cfg_scale",
|
|
"vae_use_slicing",
|
|
):
|
|
if k in sampling_params and sampling_params[k] is not None:
|
|
payload[k] = sampling_params[k]
|
|
logger.debug("img gen payload: %s", payload)
|
|
|
|
url = self.base_url + "/images/generations"
|
|
async with aiohttp.ClientSession(timeout=self.timeout) as session:
|
|
try:
|
|
async with session.post(
|
|
url,
|
|
json=payload,
|
|
headers={"Content-Type": "application/json"},
|
|
) as response:
|
|
if not response.ok:
|
|
error_text = await response.text()
|
|
raise (ValueError if response.status < 500 else RuntimeError)(
|
|
f"vLLM-Omni API returned status {response.status}: {error_text}"
|
|
)
|
|
|
|
try:
|
|
data = await response.json()
|
|
except aiohttp.ContentTypeError as e:
|
|
raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
|
|
if "data" not in data:
|
|
raise RuntimeError("API response missing 'data' field - expected OpenAI DALL-E format")
|
|
if not data["data"]:
|
|
raise RuntimeError("API returned empty data array")
|
|
|
|
image_tensors = []
|
|
for idx, img in enumerate(data["data"]):
|
|
if "b64_json" not in img:
|
|
raise RuntimeError(f"API returned image #{idx} without 'b64_json' field")
|
|
base64_str = img["b64_json"]
|
|
tensor = base64_to_image_tensor(base64_str)
|
|
image_tensors.append(tensor)
|
|
logger.debug("Image #%d has shape %s", idx, tensor.shape)
|
|
|
|
batch_tensor = torch.stack(image_tensors, dim=0)
|
|
logger.debug("batch_tensor output has shape: %s", batch_tensor.shape)
|
|
return batch_tensor
|
|
|
|
except aiohttp.ClientError as e:
|
|
raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
|
|
|
|
async def edit_image(
|
|
self,
|
|
*,
|
|
model: str,
|
|
prompt: str,
|
|
image: torch.Tensor,
|
|
width: int,
|
|
height: int,
|
|
negative_prompt: str | None = None,
|
|
mask: torch.Tensor | None = None,
|
|
sampling_params: dict | None = None,
|
|
) -> torch.Tensor:
|
|
"""Run image editing via DALLE API"""
|
|
await self._check_model_exist(model)
|
|
|
|
size = f"{width}x{height}"
|
|
image_filename = "image.png" # Required for multipart form
|
|
form = aiohttp.FormData()
|
|
form.add_field("model", model)
|
|
form.add_field(
|
|
"image",
|
|
image_tensor_to_png_bytes(image, image_filename),
|
|
filename=image_filename,
|
|
content_type="image/png",
|
|
)
|
|
form.add_field("prompt", prompt)
|
|
form.add_field("size", size)
|
|
if negative_prompt:
|
|
form.add_field("negative_prompt", negative_prompt)
|
|
if sampling_params is not None:
|
|
# Only select specific sampling params
|
|
for k in ("n", "num_inference_steps", "guidance_scale", "true_cfg_scale"):
|
|
if k in sampling_params and sampling_params[k] is not None:
|
|
form.add_field(k, str(sampling_params[k]))
|
|
if mask is not None:
|
|
mask_filename = "mask.png"
|
|
form.add_field(
|
|
"mask",
|
|
image_tensor_to_png_bytes(mask, mask_filename),
|
|
filename=mask_filename,
|
|
content_type="image/png",
|
|
)
|
|
|
|
url = self.base_url + "/images/edits"
|
|
async with aiohttp.ClientSession(timeout=self.timeout) as session:
|
|
try:
|
|
async with session.post(url, data=form) as response:
|
|
if not response.ok:
|
|
error_text = await response.text()
|
|
raise (ValueError if response.status < 500 else RuntimeError)(
|
|
f"vLLM-Omni API returned status {response.status}: {error_text}"
|
|
)
|
|
|
|
try:
|
|
data = await response.json()
|
|
except aiohttp.ContentTypeError as e:
|
|
raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
|
|
|
|
if "data" not in data:
|
|
raise RuntimeError("API response missing 'data' field - expected OpenAI DALL-E format")
|
|
if not data["data"]:
|
|
raise RuntimeError("API returned empty data array")
|
|
|
|
image_tensors = []
|
|
for idx, img in enumerate(data["data"]):
|
|
if "b64_json" not in img:
|
|
raise RuntimeError(f"API returned image #{idx} without 'b64_json' field")
|
|
base64_str = img["b64_json"]
|
|
tensor = base64_to_image_tensor(base64_str)
|
|
image_tensors.append(tensor)
|
|
|
|
return torch.stack(image_tensors, dim=0)
|
|
|
|
except aiohttp.ClientError as e:
|
|
raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
|
|
|
|
async def generate_image_chat_completion(
|
|
self,
|
|
*,
|
|
model: str,
|
|
prompt: str,
|
|
negative_prompt: str | None = None,
|
|
image: torch.Tensor | None = None,
|
|
audio: AudioInput | None = None,
|
|
video: VideoInput | None = None,
|
|
sampling_params: dict | list[dict] | None = None,
|
|
) -> torch.Tensor:
|
|
payload = VLLMOmniClient._prepare_chat_completion_messages(
|
|
model=model,
|
|
prompt=prompt,
|
|
negative_prompt=negative_prompt,
|
|
image=image,
|
|
audio=audio,
|
|
video=video,
|
|
sampling_params=sampling_params,
|
|
modalities=["image"],
|
|
)
|
|
choices = await self._generate_base_chat_completion(model, payload)
|
|
|
|
image_tensors = []
|
|
for idx, img_content in enumerate(choices[0]["message"]["content"]):
|
|
base64_str = img_content.get("image_url", {}).get("url", "")
|
|
if not base64_str:
|
|
raise RuntimeError(f"API returned image #{idx} without image url")
|
|
tensor = base64_to_image_tensor(base64_str)
|
|
image_tensors.append(tensor)
|
|
|
|
return torch.stack(image_tensors, dim=0)
|
|
|
|
async def generate_comprehension_chat_completion(
|
|
self,
|
|
*,
|
|
model: str,
|
|
prompt: str,
|
|
image: torch.Tensor | None = None,
|
|
audio: AudioInput | None = None,
|
|
video: VideoInput | None = None,
|
|
sampling_params: dict | list[dict] | None = None,
|
|
modalities: list[str] = ["text", "audio"],
|
|
**extra_body,
|
|
) -> tuple[str | None, AudioInput | None]:
|
|
# Response may contain two choices: one with text, one with audio
|
|
payload = VLLMOmniClient._prepare_chat_completion_messages(
|
|
model=model,
|
|
prompt=prompt,
|
|
negative_prompt=None,
|
|
image=image,
|
|
audio=audio,
|
|
video=video,
|
|
sampling_params=sampling_params,
|
|
modalities=modalities,
|
|
**extra_body,
|
|
)
|
|
|
|
choices = await self._generate_base_chat_completion(model, payload)
|
|
text_response = None
|
|
audio_base64 = None
|
|
for choice in choices:
|
|
try:
|
|
text_response = choice["message"]["content"]
|
|
except (KeyError, TypeError):
|
|
# Either this case (text response) or the audio response case will be hit. Checking None's later.
|
|
pass
|
|
try:
|
|
audio_base64 = choice["message"]["audio"]["data"]
|
|
except (KeyError, TypeError):
|
|
# Either this case (text response) or the audio response case will be hit. Checking None's later.
|
|
pass
|
|
if audio_base64 is None and text_response is None:
|
|
raise RuntimeError(
|
|
"API response missing both '.message.audio' and 'message.content' fields."
|
|
f"The choices object is {choices}"
|
|
)
|
|
if audio_base64 is not None:
|
|
audio = base64_to_audio(audio_base64)
|
|
logger.debug(
|
|
"audio sample rate %d, audio shape %s, duration in second %f",
|
|
audio["sample_rate"],
|
|
audio["waveform"].shape,
|
|
audio["waveform"].shape[2] / audio["sample_rate"],
|
|
)
|
|
else:
|
|
audio = None
|
|
return text_response, audio
|
|
|
|
async def generate_speech(
|
|
self,
|
|
*,
|
|
model: str,
|
|
input: str,
|
|
voice: str,
|
|
response_format: AudioFormat,
|
|
speed: float,
|
|
**extra_params,
|
|
) -> AudioInput:
|
|
await self._check_model_exist(model)
|
|
|
|
ref_audio: AudioInput | None = extra_params.pop("ref_audio", None)
|
|
|
|
payload = {
|
|
"model": model,
|
|
"input": input,
|
|
"voice": voice,
|
|
"response_format": response_format,
|
|
"speed": speed,
|
|
**extra_params,
|
|
}
|
|
|
|
if ref_audio is not None:
|
|
audio_base64 = audio_to_base64(ref_audio)
|
|
payload["ref_audio"] = audio_base64
|
|
|
|
logger.debug("Omni TTS payload: %s", pretty_printer.pformat(payload))
|
|
|
|
url = self.base_url + "/audio/speech"
|
|
async with aiohttp.ClientSession(timeout=self.timeout) as session:
|
|
try:
|
|
async with session.post(
|
|
url,
|
|
json=payload,
|
|
headers={"Content-Type": "application/json"},
|
|
) as response:
|
|
if not response.ok:
|
|
error_text = await response.text()
|
|
raise (ValueError if response.status < 500 else RuntimeError)(
|
|
f"vLLM-Omni API returned status {response.status}: {error_text}"
|
|
)
|
|
|
|
try:
|
|
audio_bytes = await response.read()
|
|
except aiohttp.ContentTypeError as e:
|
|
raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
|
|
|
|
try:
|
|
audio = bytes_to_audio(audio_bytes)
|
|
except av.error.InvalidDataError as e:
|
|
raise ValueError(
|
|
f"Invalid audio data received from vLLM-Omni: {e}"
|
|
"Check if you have input unsupported arguments (such as 'voice')"
|
|
)
|
|
return audio
|
|
|
|
except aiohttp.ClientError as e:
|
|
raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
|
|
|
|
async def _generate_base_chat_completion(self, model: str, payload: dict[str, Any]) -> list[dict[str, Any]]:
|
|
logger.debug("Omni payload: %s", pretty_printer.pformat(payload))
|
|
await self._check_model_exist(model)
|
|
|
|
url = self.base_url + "/chat/completions"
|
|
async with aiohttp.ClientSession(timeout=self.timeout) as session:
|
|
try:
|
|
async with session.post(
|
|
url,
|
|
json=payload,
|
|
headers={"Content-Type": "application/json"},
|
|
) as response:
|
|
if not response.ok:
|
|
error_text = await response.text()
|
|
raise (ValueError if response.status < 500 else RuntimeError)(
|
|
f"vLLM-Omni API returned status {response.status}: {error_text}"
|
|
)
|
|
|
|
try:
|
|
data = await response.json()
|
|
except aiohttp.ContentTypeError as e:
|
|
raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
|
|
|
|
logger.debug(
|
|
"chat completion response: %s",
|
|
pretty_printer.pformat(data),
|
|
)
|
|
|
|
try:
|
|
return data["choices"]
|
|
except (KeyError, TypeError):
|
|
raise RuntimeError("Invalid JSON response from vLLM-Omni: missing 'choices' field")
|
|
|
|
except aiohttp.ClientError as e:
|
|
raise RuntimeError(f"Network error connecting to vLLM-Omni at {self.base_url}: {e}")
|
|
|
|
async def _check_model_exist(self, model: str):
|
|
url = self.base_url + "/models"
|
|
async with aiohttp.ClientSession(timeout=self.timeout) as session:
|
|
try:
|
|
async with session.get(
|
|
url,
|
|
headers={"Content-Type": "application/json"},
|
|
) as response:
|
|
if not response.ok:
|
|
error_text = await response.text()
|
|
raise (ValueError if response.status < 500 else RuntimeError)(
|
|
f"vLLM-Omni API returned status {response.status} "
|
|
f"when getting hosted model list: {error_text}"
|
|
)
|
|
|
|
try:
|
|
data = await response.json()
|
|
except aiohttp.ContentTypeError as e:
|
|
raise RuntimeError(f"Invalid JSON response when getting hosted model list from vLLM-Omni: {e}")
|
|
|
|
except aiohttp.ClientError as e:
|
|
raise RuntimeError(f"Network error connecting to vLLM-Omni at {self.base_url}: {e}")
|
|
try:
|
|
model_list = data["data"]
|
|
model_found = next((True for m in model_list if m["id"] == model), False)
|
|
except (KeyError, TypeError):
|
|
raise RuntimeError(f"Invalid JSON response of the hosted model list: {data}")
|
|
|
|
if not model_found:
|
|
raise ValueError(f"Model {model} not served at {self.base_url}.")
|
|
|
|
@staticmethod
|
|
def _prepare_chat_completion_messages(
|
|
*,
|
|
model: str,
|
|
prompt: str,
|
|
negative_prompt: str | None,
|
|
image: torch.Tensor | None = None,
|
|
audio: AudioInput | None = None,
|
|
video: VideoInput | None = None,
|
|
sampling_params: dict | list[dict] | None = None,
|
|
modalities: list[str] | None = None, # diffusion don't have this field
|
|
**extra_body,
|
|
):
|
|
message_content: list[dict] = [{"type": "text", "text": prompt}]
|
|
if image is not None:
|
|
message_content.append(
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": image_tensor_to_base64(image)},
|
|
}
|
|
)
|
|
if audio is not None:
|
|
message_content.append({"type": "audio_url", "audio_url": {"url": audio_to_base64(audio)}})
|
|
if video is not None:
|
|
message_content.append({"type": "video_url", "video_url": {"url": video_to_base64(video)}})
|
|
messages = [{"role": "user", "content": message_content}]
|
|
|
|
combined_extra_body: dict[str, Any] = {}
|
|
if sampling_params is not None:
|
|
spec, _ = lookup_model_spec(model)
|
|
is_single_sampling_param = isinstance(sampling_params, dict) or len(sampling_params) == 1
|
|
|
|
# Exclude internal key
|
|
if isinstance(sampling_params, dict):
|
|
sampling_params = {k: v for k, v in sampling_params.items() if k != "type"}
|
|
else:
|
|
sampling_params = [{k: v for k, v in sp.items() if k != "type"} for sp in sampling_params]
|
|
|
|
if (spec is None and is_single_sampling_param) or (spec is not None and spec["stages"] == ["diffusion"]):
|
|
# Diffusion format: extra_body directly contains sampling params.
|
|
# Validation should have taken care of matching sampling params' types.
|
|
# * Use this mode if the model is a simple one-stage diffusion model.
|
|
# * Fallback to this mode if model is not registered and a single sampling param is provided.
|
|
sampling_params = sampling_params if isinstance(sampling_params, dict) else sampling_params[0]
|
|
combined_extra_body: dict[str, Any] = {**sampling_params}
|
|
if "n" in combined_extra_body:
|
|
combined_extra_body["num_outputs_per_prompt"] = combined_extra_body["n"]
|
|
del combined_extra_body["n"]
|
|
else:
|
|
# Use AR style payload, extra_body has a sampling_params_list field
|
|
combined_extra_body: dict[str, Any] = {"sampling_params_list": sampling_params}
|
|
|
|
if negative_prompt:
|
|
combined_extra_body["negative_prompt"] = negative_prompt
|
|
|
|
if extra_body:
|
|
combined_extra_body.update(extra_body)
|
|
|
|
payload: dict[str, Any] = {"messages": messages, "model": model}
|
|
if combined_extra_body:
|
|
payload["extra_body"] = combined_extra_body
|
|
if modalities:
|
|
payload["modalities"] = modalities
|
|
|
|
spec, _ = lookup_model_spec(model)
|
|
if spec:
|
|
preprocessor = spec.get("payload_preprocessor", None)
|
|
if preprocessor is not None:
|
|
payload = preprocessor(payload)
|
|
|
|
return payload
|