Benchmark and Bugfix for GLM-Image (#3024)
Signed-off-by: JaredforReal <w13431838023@gmail.com> Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Jared WenandHongsheng Liu authored and GitHub committed
1 parent
6b52db9e2b
commit
6e35bfbbae
18 files changed
+1836
-381
No files matched your search
@@ -0,0 +1,157 @@
|
||||
# GLM-Image Benchmarks
|
||||
|
||||
Benchmark GLM-Image T2I (text-to-image) and I2I (image-to-image) performance across three backends: HuggingFace baseline, vLLM-Omni offline, and vLLM-Omni online serving.
|
||||
|
||||
## Benchmarks
|
||||
|
||||
| Benchmark | Script | Description |
|
||||
|-----------|--------|-------------|
|
||||
| HuggingFace Baseline | `huggingface/inference.py` | Single-GPU transformers + diffusers pipeline |
|
||||
| vLLM-Omni Offline | `vllm-omni/inference.py` | Offline inference with continuous batching |
|
||||
| vLLM-Omni Online | `benchmark_glm_image.py` | Online serving via `/v1/chat/completions` |
|
||||
|
||||
## HuggingFace Baseline
|
||||
|
||||
Single-request sequential inference using the reference HuggingFace pipeline.
|
||||
|
||||
```bash
|
||||
# T2I
|
||||
CUDA_VISIBLE_DEVICES=0 python benchmarks/glm_image/huggingface/inference.py \
|
||||
--model-path /path/to/GLM-Image --mode t2i --num-prompts 10
|
||||
|
||||
# I2I
|
||||
CUDA_VISIBLE_DEVICES=0 python benchmarks/glm_image/huggingface/inference.py \
|
||||
--model-path /path/to/GLM-Image --mode i2i --num-prompts 10
|
||||
```
|
||||
|
||||
### Options
|
||||
|
||||
| Flag | Default | Description |
|
||||
|------|---------|-------------|
|
||||
| `--model-path` | `zai-org/GLM-Image` | Model path |
|
||||
| `--mode` | `t2i` | `t2i` or `i2i` |
|
||||
| `--dataset-path` | `prompt/prompt.json` | Path to prompt.json |
|
||||
| `--num-prompts` | `10` | Number of images to generate |
|
||||
| `--width` / `--height` | `1024` | Output image size |
|
||||
| `--num-inference-steps` | `50` | Diffusion denoising steps |
|
||||
| `--output-dir` | `benchmarks/glm_image/huggingface/outputs` | Output directory |
|
||||
| `--output-file` | - | JSON file for metrics |
|
||||
|
||||
## vLLM-Omni Offline
|
||||
|
||||
Multi-GPU offline inference with pipeline parallelism and continuous batching.
|
||||
|
||||
```bash
|
||||
# T2I
|
||||
CUDA_VISIBLE_DEVICES=0,1 python benchmarks/glm_image/vllm-omni/inference.py \
|
||||
--model-path /path/to/GLM-Image --mode t2i --num-prompts 10
|
||||
|
||||
# I2I
|
||||
CUDA_VISIBLE_DEVICES=0,1 python benchmarks/glm_image/vllm-omni/inference.py \
|
||||
--model-path /path/to/GLM-Image --mode i2i --num-prompts 10
|
||||
```
|
||||
|
||||
### Options
|
||||
|
||||
| Flag | Default | Description |
|
||||
|------|---------|-------------|
|
||||
| `--model-path` | `zai-org/GLM-Image` | Model path |
|
||||
| `--deploy-config` | - | Deploy config YAML |
|
||||
| `--mode` | `t2i` | `t2i` or `i2i` |
|
||||
| `--dataset-path` | `prompt/prompt.json` | Path to prompt.json |
|
||||
| `--num-prompts` | `10` | Number of images to generate |
|
||||
| `--width` / `--height` | `1024` | Output image size |
|
||||
| `--num-inference-steps` | `50` | Diffusion denoising steps |
|
||||
| `--output-dir` | `benchmarks/glm_image/vllm-omni/outputs` | Output directory |
|
||||
| `--output-file` | - | JSON file for metrics |
|
||||
| `--stage-init-timeout` | `600` | Stage initialization timeout (s) |
|
||||
|
||||
### Latency Computation
|
||||
|
||||
In offline mode all requests are submitted simultaneously and processed with continuous batching. The per-request latency is computed by summing the actual per-stage times (with `stage_0_gen_ms` diffed against the previous request to remove accumulated queue/scheduling wait).
|
||||
|
||||
## vLLM-Omni Online Serving
|
||||
|
||||
### Start the server
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1 vllm serve /path/to/GLM-Image \
|
||||
--omni --port 8091 --host 0.0.0.0 \
|
||||
--served-model-name glm-image
|
||||
```
|
||||
|
||||
### Run the benchmark
|
||||
|
||||
```bash
|
||||
# T2I
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode t2i --num-prompts 10 --model glm-image
|
||||
|
||||
# I2I
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode i2i --num-prompts 10 --model glm-image
|
||||
|
||||
# Custom dataset
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode i2i --dataset custom \
|
||||
--dataset-path prompts.json --num-prompts 5
|
||||
```
|
||||
|
||||
### Options
|
||||
|
||||
| Flag | Default | Description |
|
||||
|------|---------|-------------|
|
||||
| `--mode` | `t2i` | `t2i` or `i2i` |
|
||||
| `--dataset` | `prompt` | `prompt`, `random`, or `custom` |
|
||||
| `--dataset-path` | - | JSON file path (required for `custom`) |
|
||||
| `--num-prompts` | `10` | Number of benchmark requests |
|
||||
| `--max-concurrency` | `1` | Max concurrent requests |
|
||||
| `--request-rate` | `inf` | Requests per second (Poisson arrival) |
|
||||
| `--warmup-requests` | `1` | Warmup requests before measurement |
|
||||
| `--width` / `--height` | `1024` | Output image size |
|
||||
| `--num-inference-steps` | `50` | Diffusion denoising steps |
|
||||
| `--seed` | - | Random seed |
|
||||
| `--model` | `default` | Model name (must match `--served-model-name`) |
|
||||
| `--host` | `localhost` | Server host |
|
||||
| `--port` | `8091` | Server port |
|
||||
| `--output-file` | - | JSON output file for metrics |
|
||||
| `--num-input-images` | `1` | Number of input images for random I2I |
|
||||
|
||||
## Dataset
|
||||
|
||||
The default dataset is hosted on [HuggingFace](https://huggingface.co/datasets/JaredforReal/glm-image-bench) (`prompt.json`). It is automatically downloaded and cached to `prompt/prompt.json` on first run. No manual setup needed.
|
||||
|
||||
Each entry contains:
|
||||
|
||||
- `t2i_prompt`: Text prompt for text-to-image generation
|
||||
- `i2i_prompt`: Text prompt for image-to-image editing
|
||||
- `image_url`: Source image URL for I2I (downloaded and cached on first use)
|
||||
|
||||
Custom datasets use the same JSON format and can be provided via `--dataset-path`.
|
||||
|
||||
## Pipeline Timings
|
||||
|
||||
All three benchmarks report per-stage pipeline timings (in milliseconds):
|
||||
|
||||
| Key | Description |
|
||||
|-----|-------------|
|
||||
| `preprocess_ms` | Input preprocessing (tokenization, multimodal encoding) |
|
||||
| `stage_0_gen_ms` | AR (autoregressive) model generation time |
|
||||
| `ar2diffusion_ms` | AR output to diffusion input conversion |
|
||||
| `stage_1_gen_ms` | Diffusion model denoising time |
|
||||
| `queue_wait_ms` | Queue wait time before processing |
|
||||
|
||||
The stages are ordered by execution: `preprocess → stage_0 (AR) → ar2diffusion → stage_1 (Diffusion)`.
|
||||
|
||||
## Sample Results
|
||||
|
||||
Tested on 2x GPU with 10 prompts, 1024x1024, 50 denoising steps:
|
||||
|
||||
| Backend | Mode | Latency Mean (s) | Throughput (img/s) |
|
||||
|---------|------|-------------------|--------------------|
|
||||
| HuggingFace | T2I | 72.6 | 0.014 |
|
||||
| HuggingFace | I2I | 70.9 | 0.014 |
|
||||
| vLLM-Omni Offline | T2I | 35.0 | 0.044 |
|
||||
| vLLM-Omni Offline | I2I | 31.0 | 0.053 |
|
||||
| vLLM-Omni Online | T2I | 38.8 | 0.026 |
|
||||
| vLLM-Omni Online | I2I | 34.7 | 0.029 |
|
||||
Whitespace-only changes.
@@ -0,0 +1,464 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Online serving benchmark for GLM-Image (T2I and I2I modes).
|
||||
|
||||
Sends requests to the /v1/chat/completions endpoint and reports end-to-end
|
||||
latency, throughput, and per-stage durations (when the server is started with
|
||||
--enable-diffusion-pipeline-profiler and/or --enable-ar-profiler).
|
||||
|
||||
Supports three dataset types:
|
||||
- prompt: Use prompt.json (default). T2I uses t2i_prompt, I2I uses i2i_prompt
|
||||
and sends source images from image_url.
|
||||
- random: Generate synthetic prompts (and random images for I2I).
|
||||
- custom: Load from a user-specified JSON file.
|
||||
|
||||
Usage:
|
||||
# T2I with prompt.json (default)
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode t2i --num-prompts 10
|
||||
|
||||
# I2I with prompt.json (downloads source images automatically)
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode i2i --num-prompts 10
|
||||
|
||||
# Random dataset
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode t2i --dataset random --num-prompts 20
|
||||
|
||||
# Custom dataset
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode i2i --dataset custom \
|
||||
--dataset-path my_prompts.json --num-prompts 5
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
import requests as sync_requests
|
||||
from PIL import Image
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
# Import backends from the diffusion benchmark (add parent dirs to path)
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "diffusion"))
|
||||
from backends import RequestFuncOutput
|
||||
|
||||
BENCHMARK_DIR = Path(__file__).resolve().parent
|
||||
DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
|
||||
IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
|
||||
|
||||
DATASET_REPO = "JaredforReal/glm-image-bench"
|
||||
DATASET_FILE = "prompt.json"
|
||||
|
||||
|
||||
def _ensure_prompt_json(dataset_path: str | None) -> str:
|
||||
"""Return path to prompt.json, downloading from HuggingFace if needed."""
|
||||
if dataset_path:
|
||||
return dataset_path
|
||||
local = DEFAULT_PROMPT_JSON
|
||||
if local.exists():
|
||||
return str(local)
|
||||
print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
downloaded = hf_hub_download(
|
||||
repo_id=DATASET_REPO,
|
||||
filename=DATASET_FILE,
|
||||
repo_type="dataset",
|
||||
)
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
import shutil
|
||||
|
||||
shutil.copy2(downloaded, local)
|
||||
print(f"Saved to {local}")
|
||||
except ImportError:
|
||||
url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
|
||||
import urllib.request
|
||||
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
urllib.request.urlretrieve(url, local)
|
||||
print(f"Saved to {local}")
|
||||
return str(local)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class GLMImageRequest:
|
||||
prompt: str
|
||||
image_path: str | None = None # Only for I2I mode
|
||||
|
||||
|
||||
def download_image(url: str) -> str:
|
||||
"""Download an image to cache and return the local path."""
|
||||
IMAGE_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
fname = url.rsplit("/", 1)[-1]
|
||||
local_path = IMAGE_CACHE_DIR / fname
|
||||
if local_path.exists():
|
||||
return str(local_path)
|
||||
resp = sync_requests.get(url, timeout=30)
|
||||
resp.raise_for_status()
|
||||
local_path.write_bytes(resp.content)
|
||||
return str(local_path)
|
||||
|
||||
|
||||
def encode_image_as_data_url(path: str) -> str:
|
||||
"""Encode a local image file as a base64 data URL."""
|
||||
with open(path, "rb") as f:
|
||||
encoded = base64.b64encode(f.read()).decode("utf-8")
|
||||
ext = Path(path).suffix.lower()
|
||||
mime = {"png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg"}.get(ext, "image/png")
|
||||
return f"data:{mime};base64,{encoded}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Datasets
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class PromptDataset:
|
||||
"""Load from prompt.json. T2I uses t2i_prompt, I2I uses i2i_prompt + image_url."""
|
||||
|
||||
def __init__(self, args: argparse.Namespace):
|
||||
path = _ensure_prompt_json(args.dataset_path)
|
||||
with open(path, encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
|
||||
prompt_key = "t2i_prompt" if args.mode == "t2i" else "i2i_prompt"
|
||||
self.items: list[GLMImageRequest] = []
|
||||
|
||||
for entry in raw:
|
||||
prompt = entry.get(prompt_key, "").strip()
|
||||
if not prompt:
|
||||
continue
|
||||
image_path = None
|
||||
if args.mode == "i2i":
|
||||
url = entry.get("image_url", "")
|
||||
if url:
|
||||
image_path = download_image(url)
|
||||
self.items.append(GLMImageRequest(prompt=prompt, image_path=image_path))
|
||||
|
||||
if args.num_prompts and len(self.items) > args.num_prompts:
|
||||
self.items = self.items[: args.num_prompts]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.items)
|
||||
|
||||
def __getitem__(self, idx: int) -> GLMImageRequest:
|
||||
return self.items[idx]
|
||||
|
||||
def get_requests(self) -> list[GLMImageRequest]:
|
||||
return list(self.items)
|
||||
|
||||
|
||||
class RandomDataset:
|
||||
"""Generate synthetic prompts (and optional random images for I2I)."""
|
||||
|
||||
def __init__(self, args: argparse.Namespace):
|
||||
self.args = args
|
||||
self.num_prompts = args.num_prompts
|
||||
self._random_image_paths: list[str] | None = None
|
||||
if args.mode == "i2i":
|
||||
self._random_image_paths = self._generate_random_images()
|
||||
|
||||
def _generate_random_images(self) -> list[str]:
|
||||
paths: list[str] = []
|
||||
for i in range(self.args.num_input_images):
|
||||
img = Image.new("RGB", (512, 512), (128 + i * 30 % 128, 64, 192))
|
||||
path = os.path.join(tempfile.gettempdir(), f"glm_image_bench_input_{i}.png")
|
||||
img.save(path)
|
||||
paths.append(path)
|
||||
return paths
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self.num_prompts
|
||||
|
||||
def __getitem__(self, idx: int) -> GLMImageRequest:
|
||||
image_path = None
|
||||
if self._random_image_paths is not None:
|
||||
image_path = self._random_image_paths[idx % len(self._random_image_paths)]
|
||||
return GLMImageRequest(
|
||||
prompt=f"A beautiful scene with vivid colors and intricate details, prompt {idx}",
|
||||
image_path=image_path,
|
||||
)
|
||||
|
||||
def get_requests(self) -> list[GLMImageRequest]:
|
||||
return [self[i] for i in range(len(self))]
|
||||
|
||||
|
||||
class CustomDataset:
|
||||
"""Load from a user-specified JSON file.
|
||||
|
||||
Expected format:
|
||||
[
|
||||
{"prompt": "A cat sitting on a windowsill"},
|
||||
{"prompt": "Make it look like winter", "image_path": "/path/to/img.png"}
|
||||
]
|
||||
"""
|
||||
|
||||
def __init__(self, args: argparse.Namespace):
|
||||
if not args.dataset_path:
|
||||
raise ValueError("--dataset-path is required for custom dataset")
|
||||
with open(args.dataset_path, encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
self.items: list[GLMImageRequest] = []
|
||||
for item in raw:
|
||||
self.items.append(
|
||||
GLMImageRequest(
|
||||
prompt=item.get("prompt", ""),
|
||||
image_path=item.get("image_path"),
|
||||
)
|
||||
)
|
||||
if args.num_prompts and len(self.items) > args.num_prompts:
|
||||
self.items = self.items[: args.num_prompts]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.items)
|
||||
|
||||
def __getitem__(self, idx: int) -> GLMImageRequest:
|
||||
return self.items[idx]
|
||||
|
||||
def get_requests(self) -> list[GLMImageRequest]:
|
||||
return list(self.items)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Async request for GLM-Image (chat completions with image support)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def async_glm_image_request(
|
||||
req: GLMImageRequest,
|
||||
api_url: str,
|
||||
model: str,
|
||||
session: aiohttp.ClientSession,
|
||||
pbar: Any,
|
||||
args: argparse.Namespace,
|
||||
) -> RequestFuncOutput:
|
||||
"""Send a single T2I or I2I request via chat completions endpoint."""
|
||||
output = RequestFuncOutput()
|
||||
output.start_time = time.perf_counter()
|
||||
|
||||
# Build messages
|
||||
if req.image_path and args.mode == "i2i":
|
||||
data_url = encode_image_as_data_url(req.image_path)
|
||||
content = [
|
||||
{"type": "text", "text": req.prompt},
|
||||
{"type": "image_url", "image_url": {"url": data_url}},
|
||||
]
|
||||
else:
|
||||
content = req.prompt
|
||||
|
||||
messages = [{"role": "user", "content": content}]
|
||||
|
||||
extra_body: dict[str, Any] = {}
|
||||
if args.height:
|
||||
extra_body["height"] = args.height
|
||||
if args.width:
|
||||
extra_body["width"] = args.width
|
||||
if args.num_inference_steps:
|
||||
extra_body["num_inference_steps"] = args.num_inference_steps
|
||||
if args.seed is not None:
|
||||
extra_body["seed"] = args.seed
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
}
|
||||
if extra_body:
|
||||
payload["extra_body"] = extra_body
|
||||
|
||||
try:
|
||||
async with session.post(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):
|
||||
resp_content = msg.get("content", [])
|
||||
if resp_content and isinstance(resp_content, list) and len(resp_content) > 0:
|
||||
first_item = resp_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
|
||||
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 pbar:
|
||||
pbar.update(1)
|
||||
return output
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def iter_requests(n: int, request_rate: float) -> Any:
|
||||
import random as _random
|
||||
|
||||
for i in range(n):
|
||||
if request_rate != float("inf") and i > 0:
|
||||
await asyncio.sleep(_random.expovariate(request_rate))
|
||||
yield i
|
||||
|
||||
|
||||
def calculate_metrics(outputs: list[RequestFuncOutput], total_duration: float) -> dict[str, Any]:
|
||||
success = [o for o in outputs if o.success]
|
||||
errors = [o for o in outputs if not o.success]
|
||||
latencies = [o.latency for o in success]
|
||||
peak_mems = [o.peak_memory_mb for o in success if o.peak_memory_mb > 0]
|
||||
|
||||
stage_duration_lists: dict[str, list[float]] = {}
|
||||
for o in success:
|
||||
for stage, dur in (o.stage_durations or {}).items():
|
||||
stage_duration_lists.setdefault(stage, []).append(dur)
|
||||
|
||||
return {
|
||||
"duration": total_duration,
|
||||
"completed_requests": len(success),
|
||||
"failed_requests": len(errors),
|
||||
"throughput_qps": len(success) / total_duration if total_duration > 0 else 0,
|
||||
"latency_mean": float(np.mean(latencies)) if latencies else 0,
|
||||
"latency_median": float(np.median(latencies)) if latencies else 0,
|
||||
"latency_p99": float(np.percentile(latencies, 99)) if latencies else 0,
|
||||
"latency_p95": float(np.percentile(latencies, 95)) if latencies else 0,
|
||||
"peak_memory_mb_max": max(peak_mems) if peak_mems else 0,
|
||||
"stage_durations_mean": {s: float(np.mean(v)) for s, v in stage_duration_lists.items()},
|
||||
"stage_durations_p50": {s: float(np.percentile(v, 50)) for s, v in stage_duration_lists.items()},
|
||||
}
|
||||
|
||||
|
||||
async def benchmark(args: argparse.Namespace) -> None:
|
||||
api_url = f"http://{args.host}:{args.port}/v1/chat/completions"
|
||||
|
||||
# Load dataset
|
||||
if args.dataset == "prompt":
|
||||
dataset = PromptDataset(args)
|
||||
elif args.dataset == "random":
|
||||
dataset = RandomDataset(args)
|
||||
elif args.dataset == "custom":
|
||||
dataset = CustomDataset(args)
|
||||
else:
|
||||
raise ValueError(f"Unknown dataset: {args.dataset}")
|
||||
|
||||
glm_requests = dataset.get_requests()
|
||||
print(f"Prepared {len(glm_requests)} requests (mode={args.mode}, dataset={args.dataset})")
|
||||
|
||||
semaphore = asyncio.Semaphore(args.max_concurrency) if args.max_concurrency else None
|
||||
|
||||
async def limited_request(idx: int, req: GLMImageRequest, session: aiohttp.ClientSession, pbar: Any):
|
||||
if semaphore:
|
||||
async with semaphore:
|
||||
return await async_glm_image_request(req, api_url, args.model, session, pbar, args)
|
||||
return await async_glm_image_request(req, api_url, args.model, session, pbar, args)
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# Warmup
|
||||
if args.warmup_requests and glm_requests:
|
||||
print(f"Running {args.warmup_requests} warmup request(s)...")
|
||||
for i in range(args.warmup_requests):
|
||||
await limited_request(i, glm_requests[i % len(glm_requests)], session, None)
|
||||
|
||||
# Main benchmark
|
||||
pbar = tqdm(total=len(glm_requests), disable=args.disable_tqdm)
|
||||
start_time = time.perf_counter()
|
||||
tasks = []
|
||||
async for idx in iter_requests(len(glm_requests), args.request_rate):
|
||||
tasks.append(asyncio.create_task(limited_request(idx, glm_requests[idx], session, pbar)))
|
||||
outputs = await asyncio.gather(*tasks)
|
||||
total_duration = time.perf_counter() - start_time
|
||||
pbar.close()
|
||||
|
||||
# Metrics
|
||||
metrics = calculate_metrics(outputs, total_duration)
|
||||
metrics["mode"] = args.mode
|
||||
metrics["model"] = args.model
|
||||
metrics["dataset"] = args.dataset
|
||||
|
||||
print(f"\n{' GLM-Image Online Benchmark Result ':=^60}")
|
||||
print(f"{'Mode:':<40} {args.mode}")
|
||||
print(f"{'Model:':<40} {args.model}")
|
||||
print(f"{'Dataset:':<40} {args.dataset}")
|
||||
print("-" * 50)
|
||||
print(f"{'Benchmark duration (s):':<40} {metrics['duration']:.2f}")
|
||||
print(f"{'Request rate:':<40} {args.request_rate}")
|
||||
print(f"{'Max concurrency:':<40} {args.max_concurrency}")
|
||||
print(f"{'Successful requests:':<40} {metrics['completed_requests']}/{len(glm_requests)}")
|
||||
print("-" * 50)
|
||||
print(f"{'Throughput (req/s):':<40} {metrics['throughput_qps']:.2f}")
|
||||
print(f"{'Latency Mean (s):':<40} {metrics['latency_mean']:.4f}")
|
||||
print(f"{'Latency Median (s):':<40} {metrics['latency_median']:.4f}")
|
||||
print(f"{'Latency P95 (s):':<40} {metrics['latency_p95']:.4f}")
|
||||
print(f"{'Latency P99 (s):':<40} {metrics['latency_p99']:.4f}")
|
||||
|
||||
if metrics["peak_memory_mb_max"] > 0:
|
||||
print("-" * 50)
|
||||
print(f"{'Peak Memory Max (MB):':<40} {metrics['peak_memory_mb_max']:.2f}")
|
||||
|
||||
if metrics["stage_durations_mean"]:
|
||||
print("-" * 50)
|
||||
print("Stage Durations Mean:")
|
||||
for stage, val in sorted(metrics["stage_durations_mean"].items()):
|
||||
unit = "ms" if stage.endswith("_ms") else "s"
|
||||
print(f" {stage + ':':<38} {val:.4f} ({unit})")
|
||||
|
||||
print("=" * 60)
|
||||
|
||||
if args.output_file:
|
||||
with open(args.output_file, "w") as f:
|
||||
json.dump(metrics, f, indent=2)
|
||||
print(f"Metrics saved to {args.output_file}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Benchmark GLM-Image T2I/I2I online serving.")
|
||||
parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
|
||||
parser.add_argument("--dataset", type=str, default="prompt", choices=["prompt", "random", "custom"])
|
||||
parser.add_argument("--dataset-path", type=str, default=None)
|
||||
parser.add_argument("--num-prompts", type=int, default=10)
|
||||
parser.add_argument("--max-concurrency", type=int, default=1)
|
||||
parser.add_argument("--request-rate", type=float, default=float("inf"))
|
||||
parser.add_argument("--warmup-requests", type=int, default=1)
|
||||
parser.add_argument("--width", type=int, default=1024)
|
||||
parser.add_argument("--height", type=int, default=1024)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=50)
|
||||
parser.add_argument("--seed", type=int, default=None)
|
||||
parser.add_argument("--model", type=str, default="default")
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8091)
|
||||
parser.add_argument("--output-file", type=str, default=None)
|
||||
parser.add_argument("--disable-tqdm", action="store_true")
|
||||
parser.add_argument("--num-input-images", type=int, default=1, help="For random I2I dataset.")
|
||||
args = parser.parse_args()
|
||||
asyncio.run(benchmark(args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,291 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
HuggingFace (transformers + diffusers) baseline benchmark for GLM-Image.
|
||||
|
||||
Supports T2I and I2I modes with the prompt.json dataset.
|
||||
Downloads source images for I2I from image_url on first run and caches locally.
|
||||
|
||||
Usage:
|
||||
# T2I mode (text-to-image, no source images needed)
|
||||
python benchmarks/glm_image/huggingface/inference.py \
|
||||
--model-path zai-org/GLM-Image \
|
||||
--mode t2i --num-prompts 10
|
||||
|
||||
# I2I mode (image-to-image, downloads source images)
|
||||
python benchmarks/glm_image/huggingface/inference.py \
|
||||
--model-path zai-org/GLM-Image \
|
||||
--mode i2i --num-prompts 10
|
||||
|
||||
# With custom prompt.json
|
||||
python benchmarks/glm_image/huggingface/inference.py \
|
||||
--model-path zai-org/GLM-Image \
|
||||
--mode i2i --dataset-path prompts.json --num-prompts 5
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
BENCHMARK_DIR = Path(__file__).resolve().parent.parent
|
||||
DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
|
||||
IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
|
||||
|
||||
DATASET_REPO = "JaredforReal/glm-image-bench"
|
||||
DATASET_FILE = "prompt.json"
|
||||
|
||||
|
||||
def _ensure_prompt_json(dataset_path: str | None) -> str:
|
||||
"""Return path to prompt.json, downloading from HuggingFace if needed."""
|
||||
if dataset_path:
|
||||
return dataset_path
|
||||
local = DEFAULT_PROMPT_JSON
|
||||
if local.exists():
|
||||
return str(local)
|
||||
print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
downloaded = hf_hub_download(
|
||||
repo_id=DATASET_REPO,
|
||||
filename=DATASET_FILE,
|
||||
repo_type="dataset",
|
||||
)
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
import shutil
|
||||
|
||||
shutil.copy2(downloaded, local)
|
||||
print(f"Saved to {local}")
|
||||
except ImportError:
|
||||
url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
|
||||
import urllib.request
|
||||
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
urllib.request.urlretrieve(url, local)
|
||||
print(f"Saved to {local}")
|
||||
return str(local)
|
||||
|
||||
|
||||
HEIGHT = 1024
|
||||
WIDTH = 1024
|
||||
SEED = 42
|
||||
NUM_INFERENCE_STEPS = 50
|
||||
GUIDANCE_SCALE = 1.5
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataset
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_dataset(
|
||||
dataset_path: str | None,
|
||||
mode: str,
|
||||
num_prompts: int,
|
||||
) -> list[dict]:
|
||||
"""Load prompts from prompt.json and prepare per-request data."""
|
||||
path = _ensure_prompt_json(dataset_path)
|
||||
with open(path, encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
|
||||
items = []
|
||||
for entry in raw:
|
||||
if mode == "t2i":
|
||||
prompt_key = "t2i_prompt"
|
||||
else:
|
||||
prompt_key = "i2i_prompt"
|
||||
|
||||
prompt_text = entry.get(prompt_key, "").strip()
|
||||
if not prompt_text:
|
||||
continue
|
||||
|
||||
item = {"prompt": prompt_text}
|
||||
if mode == "i2i":
|
||||
item["image_url"] = entry.get("image_url", "")
|
||||
items.append(item)
|
||||
|
||||
if num_prompts and len(items) > num_prompts:
|
||||
items = items[:num_prompts]
|
||||
return items
|
||||
|
||||
|
||||
def download_image(url: str, cache_dir: Path) -> str:
|
||||
"""Download an image to cache_dir and return the local path."""
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
fname = url.rsplit("/", 1)[-1]
|
||||
local_path = cache_dir / fname
|
||||
if local_path.exists():
|
||||
return str(local_path)
|
||||
print(f" Downloading {url} ...")
|
||||
resp = requests.get(url, timeout=30)
|
||||
resp.raise_for_status()
|
||||
local_path.write_bytes(resp.content)
|
||||
return str(local_path)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def benchmark(args: argparse.Namespace) -> None:
|
||||
from diffusers.pipelines.glm_image import GlmImagePipeline
|
||||
|
||||
print("=" * 60)
|
||||
print("GLM-Image HuggingFace Baseline Benchmark")
|
||||
print(f"Mode: {args.mode} | Model: {args.model_path}")
|
||||
print(f"Size: {args.height}x{args.width} | Steps: {args.num_inference_steps}")
|
||||
print("=" * 60)
|
||||
|
||||
# Load dataset
|
||||
items = load_dataset(args.dataset_path, args.mode, args.num_prompts)
|
||||
if not items:
|
||||
print("No prompts loaded. Exiting.")
|
||||
return
|
||||
print(f"Loaded {len(items)} prompts for {args.mode} mode")
|
||||
|
||||
# Download I2I source images
|
||||
if args.mode == "i2i":
|
||||
print("Preparing source images...")
|
||||
for item in items:
|
||||
url = item.get("image_url", "")
|
||||
if url:
|
||||
item["image_path"] = download_image(url, IMAGE_CACHE_DIR)
|
||||
else:
|
||||
item["image_path"] = None
|
||||
|
||||
# Load pipeline
|
||||
print(f"\nLoading pipeline from {args.model_path} ...")
|
||||
t0 = time.perf_counter()
|
||||
pipe = GlmImagePipeline.from_pretrained(
|
||||
args.model_path,
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map="cuda",
|
||||
)
|
||||
init_time = time.perf_counter() - t0
|
||||
print(f"Pipeline loaded in {init_time:.2f}s")
|
||||
|
||||
# Create output dir
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
# Run benchmark
|
||||
generator = torch.Generator(device="cuda").manual_seed(args.seed)
|
||||
latencies = []
|
||||
success = 0
|
||||
failed = 0
|
||||
|
||||
print(f"\nRunning {len(items)} requests sequentially...")
|
||||
print("-" * 60)
|
||||
|
||||
for i, item in enumerate(items):
|
||||
prompt = item["prompt"]
|
||||
gen_kwargs: dict = {
|
||||
"prompt": prompt,
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_inference_steps": args.num_inference_steps,
|
||||
"guidance_scale": args.guidance_scale,
|
||||
"generator": generator,
|
||||
}
|
||||
|
||||
if args.mode == "i2i":
|
||||
img_path = item.get("image_path")
|
||||
if img_path and os.path.exists(img_path):
|
||||
gen_kwargs["image"] = [Image.open(img_path).convert("RGB")]
|
||||
else:
|
||||
print(f" [{i + 1}] SKIP: no source image")
|
||||
failed += 1
|
||||
continue
|
||||
|
||||
t_start = time.perf_counter()
|
||||
try:
|
||||
result = pipe(**gen_kwargs)
|
||||
image = result.images[0]
|
||||
elapsed = time.perf_counter() - t_start
|
||||
latencies.append(elapsed)
|
||||
success += 1
|
||||
|
||||
out_path = os.path.join(args.output_dir, f"{i:04d}.png")
|
||||
image.save(out_path)
|
||||
print(f" [{i + 1}/{len(items)}] {elapsed:.3f}s -> {out_path}")
|
||||
except Exception as e:
|
||||
elapsed = time.perf_counter() - t_start
|
||||
failed += 1
|
||||
print(f" [{i + 1}/{len(items)}] FAILED ({elapsed:.3f}s): {e}")
|
||||
|
||||
# Report
|
||||
total_gen_time = sum(latencies) if latencies else 0
|
||||
print("\n" + "=" * 60)
|
||||
print("HuggingFace Baseline Results")
|
||||
print("=" * 60)
|
||||
print(f"{'Mode:':<40} {args.mode}")
|
||||
print(f"{'Model:':<40} {args.model_path}")
|
||||
print(f"{'Image size:':<40} {args.height}x{args.width}")
|
||||
print(f"{'Num inference steps:':<40} {args.num_inference_steps}")
|
||||
print("-" * 50)
|
||||
print(f"{'Pipeline init time (s):':<40} {init_time:.2f}")
|
||||
print(f"{'Successful:':<40} {success}/{len(items)}")
|
||||
print(f"{'Failed:':<40} {failed}")
|
||||
print("-" * 50)
|
||||
if latencies:
|
||||
arr = np.array(latencies)
|
||||
print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
|
||||
print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
|
||||
print(f"{'Latency Mean (s):':<40} {arr.mean():.4f}")
|
||||
print(f"{'Latency Median (s):':<40} {np.median(arr):.4f}")
|
||||
print(f"{'Latency P95 (s):':<40} {np.percentile(arr, 95):.4f}")
|
||||
print(f"{'Latency P99 (s):':<40} {np.percentile(arr, 99):.4f}")
|
||||
|
||||
print(f"\n{'Output dir:':<40} {args.output_dir}")
|
||||
print("=" * 60)
|
||||
|
||||
# Save metrics JSON
|
||||
metrics = {
|
||||
"backend": "huggingface",
|
||||
"mode": args.mode,
|
||||
"model": args.model_path,
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_inference_steps": args.num_inference_steps,
|
||||
"init_time_s": init_time,
|
||||
"completed_requests": success,
|
||||
"failed_requests": failed,
|
||||
"total_gen_time_s": total_gen_time,
|
||||
"throughput_qps": success / total_gen_time if total_gen_time > 0 else 0,
|
||||
"latency_mean": float(np.mean(latencies)) if latencies else 0,
|
||||
"latency_median": float(np.median(latencies)) if latencies else 0,
|
||||
"latency_p95": float(np.percentile(latencies, 95)) if latencies else 0,
|
||||
"latency_p99": float(np.percentile(latencies, 99)) if latencies else 0,
|
||||
}
|
||||
if args.output_file:
|
||||
with open(args.output_file, "w") as f:
|
||||
json.dump(metrics, f, indent=2)
|
||||
print(f"Metrics saved to {args.output_file}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="GLM-Image HuggingFace baseline benchmark")
|
||||
parser.add_argument("--model-path", type=str, default="zai-org/GLM-Image")
|
||||
parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
|
||||
parser.add_argument("--dataset-path", type=str, default=None, help="Path to prompt.json")
|
||||
parser.add_argument("--num-prompts", type=int, default=10)
|
||||
parser.add_argument("--height", type=int, default=HEIGHT)
|
||||
parser.add_argument("--width", type=int, default=WIDTH)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=NUM_INFERENCE_STEPS)
|
||||
parser.add_argument("--guidance-scale", type=float, default=GUIDANCE_SCALE)
|
||||
parser.add_argument("--seed", type=int, default=SEED)
|
||||
parser.add_argument("--output-dir", type=str, default="benchmarks/glm_image/huggingface/outputs")
|
||||
parser.add_argument("--output-file", type=str, default=None, help="JSON file for metrics")
|
||||
args = parser.parse_args()
|
||||
benchmark(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,505 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
vLLM-Omni offline benchmark for GLM-Image.
|
||||
|
||||
Supports T2I and I2I modes with the prompt.json dataset.
|
||||
Downloads source images for I2I from image_url on first run and caches locally.
|
||||
|
||||
Usage:
|
||||
# T2I mode
|
||||
python benchmarks/glm_image/vllm-omni/inference.py \
|
||||
--model-path zai-org/GLM-Image \
|
||||
--mode t2i --num-prompts 10
|
||||
|
||||
# I2I mode (downloads source images)
|
||||
python benchmarks/glm_image/vllm-omni/inference.py \
|
||||
--model-path zai-org/GLM-Image \
|
||||
--mode i2i --num-prompts 10
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
from PIL import Image
|
||||
from vllm import SamplingParams
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
BENCHMARK_DIR = Path(__file__).resolve().parent.parent
|
||||
DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
|
||||
IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
|
||||
DEFAULT_DEPLOY_CONFIG = "vllm_omni/deploy/glm_image.yaml"
|
||||
|
||||
DATASET_REPO = "JaredforReal/glm-image-bench"
|
||||
DATASET_FILE = "prompt.json"
|
||||
|
||||
|
||||
def _ensure_prompt_json(dataset_path: str | None) -> str:
|
||||
"""Return path to prompt.json, downloading from HuggingFace if needed."""
|
||||
if dataset_path:
|
||||
return dataset_path
|
||||
local = DEFAULT_PROMPT_JSON
|
||||
if local.exists():
|
||||
return str(local)
|
||||
print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
downloaded = hf_hub_download(
|
||||
repo_id=DATASET_REPO,
|
||||
filename=DATASET_FILE,
|
||||
repo_type="dataset",
|
||||
)
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
import shutil
|
||||
|
||||
shutil.copy2(downloaded, local)
|
||||
print(f"Saved to {local}")
|
||||
except ImportError:
|
||||
url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
|
||||
import urllib.request
|
||||
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
urllib.request.urlretrieve(url, local)
|
||||
print(f"Saved to {local}")
|
||||
return str(local)
|
||||
|
||||
|
||||
SEED = 42
|
||||
HEIGHT = 1024
|
||||
WIDTH = 1024
|
||||
NUM_INFERENCE_STEPS = 50
|
||||
GUIDANCE_SCALE = 1.5
|
||||
|
||||
GLM_IMAGE_EOS_TOKEN_ID = 16385
|
||||
GLM_IMAGE_VISION_VOCAB_SIZE = 16512
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataset
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_dataset(
|
||||
dataset_path: str | None,
|
||||
mode: str,
|
||||
num_prompts: int,
|
||||
) -> list[dict]:
|
||||
path = _ensure_prompt_json(dataset_path)
|
||||
with open(path, encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
|
||||
items = []
|
||||
for entry in raw:
|
||||
prompt_key = "t2i_prompt" if mode == "t2i" else "i2i_prompt"
|
||||
prompt_text = entry.get(prompt_key, "").strip()
|
||||
if not prompt_text:
|
||||
continue
|
||||
|
||||
item = {"prompt": prompt_text}
|
||||
if mode == "i2i":
|
||||
item["image_url"] = entry.get("image_url", "")
|
||||
items.append(item)
|
||||
|
||||
if num_prompts and len(items) > num_prompts:
|
||||
items = items[:num_prompts]
|
||||
return items
|
||||
|
||||
|
||||
def download_image(url: str, cache_dir: Path) -> str:
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
fname = url.rsplit("/", 1)[-1]
|
||||
local_path = cache_dir / fname
|
||||
if local_path.exists():
|
||||
return str(local_path)
|
||||
print(f" Downloading {url} ...")
|
||||
resp = requests.get(url, timeout=30)
|
||||
resp.raise_for_status()
|
||||
local_path.write_bytes(resp.content)
|
||||
return str(local_path)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def compute_max_tokens(height: int, width: int, is_i2i: bool = False) -> int:
|
||||
factor = 32
|
||||
token_h = height // factor
|
||||
token_w = width // factor
|
||||
large_tokens = token_h * token_w
|
||||
|
||||
# Small preview tokens (half resolution in each dimension)
|
||||
|
||||
ratio = token_h / token_w if token_w > 0 else 1.0
|
||||
small_token_h = max(1, int(math.sqrt(ratio) * (factor // 2)))
|
||||
small_token_w = max(1, int(math.sqrt(1 / ratio) * (factor // 2)))
|
||||
small_tokens = small_token_h * small_token_w
|
||||
|
||||
# Mode-dependent totals:
|
||||
# - t2i: small + large + EOS
|
||||
# - i2i: large + EOS
|
||||
if is_i2i:
|
||||
return large_tokens + 1
|
||||
return small_tokens + large_tokens + 1
|
||||
|
||||
|
||||
def build_prompt_t2i(prompt: str, height: int, width: int, **gen_kw) -> dict:
|
||||
return {
|
||||
"prompt": prompt,
|
||||
"height": height,
|
||||
"width": width,
|
||||
"mm_processor_kwargs": {"target_h": height, "target_w": width},
|
||||
**gen_kw,
|
||||
}
|
||||
|
||||
|
||||
def build_prompt_i2i(prompt: str, image_path: str, height: int, width: int, **gen_kw) -> dict:
|
||||
return {
|
||||
"prompt": prompt,
|
||||
"height": height,
|
||||
"width": width,
|
||||
"mm_processor_kwargs": {"target_h": height, "target_w": width},
|
||||
"multi_modal_data": {"image": Image.open(image_path).convert("RGB")},
|
||||
**gen_kw,
|
||||
}
|
||||
|
||||
|
||||
def resolve_deploy_config(args: argparse.Namespace) -> str:
|
||||
if args.deploy_config:
|
||||
return args.deploy_config
|
||||
if os.path.exists(DEFAULT_DEPLOY_CONFIG):
|
||||
return DEFAULT_DEPLOY_CONFIG
|
||||
fallback = Path(__file__).resolve().parents[3] / DEFAULT_DEPLOY_CONFIG
|
||||
if fallback.exists():
|
||||
return str(fallback)
|
||||
raise FileNotFoundError("Deploy config not found. Specify --deploy-config.")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def benchmark(args: argparse.Namespace) -> None:
|
||||
is_i2i = args.mode == "i2i"
|
||||
|
||||
print("=" * 60)
|
||||
print("GLM-Image vLLM-Omni Benchmark")
|
||||
print(f"Mode: {args.mode} | Model: {args.model_path}")
|
||||
print(f"Size: {args.height}x{args.width} | Steps: {args.num_inference_steps}")
|
||||
print("=" * 60)
|
||||
|
||||
# Load dataset
|
||||
items = load_dataset(args.dataset_path, args.mode, args.num_prompts)
|
||||
if not items:
|
||||
print("No prompts loaded. Exiting.")
|
||||
return
|
||||
print(f"Loaded {len(items)} prompts for {args.mode} mode")
|
||||
|
||||
# Download I2I source images
|
||||
if is_i2i:
|
||||
print("Preparing source images...")
|
||||
for item in items:
|
||||
url = item.get("image_url", "")
|
||||
if url:
|
||||
item["image_path"] = download_image(url, IMAGE_CACHE_DIR)
|
||||
else:
|
||||
item["image_path"] = None
|
||||
|
||||
# Init Omni
|
||||
deploy_config = resolve_deploy_config(args)
|
||||
print(f"\nInitializing vLLM-Omni (deploy config: {deploy_config}) ...")
|
||||
t0 = time.perf_counter()
|
||||
|
||||
omni = Omni(
|
||||
model=args.model_path,
|
||||
deploy_config=deploy_config,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
enable_ar_profiler=args.enable_ar_profiler,
|
||||
)
|
||||
|
||||
init_time = time.perf_counter() - t0
|
||||
print(f"Initialized in {init_time:.2f}s")
|
||||
|
||||
# Sampling params
|
||||
max_tokens = compute_max_tokens(args.height, args.width, is_i2i=is_i2i)
|
||||
ar_params = SamplingParams(
|
||||
temperature=0.9,
|
||||
top_p=0.75,
|
||||
top_k=GLM_IMAGE_VISION_VOCAB_SIZE,
|
||||
max_tokens=max_tokens,
|
||||
stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID],
|
||||
seed=args.seed,
|
||||
detokenize=False,
|
||||
extra_args={"target_h": args.height, "target_w": args.width},
|
||||
)
|
||||
diff_params = OmniDiffusionSamplingParams(
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
seed=args.seed,
|
||||
)
|
||||
sampling_params_list = [ar_params, diff_params]
|
||||
|
||||
# Build all prompts
|
||||
gen_kw = {
|
||||
"seed": args.seed,
|
||||
"num_inference_steps": args.num_inference_steps,
|
||||
"guidance_scale": args.guidance_scale,
|
||||
}
|
||||
all_prompts = []
|
||||
for item in items:
|
||||
if is_i2i:
|
||||
img_path = item.get("image_path")
|
||||
if not img_path or not os.path.exists(img_path):
|
||||
continue
|
||||
all_prompts.append(build_prompt_i2i(item["prompt"], img_path, args.height, args.width, **gen_kw))
|
||||
else:
|
||||
all_prompts.append(build_prompt_t2i(item["prompt"], args.height, args.width, **gen_kw))
|
||||
|
||||
valid = len(all_prompts)
|
||||
print(f"Valid prompts: {valid}")
|
||||
|
||||
# Create output dir
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
# Warmup: run 1 request to prime caches, CUDA graphs, etc.
|
||||
if all_prompts:
|
||||
print("Running warmup request...")
|
||||
try:
|
||||
warmup_prompt = [all_prompts[0]]
|
||||
omni.generate(warmup_prompt, sampling_params_list, py_generator=False)
|
||||
print("Warmup done.\n")
|
||||
except Exception as e:
|
||||
print(f"Warmup failed (continuing): {e}")
|
||||
|
||||
# Run
|
||||
print(f"\nRunning {valid} requests...")
|
||||
print("-" * 60)
|
||||
|
||||
latencies = []
|
||||
all_stage_durations: list[dict[str, float]] = []
|
||||
success = 0
|
||||
failed = 0
|
||||
wall_start = time.perf_counter()
|
||||
|
||||
try:
|
||||
output_idx = 0
|
||||
for stage_outputs in omni.generate(all_prompts, sampling_params_list, py_generator=True):
|
||||
if stage_outputs.final_output_type == "image":
|
||||
request_output = stage_outputs.request_output
|
||||
request_id = getattr(request_output, "request_id", "")
|
||||
|
||||
images = getattr(request_output, "images", [])
|
||||
if not images and hasattr(request_output, "multimodal_output"):
|
||||
mm = request_output.multimodal_output
|
||||
if isinstance(mm, dict):
|
||||
images = mm.get("images", [])
|
||||
|
||||
elapsed = time.perf_counter() - wall_start
|
||||
if images:
|
||||
for img in images:
|
||||
if isinstance(img, Image.Image):
|
||||
out_path = os.path.join(args.output_dir, f"{output_idx:04d}.png")
|
||||
img.save(out_path)
|
||||
success += 1
|
||||
latencies.append(elapsed)
|
||||
stage_durations = getattr(stage_outputs, "stage_durations", {})
|
||||
if stage_durations:
|
||||
all_stage_durations.append(stage_durations)
|
||||
# Show wall-clock elapsed and pipeline breakdown if available
|
||||
preprocess_str = ""
|
||||
if "preprocess_ms" in stage_durations:
|
||||
preprocess_str = f" preprocess={stage_durations['preprocess_ms'] / 1000.0:.2f}s"
|
||||
print(f" [{success}/{valid}] id={request_id[:8]} {elapsed:.2f}s{preprocess_str}")
|
||||
output_idx += 1
|
||||
else:
|
||||
failed += 1
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
failed = valid - success
|
||||
|
||||
total_gen_time = time.perf_counter() - wall_start
|
||||
|
||||
# Diff stage_0_gen_ms with previous request to remove accumulated wait time.
|
||||
# stage_0_gen_ms is measured from submit_ts (same for all requests submitted
|
||||
# at once), so it accumulates queue/scheduling overhead across requests.
|
||||
# Other stages and pipeline timings are per-request already.
|
||||
_TIMING_ORDER = [
|
||||
"preprocess_ms",
|
||||
"stage_0_gen_ms",
|
||||
"ar2diffusion_ms",
|
||||
"stage_1_gen_ms",
|
||||
"queue_wait_ms",
|
||||
]
|
||||
|
||||
per_request_actual: list[dict[str, float]] = []
|
||||
prev_stage_0_ms = 0.0
|
||||
for sd in all_stage_durations:
|
||||
actual = dict(sd)
|
||||
s0 = sd.get("stage_0_gen_ms", 0.0)
|
||||
actual["stage_0_gen_ms"] = s0 - prev_stage_0_ms
|
||||
prev_stage_0_ms = s0
|
||||
per_request_actual.append(actual)
|
||||
|
||||
per_request_e2e_ms: list[float] = []
|
||||
for actual in per_request_actual:
|
||||
e2e_ms = sum(v for k, v in actual.items() if k in _TIMING_ORDER)
|
||||
if e2e_ms > 0:
|
||||
per_request_e2e_ms.append(e2e_ms)
|
||||
|
||||
# Report
|
||||
print("\n" + "=" * 60)
|
||||
print("vLLM-Omni Benchmark Results")
|
||||
print("=" * 60)
|
||||
print(f"{'Mode:':<40} {args.mode}")
|
||||
print(f"{'Model:':<40} {args.model_path}")
|
||||
print(f"{'Image size:':<40} {args.height}x{args.width}")
|
||||
print(f"{'Num inference steps:':<40} {args.num_inference_steps}")
|
||||
print("-" * 50)
|
||||
print(f"{'Init time (s):':<40} {init_time:.2f}")
|
||||
print(f"{'Successful:':<40} {success}/{valid}")
|
||||
print(f"{'Failed:':<40} {failed}")
|
||||
print("-" * 50)
|
||||
|
||||
if per_request_e2e_ms:
|
||||
per_request_s = np.array(per_request_e2e_ms) / 1000.0
|
||||
print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
|
||||
print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
|
||||
print(f"{'Latency Mean (s):':<40} {per_request_s.mean():.4f}")
|
||||
print(f"{'Latency Median (s):':<40} {np.median(per_request_s):.4f}")
|
||||
print(f"{'Latency P95 (s):':<40} {np.percentile(per_request_s, 95):.4f}")
|
||||
print(f"{'Latency P99 (s):':<40} {np.percentile(per_request_s, 99):.4f}")
|
||||
print(f"{'Latency Min (s):':<40} {per_request_s.min():.4f}")
|
||||
print(f"{'Latency Max (s):':<40} {per_request_s.max():.4f}")
|
||||
elif latencies:
|
||||
per_request = np.diff([0.0] + list(latencies))
|
||||
print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
|
||||
print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
|
||||
print(f"{'Latency Mean (s) [wall-clock]:':<40} {per_request.mean():.4f}")
|
||||
print(f"{'Latency Median (s) [wall-clock]:':<40} {np.median(per_request):.4f}")
|
||||
print(f"{'Latency P95 (s) [wall-clock]:':<40} {np.percentile(per_request, 95):.4f}")
|
||||
print(f"{'Latency P99 (s) [wall-clock]:':<40} {np.percentile(per_request, 99):.4f}")
|
||||
print(f"{'Latency Min (s) [wall-clock]:':<40} {per_request.min():.4f}")
|
||||
print(f"{'Latency Max (s) [wall-clock]:':<40} {per_request.max():.4f}")
|
||||
|
||||
if per_request_actual:
|
||||
print("-" * 50)
|
||||
print("Pipeline Timings Mean:")
|
||||
for key in _TIMING_ORDER:
|
||||
vals = [d.get(key, 0.0) for d in per_request_actual]
|
||||
if any(v != 0 for v in vals):
|
||||
unit = "ms" if key.endswith("_ms") else "s"
|
||||
print(f" {key + ':':<38} {np.mean(vals):.4f} ({unit})")
|
||||
# Show any extra keys not in the ordered list
|
||||
ordered_set = set(_TIMING_ORDER)
|
||||
extra_keys = sorted(k for k in per_request_actual[0].keys() if k not in ordered_set)
|
||||
for key in extra_keys:
|
||||
vals = [d.get(key, 0.0) for d in per_request_actual]
|
||||
if any(v != 0 for v in vals):
|
||||
unit = "ms" if key.endswith("_ms") else "s"
|
||||
print(f" {key + ':':<38} {np.mean(vals):.4f} ({unit})")
|
||||
|
||||
print(f"\n{'Output dir:':<40} {args.output_dir}")
|
||||
print("=" * 60)
|
||||
|
||||
# Metrics JSON
|
||||
metrics = {
|
||||
"backend": "vllm-omni",
|
||||
"mode": args.mode,
|
||||
"model": args.model_path,
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_inference_steps": args.num_inference_steps,
|
||||
"init_time_s": init_time,
|
||||
"completed_requests": success,
|
||||
"failed_requests": failed,
|
||||
"total_gen_time_s": total_gen_time,
|
||||
"throughput_qps": success / total_gen_time if total_gen_time > 0 else 0,
|
||||
}
|
||||
if per_request_e2e_ms:
|
||||
per_request_s = np.array(per_request_e2e_ms) / 1000.0
|
||||
metrics["latency_mean"] = float(per_request_s.mean())
|
||||
metrics["latency_median"] = float(np.median(per_request_s))
|
||||
metrics["latency_p95"] = float(np.percentile(per_request_s, 95))
|
||||
metrics["latency_p99"] = float(np.percentile(per_request_s, 99))
|
||||
elif latencies:
|
||||
per_request = np.diff([0.0] + list(latencies))
|
||||
metrics["latency_mean"] = float(per_request.mean())
|
||||
metrics["latency_median"] = float(np.median(per_request))
|
||||
metrics["latency_p95"] = float(np.percentile(per_request, 95))
|
||||
metrics["latency_p99"] = float(np.percentile(per_request, 99))
|
||||
else:
|
||||
metrics["latency_mean"] = 0
|
||||
metrics["latency_median"] = 0
|
||||
metrics["latency_p95"] = 0
|
||||
metrics["latency_p99"] = 0
|
||||
if per_request_actual:
|
||||
all_keys = list(_TIMING_ORDER) + sorted(k for k in per_request_actual[0].keys() if k not in set(_TIMING_ORDER))
|
||||
stage_metrics = {}
|
||||
for key in all_keys:
|
||||
vals = [d.get(key, 0.0) for d in per_request_actual]
|
||||
stage_metrics[key] = {
|
||||
"mean": float(np.mean(vals)),
|
||||
"median": float(np.median(vals)),
|
||||
"p95": float(np.percentile(vals, 95)),
|
||||
}
|
||||
metrics["stage_durations"] = stage_metrics
|
||||
if args.output_file:
|
||||
with open(args.output_file, "w") as f:
|
||||
json.dump(metrics, f, indent=2)
|
||||
print(f"Metrics saved to {args.output_file}")
|
||||
|
||||
omni.close()
|
||||
print("Done!")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="GLM-Image vLLM-Omni offline benchmark")
|
||||
parser.add_argument("--model-path", type=str, default="zai-org/GLM-Image")
|
||||
parser.add_argument("--deploy-config", type=str, default=None, help="Deploy config YAML")
|
||||
parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
|
||||
parser.add_argument("--dataset-path", type=str, default=None, help="Path to prompt.json")
|
||||
parser.add_argument("--num-prompts", type=int, default=10)
|
||||
parser.add_argument("--height", type=int, default=HEIGHT)
|
||||
parser.add_argument("--width", type=int, default=WIDTH)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=NUM_INFERENCE_STEPS)
|
||||
parser.add_argument("--guidance-scale", type=float, default=GUIDANCE_SCALE)
|
||||
parser.add_argument("--seed", type=int, default=SEED)
|
||||
parser.add_argument("--output-dir", type=str, default="benchmarks/glm_image/vllm-omni/outputs")
|
||||
parser.add_argument("--output-file", type=str, default=None, help="JSON file for metrics")
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=600)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler for stage-level timing",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-ar-profiler",
|
||||
action="store_true",
|
||||
help="Enable AR stage profiler to include AR timing in stage_durations",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
help="Enable detailed per-request pipeline stats logging",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
benchmark(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,154 +1,87 @@
|
||||
# GLM-Image Multistage End-to-End Inference
|
||||
# GLM-Image Offline Inference
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/glm_image>.
|
||||
|
||||
|
||||
This example demonstrates how to run GLM-Image with the vLLM-Omni multistage architecture.
|
||||
GLM-Image is a 2-stage image generation model (AR + Diffusion) supported by vLLM-Omni's
|
||||
declarative config system. The pipeline topology and stage structure are declared in
|
||||
`vllm_omni/model_executor/models/glm_image/pipeline.py`; deployment knobs live in
|
||||
`vllm_omni/deploy/glm_image.yaml`.
|
||||
|
||||
## Architecture
|
||||
|
||||
GLM-Image uses a 2-stage pipeline:
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ GLM-Image Pipeline │
|
||||
├─────────────────────────────────────────────────────────────┤
|
||||
│ │
|
||||
│ Stage 0 (AR Model) Stage 1 (Diffusion) │
|
||||
│ ┌─────────────────┐ ┌─────────────────────┐ │
|
||||
│ │ vLLM-optimized │ │ GlmImagePipeline │ │
|
||||
│ │ GlmImageFor │ prior │ ┌───────────────┐ │ │
|
||||
│ │ Conditional │──tokens───►│ │ DiT Denoiser │ │ │
|
||||
│ │ Generation │ │ └───────────────┘ │ │
|
||||
│ │ (9B AR model) │ │ │ │ │
|
||||
│ └─────────────────┘ │ ▼ │ │
|
||||
│ ▲ │ ┌───────────────┐ │ │
|
||||
│ │ │ │ VAE Decode │──┼──► Image
|
||||
│ Text/Image │ └───────────────┘ │ │
|
||||
│ Input └─────────────────────┘ │
|
||||
│ │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
Stage 0 (AR Model) Stage 1 (Diffusion)
|
||||
┌───────────────────┐ ┌─────────────────────┐
|
||||
│ vLLM-optimized │ prior │ GlmImagePipeline │
|
||||
│ GlmImageFor │──tokens──►│ ┌───────────────┐ │
|
||||
│ Conditional │ │ │ DiT Denoiser │ │
|
||||
│ Generation │ │ └───────┬───────┘ │
|
||||
│ (9B AR model) │ │ ▼ │
|
||||
└───────────────────┘ │ ┌───────────────┐ │
|
||||
▲ │ │ VAE Decode │──┼──► Image
|
||||
│ │ └───────────────┘ │
|
||||
Text / Image └─────────────────────┘
|
||||
Input
|
||||
```
|
||||
|
||||
## Features
|
||||
## Text-to-Image
|
||||
|
||||
- **vLLM-optimized AR**: Uses PagedAttention and tensor parallelism for faster prior token generation
|
||||
- **Flexible deployment**: AR and Diffusion stages can run on different GPUs
|
||||
- **Text-to-Image**: Generate images from text descriptions
|
||||
- **Image-to-Image**: Edit existing images with text prompts
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
## Usage
|
||||
|
||||
### Text-to-Image
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--prompt "A beautiful sunset over the ocean with sailing boats" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--output output_t2i.png
|
||||
if __name__ == "__main__":
|
||||
omni = Omni(model="zai-org/GLM-Image")
|
||||
outputs = omni.generate(
|
||||
"A photorealistic mountain landscape at sunset",
|
||||
sampling_params={
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1.5,
|
||||
"seed": 42,
|
||||
},
|
||||
)
|
||||
outputs[0].request_output.images[0].save("output.png")
|
||||
```
|
||||
|
||||
### Image-to-Image (Image Editing)
|
||||
## Image-to-Image (Image Editing)
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--prompt "Transform this scene into a winter wonderland" \
|
||||
--image input.png \
|
||||
--output output_i2i.png
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
if __name__ == "__main__":
|
||||
omni = Omni(model="zai-org/GLM-Image")
|
||||
outputs = omni.generate(
|
||||
{
|
||||
"prompt": "Convert this image to watercolor style",
|
||||
"multi_modal_data": {
|
||||
"image": "input.png",
|
||||
},
|
||||
},
|
||||
sampling_params={
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1.5,
|
||||
"seed": 42,
|
||||
},
|
||||
)
|
||||
outputs[0].request_output.images[0].save("output.png")
|
||||
```
|
||||
|
||||
### With Custom Parameters
|
||||
## Generation Parameters
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model-path /path/to/glm-image \
|
||||
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--prompt "A photorealistic cat sitting on a window sill" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--num-inference-steps 50 \
|
||||
--guidance-scale 1.5 \
|
||||
--seed 42 \
|
||||
--output output.png
|
||||
```
|
||||
| Parameter | Type | Default | Description |
|
||||
| --------------------- | ----- | ------- | ----------------------------------- |
|
||||
| `height` | int | 1024 | Image height in pixels |
|
||||
| `width` | int | 1024 | Image width in pixels |
|
||||
| `num_inference_steps` | int | 50 | Number of diffusion denoising steps |
|
||||
| `guidance_scale` | float | 1.5 | Classifier-free guidance scale |
|
||||
| `seed` | int | None | Optional random seed |
|
||||
| `negative_prompt` | str | None | Negative prompt |
|
||||
|
||||
## Shell Scripts
|
||||
## VRAM Requirements
|
||||
|
||||
### Run Text-to-Image
|
||||
|
||||
```bash
|
||||
./run_t2i.sh
|
||||
```
|
||||
|
||||
### Run Image-to-Image
|
||||
|
||||
```bash
|
||||
./run_i2i.sh --image /path/to/input.png
|
||||
```
|
||||
|
||||
## Stage Configuration
|
||||
|
||||
The stage config (`glm_image.yaml`) defines:
|
||||
|
||||
- **Stage 0 (AR)**: Uses `GPUARWorker` with vLLM engine
|
||||
|
||||
- Model: `GlmImageForConditionalGeneration`
|
||||
- Output: `token_ids` (prior tokens)
|
||||
|
||||
- **Stage 1 (Diffusion)**: Uses diffusion engine
|
||||
- Model: `GlmImagePipeline`
|
||||
- Output: Generated image
|
||||
|
||||
See `vllm_omni/model_executor/stage_configs/glm_image.yaml` for full configuration.
|
||||
|
||||
## Comparison with Single-Stage
|
||||
|
||||
| Aspect | Single-Stage (transformers) | Multistage (vLLM) |
|
||||
| ----------- | --------------------------- | ------------------- |
|
||||
| AR Model | transformers native | vLLM PagedAttention |
|
||||
| Memory | Higher (no KV cache opt) | Lower (optimized) |
|
||||
| Throughput | Lower | Higher |
|
||||
| Flexibility | Single GPU | Multi-GPU support |
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### OOM Error
|
||||
|
||||
Try reducing memory usage:
|
||||
|
||||
```bash
|
||||
# In glm_image.yaml, adjust:
|
||||
gpu_memory_utilization: 0.5 # Reduce from 0.6
|
||||
```
|
||||
|
||||
### Slow Initialization
|
||||
|
||||
The first run loads model weights. Subsequent runs are faster:
|
||||
|
||||
```bash
|
||||
--stage-init-timeout 900 # Increase timeout for slow storage
|
||||
```
|
||||
|
||||
## Requirements
|
||||
|
||||
- vLLM-Omni with GLM-Image support
|
||||
- CUDA-capable GPU (recommended: H100/A100 with 80GB)
|
||||
- GLM-Image model weights
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/glm_image/end2end.py"
|
||||
``````
|
||||
??? abstract "run_i2i.sh"
|
||||
``````sh
|
||||
--8<-- "examples/offline_inference/glm_image/run_i2i.sh"
|
||||
``````
|
||||
??? abstract "run_t2i.sh"
|
||||
``````sh
|
||||
--8<-- "examples/offline_inference/glm_image/run_t2i.sh"
|
||||
``````
|
||||
| Stage | VRAM |
|
||||
| :---------------- | :--------------------- |
|
||||
| Stage-0 (AR) | **~18 GiB + KV Cache** |
|
||||
| Stage-1 (DiT+VAE) | **~20 GiB** |
|
||||
| Total | **~38 GiB + KV Cache** |
|
||||
@@ -1,99 +1,96 @@
|
||||
# GLM-Image Online Serving
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/glm_image>.
|
||||
GLM-Image is a 2-stage image generation model (AR + Diffusion) supported by vLLM-Omni's
|
||||
declarative config system. The pipeline topology and stage structure are declared in
|
||||
`vllm_omni/model_executor/models/glm_image/pipeline.py`; deployment knobs (GPU placement,
|
||||
memory, sampling params) live in `vllm_omni/deploy/glm_image.yaml`.
|
||||
|
||||
|
||||
This example demonstrates how to deploy GLM-Image for online image generation using vLLM-Omni.
|
||||
|
||||
## 🛠️ Installation
|
||||
|
||||
Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
|
||||
|
||||
## Run examples (GLM-Image)
|
||||
|
||||
**Note**: These examples work with the default configuration on **2× NVIDIA A100 (80GB)** or equivalent. Stage 0 (AR) and Stage 1 (Diffusion) each use one GPU by default. For single-GPU setups, modify the stage configuration to share the same device.
|
||||
|
||||
### Launch the Server
|
||||
## Start Server
|
||||
|
||||
```bash
|
||||
# Use default configuration
|
||||
vllm serve zai-org/GLM-Image --omni --port 8091
|
||||
```
|
||||
|
||||
Or use the convenience script:
|
||||
The config system auto-detects the pipeline from the model's `model_index.json` — no
|
||||
manual `--stage-configs-path` or `--deploy-config` needed.
|
||||
|
||||
By default, stage 0 (AR) runs on GPU 0 and stage 1 (Diffusion) on GPU 1. To colocate
|
||||
both stages on a single GPU, override per stage:
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/glm_image
|
||||
bash run_server.sh
|
||||
vllm serve zai-org/GLM-Image --omni --port 8091 \
|
||||
--stage-0-devices 0 --stage-1-devices 0
|
||||
```
|
||||
|
||||
If you have a custom stage configs file:
|
||||
|
||||
```bash
|
||||
vllm serve zai-org/GLM-Image --omni --port 8091 --stage-configs-path /path/to/glm_image.yaml
|
||||
```
|
||||
|
||||
### Send Requests
|
||||
|
||||
Get into the glm_image folder:
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/glm_image
|
||||
```
|
||||
|
||||
Send request via Python:
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py --prompt "A cute cat sitting on a window sill"
|
||||
```
|
||||
|
||||
The Python client supports the following command-line arguments:
|
||||
|
||||
- `--prompt` (or `-p`): Text prompt for generation (default: `A beautiful sunset over the ocean with sailing boats`)
|
||||
- `--output` (or `-o`): Output file path (default: `glm_image_output.png`)
|
||||
- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
|
||||
- `--image` (or `-i`): Input image path (for image-to-image editing)
|
||||
- `--height`: Image height in pixels (default: 1024)
|
||||
- `--width`: Image width in pixels (default: 1024)
|
||||
- `--steps`: Number of inference steps (default: 50)
|
||||
- `--guidance-scale`: Classifier-free guidance scale (default: 1.5)
|
||||
- `--seed`: Random seed (default: 42)
|
||||
- `--negative`: Negative prompt
|
||||
|
||||
## Modality Control
|
||||
|
||||
GLM-Image supports **text-to-image** and **image-to-image** modes.
|
||||
|
||||
The default yaml configuration deploys AR on GPU 0 and DiT on GPU 1. You can use the default configuration file: [`glm_image.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/glm_image.yaml)
|
||||
|
||||
| Mode | Input | Output | Description |
|
||||
| -------------- | ------------ | ------ | ---------------------------------- |
|
||||
| Text-to-Image | Text | Image | Generate images from text prompts |
|
||||
| Image-to-Image | Image + Text | Image | Edit images with text instructions |
|
||||
## API Calls
|
||||
|
||||
### Text-to-Image
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "A photorealistic mountain landscape at sunset" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--output landscape.png
|
||||
|
||||
# Or use the curl script:
|
||||
bash run_curl_text_to_image.sh "A futuristic city skyline at night"
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{"role": "user", "content": "A photorealistic mountain landscape at sunset"}
|
||||
],
|
||||
"extra_body": {
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1.5,
|
||||
"seed": 42
|
||||
}
|
||||
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
|
||||
```
|
||||
|
||||
### Image-to-Image (Image Editing)
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "Convert this image to watercolor style" \
|
||||
--image input.png \
|
||||
--output watercolor.png
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Convert this image to watercolor style"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,$(base64 -w0 input.png)}"}
|
||||
]
|
||||
}
|
||||
],
|
||||
"extra_body": {
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1.5
|
||||
}
|
||||
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
|
||||
```
|
||||
|
||||
# Or use the curl script:
|
||||
bash run_curl_image_edit.sh input.png "Convert to watercolor style"
|
||||
### Using the OpenAI Python SDK
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
import base64
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="zai-org/GLM-Image",
|
||||
messages=[{"role": "user", "content": "A beautiful sunset over the ocean"}],
|
||||
extra_body={
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1.5,
|
||||
"seed": 42,
|
||||
},
|
||||
)
|
||||
|
||||
img_url = response.choices[0].message.content[0].image_url.url
|
||||
_, b64_data = img_url.split(",", 1)
|
||||
with open("output.png", "wb") as f:
|
||||
f.write(base64.b64decode(b64_data))
|
||||
```
|
||||
|
||||
For general-purpose request methods (curl, OpenAI SDK, Python `requests`), see
|
||||
@@ -116,7 +113,7 @@ dimensions and count, use `size` and `n` rather than `height` or `width`.
|
||||
| `width` | int | 1024 | Image width in pixels |
|
||||
| `num_inference_steps` | int | 50 | Number of diffusion denoising steps |
|
||||
| `guidance_scale` | float | 1.5 | Classifier-free guidance scale |
|
||||
| `seed` | int | None | Optional random seed; `/v1/images/*` generates one server-side if omitted |
|
||||
| `seed` | int | None | Optional random seed |
|
||||
| `negative_prompt` | str | None | Negative prompt |
|
||||
|
||||
## Response Format
|
||||
@@ -150,13 +147,12 @@ dimensions and count, use `size` and `n` rather than `height` or `width`.
|
||||
## Extract Image
|
||||
|
||||
```bash
|
||||
# From a saved JSON response
|
||||
cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
|
||||
```
|
||||
|
||||
## Architecture
|
||||
|
||||
GLM-Image uses a 2-stage multistage pipeline:
|
||||
GLM-Image uses a 2-stage pipeline:
|
||||
|
||||
```
|
||||
Stage 0 (AR Model) Stage 1 (Diffusion)
|
||||
@@ -181,41 +177,13 @@ Stage 0 (AR Model) Stage 1 (Diffusion)
|
||||
| Stage-1 (DiT+VAE) | **~20 GiB** |
|
||||
| Total | **~38 GiB + KV Cache** |
|
||||
|
||||
## File Description
|
||||
|
||||
| File | Description |
|
||||
| --------------------------- | ------------------------------------- |
|
||||
| `run_server.sh` | Server startup script |
|
||||
| `run_curl_text_to_image.sh` | Text-to-image curl example |
|
||||
| `run_curl_image_edit.sh` | Image-to-image (editing) curl example |
|
||||
| `openai_chat_client.py` | Python client (t2i + i2i) |
|
||||
|
||||
## FAQ
|
||||
|
||||
- If you encounter OOM errors, adjust `gpu_memory_utilization` in the stage config:
|
||||
- If you encounter OOM errors, adjust `gpu_memory_utilization` in the deploy config:
|
||||
|
||||
```yaml
|
||||
# In glm_image.yaml, reduce from default 0.6:
|
||||
# In vllm_omni/deploy/glm_image.yaml, reduce from default 0.6:
|
||||
gpu_memory_utilization: 0.5
|
||||
```
|
||||
|
||||
- The first request may be slow due to model warmup. Subsequent requests will be faster.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "openai_chat_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/glm_image/openai_chat_client.py"
|
||||
``````
|
||||
??? abstract "run_curl_image_edit.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/glm_image/run_curl_image_edit.sh"
|
||||
``````
|
||||
??? abstract "run_curl_text_to_image.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/glm_image/run_curl_text_to_image.sh"
|
||||
``````
|
||||
??? abstract "run_server.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/glm_image/run_server.sh"
|
||||
``````
|
||||
@@ -28,7 +28,6 @@ from vllm_omni.entrypoints.openai.image_api_utils import (
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.model_executor.stage_input_processors.glm_image import compute_max_tokens
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -216,7 +215,11 @@ def async_omni_test_client():
|
||||
]
|
||||
default_sampling_params_list = [
|
||||
SamplingParams(temperature=0.1),
|
||||
OmniDiffusionSamplingParams(),
|
||||
OmniDiffusionSamplingParams(
|
||||
num_inference_steps=4,
|
||||
guidance_scale=7.5,
|
||||
generator_device="cpu",
|
||||
),
|
||||
]
|
||||
self.engine = SimpleNamespace(
|
||||
stage_configs=stage_configs,
|
||||
@@ -576,7 +579,11 @@ def test_multistage_images_async_omni_construction(async_omni_test_client):
|
||||
|
||||
|
||||
def test_generate_images_async_omni_glm_image_sets_stage0_max_tokens():
|
||||
"""GLM-Image multistage requests must compute stage-0 max_tokens from size."""
|
||||
"""GLM-Image multistage: stage-0 gets target_h/w from requested size.
|
||||
|
||||
max_tokens comes from the deploy YAML default (upper-bound ceiling),
|
||||
NOT computed dynamically from height/width.
|
||||
"""
|
||||
|
||||
class FakeAsyncOmniClass(AsyncOmni):
|
||||
def __init__(self):
|
||||
@@ -584,8 +591,9 @@ def test_generate_images_async_omni_glm_image_sets_stage0_max_tokens():
|
||||
SimpleNamespace(stage_type="llm", is_comprehension=True, model_arch="GlmImageForConditionalGeneration"),
|
||||
SimpleNamespace(stage_type="diffusion", is_comprehension=False, model_arch="GlmImagePipeline"),
|
||||
]
|
||||
# YAML default max_tokens for GLM-Image AR stage (upper bound for 2048x2048 t2i)
|
||||
default_sampling_params_list = [
|
||||
SamplingParams(temperature=0.1, seed=42),
|
||||
SamplingParams(temperature=0.1, seed=42, max_tokens=4353),
|
||||
OmniDiffusionSamplingParams(height=1024, width=1024),
|
||||
]
|
||||
self.engine = SimpleNamespace(
|
||||
@@ -646,7 +654,8 @@ def test_generate_images_async_omni_glm_image_sets_stage0_max_tokens():
|
||||
captured = engine.captured_sampling_params_list
|
||||
assert captured is not None
|
||||
assert len(captured) == 2
|
||||
assert captured[0].max_tokens == compute_max_tokens(1024, 1024, is_i2i=False)
|
||||
# max_tokens comes from YAML default, not computed dynamically
|
||||
assert captured[0].max_tokens == 4353
|
||||
assert captured[0].extra_args["target_h"] == 1024
|
||||
assert captured[0].extra_args["target_w"] == 1024
|
||||
assert captured[1].height == 1024
|
||||
@@ -1161,9 +1170,10 @@ def test_image_edit_ignores_mock_like_multimodal_limit(async_omni_test_client):
|
||||
assert response.status_code == 200
|
||||
captured_prompt = engine.captured_prompt
|
||||
assert captured_prompt is not None
|
||||
processed_images = captured_prompt["multi_modal_data"]["image"]
|
||||
assert len(processed_images) == 1
|
||||
assert processed_images[0].size == (16, 16)
|
||||
# Multi-stage path uses "img2img" key for single reference image
|
||||
processed_images = captured_prompt["multi_modal_data"]["img2img"]
|
||||
assert isinstance(processed_images, Image.Image)
|
||||
assert processed_images.size == (16, 16)
|
||||
|
||||
|
||||
def test_image_edit_parameter_pass(async_omni_test_client):
|
||||
|
||||
@@ -38,7 +38,7 @@ def default_comprehension_params():
|
||||
temperature=0.4,
|
||||
top_p=0.9,
|
||||
top_k=1,
|
||||
max_tokens=2048,
|
||||
max_tokens=4353,
|
||||
seed=42,
|
||||
repetition_penalty=1.05,
|
||||
)
|
||||
@@ -146,7 +146,7 @@ def test_preserves_yaml_defaults_when_no_request_params(serving_chat, mock_reque
|
||||
assert comprehension_params.temperature == 0.4
|
||||
assert comprehension_params.top_p == 0.9
|
||||
assert comprehension_params.top_k == 1 # YAML custom param preserved
|
||||
assert comprehension_params.max_tokens == 2048
|
||||
assert comprehension_params.max_tokens == 4353
|
||||
assert comprehension_params.seed == 42
|
||||
assert comprehension_params.repetition_penalty == 1.05 # YAML custom param preserved
|
||||
|
||||
@@ -190,7 +190,7 @@ def test_max_tokens_uses_yaml_default_when_not_specified(serving_chat, mock_requ
|
||||
"""Test that max_tokens falls back to YAML default when not in request."""
|
||||
result = serving_chat._build_sampling_params_list_from_request(mock_request)
|
||||
|
||||
assert result[0].max_tokens == 2048
|
||||
assert result[0].max_tokens == 4353
|
||||
|
||||
|
||||
def test_request_seed_overrides_yaml_default(serving_chat, mock_request):
|
||||
@@ -588,12 +588,17 @@ class TestResolveHeightWidth:
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Tests for _apply_request_overrides with GLM-Image (max_tokens computation)
|
||||
# Tests for _apply_request_overrides with GLM-Image (target_h/w injection)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestApplyRequestOverridesGLMImage:
|
||||
"""Test dynamic max_tokens computation for GLM-Image AR stage."""
|
||||
"""Test target_h/w injection for GLM-Image AR stage.
|
||||
|
||||
max_tokens is NOT computed dynamically — it comes from the deploy YAML
|
||||
default (e.g. 4353). _apply_request_overrides only injects target_h/w
|
||||
into extra_args so the model can build M-RoPE position grids.
|
||||
"""
|
||||
|
||||
@pytest.fixture
|
||||
def glm_serving_chat(self, mock_engine_client, mocker: MockerFixture):
|
||||
@@ -601,7 +606,6 @@ class TestApplyRequestOverridesGLMImage:
|
||||
|
||||
instance = object.__new__(OmniOpenAIServingChat)
|
||||
instance.engine_client = mock_engine_client
|
||||
# Mock the image extraction to return no reference images (t2i by default)
|
||||
instance._extract_diffusion_prompt_and_images_from_messages = mocker.MagicMock(return_value=("a cat", []))
|
||||
return instance
|
||||
|
||||
@@ -623,40 +627,41 @@ class TestApplyRequestOverridesGLMImage:
|
||||
req.model_fields_set = set()
|
||||
return req
|
||||
|
||||
def test_t2i_computes_max_tokens(self, glm_serving_chat, glm_request, default_comprehension_params):
|
||||
"""t2i mode: max_tokens computed from height/width, no reference images."""
|
||||
def test_t2i_injects_target_h_w(self, glm_serving_chat, glm_request, default_comprehension_params):
|
||||
"""t2i mode: target_h/w injected into extra_args, max_tokens unchanged."""
|
||||
result = glm_serving_chat._apply_request_overrides(default_comprehension_params, glm_request)
|
||||
# t2i 1024x1024 = 256 + 1024 + 1 = 1281
|
||||
assert result.max_tokens == 1281
|
||||
assert result.extra_args["target_h"] == 1024
|
||||
assert result.extra_args["target_w"] == 1024
|
||||
# max_tokens stays at YAML default (not dynamically computed)
|
||||
assert result.max_tokens == 4353
|
||||
|
||||
def test_i2i_computes_fewer_tokens(
|
||||
def test_i2i_injects_target_h_w(
|
||||
self, glm_serving_chat, glm_request, default_comprehension_params, mocker: MockerFixture
|
||||
):
|
||||
"""i2i mode: max_tokens should be smaller than t2i for same dimensions."""
|
||||
# Make it detect reference images
|
||||
"""i2i mode: target_h/w injected, max_tokens unchanged."""
|
||||
glm_serving_chat._extract_diffusion_prompt_and_images_from_messages = mocker.MagicMock(
|
||||
return_value=("edit this", ["fake_image"])
|
||||
)
|
||||
|
||||
result = glm_serving_chat._apply_request_overrides(default_comprehension_params, glm_request)
|
||||
# i2i 1024x1024 = 1024 + 1 = 1025
|
||||
assert result.max_tokens == 1025
|
||||
assert result.extra_args["target_h"] == 1024
|
||||
assert result.extra_args["target_w"] == 1024
|
||||
# max_tokens stays at YAML default regardless of t2i/i2i
|
||||
assert result.max_tokens == 4353
|
||||
|
||||
def test_dynamic_max_tokens_overrides_user_value(self, glm_serving_chat, glm_request, default_comprehension_params):
|
||||
"""When height/width are provided, dynamic computation overrides user max_tokens."""
|
||||
def test_user_max_tokens_preserved(self, glm_serving_chat, glm_request, default_comprehension_params):
|
||||
"""User-provided max_tokens is respected (not overridden by dynamic computation)."""
|
||||
glm_request.max_tokens = 500
|
||||
glm_request.model_fields_set = {"max_tokens"}
|
||||
|
||||
result = glm_serving_chat._apply_request_overrides(default_comprehension_params, glm_request)
|
||||
# Dynamic computation from height/width always wins when present
|
||||
assert result.max_tokens == 1281
|
||||
assert result.max_tokens == 500
|
||||
assert result.extra_args["target_h"] == 1024
|
||||
assert result.extra_args["target_w"] == 1024
|
||||
|
||||
def test_no_height_width_preserves_default(
|
||||
self, glm_serving_chat, mocker: MockerFixture, default_comprehension_params
|
||||
):
|
||||
"""When no height/width in extra_body, keep YAML default max_tokens."""
|
||||
"""When no height/width in extra_body, keep YAML default max_tokens, no target_h/w."""
|
||||
req = mocker.MagicMock()
|
||||
req.temperature = None
|
||||
req.top_p = None
|
||||
@@ -673,7 +678,9 @@ class TestApplyRequestOverridesGLMImage:
|
||||
req.model_fields_set = set()
|
||||
|
||||
result = glm_serving_chat._apply_request_overrides(default_comprehension_params, req)
|
||||
assert result.max_tokens == 2048 # YAML default
|
||||
assert result.max_tokens == 4353 # YAML default
|
||||
# No target_h/w injected when dimensions not provided
|
||||
assert not result.extra_args or "target_h" not in (result.extra_args or {})
|
||||
|
||||
def test_size_string_parsed_for_glm_image(
|
||||
self, glm_serving_chat, mocker: MockerFixture, default_comprehension_params
|
||||
@@ -695,5 +702,7 @@ class TestApplyRequestOverridesGLMImage:
|
||||
req.model_fields_set = set()
|
||||
|
||||
result = glm_serving_chat._apply_request_overrides(default_comprehension_params, req)
|
||||
# 512x512 t2i = 256 + 256 + 1 = 513
|
||||
assert result.max_tokens == 513
|
||||
assert result.extra_args["target_h"] == 512
|
||||
assert result.extra_args["target_w"] == 512
|
||||
# max_tokens stays at YAML default (not dynamically computed)
|
||||
assert result.max_tokens == 4353
|
||||
@@ -10,6 +10,10 @@ async_chunk: false
|
||||
stages:
|
||||
# Stage 0: AR Model (GlmImageForConditionalGeneration)
|
||||
# Generates prior_token_ids for conditioning the diffusion process.
|
||||
# max_tokens is set to the theoretical maximum (2048x2048 t2i ≈ 4353
|
||||
# tokens). The AR model stops naturally at EOS (16385) well before
|
||||
# this ceiling at lower resolutions. Do NOT override max_tokens per
|
||||
# request — the actual resolution is communicated via target_h/w.
|
||||
- stage_id: 0
|
||||
max_num_seqs: 1
|
||||
gpu_memory_utilization: 0.6
|
||||
@@ -21,6 +25,7 @@ stages:
|
||||
top_p: 0.75
|
||||
top_k: 16512
|
||||
stop_token_ids: [16385]
|
||||
max_tokens: 4353
|
||||
seed: 42
|
||||
detokenize: false
|
||||
|
||||
|
||||
@@ -1032,6 +1032,7 @@ class AsyncOmniEngine:
|
||||
original_prompt = prompt
|
||||
|
||||
stage_type = self.stage_metadata[0].get("stage_type")
|
||||
_preprocess_ms = 0.0
|
||||
if stage_type != "diffusion" and not isinstance(prompt, EngineCoreRequest):
|
||||
# Inject global_request_id into the raw prompt.
|
||||
if isinstance(prompt, dict):
|
||||
@@ -1041,6 +1042,7 @@ class AsyncOmniEngine:
|
||||
_inject_global_id(item, request_id)
|
||||
|
||||
# Full input processing (tokenization, multimodal, etc.)
|
||||
_t_preprocess = time.perf_counter()
|
||||
request = self.input_processor.process_inputs(
|
||||
request_id=request_id,
|
||||
prompt=prompt,
|
||||
@@ -1054,6 +1056,7 @@ class AsyncOmniEngine:
|
||||
data_parallel_rank=data_parallel_rank,
|
||||
resumable=resumable,
|
||||
)
|
||||
_preprocess_ms = (time.perf_counter() - _t_preprocess) * 1000.0
|
||||
# TODO (Peiqi): add this for Qwen3-TTS only. Other models don't have
|
||||
# additional_information field in the prompt.
|
||||
request = _upgrade_to_omni_request(request, prompt)
|
||||
@@ -1090,6 +1093,8 @@ class AsyncOmniEngine:
|
||||
"original_prompt": original_prompt,
|
||||
"sampling_params_list": effective_sampling_params_list,
|
||||
"final_stage_id": final_stage_id,
|
||||
"preprocess_ms": _preprocess_ms,
|
||||
"enqueue_ts": time.perf_counter(),
|
||||
}
|
||||
|
||||
def _enqueue_cfg_companions(
|
||||
@@ -1316,6 +1321,7 @@ class AsyncOmniEngine:
|
||||
"num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
|
||||
"quantization": kwargs.get("quantization", None),
|
||||
"enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False),
|
||||
"enable_ar_profiler": kwargs.get("enable_ar_profiler", False),
|
||||
**(
|
||||
{
|
||||
"profiler_config": asdict(kwargs["profiler_config"])
|
||||
@@ -1458,6 +1464,17 @@ class AsyncOmniEngine:
|
||||
or cfg.engine_args.quantization_config is None
|
||||
):
|
||||
cfg.engine_args.quantization_config = quantization_config
|
||||
# Inject profiler flags for diffusion stages
|
||||
for profiler_key in (
|
||||
"enable_diffusion_pipeline_profiler",
|
||||
"enable_ar_profiler",
|
||||
):
|
||||
val = kwargs.get(profiler_key)
|
||||
if val:
|
||||
if not hasattr(cfg.engine_args, profiler_key) or not getattr(
|
||||
cfg.engine_args, profiler_key, False
|
||||
):
|
||||
setattr(cfg.engine_args, profiler_key, val)
|
||||
quantization = kwargs.get("quantization")
|
||||
if quantization is not None:
|
||||
if not hasattr(cfg.engine_args, "quantization") or cfg.engine_args.quantization is None:
|
||||
|
||||
@@ -114,6 +114,9 @@ class OrchestratorRequestState:
|
||||
|
||||
streaming: StreamingInputState = field(default_factory=lambda: StreamingInputState())
|
||||
|
||||
# Per-request pipeline timing accumulator (milliseconds)
|
||||
pipeline_timings: dict[str, float] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamingInputState:
|
||||
@@ -625,6 +628,7 @@ class Orchestrator:
|
||||
total_token=self._agg_total_tokens[stage_id],
|
||||
total_gen_time_ms=self._agg_total_gen_time_ms[stage_id],
|
||||
),
|
||||
pipeline_timings=dict(req_state.pipeline_timings),
|
||||
)
|
||||
|
||||
def _build_kv_sender_info(self, sender_stage_ids: list[int]) -> dict[int, dict[str, Any]] | None:
|
||||
@@ -682,6 +686,7 @@ class Orchestrator:
|
||||
False,
|
||||
)
|
||||
_dt_ar2d = (_time.perf_counter() - _t_ar2d) * 1000
|
||||
req_state.pipeline_timings["ar2diffusion_ms"] = _dt_ar2d
|
||||
logger.info(
|
||||
"[Orchestrator] ar2diffusion req=%s wall_time=%.3fms stage=%d->%d",
|
||||
req_id,
|
||||
@@ -911,6 +916,15 @@ class Orchestrator:
|
||||
)
|
||||
req_state.streaming.enabled = is_streaming
|
||||
req_state.stage_submit_ts[stage_id] = _time.time()
|
||||
|
||||
# Per-request pipeline timings from caller thread
|
||||
_enqueue_ts = msg.get("enqueue_ts", 0.0)
|
||||
if _enqueue_ts > 0:
|
||||
req_state.pipeline_timings["queue_wait_ms"] = (_time.perf_counter() - _enqueue_ts) * 1000.0
|
||||
_preprocess_ms = msg.get("preprocess_ms", 0.0)
|
||||
if _preprocess_ms > 0:
|
||||
req_state.pipeline_timings["preprocess_ms"] = _preprocess_ms
|
||||
|
||||
self.request_states[request_id] = req_state
|
||||
|
||||
# Stage-0 prompt is already a fully-formed OmniEngineCoreRequest
|
||||
|
||||
@@ -477,6 +477,11 @@ class OmniServeCommand(CLISubcommand):
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
omni_config_group.add_argument(
|
||||
"--enable-ar-profiler",
|
||||
action="store_true",
|
||||
help="Enable AR stage profiler to include AR stage timing in stage_durations.",
|
||||
)
|
||||
# Stash via type(self) so the docs hook (which execs this function in a
|
||||
# sandboxed globals dict via ``DummySelf``) doesn't fail on a NameError.
|
||||
type(self)._parser = serve_parser
|
||||
|
||||
@@ -124,6 +124,7 @@ class OmniBase(PDDisaggregationMixin):
|
||||
stage_init_timeout = kwargs.pop("stage_init_timeout", 300)
|
||||
init_timeout = kwargs.pop("init_timeout", 600)
|
||||
log_stats = kwargs.pop("log_stats", False)
|
||||
self._enable_ar_profiler = kwargs.pop("enable_ar_profiler", False)
|
||||
# NOTE: read-only lookup — must NOT pop. Popping here drops the key
|
||||
# before it reaches ``StageConfigFactory._create_from_registry``, so
|
||||
# ``--no-async-chunk`` (``async_chunk=False``) silently fails to
|
||||
@@ -375,6 +376,30 @@ class OmniBase(PDDisaggregationMixin):
|
||||
engine_outputs = result.get("engine_outputs")
|
||||
stage_durations = getattr(result["engine_outputs"], "stage_durations", {})
|
||||
peak_memory_mb = getattr(result["engine_outputs"], "peak_memory_mb", 0.0)
|
||||
|
||||
# Merge AR stage timing from OrchestratorAggregator.stage_events
|
||||
if self._enable_ar_profiler:
|
||||
ar_events = metrics.stage_events.get(str(req_id), [])
|
||||
for evt in ar_events:
|
||||
if evt.stage_id != stage_id:
|
||||
stage_durations[f"ar_stage_{evt.stage_id}"] = evt.stage_gen_time_ms / 1000.0
|
||||
|
||||
# Merge pipeline timings from Orchestrator into stage_durations
|
||||
_m = result.get("metrics")
|
||||
if _m is not None and hasattr(_m, "pipeline_timings") and _m.pipeline_timings:
|
||||
for key, value in _m.pipeline_timings.items():
|
||||
if key not in stage_durations:
|
||||
stage_durations[key] = value
|
||||
|
||||
# Merge per-stage gen times into stage_durations
|
||||
for evt in metrics.stage_events.get(str(req_id), []):
|
||||
key = f"stage_{evt.stage_id}_gen_ms"
|
||||
if key not in stage_durations:
|
||||
stage_durations[key] = evt.stage_gen_time_ms
|
||||
# Current stage gen time (not yet in stage_events at this point)
|
||||
if _m is not None:
|
||||
stage_durations.setdefault(f"stage_{stage_id}_gen_ms", _m.stage_gen_time_ms)
|
||||
|
||||
finished = engine_outputs.finished
|
||||
|
||||
submit_ts = result.get("stage_submit_ts")
|
||||
|
||||
@@ -1793,20 +1793,85 @@ async def edit_images(
|
||||
_update_if_not_none(gen_params, "layers", layers)
|
||||
_update_if_not_none(gen_params, "resolution", resolution)
|
||||
|
||||
# 4. Generate images using AsyncOmni (multi-stage mode)
|
||||
# 4. Generate images
|
||||
request_id = f"img_edit-{random_uuid()}"
|
||||
raw_request.state.request_metadata = RequestResponseMetadata(request_id=request_id)
|
||||
logger.info(f"Generating {n} image(s) {size_str}")
|
||||
result = await _generate_with_async_omni(
|
||||
engine_client=engine_client,
|
||||
gen_params=gen_params,
|
||||
stage_configs=stage_configs,
|
||||
prompt=prompt,
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
# 5. Extract images from result
|
||||
images = _extract_images_from_result(result)
|
||||
if len(stage_configs) > 1:
|
||||
# Multi-stage pipeline (e.g. GLM-Image AR+Diffusion): route through
|
||||
# the chat handler so the AR stage gets correct max_tokens and
|
||||
# target_h/w (same path as /v1/images/generations).
|
||||
chat_handler = getattr(raw_request.app.state, "openai_serving_chat", None)
|
||||
if chat_handler is None:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
|
||||
detail="openai_serving_chat is not initialized for multi-stage image editing.",
|
||||
)
|
||||
|
||||
# Encode input images to base64 for generate_diffusion_images.
|
||||
import base64
|
||||
import io as _io
|
||||
|
||||
ref_b64_list: list[str] = []
|
||||
for _img in pil_images:
|
||||
buf = _io.BytesIO()
|
||||
_img.save(buf, format="PNG")
|
||||
ref_b64_list.append(base64.b64encode(buf.getvalue()).decode())
|
||||
|
||||
effective_seed = seed if seed is not None else random.randint(0, MAX_UINT32_SEED)
|
||||
extra_body: dict[str, Any] = {
|
||||
"seed": effective_seed,
|
||||
"num_outputs_per_prompt": n,
|
||||
}
|
||||
if width is not None:
|
||||
extra_body["width"] = width
|
||||
if height is not None:
|
||||
extra_body["height"] = height
|
||||
if negative_prompt is not None:
|
||||
extra_body["negative_prompt"] = negative_prompt
|
||||
if num_inference_steps is not None:
|
||||
extra_body["num_inference_steps"] = num_inference_steps
|
||||
if guidance_scale is not None:
|
||||
extra_body["guidance_scale"] = guidance_scale
|
||||
if strength is not None:
|
||||
extra_body["strength"] = strength
|
||||
if true_cfg_scale is not None:
|
||||
extra_body["true_cfg_scale"] = true_cfg_scale
|
||||
if layers is not None:
|
||||
extra_body["layers"] = layers
|
||||
if resolution is not None:
|
||||
extra_body["resolution"] = resolution
|
||||
if lora is not None:
|
||||
# Validate LoRA, then pass through.
|
||||
lora_dict = _get_lora_from_json_str(lora)
|
||||
_parse_lora_request(lora_dict)
|
||||
extra_body["lora"] = lora_dict
|
||||
|
||||
prompt_text = prompt.get("prompt", "")
|
||||
generation_result = await chat_handler.generate_diffusion_images(
|
||||
prompt=prompt_text,
|
||||
extra_body=extra_body,
|
||||
reference_images=ref_b64_list,
|
||||
request_id=request_id,
|
||||
)
|
||||
if isinstance(generation_result, ErrorResponse):
|
||||
raise HTTPException(
|
||||
status_code=generation_result.error.code if generation_result.error else 400,
|
||||
detail=generation_result.message,
|
||||
)
|
||||
images, _, _ = generation_result
|
||||
else:
|
||||
# Single-stage diffusion: use the direct path.
|
||||
result = await _generate_with_async_omni(
|
||||
engine_client=engine_client,
|
||||
gen_params=gen_params,
|
||||
stage_configs=stage_configs,
|
||||
prompt=prompt,
|
||||
request_id=request_id,
|
||||
)
|
||||
images = _extract_images_from_result(result)
|
||||
|
||||
logger.info(f"Successfully generated {len(images)} image(s)")
|
||||
|
||||
# Encode images to base64
|
||||
|
||||
@@ -705,8 +705,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
Starts with YAML defaults and only overrides fields that the user
|
||||
explicitly provided (non-None values) in the request.
|
||||
|
||||
For GLM-Image AR stage, if max_tokens is not in YAML and user provides
|
||||
height/width in extra_body, computes max_tokens dynamically.
|
||||
For models needing spatial metadata (e.g. GLM-Image), target_h/w is
|
||||
injected into extra_args so the runner can build M-RoPE position grids.
|
||||
max_tokens is NOT computed dynamically — it uses the deploy YAML default.
|
||||
|
||||
Args:
|
||||
default_params: Default SamplingParams from stage config YAML.
|
||||
@@ -736,36 +737,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
extra_body = getattr(request, "extra_body", {}) or {}
|
||||
height, width = self._resolve_height_width_from_extra_body(extra_body)
|
||||
|
||||
# Best-effort mode detection from user messages.
|
||||
# i2i requests include at least one reference image in message content.
|
||||
_, reference_images = self._extract_diffusion_prompt_and_images_from_messages(request.messages)
|
||||
ref_image_count = len(reference_images)
|
||||
is_img2img = ref_image_count > 0
|
||||
|
||||
if height is not None and width is not None:
|
||||
try:
|
||||
from vllm_omni.model_executor.stage_input_processors.glm_image import compute_max_tokens
|
||||
|
||||
max_tokens = getattr(explicit_fields, "max_tokens", None)
|
||||
if max_tokens is None:
|
||||
max_tokens = compute_max_tokens(int(height), int(width), is_i2i=is_img2img)
|
||||
params.max_tokens = max_tokens
|
||||
# Keep target size in stage-0 sampling params so runner/model can
|
||||
# build deterministic M-RoPE grids for t2i (no MM features).
|
||||
extra_args = dict(getattr(params, "extra_args", {}) or {})
|
||||
extra_args["target_h"] = int(height)
|
||||
extra_args["target_w"] = int(width)
|
||||
params.extra_args = extra_args
|
||||
except (ImportError, ValueError, TypeError) as e:
|
||||
logger.warning(f"Failed to compute max_tokens: {e}, using default if available")
|
||||
else:
|
||||
logger.info(
|
||||
"[SamplingParams] Skip dynamic max_tokens (height=%s, width=%s, mode=%s, ref_images=%s)",
|
||||
height,
|
||||
width,
|
||||
"i2i" if is_img2img else "t2i",
|
||||
ref_image_count,
|
||||
)
|
||||
# Keep target size in stage-0 sampling params so runner/model can
|
||||
# build deterministic M-RoPE grids for t2i (no MM features).
|
||||
extra_args = dict(getattr(params, "extra_args", {}) or {})
|
||||
extra_args["target_h"] = int(height)
|
||||
extra_args["target_w"] = int(width)
|
||||
params.extra_args = extra_args
|
||||
|
||||
return params
|
||||
|
||||
@@ -2185,6 +2163,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
engine_prompt["mm_processor_kwargs"] = mm_processor_kwargs
|
||||
if engine_prompt_data is not None:
|
||||
engine_prompt["multi_modal_data"] = engine_prompt_data
|
||||
# Provide multi_modal_uuids so that newer vLLM versions can
|
||||
# validate multi_modal_data / multi_modal_uuids consistency.
|
||||
engine_prompt["multi_modal_uuids"] = {k: [f"img-{k}-{i}"] for i, k in enumerate(engine_prompt_data)}
|
||||
|
||||
comprehension_idx = None
|
||||
for idx, stage in enumerate(stage_configs):
|
||||
@@ -2192,24 +2173,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
comprehension_idx = idx
|
||||
break
|
||||
|
||||
is_glm_image_multistage = False
|
||||
for stage_cfg in stage_configs:
|
||||
if isinstance(stage_cfg, dict):
|
||||
engine_args = stage_cfg.get("engine_args")
|
||||
model_arch = stage_cfg.get("model_arch")
|
||||
else:
|
||||
engine_args = getattr(stage_cfg, "engine_args", None)
|
||||
model_arch = getattr(stage_cfg, "model_arch", None)
|
||||
|
||||
if model_arch is None and engine_args is not None:
|
||||
if isinstance(engine_args, dict):
|
||||
model_arch = engine_args.get("model_arch")
|
||||
else:
|
||||
model_arch = getattr(engine_args, "model_arch", None)
|
||||
if model_arch in {"GlmImageForConditionalGeneration", "GlmImagePipeline"}:
|
||||
is_glm_image_multistage = True
|
||||
break
|
||||
|
||||
sampling_params_list: list[Any] = []
|
||||
for idx, stage_cfg in enumerate(stage_configs):
|
||||
stage_type = get_stage_type(stage_cfg)
|
||||
@@ -2233,6 +2196,18 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
):
|
||||
default_stage_params.seed = seed
|
||||
|
||||
# Inject target_h/w into comprehension (AR) stage sampling params
|
||||
# for models that need M-RoPE position pre-computation (e.g.
|
||||
# GLM-Image). max_tokens is handled via the deploy YAML default
|
||||
# (upper-bound ceiling) rather than computed dynamically here.
|
||||
if comprehension_idx is not None and idx == comprehension_idx and height is not None and width is not None:
|
||||
extra_args = getattr(default_stage_params, "extra_args", None)
|
||||
if extra_args is None:
|
||||
extra_args = {}
|
||||
default_stage_params.extra_args = extra_args
|
||||
extra_args["target_h"] = int(height)
|
||||
extra_args["target_w"] = int(width)
|
||||
|
||||
if stage_type == "diffusion":
|
||||
self._set_if_supported(
|
||||
default_stage_params,
|
||||
@@ -2258,25 +2233,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
default_stage_params.lora_scale = lora_scale
|
||||
except Exception as e: # pragma: no cover - safeguard
|
||||
logger.warning("Failed to parse LoRA request: %s", e)
|
||||
elif (
|
||||
is_glm_image_multistage
|
||||
and comprehension_idx is not None
|
||||
and idx == comprehension_idx
|
||||
and isinstance(default_stage_params, SamplingParams)
|
||||
and height is not None
|
||||
and width is not None
|
||||
):
|
||||
from vllm_omni.model_executor.stage_input_processors.glm_image import compute_max_tokens
|
||||
|
||||
default_stage_params.max_tokens = compute_max_tokens(
|
||||
int(height),
|
||||
int(width),
|
||||
is_i2i=bool(reference_images),
|
||||
)
|
||||
extra_args = dict(getattr(default_stage_params, "extra_args", {}) or {})
|
||||
extra_args["target_h"] = int(height)
|
||||
extra_args["target_w"] = int(width)
|
||||
default_stage_params.extra_args = extra_args
|
||||
|
||||
sampling_params_list.append(default_stage_params)
|
||||
|
||||
@@ -2300,16 +2256,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
|
||||
engine = self._diffusion_engine if self._diffusion_engine is not None else self.engine_client
|
||||
|
||||
height = extra_body.get("height")
|
||||
width = extra_body.get("width")
|
||||
if "size" in extra_body:
|
||||
try:
|
||||
size_str = extra_body["size"]
|
||||
if isinstance(size_str, str) and "x" in size_str.lower():
|
||||
w, h = size_str.lower().split("x")
|
||||
width, height = int(w), int(h)
|
||||
except ValueError:
|
||||
logger.warning("Invalid size format: %s", extra_body.get("size"))
|
||||
height, width = self._resolve_height_width_from_extra_body(extra_body)
|
||||
|
||||
seed = extra_body.get("seed")
|
||||
generator_device = extra_body.get("generator_device")
|
||||
@@ -2341,6 +2288,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
guidance_scale_2=extra_body.get("guidance_scale_2"),
|
||||
layers=extra_body.get("layers"),
|
||||
resolution=extra_body.get("resolution"),
|
||||
strength=extra_body.get("strength"),
|
||||
)
|
||||
|
||||
if lora_body and isinstance(lora_body, dict):
|
||||
@@ -2461,16 +2409,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
extra_body = request.model_extra or {}
|
||||
|
||||
# Parse size if provided (supports "1024x1024" format)
|
||||
height = extra_body.get("height")
|
||||
width = extra_body.get("width")
|
||||
if "size" in extra_body:
|
||||
try:
|
||||
size_str = extra_body["size"]
|
||||
if isinstance(size_str, str) and "x" in size_str.lower():
|
||||
w, h = size_str.lower().split("x")
|
||||
width, height = int(w), int(h)
|
||||
except ValueError:
|
||||
logger.warning("Invalid size format: %s", extra_body.get("size"))
|
||||
height, width = self._resolve_height_width_from_extra_body(extra_body)
|
||||
|
||||
# Get request parameters from extra_body.
|
||||
# Avoid hardcoded defaults here — let each pipeline's forward()
|
||||
@@ -2787,8 +2726,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
if isinstance(size_str, str) and "x" in size_str.lower():
|
||||
w, h = size_str.lower().split("x")
|
||||
width, height = int(w), int(h)
|
||||
except Exception:
|
||||
pass
|
||||
except ValueError:
|
||||
logger.warning("Invalid size format: %s", extra_body.get("size"))
|
||||
|
||||
return height, width
|
||||
|
||||
|
||||
@@ -41,6 +41,7 @@ class StageRequestStats:
|
||||
postprocess_time_ms: float = 0.0
|
||||
diffusion_metrics: dict[str, int] = None
|
||||
audio_generated_frames: int = 0
|
||||
pipeline_timings: dict[str, float] | None = None
|
||||
|
||||
@property
|
||||
def rx_mbps(self) -> float:
|
||||
@@ -101,6 +102,7 @@ STAGE_EXCLUDE = {
|
||||
"rx_decode_time_ms",
|
||||
"rx_in_flight_time_ms",
|
||||
"final_output_type",
|
||||
"pipeline_timings",
|
||||
}
|
||||
TRANSFER_EXCLUDE = {"from_stage", "to_stage", "request_id", "used_shm"}
|
||||
E2E_EXCLUDE = {"request_id"}
|
||||
@@ -487,6 +489,15 @@ class OrchestratorAggregator:
|
||||
"e2e_avg_time_per_request_ms": float(e2e_avg_req),
|
||||
"e2e_avg_tokens_per_s": float(e2e_avg_tok),
|
||||
}
|
||||
# Add average input preprocess time across all requests
|
||||
preprocess_times = []
|
||||
for _rid, evts in self.stage_events.items():
|
||||
for evt in evts:
|
||||
if evt.pipeline_timings and "preprocess_ms" in evt.pipeline_timings:
|
||||
preprocess_times.append(evt.pipeline_timings["preprocess_ms"])
|
||||
break # only once per request
|
||||
if preprocess_times:
|
||||
overall_summary["input_preprocess_time_ms"] = sum(preprocess_times) / len(preprocess_times)
|
||||
# Add stage_wall_time_ms as separate fields for each stage
|
||||
for idx, wall_time in enumerate(stage_wall_time_ms):
|
||||
overall_summary[f"e2e_stage_{idx}_wall_time_ms"] = wall_time
|
||||
@@ -534,11 +545,41 @@ class OrchestratorAggregator:
|
||||
),
|
||||
)
|
||||
|
||||
# === Stage table (columns = stage_id) ===
|
||||
# === [OmniTiming] concise per-request summary ===
|
||||
stage_evts = sorted(
|
||||
self.stage_events.get(rid, []),
|
||||
key=lambda e: e.stage_id if e.stage_id is not None else -1,
|
||||
)
|
||||
pt = {}
|
||||
if stage_evts:
|
||||
pt = stage_evts[-1].pipeline_timings or {}
|
||||
if pt or e2e_evt:
|
||||
parts = [f"req={rid}"]
|
||||
if e2e_evt:
|
||||
parts.append(f"total={e2e_evt.e2e_total_ms / 1000.0:.2f}s")
|
||||
if "preprocess_ms" in pt:
|
||||
parts.append(f"preprocess={pt['preprocess_ms'] / 1000.0:.2f}s")
|
||||
if e2e_evt:
|
||||
engine_ms = e2e_evt.e2e_total_ms - pt.get("preprocess_ms", 0.0)
|
||||
parts.append(f"engine={engine_ms / 1000.0:.2f}s")
|
||||
stage_parts = []
|
||||
for evt in stage_evts:
|
||||
sid = evt.stage_id if evt.stage_id is not None else "?"
|
||||
t = evt.stage_gen_time_ms / 1000.0
|
||||
stage_parts.append(f"{sid}:{t:.2f}s")
|
||||
if stage_parts:
|
||||
parts.append(f"stages=[{','.join(stage_parts)}]")
|
||||
transfer_parts = []
|
||||
for te in self.transfer_events.values():
|
||||
if te.request_id == rid:
|
||||
transfer_parts.append(f"{te.from_stage}->{te.to_stage}={te.tx_time_ms:.2f}ms")
|
||||
if transfer_parts:
|
||||
parts.append(f"transfers=[{','.join(transfer_parts)}]")
|
||||
if "ar2diffusion_ms" in pt:
|
||||
parts.append(f"ar2diffusion={pt['ar2diffusion_ms']:.2f}ms")
|
||||
logger.info("[OmniTiming] %s", " ".join(parts))
|
||||
|
||||
# === Stage table (columns = stage_id) ===
|
||||
# if any stage has diffusion_metrics, remove postprocess_time_ms field
|
||||
# because it is already included in diffusion_metrics
|
||||
local_exclude = STAGE_EXCLUDE.copy()
|
||||
|
||||
@@ -330,6 +330,13 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo
|
||||
- Grid dimension calculation for M-RoPE position encoding
|
||||
"""
|
||||
|
||||
def _cached_apply_hf_processor(self, inputs, timing_ctx):
|
||||
# i2i: prompt text must be modified based on mm data presence,
|
||||
# and grid computation requires all images together — bypass cache.
|
||||
if inputs.mm_data_items.get_all_counts().get("image", 0) > 0:
|
||||
return self._apply_hf_processor(inputs, timing_ctx)
|
||||
return super()._cached_apply_hf_processor(inputs, timing_ctx)
|
||||
|
||||
def _call_hf_processor(
|
||||
self,
|
||||
prompt: str,
|
||||
@@ -678,7 +685,7 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo
|
||||
try:
|
||||
mrope_grid_thw = self._build_generation_grids(hf_processor_mm_kwargs)
|
||||
mm_processed_data["mrope_image_grid_thw"] = mrope_grid_thw
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"_apply_hf_processor_main t2i: mrope_image_grid_thw=%s",
|
||||
mrope_grid_thw.tolist(),
|
||||
)
|
||||
@@ -786,7 +793,7 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo
|
||||
hf_config = self.info.get_hf_config()
|
||||
image_token_id = getattr(hf_config, "image_token_id", 167855)
|
||||
image_token_count = prompt_ids.count(image_token_id)
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"_apply_hf_processor_main i2i(HF): num_images=%s, prompt_len=%s, image_token_count=%s, "
|
||||
"source_grid_shape=%s, mrope_grid_shape=%s",
|
||||
num_images,
|
||||
@@ -2327,7 +2334,7 @@ class GlmImageModel(nn.Module):
|
||||
|
||||
# Debug: log prior_token_image_ids_info
|
||||
shapes = [t.shape for t in upsampled_token_ids]
|
||||
logger.info(
|
||||
logger.debug(
|
||||
f"[GlmImageModel.forward] Built prior_token_image_ids_info: "
|
||||
f"num_images={len(upsampled_token_ids)}, shapes={shapes}, "
|
||||
f"image_grid_thw={image_grid_thw.tolist()}"
|
||||
@@ -2537,7 +2544,7 @@ class GlmImageForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP
|
||||
|
||||
# Debug: log prior_token_info
|
||||
shapes = [t.shape for t in upsampled_token_ids]
|
||||
logger.info(
|
||||
logger.debug(
|
||||
f"[_process_image_input] Built prior_token_info: "
|
||||
f"num_images={len(upsampled_token_ids)}, shapes={shapes}, "
|
||||
f"image_grid_thw={image_grid_thw.tolist()}"
|
||||
|
||||
Reference in new issue
Block a user