add dc-sampler

CHANGELOG.md

@@ -1,8 +1,8 @@
 # Change Log for SD.Next
 
-## Update for 2024-09-08
+## Update for 2024-09-09
 
-### Highlights for 2024-09-08
+### Highlights for 2024-09-09
 
 Major refactor of [FLUX.1](https://blackforestlabs.ai/announcing-black-forest-labs/) support:  
 - Full **ControlNet** support, better **LoRA** support, full **prompt attention** implementation  
@@ -25,7 +25,7 @@ And few video related goodies...
 
 Plus tons of minor items and fixes - see [changelog](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md) for details!  
 
-### Details for 2024-09-08
+### Details for 2024-09-09
 
 **Major refactor of FLUX.1 support:**
 - allow configuration of individual FLUX.1 model components: *transformer, text-encoder, vae*  
@@ -92,6 +92,7 @@ Plus tons of minor items and fixes - see [changelog](https://github.com/vladmand
   ui result is always 8bit/channel hdr-effect image plus grid of original images used to create hdr  
   grid image can be disabled via settings -> user interface -> show grid  
   actual full-hdr image is not displayed in ui, only optionally saved to disk  
+- new scheduler: [DC Solver](https://github.com/wl-zhao/DC-Solver)  
 - **color grading** apply professional color grading to your images  
   using industry-standard *.cube* LUTs!
   enable via *scripts -> color-grading*  

modules/processing_vae.py

@@ -5,8 +5,9 @@ import torchvision.transforms.functional as TF
 from modules import shared, devices, sd_models, sd_vae, sd_vae_taesd
 
 
-debug = shared.log.trace if os.environ.get('SD_VAE_DEBUG', None) is not None else lambda *args, **kwargs: None
-debug('Trace: VAE')
+debug = os.environ.get('SD_VAE_DEBUG', None) is not None
+log_debug = shared.log.trace if debug else lambda *args, **kwargs: None
+log_debug('Trace: VAE')
 
 
 def create_latents(image, p, dtype=None, device=None):
@@ -33,6 +34,8 @@ def create_latents(image, p, dtype=None, device=None):
 
 def full_vae_decode(latents, model):
     t0 = time.time()
+    if debug:
+        shared.mem_mon.reset()
     base_device = None
     if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False):
         base_device = sd_models.move_base(model, devices.cpu)
@@ -78,14 +81,16 @@ def full_vae_decode(latents, model):
     elif shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and base_device is not None:
         sd_models.move_base(model, base_device)
     t1 = time.time()
-    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={upcast} images={latents.shape[0]} latents={latents.shape} time={round(t1-t0, 3)}')
+    if debug:
+        log_debug(f'VAE memory: {shared.mem_mon.read()}')
+    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={upcast} images={latents.shape[0]} latents={latents.shape} time={round(t1-t0, 3)}')
     return decoded
 
 
 def full_vae_encode(image, model):
-    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)}')
+    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 getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'):
-        debug('Moving to CPU: model=UNet')
+        log_debug('Moving to CPU: model=UNet')
         unet_device = model.unet.device
         sd_models.move_model(model.unet, devices.cpu)
     if not shared.opts.diffusers_offload_mode == "sequential" and hasattr(model, 'vae'):
@@ -97,7 +102,7 @@ def full_vae_encode(image, model):
 
 
 def taesd_vae_decode(latents):
-    debug(f'VAE decode: name=TAESD images={len(latents)} latents={latents.shape} slicing={shared.opts.diffusers_vae_slicing}')
+    log_debug(f'VAE decode: name=TAESD images={len(latents)} latents={latents.shape} slicing={shared.opts.diffusers_vae_slicing}')
     if len(latents) == 0:
         return []
     if shared.opts.diffusers_vae_slicing and len(latents) > 1:
@@ -110,7 +115,7 @@ def taesd_vae_decode(latents):
 
 
 def taesd_vae_encode(image):
-    debug(f'VAE encode: name=TAESD image={image.shape}')
+    log_debug(f'VAE encode: name=TAESD image={image.shape}')
     encoded = sd_vae_taesd.encode(image)
     return encoded
 
@@ -148,7 +153,7 @@ def vae_decode(latents, model, output_type='np', full_quality=True, width=None,
         image_processor = diffusers.image_processor.VaeImageProcessor()
         imgs = image_processor.postprocess(decoded, output_type=output_type)
     shared.state.job = prev_job
-    if shared.cmd_opts.profile:
+    if shared.cmd_opts.profile or debug:
         t1 = time.time()
         shared.log.debug(f'Profile: VAE decode: {t1-t0:.2f}')
     devices.torch_gc()

modules/sd_samplers_diffusers.py

@@ -5,6 +5,7 @@ import inspect
 from modules import shared
 from modules import sd_samplers_common
 from modules.tcd import TCDScheduler
+from modules.dcsolver import DCSolverMultistepScheduler #https://github.com/wl-zhao/DC-Solver
 
 
 debug = shared.log.trace if os.environ.get('SD_SAMPLER_DEBUG', None) is not None else lambda *args, **kwargs: None
@@ -62,6 +63,7 @@ config = {
     'LMSD': { 'use_karras_sigmas': False, 'timestep_spacing': 'linspace', 'steps_offset': 0 },
     'PNDM': { 'skip_prk_steps': False, 'set_alpha_to_one': False, 'steps_offset': 0, 'timestep_spacing': 'linspace' },
     'SA Solver': {'predictor_order': 2, 'corrector_order': 2, 'thresholding': False, 'lower_order_final': True, 'use_karras_sigmas': False, 'timestep_spacing': 'linspace'},
+    'DC Solver': { 'beta_start': 0.0001, 'beta_end': 0.02, 'solver_order': 2, 'prediction_type': "epsilon", 'thresholding': False, 'solver_type': 'bh2', 'lower_order_final': True, 'dc_order': 2, 'disable_corrector': [0] },
     'LCM': { 'beta_start': 0.00085, 'beta_end': 0.012, 'beta_schedule': "scaled_linear", 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False, 'thresholding': False, 'timestep_spacing': 'linspace' },
     'TCD': { 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False, 'beta_schedule': 'scaled_linear' },
     'Euler SGM': { 'timestep_spacing': "trailing", 'prediction_type': "sample" },
@@ -79,6 +81,7 @@ samplers_data_diffusers = [
     sd_samplers_common.SamplerData('UniPC', lambda model: DiffusionSampler('UniPC', UniPCMultistepScheduler, model), [], {}),
     sd_samplers_common.SamplerData('DEIS', lambda model: DiffusionSampler('DEIS', DEISMultistepScheduler, model), [], {}),
     sd_samplers_common.SamplerData('SA Solver', lambda model: DiffusionSampler('SA Solver', SASolverScheduler, model), [], {}),
+    sd_samplers_common.SamplerData('DC Solver', lambda model: DiffusionSampler('DC Solver', DCSolverMultistepScheduler, model), [], {}),
     sd_samplers_common.SamplerData('DDIM', lambda model: DiffusionSampler('DDIM', DDIMScheduler, model), [], {}),
     sd_samplers_common.SamplerData('Heun', lambda model: DiffusionSampler('Heun', HeunDiscreteScheduler, model), [], {}),
     sd_samplers_common.SamplerData('Euler', lambda model: DiffusionSampler('Euler', EulerDiscreteScheduler, model), [], {}),
@@ -194,5 +197,9 @@ class DiffusionSampler:
         debug(f'Sampler: signature={possible}')
         # shared.log.debug(f'Sampler: sampler="{name}" config={self.config}')
         self.sampler = constructor(**self.config)
+        if name == 'DC Solver':
+            if not hasattr(self.sampler, 'dc_ratios'):
+                pass
+                # self.sampler.dc_ratios = self.sampler.cascade_polynomial_regression(test_CFG=6.0, test_NFE=10, cpr_path='tmp/sd2.1.npy')
         # shared.log.debug(f'Sampler: class="{self.sampler.__class__.__name__}" config={self.sampler.config}')
         self.sampler.name = name

modules/sd_vae_natten.py

@@ -0,0 +1,90 @@
+# copied from https://github.com/Birch-san/sdxl-play/blob/main/src/attn/natten_attn_processor.py
+
+import os
+from typing import Optional
+from diffusers.models.attention import Attention
+import torch
+from torch.nn import Linear
+from einops import rearrange
+from installer import install, log
+
+
+def init():
+    try:
+        os.environ['NATTEN_CUDA_ARCH'] = '8.0;8.6'
+        install('natten')
+        import natten
+        return natten
+    except Exception as e:
+        log.error(f'Init natten: {e}')
+        return None
+
+
+def fuse_qkv(attn: Attention) -> None:
+    has_bias = attn.to_q.bias is not None
+    qkv = Linear(in_features=attn.to_q.in_features, out_features=attn.to_q.out_features*3, bias=has_bias, dtype=attn.to_q.weight.dtype, device=attn.to_q.weight.device)
+    qkv.weight.data.copy_(torch.cat([attn.to_q.weight.data * attn.scale, attn.to_k.weight.data, attn.to_v.weight.data]))
+    if has_bias:
+        qkv.bias.data.copy_(torch.cat([attn.to_q.bias.data * attn.scale, attn.to_k.bias.data, attn.to_v.bias.data]))
+    setattr(attn, 'qkv', qkv) # noqa: B010
+    del attn.to_q, attn.to_k, attn.to_v
+
+
+def fuse_vae_qkv(vae) -> None:
+    for attn in [*vae.encoder.mid_block.attentions, *vae.decoder.mid_block.attentions]:
+        fuse_qkv(attn)
+
+
+class NattenAttnProcessor:
+    kernel_size: int
+
+    def __init__(self, kernel_size: int):
+        self.kernel_size = kernel_size
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+        temb: Optional[torch.FloatTensor] = None,
+    ):
+        import natten
+        assert hasattr(attn, 'qkv'), "Did not find property qkv on attn. Expected you to fuse its q_proj, k_proj, v_proj weights and biases beforehand, and multiply attn.scale into the q weights and bias."
+        residual = hidden_states
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+        # assumes MHA (as opposed to GQA)
+        inner_dim: int = attn.qkv.out_features // 3
+        if attention_mask is not None:
+            raise ValueError("No mask customization for neighbourhood attention; the mask is already complicated enough as it is")
+        if encoder_hidden_states is not None:
+            raise ValueError("NATTEN cannot be used for cross-attention. I think.")
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states)
+            hidden_states = rearrange(hidden_states, '... c h w -> ... h w c')
+        qkv = attn.qkv(hidden_states)
+        # assumes MHA (as opposed to GQA)
+        q, k, v = rearrange(qkv, "n h w (t nh e) -> t n nh h w e", t=3, e=inner_dim)
+        qk = natten.functional.na2d_qk(q, k, self.kernel_size, 1) # natten2dqk
+        a = torch.softmax(qk, dim=-1)
+        hidden_states = natten.functional.na2d_av(a, v, self.kernel_size, 1) # natten2dav
+        hidden_states = rearrange(hidden_states, "n nh h w e -> n h w (nh e)")
+        linear_proj, dropout = attn.to_out
+        hidden_states = linear_proj(hidden_states)
+        hidden_states = dropout(hidden_states)
+        hidden_states = rearrange(hidden_states, '... h w c -> ... c h w')
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+        return hidden_states
+
+
+def enable_natten(pipe):
+    if not hasattr(pipe, 'vae'):
+        return
+    natten = init()
+    kernel_size = 17
+    if natten is not None:
+        log.info(f'VAE natten: version={natten.__version__} kernel={kernel_size}')
+        fuse_vae_qkv(pipe.vae)
+        pipe.vae.set_attn_processor(NattenAttnProcessor(kernel_size=kernel_size))