[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:
125 files changed
+5210
-9322
No files matched your search
@@ -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)
|
||||
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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"]
|
||||
|
||||
+6
-8
@@ -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"]
|
||||
|
||||
Reference in new issue
Block a user