device_map in load_model_dict_into_meta

Author: hlky

src/diffusers/loaders/transformer_flux.py

@@ -82,7 +82,8 @@ def _convert_ip_adapter_image_proj_to_diffusers(self, state_dict, low_cpu_mem_us
         if not low_cpu_mem_usage:
             image_projection.load_state_dict(updated_state_dict, strict=True)
         else:
-            load_model_dict_into_meta(image_projection, updated_state_dict, device=self.device, dtype=self.dtype)
+            device_map = {"": self.device}
+            load_model_dict_into_meta(image_projection, updated_state_dict, device_map=device_map, dtype=self.dtype)
 
         return image_projection
 
@@ -151,9 +152,9 @@ def _convert_ip_adapter_attn_to_diffusers(self, state_dicts, low_cpu_mem_usage=F
                 if not low_cpu_mem_usage:
                     attn_procs[name].load_state_dict(value_dict)
                 else:
-                    device = self.device
+                    device_map = {"": self.device}
                     dtype = self.dtype
-                    load_model_dict_into_meta(attn_procs[name], value_dict, device=device, dtype=dtype)
+                    load_model_dict_into_meta(attn_procs[name], value_dict, device_map=device_map, dtype=dtype)
 
                 key_id += 1
 

src/diffusers/loaders/transformer_sd3.py

@@ -75,8 +75,9 @@ def _convert_ip_adapter_attn_to_diffusers(
             if not low_cpu_mem_usage:
                 attn_procs[name].load_state_dict(layer_state_dict[idx], strict=True)
             else:
+                device_map = {"": self.device}
                 load_model_dict_into_meta(
-                    attn_procs[name], layer_state_dict[idx], device=self.device, dtype=self.dtype
+                    attn_procs[name], layer_state_dict[idx], device_map=device_map, dtype=self.dtype
                 )
 
         return attn_procs
@@ -144,7 +145,8 @@ def _convert_ip_adapter_image_proj_to_diffusers(
         if not low_cpu_mem_usage:
             image_proj.load_state_dict(updated_state_dict, strict=True)
         else:
-            load_model_dict_into_meta(image_proj, updated_state_dict, device=self.device, dtype=self.dtype)
+            device_map = {"": self.device}
+            load_model_dict_into_meta(image_proj, updated_state_dict, device_map=device_map, dtype=self.dtype)
 
         return image_proj
 

src/diffusers/loaders/unet.py

@@ -753,7 +753,8 @@ def _convert_ip_adapter_image_proj_to_diffusers(self, state_dict, low_cpu_mem_us
         if not low_cpu_mem_usage:
             image_projection.load_state_dict(updated_state_dict, strict=True)
         else:
-            load_model_dict_into_meta(image_projection, updated_state_dict, device=self.device, dtype=self.dtype)
+            device_map = {"": self.device}
+            load_model_dict_into_meta(image_projection, updated_state_dict, device_map=device_map, dtype=self.dtype)
 
         return image_projection
 
@@ -846,7 +847,8 @@ def _convert_ip_adapter_attn_to_diffusers(self, state_dicts, low_cpu_mem_usage=F
                 else:
                     device = next(iter(value_dict.values())).device
                     dtype = next(iter(value_dict.values())).dtype
-                    load_model_dict_into_meta(attn_procs[name], value_dict, device=device, dtype=dtype)
+                    device_map = {"": device}
+                    load_model_dict_into_meta(attn_procs[name], value_dict, device_map=device_map, dtype=dtype)
 
                 key_id += 2