[Entrypoint][Refactor] vLLM-Omni Entrypoint Refactoring (#1908)

Signed-off-by: yinpe <11810305@mail.sustech.edu.cn>
Signed-off-by: Peiqi Yin <yinpeiqi809@gmail.com>
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Signed-off-by: wuhang <whlbx@hotmail.com>
Signed-off-by: linyueqian <linyueqian@outlook.com>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: lishunyang <lishunyang12@163.com>
Signed-off-by: Chenguang ZHENG <645327136@qq.com>
Signed-off-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Co-authored-by: yinpe <11810305@mail.sustech.edu.cn>
Co-authored-by: Peiqi Yin <yinpeiqi809@gmail.com>
Co-authored-by: wuhang <whlbx@hotmail.com>
Co-authored-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Co-authored-by: linyueqian <linyueqian@outlook.com>
Co-authored-by: princepride <wangzhipeng628@gmail.com>
Co-authored-by: lishunyang <lishunyang12@163.com>
This commit is contained in:
authored and GitHub committed 2026-03-18 17:33:53 +08:00
1 parent 6a66a1fe36
commit 3158e9d28f
125 files changed
+5210 -9322

No files matched your search

+10 -18
View File
@@ -34,7 +34,7 @@ def parse_args():
help="Path to input image for img2img.",
)
# OmniLLM init args
# Omni runtime init args
parser.add_argument("--log-stats", action="store_true", default=False)
parser.add_argument("--init-sleep-seconds", type=int, default=20)
parser.add_argument("--batch-timeout", type=int, default=5)
@@ -163,25 +163,17 @@ def main():
omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
for i, req_output in enumerate(omni_outputs):
img_idx = 0
for req_output in omni_outputs:
images = getattr(req_output, "images", None)
if not images and hasattr(req_output, "output"):
if isinstance(req_output.output, list):
images = req_output.output
else:
images = [req_output.output]
if not images:
continue
if images:
for j, img in enumerate(images):
img.save(f"output_{i}_{j}.png")
if hasattr(req_output, "request_output") and req_output.request_output:
for stage_out in req_output.request_output:
if hasattr(stage_out, "images") and stage_out.images:
for k, img in enumerate(stage_out.images):
save_path = f"output_{i}_stage_{getattr(stage_out, 'stage_id', '?')}_{k}.png"
img.save(save_path)
print(f"[Info] Saved stage output image to {save_path}")
for j, img in enumerate(images):
save_path = f"output_{img_idx}_{j}.png"
img.save(save_path)
print(f"[Info] Saved image to {save_path}")
img_idx += 1
print(omni_outputs)
@@ -167,27 +167,26 @@ def run_e2e():
print(f"Received {len(outputs)} outputs.")
for i, output in enumerate(outputs):
try:
ro_list = output.request_output or []
if not ro_list:
ro = output.request_output
if ro is None:
print("No request_output found.")
continue
for ro in ro_list:
# Multimodal output may be attached to RequestOutput or CompletionOutput.
mm = getattr(ro, "multimodal_output", None)
if not mm and ro.outputs:
mm = getattr(ro.outputs[0], "multimodal_output", None)
# Multimodal output may be attached to RequestOutput or CompletionOutput.
mm = getattr(ro, "multimodal_output", None)
if not mm and ro.outputs:
mm = getattr(ro.outputs[0], "multimodal_output", None)
if mm:
print(f"Multimodal output keys: {mm.keys()}")
if "audio" in mm:
audio_out = mm["audio"]
print(f"Generated Audio Shape: {audio_out.shape}")
out_path = f"output_{i}.wav"
sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050)
print(f"Saved audio to {out_path}")
else:
print("No multimodal output found.")
if mm:
print(f"Multimodal output keys: {mm.keys()}")
if "audio" in mm:
audio_out = mm["audio"]
print(f"Generated Audio Shape: {audio_out.shape}")
out_path = f"output_{i}.wav"
sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050)
print(f"Saved audio to {out_path}")
else:
print("No multimodal output found.")
except Exception as e:
print(f"Error inspecting output: {e}")
omni.close()
@@ -209,7 +209,7 @@ async def main():
if not outputs:
raise ValueError("No output produced from omni.generate()")
first_out = outputs[0].request_output[0]
first_out = outputs[0].request_output
req_out: OmniRequestOutput = first_out
# Verify trajectory data (from custom pipeline)
@@ -129,8 +129,14 @@ def main(args):
t_start = time.perf_counter()
for stage_outputs in omni.generate(inputs):
for output in stage_outputs.request_output:
_save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
request_output = stage_outputs.request_output
if request_output is None or not request_output.outputs:
continue
_save_wav(
output_dir,
request_output.request_id,
request_output.outputs[0].multimodal_output,
)
t_end = time.perf_counter()
logger.info("Total inference time: %.1f ms", (t_end - t_start) * 1000)
+20 -21
View File
@@ -310,36 +310,35 @@ def main(args: argparse.Namespace) -> None:
output_count = 0
for stage_outputs in omni.generate(prompts, sampling_params_list, py_generator=True):
output = stage_outputs.request_output
if stage_outputs.final_output_type == "image":
for output in stage_outputs.request_output:
request_id = output.request_id
request_id = output.request_id
# Get generated images
images = output.images if hasattr(output, "images") else []
if not images and hasattr(output, "multimodal_output"):
images = output.multimodal_output.get("images", [])
# Get generated images
images = output.images if hasattr(output, "images") else []
if not images and hasattr(output, "multimodal_output"):
images = output.multimodal_output.get("images", [])
# Save each generated image
for idx, img in enumerate(images):
if args.num_prompts == 1 and len(images) == 1:
output_path = args.output
else:
base, ext = os.path.splitext(args.output)
output_path = f"{base}_{request_id}_{idx}{ext}"
# Save each generated image
for idx, img in enumerate(images):
if args.num_prompts == 1 and len(images) == 1:
output_path = args.output
else:
base, ext = os.path.splitext(args.output)
output_path = f"{base}_{request_id}_{idx}{ext}"
if isinstance(img, Image.Image):
save_image(img, output_path)
else:
print(f"Warning: Unexpected image type for request {request_id}: {type(img)}")
if isinstance(img, Image.Image):
save_image(img, output_path)
else:
print(f"Warning: Unexpected image type for request {request_id}: {type(img)}")
output_count += 1
output_count += 1
elif stage_outputs.final_output_type == "text":
# AR stage output (intermediate, for debugging)
if args.verbose:
for output in stage_outputs.request_output:
print(f"AR output for request {output.request_id}:")
print(f" Token count: {len(output.outputs[0].token_ids)}")
print(f"AR output for request {output.request_id}:")
print(f" Token count: {len(output.outputs[0].token_ids)}")
gen_time = time.time() - gen_start_time
print(f"\nGeneration completed in {gen_time:.2f}s")
+5 -6
View File
@@ -293,12 +293,11 @@ def main():
)
if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
inner_output = first_item.request_output[0]
if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
frames = inner_output.images[0] if inner_output.images else None
if frames is None:
raise ValueError("No video frames found in output.")
inner_output = first_item.request_output
if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
frames = inner_output.images[0] if inner_output.images else None
if frames is None:
raise ValueError("No video frames found in output.")
elif hasattr(first_item, "images") and first_item.images:
frames = first_item.images
else:
@@ -9,7 +9,7 @@ from PIL import Image
from vllm_omni.entrypoints.omni import Omni
"""
The tencent/HunyuanImage-3.0-Instruct base model is built on the Hunyuan v1 architecture, specifically the tencent/Hunyuan-A13B-Instruct model. It utilizes two tokenizer delimiter templates:
The tencent/HunyuanImage-3.0-Instruct base model uses the tencent/Hunyuan-A13B-Instruct backbone. It utilizes two tokenizer delimiter templates:
1) Pretrained template (default for gen_text mode), which concatenates system, image
tokens, and user question WITHOUT role delimiters:
@@ -70,8 +70,8 @@ def main(args: argparse.Namespace) -> None:
prompts = [prompt_dict]
omni_outputs = omni.generate(prompts=prompts)
prompt_text = omni_outputs[0].request_output[0].prompt
generated_text = omni_outputs[0].request_output[0].outputs[0].text
prompt_text = omni_outputs[0].request_output.prompt
generated_text = omni_outputs[0].request_output.outputs[0].text
print(f"Prompt: {prompt_text}")
print(f"Text: {generated_text}")
@@ -463,12 +463,12 @@ def main():
raise ValueError("No output generated from omni.generate()")
# Extract images from OmniRequestOutput
# omni.generate() returns list[OmniRequestOutput], extract images from request_output[0].images
# omni.generate() returns list[OmniRequestOutput], extract images from request_output.images
first_output = outputs[0]
if not hasattr(first_output, "request_output") or not first_output.request_output:
raise ValueError("No request_output found in OmniRequestOutput")
req_out = first_output.request_output[0]
req_out = first_output.request_output
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
@@ -338,8 +338,6 @@ def main():
audio = frames.multimodal_output["audio"]
if frames.is_pipeline_output and frames.request_output is not None:
inner_output = frames.request_output
if isinstance(inner_output, list):
inner_output = inner_output[0] if inner_output else None
if isinstance(inner_output, OmniRequestOutput):
if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
audio = inner_output.multimodal_output["audio"]
@@ -99,14 +99,12 @@ def main() -> None:
lines: list[str] = []
for stage_outputs in outputs:
req_outputs = getattr(stage_outputs, "request_output", stage_outputs)
req_outputs = req_outputs if isinstance(req_outputs, list) else [req_outputs]
for ro in req_outputs:
text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
lines.append("answer:\n")
lines.append(text.strip() + "\n")
lines.append("\n")
ro = getattr(stage_outputs, "request_output", stage_outputs)
text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
lines.append("answer:\n")
lines.append(text.strip() + "\n")
lines.append("\n")
print("\n".join(lines))
@@ -148,19 +148,16 @@ def _collect_images(outputs: list) -> list[torch.Tensor]:
"""Extract all image tensors produced by the final (DiT) stage."""
images: list[torch.Tensor] = []
for out in outputs:
ro_list = getattr(out, "request_output", out)
if not isinstance(ro_list, list):
ro_list = [ro_list]
for ro_item in ro_list:
for completion in getattr(ro_item, "outputs", None) or []:
mm = getattr(completion, "multimodal_output", None)
if not isinstance(mm, dict) or "image" not in mm:
raise RuntimeError(f"Missing image in multimodal output: {mm}")
payload = mm["image"]
for tensor in payload if isinstance(payload, list) else [payload]:
if not isinstance(tensor, torch.Tensor):
raise TypeError(f"Expected image tensor, got {type(tensor)}")
images.append(tensor)
ro_item = getattr(out, "request_output", out)
for completion in getattr(ro_item, "outputs", None) or []:
mm = getattr(completion, "multimodal_output", None)
if not isinstance(mm, dict) or "image" not in mm:
raise RuntimeError(f"Missing image in multimodal output: {mm}")
payload = mm["image"]
for tensor in payload if isinstance(payload, list) else [payload]:
if not isinstance(tensor, torch.Tensor):
raise TypeError(f"Expected image tensor, got {type(tensor)}")
images.append(tensor)
return images
@@ -180,11 +180,11 @@ def main(args):
# Get the query function and call it with appropriate parameters
query_func = query_map[args.query_type]
omni_llm = Omni(
omni = Omni(
model=model_name,
stage_configs_path=args.stage_configs_path,
log_stats=args.enable_stats,
log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
log_file=("omni_pipeline.log" if args.enable_stats else None),
init_sleep_seconds=args.init_sleep_seconds,
batch_timeout=args.batch_timeout,
init_timeout=args.init_timeout,
@@ -290,7 +290,7 @@ def main(args):
prompts = [copy.deepcopy(query_result) for _ in range(args.num_prompts)]
print("prompts", prompts)
omni_outputs = omni_llm.generate(prompts, sampling_params_list)
omni_outputs = omni.generate(prompts, sampling_params_list)
output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
if args.query_type is not None:
@@ -298,45 +298,44 @@ def main(args):
os.makedirs(output_dir, exist_ok=True)
for stage_outputs in omni_outputs:
output = stage_outputs.request_output
if stage_outputs.final_output_type == "text":
for output in stage_outputs.request_output:
request_id = output.request_id
text_output = output.outputs[0].text
# Save aligned text file per request
prompt_text = output.prompt
out_txt = os.path.join(output_dir, f"{request_id}.txt")
lines = []
lines.append("Prompt:\n")
lines.append(str(prompt_text) + "\n")
lines.append("vllm_text_output:\n")
lines.append(str(text_output).strip() + "\n")
try:
with open(out_txt, "w", encoding="utf-8") as f:
print("lines", lines)
f.writelines(lines)
except Exception as e:
print(f"[Warn] Failed writing text file {out_txt}: {e}")
print(f"Request ID: {request_id}, Text saved to {out_txt}\n")
request_id = output.request_id
text_output = output.outputs[0].text
# Save aligned text file per request
prompt_text = output.prompt
out_txt = os.path.join(output_dir, f"{request_id}.txt")
lines = []
lines.append("Prompt:\n")
lines.append(str(prompt_text) + "\n")
lines.append("vllm_text_output:\n")
lines.append(str(text_output).strip() + "\n")
try:
with open(out_txt, "w", encoding="utf-8") as f:
print("lines", lines)
f.writelines(lines)
except Exception as e:
print(f"[Warn] Failed writing text file {out_txt}: {e}")
print(f"Request ID: {request_id}, Text saved to {out_txt}\n")
elif stage_outputs.final_output_type == "audio":
for output in stage_outputs.request_output:
request_id = output.request_id
audio_tensor = output.outputs[0].multimodal_output.get("audio")
request_id = output.request_id
audio_tensor = output.outputs[0].multimodal_output.get("audio")
if audio_tensor is None:
continue
if audio_tensor is None:
continue
output_wav = os.path.join(output_dir, f"{request_id}.wav")
output_wav = os.path.join(output_dir, f"{request_id}.wav")
# Convert to numpy array and ensure correct format
audio_numpy = audio_tensor.float().detach().cpu().numpy()
# Convert to numpy array and ensure correct format
audio_numpy = audio_tensor.float().detach().cpu().numpy()
# Ensure audio is 1D (flatten if needed)
if audio_numpy.ndim > 1:
audio_numpy = audio_numpy.flatten()
# Ensure audio is 1D (flatten if needed)
if audio_numpy.ndim > 1:
audio_numpy = audio_numpy.flatten()
# Save audio file with explicit WAV format
sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
print(f"Request ID: {request_id}, Audio saved to {output_wav}")
# Save audio file with explicit WAV format
sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
print(f"Request ID: {request_id}, Audio saved to {output_wav}")
def parse_args():
@@ -320,7 +320,7 @@ def main(args):
query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate)
else:
query_result = query_func()
omni_llm = Omni(
omni = Omni(
model=model_name,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
@@ -378,9 +378,11 @@ def main(args):
prompt["modalities"] = output_modalities
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
if profiler_enabled:
omni_llm.start_profile(stages=[0])
omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
if profiler_enabled and hasattr(omni, "start_profile"):
omni.start_profile(stages=[0])
elif profiler_enabled:
print("[Warn] VLLM_TORCH_PROFILER_DIR is set, but current engine does not support profiler controls.")
omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
# Determine output directory: prefer --output-dir; fallback to --output-wav
output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
@@ -389,43 +391,42 @@ def main(args):
total_requests = len(prompts)
processed_count = 0
for stage_outputs in omni_generator:
output = stage_outputs.request_output
if stage_outputs.final_output_type == "text":
for output in stage_outputs.request_output:
request_id = output.request_id
text_output = output.outputs[0].text
# Save aligned text file per request
prompt_text = output.prompt
out_txt = os.path.join(output_dir, f"{request_id}.txt")
lines = []
lines.append("Prompt:\n")
lines.append(str(prompt_text) + "\n")
lines.append("vllm_text_output:\n")
lines.append(str(text_output).strip() + "\n")
try:
with open(out_txt, "w", encoding="utf-8") as f:
f.writelines(lines)
except Exception as e:
print(f"[Warn] Failed writing text file {out_txt}: {e}")
print(f"Request ID: {request_id}, Text saved to {out_txt}")
request_id = output.request_id
text_output = output.outputs[0].text
# Save aligned text file per request
prompt_text = output.prompt
out_txt = os.path.join(output_dir, f"{request_id}.txt")
lines = []
lines.append("Prompt:\n")
lines.append(str(prompt_text) + "\n")
lines.append("vllm_text_output:\n")
lines.append(str(text_output).strip() + "\n")
try:
with open(out_txt, "w", encoding="utf-8") as f:
f.writelines(lines)
except Exception as e:
print(f"[Warn] Failed writing text file {out_txt}: {e}")
print(f"Request ID: {request_id}, Text saved to {out_txt}")
elif stage_outputs.final_output_type == "audio":
for output in stage_outputs.request_output:
request_id = output.request_id
audio_tensor = output.outputs[0].multimodal_output["audio"]
output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
print(f"Request ID: {request_id}, Saved audio to {output_wav}")
request_id = output.request_id
audio_tensor = output.outputs[0].multimodal_output["audio"]
output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
print(f"Request ID: {request_id}, Saved audio to {output_wav}")
processed_count += len(stage_outputs.request_output)
if profiler_enabled and processed_count >= total_requests:
processed_count += 1
if profiler_enabled and hasattr(omni, "stop_profile") and processed_count >= total_requests:
print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
# Stop the profiler while workers are still alive
omni_llm.stop_profile()
omni.stop_profile()
print("[Info] Waiting 30s for workers to write massive trace files to disk...")
time.sleep(30)
print("[Info] Trace export wait finished.")
omni_llm.close()
omni.close()
def parse_args():
@@ -294,7 +294,7 @@ def main(args):
else:
query_result = query_func()
omni_llm = Omni(
omni = Omni(
model=model_name,
stage_configs_path=args.stage_configs_path,
log_stats=args.log_stats,
@@ -354,8 +354,8 @@ def main(args):
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
if profiler_enabled:
omni_llm.start_profile(stages=[0])
omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
omni.start_profile(stages=[0])
omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
# Determine output directory: prefer --output-dir; fallback to --output-wav
output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
os.makedirs(output_dir, exist_ok=True)
@@ -366,51 +366,50 @@ def main(args):
print(f"query type: {args.query_type}")
for stage_outputs in omni_generator:
output = stage_outputs.request_output
if stage_outputs.final_output_type == "text":
for output in stage_outputs.request_output:
request_id = output.request_id
text_output = output.outputs[0].text
# Save aligned text file per request
prompt_text = output.prompt
out_txt = os.path.join(output_dir, f"{request_id}.txt")
lines = []
lines.append("Prompt:\n")
lines.append(str(prompt_text) + "\n")
lines.append("vllm_text_output:\n")
lines.append(str(text_output).strip() + "\n")
try:
with open(out_txt, "w", encoding="utf-8") as f:
f.writelines(lines)
except Exception as e:
print(f"[Warn] Failed writing text file {out_txt}: {e}")
print(f"Request ID: {request_id}, Text saved to {out_txt}")
request_id = output.request_id
text_output = output.outputs[0].text
# Save aligned text file per request
prompt_text = output.prompt
out_txt = os.path.join(output_dir, f"{request_id}.txt")
lines = []
lines.append("Prompt:\n")
lines.append(str(prompt_text) + "\n")
lines.append("vllm_text_output:\n")
lines.append(str(text_output).strip() + "\n")
try:
with open(out_txt, "w", encoding="utf-8") as f:
f.writelines(lines)
except Exception as e:
print(f"[Warn] Failed writing text file {out_txt}: {e}")
print(f"Request ID: {request_id}, Text saved to {out_txt}")
elif stage_outputs.final_output_type == "audio":
for output in stage_outputs.request_output:
request_id = output.request_id
audio_tensor = output.outputs[0].multimodal_output["audio"]
output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
request_id = output.request_id
audio_tensor = output.outputs[0].multimodal_output["audio"]
output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
# Convert to numpy array and ensure correct format
audio_numpy = audio_tensor.float().detach().cpu().numpy()
# Convert to numpy array and ensure correct format
audio_numpy = audio_tensor.float().detach().cpu().numpy()
# Ensure audio is 1D (flatten if needed)
if audio_numpy.ndim > 1:
audio_numpy = audio_numpy.flatten()
# Ensure audio is 1D (flatten if needed)
if audio_numpy.ndim > 1:
audio_numpy = audio_numpy.flatten()
# Save audio file with explicit WAV format
sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
print(f"Request ID: {request_id}, Saved audio to {output_wav}")
# Save audio file with explicit WAV format
sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
print(f"Request ID: {request_id}, Saved audio to {output_wav}")
processed_count += len(stage_outputs.request_output)
processed_count += 1
if profiler_enabled and processed_count >= total_requests:
print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
# Stop the profiler while workers are still alive
omni_llm.stop_profile()
omni.stop_profile()
print("[Info] Waiting 30s for workers to write trace files to disk...")
time.sleep(30)
print("[Info] Trace export wait time finished.")
omni_llm.close()
omni.close()
def parse_args():
@@ -231,52 +231,47 @@ async def run_single_request(
sampling_params_list=sampling_params_list,
output_modalities=output_modalities,
):
if not isinstance(omni_output.request_output, list):
outputs_list = [omni_output.request_output]
else:
outputs_list = omni_output.request_output
output = omni_output.request_output
if omni_output.final_output_type == "text":
if stage_0_first_output_ts is None:
stage_0_first_output_ts = time.perf_counter()
text_output = output.outputs[0].text
if output.finished:
text_parts.append(text_output)
elif omni_output.final_output_type == "audio":
mm_out = output.outputs[0].multimodal_output
if mm_out and "audio" in mm_out:
if first_audio_ts is None:
first_audio_ts = time.perf_counter()
if audio_sr is None and "sr" in mm_out:
sr_val = mm_out["sr"]
audio_sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
samplerate = audio_sr
audio_data = mm_out["audio"]
if isinstance(audio_data, list):
new_chunks = audio_data[audio_list_consumed:]
audio_list_consumed = len(audio_data)
elif isinstance(audio_data, torch.Tensor):
new_chunks = [audio_data]
audio_last_tensor = audio_data
else:
new_chunks = []
for output in outputs_list:
if omni_output.final_output_type == "text":
if stage_0_first_output_ts is None:
stage_0_first_output_ts = time.perf_counter()
text_output = output.outputs[0].text
if output.finished:
text_parts.append(text_output)
elif omni_output.final_output_type == "audio":
mm_out = output.outputs[0].multimodal_output
if mm_out and "audio" in mm_out:
if first_audio_ts is None:
first_audio_ts = time.perf_counter()
if audio_sr is None and "sr" in mm_out:
sr_val = mm_out["sr"]
audio_sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
samplerate = audio_sr
audio_data = mm_out["audio"]
if isinstance(audio_data, list):
new_chunks = audio_data[audio_list_consumed:]
audio_list_consumed = len(audio_data)
elif isinstance(audio_data, torch.Tensor):
new_chunks = [audio_data]
audio_last_tensor = audio_data
else:
new_chunks = []
if stream_audio_to_disk and new_chunks:
if sf_writer is None:
sf_writer = sf.SoundFile(
wav_file,
mode="w",
samplerate=samplerate,
channels=1,
subtype="FLOAT",
)
for chunk in new_chunks:
chunk_np = chunk.float().detach().cpu().numpy().flatten()
sf_writer.write(chunk_np)
audio_samples_written += len(chunk_np)
else:
audio_chunks.extend(new_chunks)
if stream_audio_to_disk and new_chunks:
if sf_writer is None:
sf_writer = sf.SoundFile(
wav_file,
mode="w",
samplerate=samplerate,
channels=1,
subtype="FLOAT",
)
for chunk in new_chunks:
chunk_np = chunk.float().detach().cpu().numpy().flatten()
sf_writer.write(chunk_np)
audio_samples_written += len(chunk_np)
else:
audio_chunks.extend(new_chunks)
finally:
if sf_writer is not None:
sf_writer.close()
@@ -51,16 +51,68 @@ def _estimate_prompt_len(
tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left")
cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True)
_cache[model_name] = (tok, getattr(cfg, "talker_config", None))
tok, tcfg = _cache[model_name]
# Load speech tokenizer (codec encoder) for exact ref_code_len.
speech_tok = None
try:
import os
from transformers.utils import cached_file
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer
st_cfg_path = cached_file(model_name, "speech_tokenizer/config.json")
if st_cfg_path:
speech_tok = Qwen3TTSTokenizer.from_pretrained(
os.path.dirname(st_cfg_path), torch_dtype=torch.bfloat16
)
logger.info("Loaded speech tokenizer for exact ref_code_len estimation")
except Exception as e:
logger.debug("Could not load speech tokenizer: %s", e)
_cache[model_name] = (tok, getattr(cfg, "talker_config", None), speech_tok)
tok, tcfg, speech_tok = _cache[model_name]
task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]
def _estimate_ref_code_len(ref_audio: object) -> int | None:
"""Encode ref_audio with the actual codec to get exact frame count."""
if not isinstance(ref_audio, (str, list)):
return None
audio_path = ref_audio[0] if isinstance(ref_audio, list) else ref_audio
if not isinstance(audio_path, str) or not audio_path.strip():
return None
try:
from vllm.multimodal.media import MediaConnector
connector = MediaConnector(allowed_local_media_path="/")
audio, sr = connector.fetch_audio(audio_path)
import numpy as np
wav_np = np.asarray(audio, dtype=np.float32)
if speech_tok is not None:
enc = speech_tok.encode(wav_np, sr=int(sr), return_dict=True)
ref_code = getattr(enc, "audio_codes", None)
if isinstance(ref_code, list):
ref_code = ref_code[0] if ref_code else None
if ref_code is not None and hasattr(ref_code, "shape"):
shape = ref_code.shape
return int(shape[0]) if len(shape) == 2 else int(shape[1]) if len(shape) == 3 else None
# Fallback: estimate from duration
codec_hz = getattr(tcfg, "codec_frame_rate", None) or 12
return int(len(audio) / sr * codec_hz)
except Exception:
return None
return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
additional_information=additional_information,
task_type=task_type,
tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
codec_language_id=getattr(tcfg, "codec_language_id", None),
spk_is_dialect=getattr(tcfg, "spk_is_dialect", None),
estimate_ref_code_len=_estimate_ref_code_len,
)
except Exception as exc:
logger.warning("Failed to estimate prompt length, using fallback 2048: %s", exc)
@@ -325,8 +377,8 @@ def main(args):
for batch_start in range(0, len(inputs), batch_size):
batch = inputs[batch_start : batch_start + batch_size]
for stage_outputs in omni.generate(batch):
for output in stage_outputs.request_output:
_save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
output = stage_outputs.request_output
_save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
async def main_streaming(args):
@@ -177,7 +177,7 @@ def main():
output = outputs[0]
if not hasattr(output, "request_output") or not output.request_output:
raise ValueError("No request_output found in OmniRequestOutput")
request_output = output.request_output[0]
request_output = output.request_output
if not hasattr(request_output, "multimodal_output"):
raise ValueError("No multimodal_output found in request_output")
@@ -16,7 +16,7 @@ if __name__ == "__main__":
omni = Omni(model="Qwen/Qwen-Image")
prompt = "a cup of coffee on the table"
outputs = omni.generate(prompt)
images = outputs[0].request_output[0].images
images = outputs[0].request_output.images
images[0].save("coffee.png")
```
@@ -34,7 +34,7 @@ if __name__ == "__main__":
]
outputs = omni.generate(prompts)
for i, output in enumerate(outputs):
image = output.request_output[0].images[0].save(f"{i}.jpg")
image = output.request_output.images[0].save(f"{i}.jpg")
```
!!! info
@@ -69,7 +69,7 @@ if __name__ == "__main__":
}
])
for i, output in enumerate(outputs):
image = output.request_output[0].images[0].save(f"{i}.jpg")
image = output.request_output.images[0].save(f"{i}.jpg")
```
## Local CLI Usage
@@ -112,7 +112,7 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
)
images_outputs = []
for output in outputs:
req_out = output.request_output[0]
req_out = output.request_output
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
images = req_out.images
@@ -14,7 +14,6 @@ from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.lora.request import LoRARequest
from vllm_omni.lora.utils import stable_lora_int_id
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
@@ -400,20 +399,18 @@ def main():
else:
print("[Profiler] No valid profiling data returned.")
# Extract images from OmniRequestOutput
# omni.generate() returns list[OmniRequestOutput], extract images from the first output
# omni.generate() returns list[OmniRequestOutput]
if not outputs or len(outputs) == 0:
raise ValueError("No output generated from omni.generate()")
logger.info(f"Outputs: {outputs}")
# Extract images from request_output[0]['images']
first_output = outputs[0]
if not hasattr(first_output, "request_output") or not first_output.request_output:
raise ValueError("No request_output found in OmniRequestOutput")
req_out = first_output.request_output[0]
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
req_out = first_output.request_output
if not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images'.")
images = req_out.images
if not images:
@@ -240,8 +240,6 @@ def main():
audio = frames.multimodal_output["audio"]
if frames.is_pipeline_output and frames.request_output is not None:
inner_output = frames.request_output
if isinstance(inner_output, list):
inner_output = inner_output[0] if inner_output else None
if isinstance(inner_output, OmniRequestOutput):
if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
audio = inner_output.multimodal_output["audio"]