From c12d00dca710e69c8f7e692af2726c38e6a9ce4a Mon Sep 17 00:00:00 2001 From: Feng Wang Date: Mon, 9 Jun 2025 17:43:19 +0800 Subject: [PATCH] fix device --- pipeline.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pipeline.py b/pipeline.py index b133012..21f9ccb 100644 --- a/pipeline.py +++ b/pipeline.py @@ -109,7 +109,7 @@ def encode_image_emb(self, siglip_image, device, dtype): def init_ccp_and_attn_processor(self, *args, **kwargs): subject_ip_adapter_path = kwargs['subject_ip_adapter_path'] nb_token = kwargs['nb_token'] - device=kwargs.get('nb_token', torch.device('cuda')) + device=kwargs.get('device', torch.device('cuda')) state_dict = torch.load(subject_ip_adapter_path, map_location="cpu") dtype = self.transformer.dtype @@ -627,4 +627,4 @@ def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Un # make sure to offload buffers if not all high level weights # are of type nn.Module offload_buffers = len(model._parameters) > 0 - cpu_offload(model, device, offload_buffers=offload_buffers) \ No newline at end of file + cpu_offload(model, device, offload_buffers=offload_buffers)