From aa780a6258e506a0cbbe8d069b04a62902693289 Mon Sep 17 00:00:00 2001
From: CHEN Xi <78632976+RubiaCx@users.noreply.github.com>
Date: Mon, 2 Feb 2026 17:24:19 +0800
Subject: [PATCH] [diffusion] fix: remove accelerate dependency for device
 mapping (#18026)

Co-authored-by: Kangyan-Zhou <zky314343421@gmail.com>
---
 .../runtime/pipelines/diffusers_pipeline.py   | 32 ++-----------------
 1 file changed, 3 insertions(+), 29 deletions(-)

diff --git a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py
index b408790d7..73f44f6b3 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py
@@ -21,6 +21,7 @@ from diffusers import DiffusionPipeline
 from PIL import Image
 
 from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
 )
@@ -391,12 +392,7 @@ class DiffusersPipeline(ComposedPipelineBase):
         self.model_path = model_path
 
         dtype = self._get_dtype(server_args)
-        device_map = self._get_device_map(server_args)
-        logger.info(
-            "Loading diffusers pipeline with dtype=%s, device_map=%s",
-            dtype,
-            device_map,
-        )
+        logger.info("Loading diffusers pipeline with dtype=%s", dtype)
 
         # Build common kwargs for from_pretrained
         load_kwargs = {
@@ -405,11 +401,6 @@ class DiffusersPipeline(ComposedPipelineBase):
             "revision": server_args.revision,
         }
 
-        # Add device_map for direct GPU loading and parallel shard loading
-        # This warms up CUDA caching allocator and enables parallel loading via accelerate
-        if device_map is not None:
-            load_kwargs["device_map"] = device_map
-
         # Add quantization config if provided (e.g., BitsAndBytesConfig for 4/8-bit)
         config = server_args.pipeline_config
         if config is not None:
@@ -458,22 +449,13 @@ class DiffusersPipeline(ComposedPipelineBase):
             else:
                 raise
 
-        # Only move to device if device_map wasn't used (already on device)
-        if device_map is None:
-            if torch.cuda.is_available():
-                pipe = pipe.to("cuda")
-            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
-                pipe = pipe.to("mps")
-
+        pipe = pipe.to(get_local_torch_device())
         # Apply VAE memory optimizations from pipeline config
         self._apply_vae_optimizations(pipe, server_args)
-
         # Apply attention backend if specified
         self._apply_attention_backend(pipe, server_args)
-
         # Apply cache-dit acceleration if configured
         pipe = self._apply_cache_dit(pipe, server_args)
-
         logger.info("Loaded diffusers pipeline: %s", pipe.__class__.__name__)
         return pipe
 
@@ -581,14 +563,6 @@ class DiffusersPipeline(ComposedPipelineBase):
         logger.info("Enabled cache-dit for diffusers pipeline")
         return pipe
 
-    def _get_device_map(self, server_args: ServerArgs) -> str | None:
-        """
-        Determine device_map for pipeline loading.
-        """
-        if not torch.cuda.is_available():
-            return None
-        return "cuda"
-
     def _get_dtype(self, server_args: ServerArgs) -> torch.dtype:
         """
         Determine the dtype to use for model loading.