add new hires with refiner and non-latent modes

CHANGELOG.md

@@ -1,8 +1,9 @@
 # Change Log for SD.Next
 
-## Update for 2023-09-10
+## Update for 2023-09-12
 
-Mostly a service release  
+Mostly a service release, but with some changes in behavior, especially in HiRes area of the code...
+  
 - tons of fixes  
 - changes to **hires**
   - enable non-latent upscale modes (standard upscalers)  
@@ -11,6 +12,12 @@ Mostly a service release
     enabled using **force hires** option in ui  
     hires was not designed to work with standard upscalers, but i understand this is a common workflow  
   - when using refiner, upscale/hires runs before refiner pass  
+  - second pass can now also utilize full/quick vae quality  
+  - note that when combining non-latent upscale, hires and refiner output quality is maximum,  
+    but operations are really resource intensive as it includes: *base->decode->upscale->encode->hires->refine*
+  - all combinations of: decode full/quick + upscale none/latent/non-latent + hires on/off + refiner on/off  
+    should be supported, but given the number of combinations, issues are possible  
+  - all operations are captured in image medata
 - update **ui hints**
 - updated **models -> civitai**
   - search and download loras  
@@ -29,6 +36,7 @@ Mostly a service release
   - capture extension output  
   - capture ldm output  
   - cleaner server restart  
+  - custom exception handling
 
 
 ## Update for 2023-09-06

installer.py

@@ -108,6 +108,18 @@ def setup_logging():
     # logging.getLogger("DeepSpeed").handlers = log.handlers
 
 
+def custom_excepthook(exc_type, exc_value, exc_traceback):
+    import traceback
+    if issubclass(exc_type, KeyboardInterrupt):
+        sys.__excepthook__(exc_type, exc_value, exc_traceback)
+        return
+    log.error(f"Uncaught exception occurred: type={exc_type} value={exc_value}")
+    if exc_traceback:
+        format_exception = traceback.format_tb(exc_traceback)
+        for line in format_exception:
+            log.error(repr(line))
+
+
 def print_dict(d):
     return ' '.join([f'{k}={v}' for k, v in d.items()])
 

javascript/style.css

@@ -119,7 +119,7 @@ div#extras_scale_to_tab div.form{ flex-direction: row; }
 #quicksettings > button { padding: 0 1em 0 0 }
 
 #settings { display: flex; gap: var(--layout-gap); }
-#settings div { border: none; gap: 0.5em; width: fit-content; display: inline-flex; }
+#settings div { border: none; gap: 0.5em; }
 #settings > div.tab-content { flex: 10 0 75%; display: grid; }
 
 #settings > div.tab-content > div { border: none; padding: 0; }

launch.py

@@ -165,6 +165,7 @@ if __name__ == "__main__":
     installer.args = args
     installer.setup_logging()
     installer.log.info('Starting SD.Next')
+    sys.excepthook = installer.custom_excepthook
     installer.read_options()
     if args.skip_all:
         args.quick = True

modules/images.py

@@ -201,7 +201,7 @@ def draw_prompt_matrix(im, width, height, all_prompts, margin=0):
     return draw_grid_annotations(im, width, height, hor_texts, ver_texts, margin)
 
 
-def resize_image(resize_mode, im, width, height, upscaler_name=None):
+def resize_image(resize_mode, im, width, height, upscaler_name=None, output_type='image'):
     """
     Resizes an image with the specified resize_mode, width, and height.
     Args:
@@ -261,6 +261,8 @@ def resize_image(resize_mode, im, width, height, upscaler_name=None):
             fill_width = width // 2 - src_w // 2
             res.paste(resized.resize((fill_width, height), box=(0, 0, 0, height)), box=(0, 0))
             res.paste(resized.resize((fill_width, height), box=(resized.width, 0, resized.width, height)), box=(fill_width + src_w, 0))
+    if output_type == 'np':
+        return np.array(res)
     return res
 
 

modules/processing.py

@@ -488,7 +488,7 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts, all_seeds, all_su
         "Backend": 'Diffusers' if shared.backend == shared.Backend.DIFFUSERS else 'Original',
         "Version": git_commit,
         "Comment": comment,
-        "Operations": ', '.join(list(set(p.ops))).replace('"', '') if len(p.ops) > 0 else None,
+        "Operations": '; '.join(p.ops).replace('"', '') if len(p.ops) > 0 else 'none',
     }
     if 'txt2img' in p.ops:
         pass
@@ -800,8 +800,13 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
 
             for i, x_sample in enumerate(x_samples_ddim):
                 p.batch_index = i
-                x_sample = 255. * (np.moveaxis(x_sample.cpu().numpy(), 0, 2) if shared.backend == shared.Backend.ORIGINAL else x_sample)
-                x_sample = validate_sample(x_sample)
+                if type(x_sample) == Image.Image:
+                    image = x_sample
+                    x_sample = np.array(x_sample)
+                else:
+                    x_sample = 255. * (np.moveaxis(x_sample.cpu().numpy(), 0, 2) if shared.backend == shared.Backend.ORIGINAL else x_sample)
+                    x_sample = validate_sample(x_sample)
+                    image = Image.fromarray(x_sample)
                 if p.restore_faces:
                     if shared.opts.save and not p.do_not_save_samples and shared.opts.save_images_before_face_restoration:
                         orig = p.restore_faces
@@ -811,7 +816,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                         images.save_image(Image.fromarray(x_sample), path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-face-restoration")
                     p.ops.append('face')
                     x_sample = modules.face_restoration.restore_faces(x_sample)
-                image = Image.fromarray(x_sample)
+                    image = Image.fromarray(x_sample)
                 if p.scripts is not None:
                     pp = modules.scripts.PostprocessImageArgs(image)
                     p.scripts.postprocess_image(p, pp)

modules/processing_diffusers.py

@@ -2,6 +2,7 @@ import time
 import inspect
 import typing
 import torch
+import torchvision.transforms.functional as TF
 import modules.devices as devices
 import modules.shared as shared
 import modules.sd_samplers as sd_samplers
@@ -31,14 +32,17 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
         shared.log.info(f'Hires: upscaler={p.hr_upscaler} width={p.hr_upscale_to_x} height={p.hr_upscale_to_y} images={latents.shape[0]}')
         if latent_upscaler is not None:
             latents = torch.nn.functional.interpolate(latents, size=(p.hr_upscale_to_y // 8, p.hr_upscale_to_x // 8), mode=latent_upscaler["mode"], antialias=latent_upscaler["antialias"])
-        first_pass_images = vae_decode(latents=latents, model=shared.sd_model, full_quality=True, output_type='pil')
+        first_pass_images = vae_decode(latents=latents, model=shared.sd_model, full_quality=p.full_quality, output_type='pil')
         p.init_images = []
         for first_pass_image in first_pass_images:
             if latent_upscaler is None:
                 init_image = images.resize_image(1, first_pass_image, p.hr_upscale_to_x, p.hr_upscale_to_y, upscaler_name=p.hr_upscaler)
             else:
                 init_image = first_pass_image
+            # if is_refiner_enabled:
+            #    init_image = vae_encode(init_image, model=shared.sd_model, full_quality=p.full_quality)
             p.init_images.append(init_image)
+        return p.init_images
 
     def save_intermediate(latents, suffix):
         for i in range(len(latents)):
@@ -64,7 +68,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
                 time.sleep(0.1)
 
     def full_vae_decode(latents, model):
-        shared.log.debug(f'VAE decode: name={sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "baked"} dtype={model.vae.dtype} upcast={model.vae.config.get("force_upcast", None)} images={latents.shape[0]}')
+        shared.log.debug(f'VAE decode: name={sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "baked"} dtype={model.vae.dtype} upcast={model.vae.config.get("force_upcast", None)} images={latents.shape[0]} latents={latents.shape}')
         if shared.opts.diffusers_move_unet and not model.has_accelerate:
             shared.log.debug('Moving to CPU: model=UNet')
             unet_device = model.unet.device
@@ -78,13 +82,34 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
             model.unet.to(unet_device)
         return decoded
 
+    def full_vae_encode(image, model):
+        shared.log.debug(f'VAE encode: name={sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "baked"} dtype={model.vae.dtype} upcast={model.vae.config.get("force_upcast", None)}')
+        if shared.opts.diffusers_move_unet and not model.has_accelerate:
+            shared.log.debug('Moving to CPU: model=UNet')
+            unet_device = model.unet.device
+            model.unet.to(devices.cpu)
+            devices.torch_gc()
+        if not shared.cmd_opts.lowvram and not shared.opts.diffusers_seq_cpu_offload:
+            model.vae.to(devices.device)
+        encoded = model.vae.encode(image.to(model.vae.device, model.vae.dtype))
+        if shared.opts.diffusers_move_unet and not model.has_accelerate:
+            model.unet.to(unet_device)
+        return encoded
+
     def taesd_vae_decode(latents):
-        shared.log.debug(f'VAE decode: name=TAESD images={latents.shape[0]}')
-        decoded = torch.zeros((len(latents), 3, p.height, p.width), dtype=devices.dtype_vae, device=devices.device)
+        shared.log.debug(f'VAE decode: name=TAESD images={len(latents)} latents={latents.shape}')
+        if len(latents) == 0:
+            return []
+        decoded = torch.zeros((len(latents), 3, latents.shape[2] * 8, latents.shape[3] * 8), dtype=devices.dtype_vae, device=devices.device)
         for i in range(len(output.images)):
             decoded[i] = (sd_vae_taesd.decode(latents[i]) * 2.0) - 1.0
         return decoded
 
+    def taesd_vae_encode(image):
+        shared.log.debug(f'VAE encode: name=TAESD image={image.shape}')
+        encoded = sd_vae_taesd.encode(image)
+        return encoded
+
     def vae_decode(latents, model, output_type='np', full_quality=True):
         if not torch.is_tensor(latents): # already decoded
             return latents
@@ -105,6 +130,19 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
         imgs = model.image_processor.postprocess(decoded, output_type=output_type)
         return imgs
 
+    def vae_encode(image, model, full_quality=True): # pylint: disable=unused-variable
+        if shared.state.interrupted or shared.state.skipped:
+            return []
+        if not hasattr(model, 'vae'):
+            shared.log.error('VAE not found in model')
+            return []
+        tensor = TF.to_tensor(image.convert("RGB")).unsqueeze(0).to(devices.device, devices.dtype_vae)
+        if full_quality:
+            latents = full_vae_encode(image=tensor, model=shared.sd_model)
+        else:
+            latents = taesd_vae_encode(image=tensor)
+        return latents
+
     def fix_prompts(prompts, negative_prompts, prompts_2, negative_prompts_2):
         if type(prompts) is str:
             prompts = [prompts]
@@ -323,8 +361,9 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
             p.ops.append('upscale')
             if shared.opts.save and not p.do_not_save_samples and shared.opts.save_images_before_highres_fix and hasattr(shared.sd_model, 'vae'):
                 save_intermediate(latents=output.images, suffix="-before-hires")
-            hires_resize(latents=output.images)
+            output.images = hires_resize(latents=output.images)
             if latent_scale_mode is not None or p.hr_force:
+                p.ops.append('hires')
                 recompile_model(hires=True)
                 sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE)
                 hires_args = set_pipeline_args(
@@ -372,10 +411,10 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
         p.ops.append('refine')
         for i in range(len(output.images)):
             image = output.images[i]
-            if (image.shape[2] == 3) and (image.shape[0] % 8 != 0 or image.shape[1] % 8 != 0):
-                shared.log.warning(f'Refiner requires image size to be divisible by 8: {image.shape}')
-                results.append(image)
-                return results
+            # if (image.shape[2] == 3) and (image.shape[0] % 8 != 0 or image.shape[1] % 8 != 0):
+            #    shared.log.warning(f'Refiner requires image size to be divisible by 8: {image.shape}')
+            #    results.append(image)
+            #    return results
             refiner_args = set_pipeline_args(
                 model=shared.sd_refiner,
                 prompts=[p.refiner_prompt] if len(p.refiner_prompt) > 0 else prompts[i],

modules/taesd/sd_vae_taesd.py

@@ -61,3 +61,22 @@ def decode(latents):
     enc = latents.unsqueeze(0).to(devices.device, devices.dtype_vae)
     image = vae.decoder(enc).clamp(0, 1).detach()
     return image[0]
+
+def encode(image):
+    from modules import shared
+    model_class = shared.sd_model_type
+    if model_class == 'ldm':
+        model_class = 'sd'
+    if 'sd' not in model_class:
+        shared.log.warning(f'TAESD unsupported model type: {model_class}')
+        return Image.new('RGB', (8, 8), color = (0, 0, 0))
+    vae = taesd_models[f'{model_class}-encoder']
+    if vae is None:
+        model_path = os.path.join(paths_internal.models_path, "TAESD", f"tae{model_class}_encoder.pth")
+        download_model(model_path)
+        if os.path.exists(model_path):
+            taesd_models[f'{model_class}-encoder'] = TAESD(encoder_path=model_path, decoder_path=None)
+            vae = taesd_models[f'{model_class}-encoder']
+            vae.to(devices.device, devices.dtype_vae)
+    latents = vae.encoder(image).detach()
+    return latents

modules/ui.py

@@ -466,13 +466,14 @@ def create_ui(startup_timer = None):
             txt2img_prompt.submit(**txt2img_args)
             submit.click(**txt2img_args)
 
-            def enable_hr_change(visible: bool):
-                return {"visible": visible, "__type__": "update"}, f'Refiner: {"disabled" if modules.shared.opts.sd_model_refiner == "None" else "enabled"}'
+            def enable_hr_change(visible: bool, refiner_start):
+                enabled = modules.shared.opts.sd_model_refiner != "None" and refiner_start > 0 and refiner_start < 1
+                return {"visible": visible, "__type__": "update"}, f'Refiner: {"enabled" if enabled else "disabled"}'
 
             res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height], show_progress=False)
             batch_switch_btn.click(lambda w, h: (h, w), inputs=[batch_count, batch_size], outputs=[batch_count, batch_size], show_progress=False)
             txt_prompt_img.change(fn=modules.images.image_data, inputs=[txt_prompt_img], outputs=[txt2img_prompt, txt_prompt_img])
-            show_second_pass.change(enable_hr_change, inputs=[show_second_pass], outputs=[second_pass_group, hr_refiner], show_progress = False)
+            show_second_pass.change(enable_hr_change, inputs=[show_second_pass, refiner_start], outputs=[second_pass_group, hr_refiner], show_progress = False)
             show_seed.change(gr_show, inputs=[show_seed], outputs=[seed_group], show_progress = False)
             show_batch.change(gr_show, inputs=[show_batch], outputs=[batch_group], show_progress = False)
             show_advanced.change(gr_show, inputs=[show_advanced], outputs=[advanced_group], show_progress = False)

webui.py

@@ -11,9 +11,8 @@ from threading import Thread
 import modules.loader
 import torch # pylint: disable=wrong-import-order
 from modules import timer, errors, paths # pylint: disable=unused-import
-
 local_url = None
-from installer import log, git_commit
+from installer import log, git_commit, custom_excepthook
 import ldm.modules.encoders.modules # pylint: disable=W0611,C0411,E0401
 from modules import shared, extensions, extra_networks, ui_tempdir, ui_extra_networks, modelloader # pylint: disable=ungrouped-imports
 from modules.paths import create_paths
@@ -39,6 +38,8 @@ from modules.shared import cmd_opts, opts
 import modules.hypernetworks.hypernetwork
 from modules.middleware import setup_middleware
 
+
+sys.excepthook = custom_excepthook
 state = shared.state
 if not modules.loader.initialized:
     timer.startup.record("libraries")
@@ -63,6 +64,7 @@ fastapi_args = {
 }
 modules.loader.initialized = True
 
+
 def check_rollback_vae():
     if shared.cmd_opts.rollback_vae:
         if not torch.cuda.is_available():