add hdr

CHANGELOG.md

@@ -8,6 +8,8 @@ Major refactor of [FLUX.1](https://blackforestlabs.ai/announcing-black-forest-la
 - Full **ControlNet** support, better **LoRA** support, full **prompt attention** support,
 - faster, more flexible loading, with additional quantization options, and more...
 
+Oh, as a sidenote, and also new auto **HDR** image create for SD and SDXL ;)
+
 ### Details
 
 **Major refactor of FLUX.1 support:**
@@ -40,6 +42,10 @@ Major refactor of [FLUX.1](https://blackforestlabs.ai/announcing-black-forest-la
   enable via *settings -> compute -> fused projections*  
 
 **Other improvements:**
+- **HDR** high-dynamic-range image create for SD and SDXL  
+  create HDR images from in multiple exposures by latent-space modifications during generation  
+  use via *scripts -> hdr*  
+  *note*: save hdr saves image in standard 8bit/channel *and* 16bit/channel PNG format  
 - **taesd** configurable number of layers  
   can be used to speed-up taesd decoding by reducing number of ops  
   e.g. if generating 1024px image, reducing layers by 1 will result in preview being 512px  

modules/processing_correction.py

@@ -85,8 +85,7 @@ def correction(p, timestep, latent):
     if timestep > 950 and p.hdr_clamp:
         p.extra_generation_params["HDR clamp"] = f'{p.hdr_threshold}/{p.hdr_boundary}'
         latent = soft_clamp_tensor(latent, threshold=p.hdr_threshold, boundary=p.hdr_boundary)
-    if 500 < timestep < 800 and (p.hdr_brightness != 0 or p.hdr_color != 0 or p.hdr_tint_ratio != 0):
-        p.extra_generation_params["HDR center"] = f'{p.hdr_color}/{p.hdr_brightness}'
+    if 600 < timestep < 900 and (p.hdr_color != 0 or p.hdr_tint_ratio != 0):
         if p.hdr_brightness != 0:
             latent[0:1] = center_tensor(latent[0:1], full_shift=float(p.hdr_mode), offset=2*p.hdr_brightness)  # Brightness
             p.extra_generation_params["HDR brightness"] = f'{p.hdr_brightness}'
@@ -98,6 +97,11 @@ def correction(p, timestep, latent):
         if p.hdr_tint_ratio != 0:
             latent = color_adjust(latent, p.hdr_color_picker, p.hdr_tint_ratio)
             p.hdr_tint_ratio = 0
+    if timestep < 200 and (p.hdr_brightness != 0): # do it late so it doesn't change the composition
+        if p.hdr_brightness != 0:
+            latent[0:1] = center_tensor(latent[0:1], full_shift=float(p.hdr_mode), offset=2*p.hdr_brightness)  # Brightness
+            p.extra_generation_params["HDR brightness"] = f'{p.hdr_brightness}'
+            p.hdr_brightness = 0
     if timestep < 350 and p.hdr_sharpen != 0:
         p.extra_generation_params["HDR sharpen"] = f'{p.hdr_sharpen}'
         per_step_ratio = 2 ** (timestep / 250) * p.hdr_sharpen / 16

scripts/hdr.py

@@ -0,0 +1,101 @@
+import os
+import cv2
+import numpy as np
+import gradio as gr
+from PIL import Image
+import modules.scripts as scripts
+from modules import images, processing, shared
+from modules.processing import Processed
+from modules.shared import opts, state
+
+
+class Script(scripts.Script):
+    def title(self):
+        return "HDR"
+
+    def show(self, is_img2img):
+        return True
+
+    def ui(self, is_img2img):
+        with gr.Row():
+            gr.HTML("<span>&nbsp High Dynamic Range</span><br>")
+        with gr.Row():
+            save_hdr = gr.Checkbox(label="Save HDR image", value=True)
+            hdr_range = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.65, label='HDR range')
+        with gr.Row():
+            is_tonemap = gr.Checkbox(label="Enable tonemap", value=False)
+            gamma = gr.Slider(minimum=0, maximum=2, step=0.05, value=1.0, label='Gamma', visible=False)
+        with gr.Row():
+            scale = gr.Slider(minimum=0, maximum=2, step=0.05, value=1.0, label='Scale', visible=False)
+            saturation = gr.Slider(minimum=0, maximum=2, step=0.05, value=1.0, label='Saturation', visible=False)
+        is_tonemap.change(fn=self.change_tonemap, inputs=[is_tonemap], outputs=[gamma, scale, saturation])
+        return [hdr_range, save_hdr, is_tonemap, gamma, scale, saturation]
+
+    def change_tonemap(self, is_tonemap):
+        return [gr.update(visible=is_tonemap), gr.update(visible=is_tonemap), gr.update(visible=is_tonemap)]
+
+    def merge(self, imgs: list, is_tonemap: bool, gamma, scale, saturation):
+        shared.log.info(f'HDR: merge images={len(imgs)} tonemap={is_tonemap} sgamma={gamma} scale={scale} saturation={saturation}')
+        imgs_np = [np.asarray(img).astype(np.uint8) for img in imgs]
+
+        align = cv2.createAlignMTB()
+        align.process(imgs_np, imgs_np)
+
+        # cv2.createMergeRobertson()
+        # cv2.createMergeDebevec()
+        merge = cv2.createMergeMertens()
+        hdr = merge.process(imgs_np)
+
+        # cv2.createTonemapDrago()
+        # cv2.createTonemapReinhard()
+        if is_tonemap:
+            tonemap = cv2.createTonemapMantiuk(gamma, scale, saturation)
+            hdr = tonemap.process(hdr)
+
+        ldr = np.clip(hdr * 255, 0, 255).astype(np.uint8)
+        hdr = np.clip(hdr * 65535, 0, 65535).astype(np.uint16)
+        hdr = cv2.cvtColor(hdr, cv2.COLOR_BGR2RGB)
+        return hdr, ldr
+
+    def run(self, p, hdr_range, save_hdr, is_tonemap, gamma, scale, saturation): # pylint: disable=arguments-differ
+        if shared.sd_model_type != 'sd' and shared.sd_model_type != 'sdxl':
+            shared.log.error(f'HDR: incorrect base model: {shared.sd_model.__class__.__name__}')
+            return
+        p.extra_generation_params = {
+            "HDR range": hdr_range,
+        }
+        shared.log.info(f'HDR: range={hdr_range}')
+        processing.fix_seed(p)
+        imgs = []
+        info = ''
+        for i in range(3):
+            p.n_iter = 1
+            p.batch_size = 1
+            p.do_not_save_grid = True
+            p.hdr_brightness = (i - 1) * (2.0 * hdr_range)
+            p.hdr_mode = 0
+            p.task_args['seed'] = p.seed
+            processed: processing.Processed = processing.process_images(p)
+            imgs += processed.images
+            if i == 1:
+                info = processed.info
+            if state.interrupted:
+                break
+
+        if len(imgs) > 1:
+            hdr, ldr = self.merge(imgs, is_tonemap, gamma, scale, saturation)
+            img = Image.fromarray(ldr)
+            imgs.insert(0, img)
+            if save_hdr:
+                fn, _txt = images.save_image(img, shared.opts.outdir_save, "", p.seed, p.prompt, opts.grid_format, info=processed.info, p=p)
+                fn = os.path.splitext(fn)[0] + '-hdr.png'
+                shared.log.debug(f'Save: image="{fn}" type=PNG channels=16')
+                cv2.imwrite(fn, hdr)
+            # if opts.grid_save:
+            #    images.save_image(grid, p.outpath_grids, "grid", p.seed, p.prompt, opts.grid_format, info=processed.info, grid=True, p=p)
+            if opts.return_grid:
+                grid = images.image_grid(imgs, rows=1)
+                imgs.append(grid)
+
+        processed = Processed(p, images_list=imgs, seed=p.seed, info=info)
+        return processed