File tree Expand file tree Collapse file tree
Expand file tree Collapse file tree Original file line number Diff line number Diff line change @@ -601,12 +601,20 @@ def module_is_offloaded(module):
601601 def device (self ) -> torch .device :
602602 r"""
603603 Returns:
604- `torch.device`: The torch device on which the pipeline is located.
604+ `torch.device`: The torch device on which the pipeline is located. When components are split across devices
605+ (for example, text encoders on CPU while the denoising backbone runs on an accelerator), the accelerator
606+ device is returned.
605607 """
606608 module_names , _ = self ._get_signature_keys (self )
607609 modules = [getattr (self , n , None ) for n in module_names ]
608610 modules = [m for m in modules if isinstance (m , torch .nn .Module )]
609611
612+ # Prefer a non-CPU, non-meta component so a split pipeline reports the accelerator it computes on,
613+ # rather than whichever component happens to sort first.
614+ for module in modules :
615+ if module .device .type not in ("cpu" , "meta" ):
616+ return module .device
617+
610618 for module in modules :
611619 return module .device
612620
Original file line number Diff line number Diff line change @@ -1944,6 +1944,34 @@ def test_pipe_to(self):
19441944 assert sd1 .device .type == device_type
19451945 assert sd2 .device .type == device_type
19461946
1947+ @require_torch_accelerator
1948+ def test_pipe_device_split_across_devices (self ):
1949+ unet = self .dummy_cond_unet ()
1950+ scheduler = PNDMScheduler (skip_prk_steps = True )
1951+ vae = self .dummy_vae
1952+ bert = self .dummy_text_encoder
1953+ tokenizer = CLIPTokenizer .from_pretrained ("hf-internal-testing/tiny-random-clip" )
1954+
1955+ sd = StableDiffusionPipeline (
1956+ unet = unet ,
1957+ scheduler = scheduler ,
1958+ vae = vae ,
1959+ text_encoder = bert ,
1960+ tokenizer = tokenizer ,
1961+ safety_checker = None ,
1962+ feature_extractor = self .dummy_extractor ,
1963+ )
1964+
1965+ device_type = torch .device (torch_device ).type
1966+
1967+ # Text encoder stays on CPU while the denoising backbone runs on the accelerator. `text_encoder` sorts
1968+ # before `unet`/`vae`, so a first-component rule would report `cpu` here.
1969+ sd .unet .to (torch_device )
1970+ sd .vae .to (torch_device )
1971+
1972+ assert sd .text_encoder .device .type == "cpu"
1973+ assert sd .device .type == device_type
1974+
19471975 def test_pipe_same_device_id_offload (self ):
19481976 unet = self .dummy_cond_unet ()
19491977 scheduler = PNDMScheduler (skip_prk_steps = True )
You can’t perform that action at this time.
0 commit comments