cleanup corrections

This commit is contained in:
AI-Casanova
2024-01-27 19:16:30 -06:00
committed by Vladimir Mandic
parent 4b30d6aeb7
commit dc70246c78
+14 -12
View File
@@ -29,29 +29,25 @@ def soft_clamp_tensor(tensor, threshold=0.8, boundary=4):
return tensor
def center_tensor(tensor, channel_shift=1.0, full_shift=1.0, channels=[0, 1, 2, 3]): # pylint: disable=dangerous-default-value # noqa: B006
def center_tensor(tensor, channel_shift=1.0, full_shift=1.0):
if channel_shift == 0 and full_shift == 0:
return tensor
means = []
for channel in channels:
means.append(tensor[0, channel].mean())
# 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}')
tensor -= tensor.mean(dim=(-1, -2), keepdim=True) * channel_shift
tensor -= tensor.mean() * full_shift
debug(f'HDR center: channel-shift={channel_shift} full-shift={full_shift}')
debug(f'HDR center: After Adjustment: Full mean={tensor.mean().item()} Channel means={tensor.mean(dim=(-1, -2)).cpu().numpy()}')
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):
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
tensor *= normalization_factor
debug(f'HDR maximize: boundary={boundary} min={min_val} max={max_val} factor={normalization_factor} shape={tensor.shape}')
debug(f'HDR maximize: boundary={boundary} min={min_val} max={max_val} factor={normalization_factor}')
return tensor
@@ -62,7 +58,7 @@ def correction(p, timestep, latent):
if timestep > 700 and p.hdr_center:
p.extra_generation_params["HDR center"] = f'{p.hdr_channel_shift}/{p.hdr_full_shift}'
latent = center_tensor(latent, channel_shift=p.hdr_channel_shift, full_shift=p.hdr_full_shift)
if timestep > 1 and timestep < 100 and p.hdr_maximize:
if 1 < timestep < 100 and p.hdr_maximize:
p.extra_generation_params["HDR max"] = f'{p.hdr_max_center}/{p.hdr_max_boundry}'
latent = center_tensor(latent, channel_shift=p.hdr_max_center, full_shift=1.0)
latent = maximize_tensor(latent, boundary=p.hdr_max_boundry)
@@ -73,10 +69,16 @@ 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('')
debug(f' Timestep: {timestep}')
# 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])
debug(f"Full Mean: {latents[i].mean().item()}")
debug(f"Channel Means: {latents[i].mean(dim=(-1, -2), keepdim=True).flatten().cpu().numpy()}")
debug(f"Channel Mins: {latents[i].min(-1, keepdim=True)[0].min(-2, keepdim=True)[0].flatten().cpu().numpy()}")
debug(f"Channel Maxes: {latents[i].max(-1, keepdim=True)[0].min(-2, keepdim=True)[0].flatten().cpu().numpy()}")
elif len(latents.shape) == 5 and latents.shape[0] == 1: # probably animatediff
latents = latents.squeeze(0).permute(1, 0, 2, 3)
for i in range(latents.shape[0]):