fix pre-commit

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

tests/diffusion/models/qwen_image_21/test_decode_graph.py

@@ -205,9 +205,7 @@ def test_replay_uses_active_request_and_layout(graph_case):
         register(manager, second)
         assert decode(manager, second) is None
         output = decode(manager, first)
-        expected = (
+        expected = torch.ones_like(output).sin() + first[0]["cond"]["key"].mean() + first[0]["cond"]["value"].mean() + 2
-            torch.ones_like(output).sin() + first[0]["cond"]["key"].mean() + first[0]["cond"]["value"].mean() + 2
-        )
         torch.testing.assert_close(output, expected)
         # Once the owner's cache is released, the key is free to adopt.
         del first

tests/diffusion/offloader/test_sequential_backend.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Unit tests for SequentialOffloadBackend."""
 

vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py

@@ -920,8 +920,7 @@ class QwenImage21Transformer2DModel(CachedTransformer):
                 self,
                 max_entries=cuda_graph_max_decode_graphs,
                 model_level_offload=(
-                    od_config.enable_cpu_offload
+                    od_config.enable_cpu_offload or getattr(od_config, "enable_distributed_layerwise_offload", False)
-                    or getattr(od_config, "enable_distributed_layerwise_offload", False)
                 ),
             )
             if self.enable_cuda_graph_decode

vllm_omni/diffusion/offloader/sequential_backend.py

@@ -420,8 +420,7 @@ class ModelLevelOffloadBackend(OffloadBackend):
                 # plain move semantics so offloading actually frees DiT VRAM
                 # while the encoders run.
                 persistent_dit_staging=(
-                    self.device.type == "cuda"
+                    self.device.type == "cuda" and any(getattr(dit, "enable_cuda_graph_decode", False) for dit in dits)
-                    and any(getattr(dit, "enable_cuda_graph_decode", False) for dit in dits)
                 ),
             )
         except BaseException: