[CI][Bugfix] Refactor the test case to add support for increasing init timeout and stage init timeout in order to resolve the CI timeout error. (#2711)

Signed-off-by: wangyu <410167048@qq.com>

tests/e2e/offline_inference/test_diffusion_cpu_offload.py

@@ -1,22 +1,14 @@
 import gc
-import sys
-from pathlib import Path
 
 import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
+from tests.conftest import OmniRunner
 from tests.utils import DeviceMemoryMonitor, hardware_test
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
-# ruff: noqa: E402
-REPO_ROOT = Path(__file__).resolve().parents[2]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
-
-from vllm_omni import Omni
-
 models = ["riverclouds/qwen_image_random"]
 
 
@@ -27,30 +19,29 @@ def inference(model_name: str, offload: bool = True):
     current_omni_platform.reset_peak_memory_stats()
     monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
-    m = Omni(
-        model=model_name,
+    with OmniRunner(
+        model_name,
         # TODO: we might want to add overlapped feature e2e tests
         # cache_backend="cache_dit",
         enable_cpu_offload=offload,
-    )
-    current_omni_platform.reset_peak_memory_stats()
-    height = 256
-    width = 256
+    ) as runner:
+        current_omni_platform.reset_peak_memory_stats()
+        height = 256
+        width = 256
 
-    m.generate(
-        "a photo of a cat sitting on a laptop keyboard",
-        OmniDiffusionSamplingParams(
-            height=height,
-            width=width,
-            num_inference_steps=9,
-            guidance_scale=0.0,
-            generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
-        ),
-    )
+        runner.omni.generate(
+            "a photo of a cat sitting on a laptop keyboard",
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                num_inference_steps=9,
+                guidance_scale=0.0,
+                generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
+            ),
+        )
     peak = monitor.peak_used_mb
     monitor.stop()
 
-    del m
     gc.collect()
     current_omni_platform.empty_cache()