[Misc] Add per-request generator_device to online image gen and edit (#1183)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -158,7 +158,13 @@ class DiffusionModelRunner:
         self.kv_transfer_manager.receive_kv_cache(req, target_device=getattr(self.pipeline, "device", None))
 
         if req.sampling_params.generator is None and req.sampling_params.seed is not None:
-            req.sampling_params.generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed)
+            if req.sampling_params.generator_device is not None:
+                gen_device = req.sampling_params.generator_device
+            elif self.device.type == "cpu":
+                gen_device = "cpu"
+            else:
+                gen_device = self.device
+            req.sampling_params.generator = torch.Generator(device=gen_device).manual_seed(req.sampling_params.seed)
 
         # Refresh cache context if needed
         if (

vllm_omni/entrypoints/openai/api_server.py

@@ -971,6 +971,7 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
         _update_if_not_none(
             gen_params, "seed", request.seed if request.seed is not None else random.randint(0, 2**32 - 1)
         )
+        _update_if_not_none(gen_params, "generator_device", request.generator_device)
 
         request_id = f"img_gen_{uuid.uuid4().hex}"
 
@@ -1047,6 +1048,7 @@ async def edit_images(
     guidance_scale: float | None = Form(None),
     true_cfg_scale: float | None = Form(None),
     seed: int | None = Form(None),
+    generator_device: str | None = Form(None),
     # vllm-omni extension for per-request LoRA.
     lora: str | None = Form(None),  # Json string
 ) -> ImageGenerationResponse:
@@ -1129,6 +1131,7 @@ async def edit_images(
         # This fixes issues where using the default global generator
         # might produce blurry images in some environments.
         _update_if_not_none(gen_params, "seed", seed if seed is not None else random.randint(0, 2**32 - 1))
+        _update_if_not_none(gen_params, "generator_device", generator_device)
 
         # 4. Generate images using AsyncOmni (multi-stage mode)
         request_id = f"img_edit_{int(time.time())}"

vllm_omni/entrypoints/openai/protocol/images.py

@@ -88,6 +88,10 @@ class ImageGenerationRequest(BaseModel):
         description="True CFG scale (model-specific parameter, may be ignored if not supported)",
     )
     seed: int | None = Field(default=None, description="Random seed for reproducibility")
+    generator_device: str | None = Field(
+        default=None,
+        description="Device for the seeded torch.Generator (e.g. 'cpu', 'cuda'). Defaults to the runner's device.",
+    )
 
     # vllm-omni extension for per-request LoRA.
     # This mirrors the `extra_body.lora` convention in /v1/chat/completions.

vllm_omni/inputs/data.py

@@ -166,6 +166,7 @@ class OmniDiffusionSamplingParams:
     num_outputs_per_prompt: int = 1
     seed: int | None = None
     generator: torch.Generator | list[torch.Generator] | None = None
+    generator_device: str | None = None
 
     # layered info
     layers: int = 4