From 4609afa613f14e6324ada7fdb39b330dca86cca3 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Wed, 22 Oct 2025 17:42:29 +0800 Subject: [PATCH] Refactor sampled_token_ids assignment in OmniGPUModelRunner - Simplified the assignment of sampled_token_ids by removing the conditional check for "model_stage". - Ensured consistent handling of sampled_token_ids for improved clarity and maintainability. --- vllm_omni/worker/gpu_model_runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index acbda8648..df021e4a5 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -542,7 +542,7 @@ class OmniGPUModelRunner(GPUModelRunner): # Get the valid generated tokens. import os - sampled_token_ids = sampler_output.sampled_token_ids if os.environ.get("model_stage") != "code2wav" else torch.tensor([[8294]]).to(torch.int32).cuda() + sampled_token_ids = sampler_output.sampled_token_ids max_gen_len = sampled_token_ids.shape[-1] if max_gen_len == 1: # No spec decode tokens.