From 6cd178541b527021371759a33d8632f319a4c762 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Sat, 9 Dec 2023 10:00:29 -0500
Subject: [PATCH] fix hdr

---
 modules/processing_correction.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/modules/processing_correction.py b/modules/processing_correction.py
index d0f189dcc..a421da052 100644
--- a/modules/processing_correction.py
+++ b/modules/processing_correction.py
@@ -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])