wehub-resource-sync
479 lines
16 KiB
Python
479 lines
16 KiB
Python
import asyncio
|
|
import base64
|
|
import mimetypes
|
|
import os
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
import aiohttp
|
|
from tqdm import tqdm
|
|
|
|
DEFAULT_EDITS_BOT_TASK = "think"
|
|
|
|
|
|
@dataclass
|
|
class RequestFuncInput:
|
|
prompt: str
|
|
api_url: str
|
|
model: str
|
|
width: int | None = None
|
|
height: int | None = None
|
|
num_frames: int | None = None
|
|
num_inference_steps: int | None = None
|
|
seed: int | None = None
|
|
fps: int | None = None
|
|
timestamp: float | None = None
|
|
slo_ms: float | None = None
|
|
extra_body: dict[str, Any] = field(default_factory=dict)
|
|
image_paths: list[str] | None = None
|
|
request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
|
default_bot_task: str | None = DEFAULT_EDITS_BOT_TASK
|
|
|
|
|
|
@dataclass
|
|
class RequestFuncOutput:
|
|
success: bool = False
|
|
latency: float = 0.0
|
|
error: str = ""
|
|
start_time: float = 0.0
|
|
response_body: dict[str, Any] = field(default_factory=dict)
|
|
stage_durations: dict[str, float] = field(default_factory=dict)
|
|
peak_memory_mb: float = 0.0
|
|
slo_achieved: bool | None = None
|
|
|
|
|
|
def _guess_mime_type(path: str) -> str:
|
|
mime, _ = mimetypes.guess_type(path)
|
|
return mime or "application/octet-stream"
|
|
|
|
|
|
def _encode_image_as_data_url(path: str) -> str:
|
|
with open(path, "rb") as f:
|
|
encoded = base64.b64encode(f.read()).decode("utf-8")
|
|
mime = _guess_mime_type(path)
|
|
return f"data:{mime};base64,{encoded}"
|
|
|
|
|
|
async def async_request_image_edits(
|
|
input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
enable_diffusion_pipeline_profiler: bool = False,
|
|
) -> RequestFuncOutput:
|
|
"""POST /v1/images/edits (multipart)."""
|
|
del enable_diffusion_pipeline_profiler
|
|
output = RequestFuncOutput()
|
|
output.start_time = time.perf_counter()
|
|
|
|
extra_body = dict(input.extra_body)
|
|
width = input.width or extra_body.get("width") or 1024
|
|
height = input.height or extra_body.get("height") or 1024
|
|
edits_url = input.api_url
|
|
|
|
form = aiohttp.FormData()
|
|
form.add_field("model", input.model)
|
|
form.add_field("prompt", input.prompt)
|
|
form.add_field("size", f"{width}x{height}")
|
|
form.add_field("response_format", "b64_json")
|
|
|
|
if input.num_inference_steps is not None:
|
|
form.add_field("num_inference_steps", str(input.num_inference_steps))
|
|
elif extra_body.get("num_inference_steps") is not None:
|
|
form.add_field("num_inference_steps", str(extra_body["num_inference_steps"]))
|
|
|
|
if input.seed is not None:
|
|
form.add_field("seed", str(input.seed))
|
|
elif extra_body.get("seed") is not None:
|
|
form.add_field("seed", str(extra_body["seed"]))
|
|
|
|
if extra_body.get("guidance_scale") is not None:
|
|
form.add_field("guidance_scale", str(extra_body["guidance_scale"]))
|
|
if extra_body.get("negative_prompt") is not None:
|
|
form.add_field("negative_prompt", str(extra_body["negative_prompt"]))
|
|
if extra_body.get("true_cfg_scale") is not None:
|
|
form.add_field("true_cfg_scale", str(extra_body["true_cfg_scale"]))
|
|
if extra_body.get("sys_type") is not None:
|
|
form.add_field("sys_type", str(extra_body["sys_type"]))
|
|
if extra_body.get("system_prompt") is not None:
|
|
form.add_field("system_prompt", str(extra_body["system_prompt"]))
|
|
|
|
bot_task = extra_body.get("bot_task")
|
|
if bot_task is None and input.default_bot_task is not None:
|
|
bot_task = input.default_bot_task
|
|
if bot_task is not None:
|
|
form.add_field("bot_task", str(bot_task))
|
|
|
|
assert input.image_paths is not None
|
|
for img_path in input.image_paths:
|
|
if not os.path.exists(img_path):
|
|
output.error = f"Image file not found: {img_path}"
|
|
output.success = False
|
|
if pbar:
|
|
pbar.update(1)
|
|
return output
|
|
with open(img_path, "rb") as img_f:
|
|
image_bytes = img_f.read()
|
|
form.add_field(
|
|
"image",
|
|
image_bytes,
|
|
filename=os.path.basename(img_path),
|
|
content_type=_guess_mime_type(img_path),
|
|
)
|
|
|
|
try:
|
|
async with session.post(edits_url, data=form) as response:
|
|
if response.status == 200:
|
|
resp_json = await response.json()
|
|
output.response_body = resp_json
|
|
output.success = True
|
|
else:
|
|
output.error = f"HTTP {response.status}: {await response.text()}"
|
|
output.success = False
|
|
except Exception as e:
|
|
output.error = str(e)
|
|
output.success = False
|
|
|
|
output.latency = time.perf_counter() - output.start_time
|
|
|
|
if output.success and input.slo_ms is not None:
|
|
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
|
|
|
|
if pbar:
|
|
pbar.update(1)
|
|
return output
|
|
|
|
|
|
async def async_request_chat_completions(
|
|
input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
enable_diffusion_pipeline_profiler: bool = False,
|
|
) -> RequestFuncOutput:
|
|
output = RequestFuncOutput()
|
|
output.start_time = time.perf_counter()
|
|
|
|
extra_body = dict(input.extra_body)
|
|
if input.width and input.height:
|
|
extra_body.setdefault("height", input.height)
|
|
extra_body.setdefault("width", input.width)
|
|
if input.num_frames:
|
|
extra_body.setdefault("num_frames", input.num_frames)
|
|
if input.num_inference_steps:
|
|
extra_body.setdefault("num_inference_steps", input.num_inference_steps)
|
|
if input.seed is not None:
|
|
extra_body.setdefault("seed", input.seed)
|
|
if input.fps:
|
|
extra_body.setdefault("fps", input.fps)
|
|
|
|
if input.image_paths and len(input.image_paths) > 0:
|
|
content = []
|
|
if input.prompt:
|
|
content.append({"type": "text", "text": input.prompt})
|
|
for img_path in input.image_paths:
|
|
if not os.path.exists(img_path):
|
|
output.error = f"Image file not found: {img_path}"
|
|
output.success = False
|
|
if pbar:
|
|
pbar.update(1)
|
|
return output
|
|
content.append(
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": _encode_image_as_data_url(img_path)},
|
|
}
|
|
)
|
|
messages = [{"role": "user", "content": content}]
|
|
else:
|
|
messages = [{"role": "user", "content": input.prompt}]
|
|
|
|
payload = {
|
|
"model": input.model,
|
|
"messages": messages,
|
|
}
|
|
if extra_body:
|
|
payload["extra_body"] = extra_body
|
|
|
|
try:
|
|
async with session.post(input.api_url, json=payload) as response:
|
|
if response.status == 200:
|
|
resp_json = await response.json()
|
|
output.response_body = resp_json
|
|
output.success = True
|
|
try:
|
|
choices = resp_json.get("choices", [])
|
|
if choices and isinstance(choices, list):
|
|
msg = choices[0].get("message", {})
|
|
if isinstance(msg, dict):
|
|
content = msg.get("content", [])
|
|
if content and isinstance(content, list) and len(content) > 0:
|
|
first_item = content[0]
|
|
if isinstance(first_item, dict):
|
|
output.stage_durations = first_item.get("stage_durations") or {}
|
|
output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0)
|
|
except (IndexError, TypeError, AttributeError):
|
|
pass
|
|
|
|
if (not output.stage_durations or output.peak_memory_mb == 0.0) and isinstance(
|
|
resp_json.get("metrics"), dict
|
|
):
|
|
m = resp_json["metrics"]
|
|
if not output.stage_durations and isinstance(m.get("stage_durations"), dict):
|
|
output.stage_durations = m.get("stage_durations") or {}
|
|
if output.peak_memory_mb == 0.0 and m.get("peak_memory_mb") is not None:
|
|
try:
|
|
output.peak_memory_mb = float(m.get("peak_memory_mb") or 0.0)
|
|
except (TypeError, ValueError):
|
|
pass
|
|
else:
|
|
output.error = f"HTTP {response.status}: {await response.text()}"
|
|
output.success = False
|
|
except Exception as e:
|
|
output.error = str(e)
|
|
output.success = False
|
|
|
|
output.latency = time.perf_counter() - output.start_time
|
|
|
|
if output.success and input.slo_ms is not None:
|
|
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
|
|
|
|
if pbar:
|
|
pbar.update(1)
|
|
return output
|
|
|
|
|
|
async def async_request_openai_image_generations(
|
|
input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
) -> RequestFuncOutput:
|
|
"""
|
|
Send request to OpenAI's /v1/images/generations endpoint.
|
|
"""
|
|
output = RequestFuncOutput()
|
|
output.start_time = time.perf_counter()
|
|
|
|
# Build size string from width/height
|
|
width = input.width or 1024
|
|
height = input.height or 1024
|
|
size = f"{width}x{height}"
|
|
|
|
payload: dict[str, Any] = {
|
|
"model": input.model,
|
|
"prompt": input.prompt,
|
|
"n": 1,
|
|
"size": size,
|
|
"response_format": "b64_json",
|
|
}
|
|
|
|
# Add optional parameters
|
|
if input.seed is not None:
|
|
payload["seed"] = input.seed
|
|
if input.num_inference_steps is not None:
|
|
payload["num_inference_steps"] = input.num_inference_steps
|
|
|
|
# Add any extra body parameters
|
|
if input.extra_body:
|
|
for key, value in input.extra_body.items():
|
|
if key not in payload:
|
|
payload[key] = value
|
|
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": "Bearer EMPTY",
|
|
}
|
|
|
|
try:
|
|
async with session.post(input.api_url, json=payload, headers=headers) as response:
|
|
if response.status == 200:
|
|
resp_json = await response.json()
|
|
output.response_body = resp_json
|
|
output.success = True
|
|
# Check for usage/memory info if available
|
|
if "usage" in resp_json and "peak_memory_mb" in resp_json.get("usage", {}):
|
|
output.peak_memory_mb = resp_json["usage"]["peak_memory_mb"]
|
|
else:
|
|
output.error = f"HTTP {response.status}: {await response.text()}"
|
|
output.success = False
|
|
except Exception as e:
|
|
output.error = str(e)
|
|
output.success = False
|
|
|
|
output.latency = time.perf_counter() - output.start_time
|
|
|
|
if output.success and input.slo_ms is not None:
|
|
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
|
|
|
|
if pbar:
|
|
pbar.update(1)
|
|
return output
|
|
|
|
|
|
async def async_request_v1_videos(
|
|
input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
) -> RequestFuncOutput:
|
|
output = RequestFuncOutput()
|
|
output.start_time = time.perf_counter()
|
|
|
|
files = dict(input.extra_body)
|
|
if input.prompt:
|
|
files.setdefault("prompt", input.prompt)
|
|
if input.width and input.height:
|
|
files.setdefault("height", input.height)
|
|
files.setdefault("width", input.width)
|
|
if input.num_frames:
|
|
files.setdefault("num_frames", input.num_frames)
|
|
if input.num_inference_steps:
|
|
files.setdefault("num_inference_steps", input.num_inference_steps)
|
|
if input.seed is not None:
|
|
files.setdefault("seed", input.seed)
|
|
if input.fps:
|
|
files.setdefault("fps", input.fps)
|
|
|
|
form = aiohttp.FormData()
|
|
for k, v in files.items():
|
|
form.add_field(k, str(v))
|
|
|
|
image_file = None
|
|
if input.image_paths and len(input.image_paths) > 0:
|
|
image_path = input.image_paths[0]
|
|
image_file = open(image_path, "rb")
|
|
form.add_field(
|
|
"input_reference",
|
|
image_file,
|
|
filename=os.path.basename(image_path),
|
|
content_type="application/octet-stream",
|
|
)
|
|
|
|
job_id = None
|
|
job_status = None
|
|
poll_json = {}
|
|
resp_json = {}
|
|
|
|
try:
|
|
# invoke a post request (POST /v1/videos)
|
|
async with session.post(input.api_url, data=form) as response:
|
|
if response.status == 200:
|
|
resp_json = await response.json()
|
|
job_id = resp_json.get("id")
|
|
job_status = resp_json.get("status")
|
|
if not job_id or not job_status:
|
|
output.error = "API response missing job 'id' or 'status' field."
|
|
output.success = False
|
|
return output
|
|
else:
|
|
output.error = f"HTTP {response.status}: {await response.text()}"
|
|
output.success = False
|
|
return output
|
|
|
|
# invoke a poll request (GET /v1/videos/{video_id})
|
|
poll_interval = 2.0 # Unit(s)
|
|
timeout_seconds = 600.0
|
|
deadline = time.perf_counter() + timeout_seconds
|
|
job_url = f"{input.api_url}/{job_id}"
|
|
|
|
while job_status not in {"completed", "failed"}:
|
|
await asyncio.sleep(poll_interval)
|
|
|
|
async with session.get(job_url) as poll_response:
|
|
if poll_response.status != 200:
|
|
output.error = f"Polling failed HTTP {poll_response.status}: {await poll_response.text()}"
|
|
output.success = False
|
|
return output
|
|
|
|
poll_json = await poll_response.json()
|
|
job_status = poll_json.get("status")
|
|
|
|
if time.perf_counter() >= deadline:
|
|
output.error = f"Timed out waiting for video job {job_id} to complete."
|
|
output.success = False
|
|
return output
|
|
|
|
if job_status == "failed":
|
|
output.error = f"Video job failed: {poll_json}"
|
|
output.success = False
|
|
return output
|
|
|
|
# invoke a get request (GET /v1/videos/{video_id}/content)
|
|
content_url = f"{job_url}/content"
|
|
async with session.get(content_url) as content_response:
|
|
if content_response.status != 200:
|
|
output.error = (
|
|
f"Content retrieval failed HTTP {content_response.status}: {await content_response.text()}"
|
|
)
|
|
output.success = False
|
|
return output
|
|
|
|
video_bytes = await content_response.read()
|
|
output.response_body = video_bytes
|
|
output.success = True
|
|
if "stage_durations" in poll_json:
|
|
output.stage_durations = poll_json["stage_durations"] or {}
|
|
if "peak_memory_mb" in poll_json:
|
|
output.peak_memory_mb = poll_json["peak_memory_mb"]
|
|
elif "peak_memory_mb" in resp_json:
|
|
output.peak_memory_mb = resp_json["peak_memory_mb"]
|
|
except Exception as e:
|
|
output.error = str(e)
|
|
output.success = False
|
|
finally:
|
|
if image_file is not None:
|
|
image_file.close()
|
|
|
|
if job_id is not None:
|
|
try:
|
|
async with session.delete(f"{input.api_url}/{job_id}") as _:
|
|
pass
|
|
except Exception as e:
|
|
print(f"Failed to clean up video job {job_id}: {e}")
|
|
|
|
output.latency = time.perf_counter() - output.start_time
|
|
|
|
if output.success and input.slo_ms is not None:
|
|
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
|
|
|
|
if pbar:
|
|
pbar.update(1)
|
|
return output
|
|
|
|
|
|
LEGACY_BACKEND_ENDPOINT_ALIASES = {
|
|
"vllm-omni": "/v1/chat/completions",
|
|
"openai": "/v1/images/generations",
|
|
}
|
|
|
|
|
|
def normalize_endpoint(value: str) -> str:
|
|
endpoint = str(value).strip()
|
|
if not endpoint:
|
|
raise ValueError("endpoint must not be empty.")
|
|
endpoint = LEGACY_BACKEND_ENDPOINT_ALIASES.get(
|
|
endpoint,
|
|
LEGACY_BACKEND_ENDPOINT_ALIASES.get(endpoint.lstrip("/"), endpoint),
|
|
)
|
|
if not endpoint.startswith("/"):
|
|
endpoint = f"/{endpoint}"
|
|
return endpoint
|
|
|
|
|
|
def endpoint_filename_token(value: str) -> str:
|
|
token = normalize_endpoint(value).lstrip("/")
|
|
for bad in ("/", "\\", ":", "*", "?", '"', "<", ">", "|"):
|
|
token = token.replace(bad, "_")
|
|
return token or "endpoint"
|
|
|
|
|
|
backends_function_mapping = {
|
|
"2i": {
|
|
"/v1/chat/completions": (async_request_chat_completions, "/v1/chat/completions"),
|
|
"/v1/images/generations": (async_request_openai_image_generations, "/v1/images/generations"),
|
|
"/v1/images/edits": (async_request_image_edits, "/v1/images/edits"),
|
|
},
|
|
"2v": {
|
|
"/v1/videos": (async_request_v1_videos, "/v1/videos"),
|
|
},
|
|
}
|