[diffusion] feat: Add Configurable Generator Device and Seed Support via API (#14366)

Co-authored-by: niehen6174 <niehen.6174@gmail.com>
Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -101,6 +101,7 @@ class SamplingParams:
     # Batch info
     num_outputs_per_prompt: int = 1
     seed: int = 1024
+    generator_device: str = "cuda"  # Device for random generator: "cuda" or "cpu"
 
     # Original dimensions (before VAE scaling)
     num_frames: int = 125
@@ -393,6 +394,13 @@ class SamplingParams:
             default=SamplingParams.seed,
             help="Random seed for generation",
         )
+        parser.add_argument(
+            "--generator-device",
+            type=str,
+            default=SamplingParams.generator_device,
+            choices=["cuda", "cpu"],
+            help="Device for random generator (cuda or cpu). Default: cuda",
+        )
         parser.add_argument(
             "--num-frames",
             type=int,

python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py

@@ -53,6 +53,8 @@ def _build_sampling_params_from_request(
     output_format: Optional[str],
     background: Optional[str],
     image_path: Optional[str] = None,
+    seed: Optional[int] = None,
+    generator_device: Optional[str] = None,
 ) -> SamplingParams:
     if size is None:
         width, height = None, None
@@ -73,6 +75,8 @@ def _build_sampling_params_from_request(
         save_output=True,
         server_args=server_args,
         output_file_name=f"{request_id}.{ext}",
+        seed=seed,
+        generator_device=generator_device,
     )
     return sampling_params
 
@@ -88,6 +92,7 @@ def _build_req_from_sampling(s: SamplingParams) -> Req:
         fps=1,
         num_frames=s.num_frames,
         seed=s.seed,
+        generator_device=s.generator_device,
         output_path=s.output_path,
         output_file_name=s.output_file_name,
         num_outputs_per_prompt=s.num_outputs_per_prompt,
@@ -107,6 +112,8 @@ async def generations(
         size=request.size,
         output_format=request.output_format,
         background=request.background,
+        seed=request.seed,
+        generator_device=request.generator_device,
     )
     batch = prepare_request(
         server_args=get_global_server_args(),
@@ -155,6 +162,8 @@ async def edits(
     size: Optional[str] = Form(None),
     output_format: Optional[str] = Form(None),
     background: Optional[str] = Form("auto"),
+    seed: Optional[int] = Form(1024),
+    generator_device: Optional[str] = Form("cuda"),
     user: Optional[str] = Form(None),
 ):
     request_id = generate_request_id()
@@ -178,6 +187,8 @@ async def edits(
         output_format=output_format,
         background=background,
         image_path=input_path,
+        seed=seed,
+        generator_device=generator_device,
     )
     batch = _build_req_from_sampling(sampling)
 

python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py

@@ -26,6 +26,8 @@ class ImageGenerationsRequest(BaseModel):
     style: Optional[str] = "vivid"
     background: Optional[str] = "auto"  # transparent | opaque | auto
     output_format: Optional[str] = None  # png | jpeg | webp
+    seed: Optional[int] = 1024
+    generator_device: Optional[str] = "cuda"
     user: Optional[str] = None
 
 
@@ -54,6 +56,8 @@ class VideoGenerationsRequest(BaseModel):
     size: Optional[str] = "720x1280"
     fps: Optional[int] = None
     num_frames: Optional[int] = None
+    seed: Optional[int] = 1024
+    generator_device: Optional[str] = "cuda"
 
 
 class VideoListResponse(BaseModel):

python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py

@@ -73,6 +73,8 @@ def _build_sampling_params_from_request(
         save_output=True,
         server_args=server_args,
         output_file_name=request_id,
+        seed=request.seed,
+        generator_device=request.generator_device,
     )
 
     return sampling_params

python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py

@@ -88,6 +88,7 @@ class Req:
     num_outputs_per_prompt: int = 1
     seed: int | None = None
     seeds: list[int] | None = None
+    generator_device: str = "cuda"  # Device for random generator: "cuda" or "cpu"
 
     # Tracking if embeddings are already processed
     is_prompt_processed: bool = False

python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py

@@ -344,7 +344,8 @@ class CausalDMDDenoisingStage(DenoisingStage):
                                 if isinstance(batch.generator, list)
                                 else batch.generator
                             ),
-                        ).to(self.device)
+                            device=self.device,
+                        )
                         noise_btchw = noise
                         noise_latents_btchw = self.scheduler.add_noise(
                             pred_video_btchw.flatten(0, 1),

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py

@@ -166,7 +166,8 @@ class DmdDenoisingStage(DenoisingStage):
                                 video_raw_latent_shape,
                                 dtype=pred_video.dtype,
                                 generator=batch.generator[0],
-                            ).to(self.device)
+                                device=self.device,
+                            )
                             latents = self.scheduler.add_noise(
                                 pred_video.flatten(0, 1),
                                 noise.flatten(0, 1),

python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py

@@ -53,9 +53,19 @@ class InputValidationStage(PipelineStage):
         assert seed is not None
         seeds = [seed + i for i in range(num_videos_per_prompt)]
         batch.seeds = seeds
-        # Peiyuan: using GPU seed will cause A100 and H100 to generate different results...
-        # FIXME: the generator's in latent preparation stage seems to be different from seeds
-        batch.generator = [torch.Generator("cpu").manual_seed(seed) for seed in seeds]
+
+        # Create generators based on generator_device parameter
+        # Note: This will overwrite any existing batch.generator
+        generator_device = batch.generator_device
+
+        if generator_device == "cpu":
+            device_str = "cpu"
+        else:
+            device_str = "cuda" if torch.cuda.is_available() else "cpu"
+
+        batch.generator = [
+            torch.Generator(device_str).manual_seed(seed) for seed in seeds
+        ]
 
     def preprocess_condition_image(
         self,