diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index ef691b1ae..d9610234f 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index 7ac86c2c4..7df023243 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py index 00800ab15..d8fdd7c42 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py +++ b/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): diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py index c346fc528..6f6743a47 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index 781d31e5d..26bb392d6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py index d06dc142b..e06ee9bed 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py +++ b/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), diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py index 1ee183b0a..ec9dad38a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py +++ b/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), diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py index 8ad061d95..579005b48 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py +++ b/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,