This commit is contained in:
Vladimir Mandic
2023-12-09 10:00:29 -05:00
parent b96131fad2
commit 6cd178541b
+10 -7
View File
@@ -15,7 +15,7 @@ def soft_clamp_tensor(tensor, threshold=0.8, boundary=4):
# shrinking towards the mean; will also remove outliers
if max(abs(tensor.max()), abs(tensor.min())) < boundary or threshold == 0:
return tensor
channel_dim = 1
channel_dim = 0
threshold *= boundary
max_vals = tensor.max(channel_dim, keepdim=True)[0]
max_replace = ((tensor - threshold) / (max_vals - threshold)) * (boundary - threshold) + threshold
@@ -23,8 +23,8 @@ def soft_clamp_tensor(tensor, threshold=0.8, boundary=4):
min_vals = tensor.min(channel_dim, keepdim=True)[0]
min_replace = ((tensor + threshold) / (min_vals + threshold)) * (-boundary + threshold) - threshold
under_mask = tensor < -threshold
debug(f'HDE soft clamp: threshold={threshold} boundary={boundary}')
tensor = torch.where(over_mask, max_replace, torch.where(under_mask, min_replace, tensor))
debug(f'HDR soft clamp: threshold={threshold} boundary={boundary} shape={tensor.shape}')
return tensor
@@ -34,21 +34,23 @@ def center_tensor(tensor, channel_shift=1.0, full_shift=1.0, channels=[0, 1, 2,
means = []
for channel in channels:
means.append(tensor[0, channel].mean())
tensor[0, channel] -= means[-1] * channel_shift
debug(f'HDR center: channel-shift{channel_shift} full-shift={full_shift} means={torch.stack(means)}')
# tensor[0, channel] -= means[-1] * channel_shift
tensor[channel] -= means[-1] * channel_shift
tensor = tensor - tensor.mean() * full_shift
debug(f'HDR center: channel-shift={channel_shift} full-shift={full_shift} means={torch.stack(means)} shape={tensor.shape}')
return tensor
def maximize_tensor(tensor, boundary=1.0, channels=[0, 1, 2]): # pylint: disable=dangerous-default-value # noqa: B006
def maximize_tensor(tensor, boundary=1.0, _channels=[0, 1, 2]): # pylint: disable=dangerous-default-value # noqa: B006
if boundary == 1.0:
return tensor
boundary *= 4
min_val = tensor.min()
max_val = tensor.max()
normalization_factor = boundary / max(abs(min_val), abs(max_val))
tensor[0, channels] *= normalization_factor
debug(f'HDR maximize: boundary={boundary} min={min_val} max={max_val} factor={normalization_factor}')
# tensor[0, channels] *= normalization_factor
tensor *= normalization_factor
debug(f'HDR maximize: boundary={boundary} min={min_val} max={max_val} factor={normalization_factor} shape={tensor.shape}')
return tensor
@@ -70,6 +72,7 @@ def correction_callback(p, timestep, kwargs):
if not p.hdr_clamp and not p.hdr_center and not p.hdr_maximize:
return kwargs
latents = kwargs["latents"]
# debug(f'HDR correction: latents={latents.shape}')
if len(latents.shape) == 4: # standard batched latent
for i in range(latents.shape[0]):
latents[i] = correction(p, timestep, latents[i])