From 41e22e27701bab010b22677c33da6637441dfef3 Mon Sep 17 00:00:00 2001
From: tzhouam <tzhouam@connect.ust.hk>
Date: Thu, 5 Feb 2026 11:02:03 +0000
Subject: [PATCH] [Refactor] Clean up code by removing unnecessary lines and
 adding missing imports in async_omni_llm.py, omni_stage.py, omni.py, and
 api_server.py

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
---
 vllm_omni/entrypoints/async_omni_llm.py    | 5 +----
 vllm_omni/entrypoints/omni.py              | 1 +
 vllm_omni/entrypoints/omni_stage.py        | 5 -----
 vllm_omni/entrypoints/openai/api_server.py | 1 +
 4 files changed, 3 insertions(+), 9 deletions(-)

diff --git a/vllm_omni/entrypoints/async_omni_llm.py b/vllm_omni/entrypoints/async_omni_llm.py
index aadbe5756..02d175ad0 100644
--- a/vllm_omni/entrypoints/async_omni_llm.py
+++ b/vllm_omni/entrypoints/async_omni_llm.py
@@ -164,10 +164,7 @@ class AsyncOmniLLM(AsyncLLM):
             pass
 
         # Use profiler_config from vllm_config (new way, aligned with vllm v1)
-        if (
-            vllm_config.profiler_config.profiler == "torch"
-            and not vllm_config.profiler_config.ignore_frontend
-        ):
+        if vllm_config.profiler_config.profiler == "torch" and not vllm_config.profiler_config.ignore_frontend:
             profiler_dir = vllm_config.profiler_config.torch_profiler_dir
             logger.info(
                 "Torch profiler enabled. AsyncOmniLLM CPU traces will be collected under %s",
diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py
index 33355ed3f..97357dc3b 100644
--- a/vllm_omni/entrypoints/omni.py
+++ b/vllm_omni/entrypoints/omni.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
 import multiprocessing as mp
 import os
 import time
diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py
index 8bcdd24ea..2fb9eddec 100644
--- a/vllm_omni/entrypoints/omni_stage.py
+++ b/vllm_omni/entrypoints/omni_stage.py
@@ -626,7 +626,6 @@ def _stage_worker(
         for device_id in devices_to_lock:
             lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
             lock_acquired = False
-            next_wait_log = wait_start + 5.0
 
             while not lock_acquired:
                 try:
@@ -655,7 +654,6 @@ def _stage_worker(
                             )
                             break
 
-
                         # Wait a bit before retrying
                         _time.sleep(0.1)
                 except OSError as e:
@@ -1173,8 +1171,6 @@ async def _stage_worker_async(
         for device_id in devices_to_lock:
             lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
             lock_acquired = False
-            next_wait_log = wait_start + 5.0
-
             while not lock_acquired:
                 try:
                     # Open or create the lock file
@@ -1203,7 +1199,6 @@ async def _stage_worker_async(
                             )
                             break
 
-
                         # Wait a bit before retrying
                         _time.sleep(0.1)
                 except OSError as e:
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index 4ceea5516..798cb9229 100644
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import base64
 import io
+import json
 import multiprocessing
 import multiprocessing.forkserver as forkserver
 import os