update samplers and callbacks

TODO.md

@@ -4,13 +4,13 @@
 
 Stuff to be fixed, in no particular order...
 
-- SD-XL VAE `AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)`
+- SD-XL VAE
 - SD-XL Lora
 - SD-XL Sketch/Inpaint
 - Kandinsky 2.2 (2.1 is working)
 - Misterious Extensions auto-enabling
 - Misterious Extra network corruptions
-- script_callbacks.on_model_loaded 
+- script_callbacks.on_model_loaded
 
 ## Features
 

extensions-builtin/Lora/lora.py

@@ -139,6 +139,8 @@ class LoraUpDownModule:
 
 def assign_lora_names_to_compvis_modules(sd_model):
     lora_layer_mapping = {}
+    if not hasattr(shared.sd_model, 'cond_stage_model'):
+        return
 
     for name, module in shared.sd_model.cond_stage_model.wrapped.named_modules():
         lora_name = name.replace(".", "_")

modules/sd_models.py

@@ -778,7 +778,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
 
         if sd_model is None:
             shared.log.error('Diffuser model not loaded')
-            return
+            return        
         sd_model.sd_checkpoint_info = checkpoint_info # pylint: disable=attribute-defined-outside-init
         sd_model.sd_model_checkpoint = checkpoint_info.filename # pylint: disable=attribute-defined-outside-init
         sd_model.sd_model_hash = checkpoint_info.hash # pylint: disable=attribute-defined-outside-init
@@ -809,6 +809,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
     timer.record("load")
     shared.log.info(f"Model loaded in {timer.summary()}")
     devices.torch_gc(force=True)
+    script_callbacks.model_loaded_callback(sd_model)
     shared.log.info(f'Model load finished: {memory_stats()}')
 
 
@@ -1008,6 +1009,7 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model')
             shared.opts.data["sd_model_checkpoint"] = next_checkpoint_info.title
             reload_model_weights(reuse_dict=True) # ok we loaded dict now lets redo and load model on top of it
         return model_data.sd_model if op == 'model' or op == 'dict' else model_data.sd_refiner
+
     try:
         load_model_weights(sd_model, checkpoint_info, state_dict, timer)
     except Exception:

modules/sd_samplers_diffusers.py

@@ -21,23 +21,16 @@ except Exception as e:
     import diffusers
     log.error(f'Diffusers import error: version={diffusers.__version__} error: {e}')
 
-
 config = {
-    #removed beta start and end from all into each sampler
     'All': { 'num_train_timesteps': 1000, 'beta_schedule': 'linear', 'prediction_type': 'epsilon' },
     'UniPC': { 'beta_start': 0.0001, 'beta_end': 0.02,'solver_order': 2, 'thresholding': False, 'sample_max_value': 1.0, 'predict_x0': 'bh2', 'lower_order_final': True },
-    #DDIM FIX - changed timestep spacing
     'DDIM': { 'beta_start': 0.0001, 'beta_end': 0.02,'clip_sample': True, 'set_alpha_to_one': True, 'steps_offset': 0, 'thresholding': False, 'clip_sample_range': 1.0, 'sample_max_value': 1.0, 'timestep_spacing': 'linspace', 'rescale_betas_zero_snr': False },
-    #DDPM FIX - changed timestep spacing
     'DDPM': { 'beta_start': 0.0001, 'beta_end': 0.02,'variance_type': "fixed_small", 'clip_sample': True, 'thresholding': False, 'clip_sample_range': 1.0, 'sample_max_value': 1.0, 'timestep_spacing': 'linspace'},
-    #ADDED KDPM2
     'KDPM2': {'beta_start': 0.00085, 'beta_end': 0.012, 'steps_offset': 0 },
-    #ADDED KDPM2 A
-    'KDPM2 A': {'beta_start': 0.00085, 'beta_end': 0.012, 'steps_offset': 0 },
-
+    'KDPM2 a': {'beta_start': 0.00085, 'beta_end': 0.012, 'steps_offset': 0 },
     'DEIS': { 'beta_start': 0.0001, 'beta_end': 0.02,'solver_order': 2, 'thresholding': False, 'sample_max_value': 1.0, 'algorithm_type': "deis", 'solver_type': "logrho", 'lower_order_final': True },
     'Euler': {'beta_start': 0.0001, 'beta_end': 0.02, 'interpolation_type': "linear", 'use_karras_sigmas': False },
-    'Euler a': {},
+    'Euler a': { 'beta_start': 0.0001, 'beta_end': 0.02 },
     'Heun': { 'beta_start': 0.0001, 'beta_end': 0.02,'use_karras_sigmas': False },
     'PNDM': { 'beta_start': 0.0001, 'beta_end': 0.02,'skip_prk_steps': False, 'set_alpha_to_one': False, 'steps_offset': 0 },
     'DPM 1S': { 'beta_start': 0.0001, 'beta_end': 0.02,'solver_order': 2, 'thresholding': False, 'sample_max_value': 1.0, 'algorithm_type': "dpmsolver++", 'solver_type': "midpoint", 'lower_order_final': True, 'use_karras_sigmas': False },
@@ -49,10 +42,8 @@ samplers_data_diffusers = [
     sd_samplers_common.SamplerData('UniPC', lambda model: DiffusionSampler('UniPC', UniPCMultistepScheduler, model), [], {}),
     sd_samplers_common.SamplerData('DDIM', lambda model: DiffusionSampler('DDIM', DDIMScheduler, model), [], {}),
     sd_samplers_common.SamplerData('DDPM', lambda model: DiffusionSampler('DDPM', DDPMScheduler, model), [], {}),
-    #ADDED KDPM2
     sd_samplers_common.SamplerData('KDPM2', lambda model: DiffusionSampler('KDPM2', KDPM2DiscreteScheduler, model), [], {}),
-    #ADDED KDPM2 A
-    sd_samplers_common.SamplerData('KDPM2 A', lambda model: DiffusionSampler('KDPM2 A', KDPM2AncestralDiscreteScheduler, model), [], {}),
+    sd_samplers_common.SamplerData('KDPM2 a', lambda model: DiffusionSampler('KDPM2 a', KDPM2AncestralDiscreteScheduler, model), [], {}),
     sd_samplers_common.SamplerData('DEIS', lambda model: DiffusionSampler('DEIS', DEISMultistepScheduler, model), [], {}),
     sd_samplers_common.SamplerData('DPM 1S', lambda model: DiffusionSampler('DPM++ 1S', DPMSolverSinglestepScheduler, model), [], {}),
     sd_samplers_common.SamplerData('DPM 2M', lambda model: DiffusionSampler('DPM++ 2M', DPMSolverMultistepScheduler, model), [], {}),
@@ -91,10 +82,8 @@ class DiffusionSampler:
             self.config['solver_order'] = opts.schedulers_solver_order
         if 'predict_x0' in self.config:
             self.config['predict_x0'] = opts.uni_pc_variant
-##disbaled this for now until full fix is working            
-        #if name.startswith('DPM'):
-        #    self.config['algorithm_type'] = opts.schedulers_dpm_solver
-
+        if name == 'DPM 2M':
+            self.config['algorithm_type'] = opts.schedulers_dpm_solver
         self.sampler = constructor(**self.config)
         self.sampler.name = name
         log.debug(f'Diffusers sampler: {name} {self.config}')