From 590ab5b3e98c534f888ce8c545c7dde86e8a4579 Mon Sep 17 00:00:00 2001
From: Cheung Ka Wai <zhtmike@gmail.com>
Date: Fri, 27 Mar 2026 00:56:46 +0800
Subject: [PATCH] [Bug Fix] Resolve broken image issue when TP is enabled and
 no seed is provided. (#2176)

Signed-off-by: Mike Cheung <zhtmike@gmail.com>
---
 tests/diffusion/test_diffusion_request.py | 30 +++++++++++++++++++++++
 vllm_omni/diffusion/request.py            |  6 +++++
 2 files changed, 36 insertions(+)
 create mode 100644 tests/diffusion/test_diffusion_request.py

diff --git a/tests/diffusion/test_diffusion_request.py b/tests/diffusion/test_diffusion_request.py
new file mode 100644
index 000000000..56ec95d0c
--- /dev/null
+++ b/tests/diffusion/test_diffusion_request.py
@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import random
+
+import pytest
+
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _make_request() -> OmniDiffusionRequest:
+    return OmniDiffusionRequest(
+        prompts=[{"prompt": "a cup of coffee on a table"}],
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+    )
+
+
+def test_tp_seed_same_across_ranks_and_varies_across_requests():
+    random.seed(0)
+    n_requests = 5
+    seeds = [_make_request().sampling_params.seed for _ in range(n_requests)]
+
+    # Seed must be auto-assigned (not None) so every TP rank can use it.
+    assert all(s is not None for s in seeds)
+
+    # Seeds must vary across requests (non-determinism preserved).
+    assert len(set(seeds)) == n_requests, f"Expected {n_requests} unique seeds but got {len(set(seeds))}: {seeds}"
diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py
index 56a770461..1d6d64905 100644
--- a/vllm_omni/diffusion/request.py
+++ b/vllm_omni/diffusion/request.py
@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import random
 from dataclasses import dataclass, field
 
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
@@ -28,6 +29,11 @@ class OmniDiffusionRequest:
 
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""
+        # When neither a generator nor a seed is provided, assign a random seed
+        # so that all ranks derive the same generator state.
+        if self.sampling_params.generator is None and self.sampling_params.seed is None:
+            self.sampling_params.seed = random.randint(0, 2**31 - 1)
+
         # Detect whether user explicitly provided guidance_scale.
         # The sentinel default is 0.0 (false-like); any truthy value means
         # the caller set it intentionally.  We must resolve this BEFORE