Update merged_files3.py

merged_files3.py

@@ -487,8 +487,12 @@ def hooked_unet_forward(sample, timestep, encoder_hidden_states, **kwargs):
 
 unet.forward = hooked_unet_forward
 
+# Device and dtype setup
+device = torch.device('cuda')
+dtype = torch.float16  # Use float16 consistently for all models
+
 
-sd_offset = sf.load_file(model_path)
+sd_offset = sf.load_file(model_path, device=device) # Use device variable
 sd_origin = unet.state_dict()
 keys = sd_origin.keys()
 sd_merged = {k: sd_origin[k] + sd_offset[k] for k in sd_origin.keys()}
@@ -496,9 +500,6 @@ unet.load_state_dict(sd_merged, strict=True)
 del sd_offset, sd_origin, sd_merged, keys
 
 
-# Device and dtype setup
-device = torch.device('cuda')
-dtype = torch.float16  # Use float16 consistently for all models
 
 pipe = prepare_pipeline(
     base_model="stabilityai/stable-diffusion-xl-base-1.0",