remove non-blocking

Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2024-12-06 09:15:37 -05:00
parent 2a1fbd904f
commit 28eeb00432
6 changed files with 16 additions and 16 deletions
+2 -2
View File
@@ -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
+6 -6
View File
@@ -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)
+1 -1
View File
@@ -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)
+2 -2
View File
@@ -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])
+4 -4
View File
@@ -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)
+1 -1
View File
@@ -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: