diff --git a/cli/load-unet.py b/cli/load-unet.py index 2398cdb64..c910101b0 100644 --- a/cli/load-unet.py +++ b/cli/load-unet.py @@ -33,13 +33,13 @@ def set_module_tensor( stats.dtypes[value.dtype] = 0 stats.dtypes[value.dtype] += 1 if name in module._buffers: # pylint: disable=protected-access - module._buffers[name] = value.to(device=device, dtype=dtype, non_blocking=True) # pylint: disable=protected-access + module._buffers[name] = value.to(device=device, dtype=dtype) # pylint: disable=protected-access if 'buffers' not in stats.weights: stats.weights['buffers'] = 0 stats.weights['buffers'] += 1 elif value is not None: param_cls = type(module._parameters[name]) # pylint: disable=protected-access - module._parameters[name] = param_cls(value, requires_grad=old_value.requires_grad).to(device, dtype=dtype, non_blocking=True) # pylint: disable=protected-access + module._parameters[name] = param_cls(value, requires_grad=old_value.requires_grad).to(device, dtype=dtype) # pylint: disable=protected-access if 'parameters' not in stats.weights: stats.weights['parameters'] = 0 stats.weights['parameters'] += 1 diff --git a/modules/lora/networks.py b/modules/lora/networks.py index f1fdb0c45..5a093370c 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -359,7 +359,7 @@ def network_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.n def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], weight: torch.Tensor, network_layer_name: str): if shared.opts.diffusers_offload_mode == "none": - self.to(devices.device, non_blocking=True) + self.to(devices.device) batch_updown = None batch_ex_bias = None for net in loaded_networks: @@ -370,11 +370,11 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn. updown, ex_bias = module.calc_updown(weight) t1 = time.time() if batch_updown is not None and updown is not None: - batch_updown += updown.to(batch_updown.device, non_blocking=True) + batch_updown += updown.to(batch_updown.device) else: batch_updown = updown if batch_ex_bias is not None and ex_bias is not None: - batch_ex_bias += ex_bias.to(batch_ex_bias.device, non_blocking=True) + batch_ex_bias += ex_bias.to(batch_ex_bias.device) else: batch_ex_bias = ex_bias timer['calc'] += t1 - t0 @@ -412,7 +412,7 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn if updown is not None and len(weights_backup.shape) == 4 and weights_backup.shape[1] == 9: # inpainting model. zero pad updown to make channel[1] 4 to 9 updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5)) # pylint: disable=not-callable if updown is not None: - new_weight = weights_backup.to(devices.device, non_blocking=True) + updown.to(devices.device, non_blocking=True) + new_weight = weights_backup.to(devices.device) + updown.to(devices.device) if getattr(self, "quant_type", None) in ['nf4', 'fp4'] and bnb is not None: self.weight = bnb.nn.Params4bit(new_weight, quant_state=self.quant_state, quant_type=self.quant_type, blocksize=self.blocksize) else: @@ -429,7 +429,7 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn else: self.bias = None if ex_bias is not None: - new_weight = bias_backup.to(devices.device, non_blocking=True) + ex_bias.to(devices.device, non_blocking=True) + new_weight = bias_backup.to(devices.device) + ex_bias.to(devices.device) self.bias = torch.nn.Parameter(new_weight.to(device=orig_device), requires_grad=False) del new_weight else: @@ -482,7 +482,7 @@ def network_activate(): pbar.update(task, advance=1, description=f'networks={len(loaded_networks)} skip') continue weight = getattr(module, 'weight', None) - weight = weight.to(devices.device, non_blocking=True) if weight is not None else None + weight = weight.to(devices.device) if weight is not None else None backup_size += network_backup_weights(module, weight, network_layer_name, wanted_names) batch_updown, batch_ex_bias = network_calc_weights(module, weight, network_layer_name) weights_device, weights_dtype = network_apply_weights(module, batch_updown, batch_ex_bias, orig_device) diff --git a/modules/processing_vae.py b/modules/processing_vae.py index b114e01d3..1c4a45f07 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -117,7 +117,7 @@ def full_vae_decode(latents, model): model.vae.orig_dtype = model.vae.dtype model.vae = model.vae.to(dtype=torch.float32) latents = latents.to(torch.float32) - latents = latents.to(devices.device, non_blocking=True) + latents = latents.to(devices.device) if getattr(model.vae, "post_quant_conv", None) is not None: latents = latents.to(next(iter(model.vae.post_quant_conv.parameters())).dtype) diff --git a/modules/rife/__init__.py b/modules/rife/__init__.py index f74f3d984..2a636eb2f 100644 --- a/modules/rife/__init__.py +++ b/modules/rife/__init__.py @@ -82,13 +82,13 @@ def interpolate(images: list, count: int = 2, scale: float = 1.0, pad: int = 1, for _i in range(pad): # fill starting frames buffer.put(frame) - I1 = f_pad(torch.from_numpy(np.transpose(frame, (2,0,1))).to(devices.device, non_blocking=True).unsqueeze(0).float() / 255.) + I1 = f_pad(torch.from_numpy(np.transpose(frame, (2,0,1))).to(devices.device).unsqueeze(0).float() / 255.) with torch.no_grad(): with tqdm(total=len(images), desc='Interpolate', unit='frame') as pbar: for image in images: frame = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) I0 = I1 - I1 = f_pad(torch.from_numpy(np.transpose(frame, (2,0,1))).to(devices.device, non_blocking=True).unsqueeze(0).float() / 255.) + I1 = f_pad(torch.from_numpy(np.transpose(frame, (2,0,1))).to(devices.device).unsqueeze(0).float() / 255.) I0_small = F.interpolate(I0, (32, 32), mode='bilinear', align_corners=False).to(torch.float32) I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False).to(torch.float32) ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3]) diff --git a/modules/sd_hijack_accelerate.py b/modules/sd_hijack_accelerate.py index 90eac5c4e..f8cf8983f 100644 --- a/modules/sd_hijack_accelerate.py +++ b/modules/sd_hijack_accelerate.py @@ -35,10 +35,10 @@ def hijack_set_module_tensor( with devices.inference_context(): # note: majority of time is spent on .to(old_value.dtype) if tensor_name in module._buffers: # pylint: disable=protected-access - module._buffers[tensor_name] = value.to(device, old_value.dtype, non_blocking=True) # pylint: disable=protected-access + module._buffers[tensor_name] = value.to(device, old_value.dtype) # pylint: disable=protected-access elif value is not None or not devices.same_device(torch.device(device), module._parameters[tensor_name].device): # pylint: disable=protected-access param_cls = type(module._parameters[tensor_name]) # pylint: disable=protected-access - module._parameters[tensor_name] = param_cls(value, requires_grad=old_value.requires_grad).to(device, old_value.dtype, non_blocking=True) # pylint: disable=protected-access + module._parameters[tensor_name] = param_cls(value, requires_grad=old_value.requires_grad).to(device, old_value.dtype) # pylint: disable=protected-access t1 = time.time() tensor_to_timer += (t1 - t0) @@ -63,10 +63,10 @@ def hijack_set_module_tensor_simple( old_value = getattr(module, tensor_name) with devices.inference_context(): if tensor_name in module._buffers: # pylint: disable=protected-access - module._buffers[tensor_name] = value.to(device, non_blocking=True) # pylint: disable=protected-access + module._buffers[tensor_name] = value.to(device) # pylint: disable=protected-access elif value is not None or not devices.same_device(torch.device(device), module._parameters[tensor_name].device): # pylint: disable=protected-access param_cls = type(module._parameters[tensor_name]) # pylint: disable=protected-access - module._parameters[tensor_name] = param_cls(value, requires_grad=old_value.requires_grad).to(device, non_blocking=True) # pylint: disable=protected-access + module._parameters[tensor_name] = param_cls(value, requires_grad=old_value.requires_grad).to(device) # pylint: disable=protected-access t1 = time.time() tensor_to_timer += (t1 - t0) diff --git a/modules/sd_models.py b/modules/sd_models.py index 63ec6b327..8853916e4 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -529,7 +529,7 @@ def move_model(model, device=None, force=False): t0 = time.time() try: if hasattr(model, 'to'): - model.to(device, non_blocking=True) + model.to(device) if hasattr(model, "prior_pipe"): model.prior_pipe.to(device) except Exception as e0: