diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 0e12460e8..1e00e4b0c 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/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 ( diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index b5e48ab95..c4fc970c4 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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())}" diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index 63514b191..aed6536a6 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/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. diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 3684744ba..bc9434817 100644 --- a/vllm_omni/inputs/data.py +++ b/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