diffusers main code refactor

modules/processing_helpers.py

@@ -398,6 +398,7 @@ def resize_hires(p, latents): # input=latents output=pil if not latent_upscaler
         resized_images.append(resized_image)
     return resized_images
 
+
 def fix_prompts(prompts, negative_prompts, prompts_2, negative_prompts_2):
     if type(prompts) is str:
         prompts = [prompts]
@@ -419,6 +420,7 @@ def fix_prompts(prompts, negative_prompts, prompts_2, negative_prompts_2):
             negative_prompts_2.append(negative_prompts_2[-1])
     return prompts, negative_prompts, prompts_2, negative_prompts_2
 
+
 def calculate_base_steps(p, use_denoise_start, use_refiner_start):
     if len(getattr(p, 'timesteps', [])) > 0:
         return None
@@ -437,6 +439,7 @@ def calculate_base_steps(p, use_denoise_start, use_refiner_start):
     debug_steps(f'Steps: type=base input={p.steps} output={steps} task={sd_models.get_diffusers_task(shared.sd_model)} refiner={use_refiner_start} denoise={p.denoising_strength} model={shared.sd_model_type}')
     return max(1, int(steps))
 
+
 def calculate_hires_steps(p):
     # if len(getattr(p, 'timesteps', [])) > 0:
     #    return None
@@ -449,6 +452,7 @@ def calculate_hires_steps(p):
     debug_steps(f'Steps: type=hires input={p.hr_second_pass_steps} output={steps} denoise={p.denoising_strength} model={shared.sd_model_type}')
     return max(1, int(steps))
 
+
 def calculate_refiner_steps(p):
     # if len(getattr(p, 'timesteps', [])) > 0:
     #    return None
@@ -465,3 +469,84 @@ def calculate_refiner_steps(p):
         steps = (p.refiner_steps * 1.25) + 1
     debug_steps(f'Steps: type=refiner input={p.refiner_steps} output={steps} start={p.refiner_start} denoise={p.denoising_strength}')
     return max(1, int(steps))
+
+
+def get_generator(p):
+    if shared.opts.diffusers_generator_device == "Unset":
+        generator_device = None
+        generator = None
+    elif getattr(p, "generator", None) is not None:
+        generator_device = devices.cpu if shared.opts.diffusers_generator_device == "CPU" else shared.device
+        generator = p.generator
+    else:
+        generator_device = devices.cpu if shared.opts.diffusers_generator_device == "CPU" else shared.device
+        try:
+            generator = [torch.Generator(generator_device).manual_seed(s) for s in p.seeds]
+        except Exception as e:
+            shared.log.error(f'Torch generator: seeds={p.seeds} device={generator_device} {e}')
+            generator = None
+    return generator
+
+
+def set_latents(p):
+    def dummy_prepare_latents(*args, **_kwargs):
+        return args[0] # just return image to skip re-processing it
+
+    from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
+    image = shared.sd_model.image_processor.preprocess(p.init_images) # resize to mod8, normalize, transpose, to tensor
+    timesteps, steps = retrieve_timesteps(shared.sd_model.scheduler, p.steps, devices.device)
+    timesteps, steps = shared.sd_model.get_timesteps(steps, p.denoising_strength, devices.device)
+    timestep = timesteps[:1].repeat(p.batch_size) # need to determine level of added noise
+    latents = shared.sd_model.prepare_latents(image, timestep, batch_size=p.batch_size, num_images_per_prompt=1, dtype=devices.dtype, device=devices.device, generator=get_generator(p))
+    shared.sd_model.prepare_latents = dummy_prepare_latents # stop diffusers processing latents again
+    return latents
+
+def apply_circular(enable, model):
+    try:
+        for layer in [layer for layer in model.unet.modules() if type(layer) is torch.nn.Conv2d]:
+            layer.padding_mode = 'circular' if enable else 'zeros'
+        for layer in [layer for layer in model.vae.modules() if type(layer) is torch.nn.Conv2d]:
+            layer.padding_mode = 'circular' if enable else 'zeros'
+    except Exception as e:
+        debug(f"Diffusers tiling failed: {e}")
+
+def save_intermediate(p, latents, suffix):
+    for i in range(len(latents)):
+        from modules.processing import create_infotext
+        info=create_infotext(p, p.all_prompts, p.all_seeds, p.all_subseeds, [], iteration=p.iteration, position_in_batch=i)
+        decoded = processing_vae.vae_decode(latents=latents, model=shared.sd_model, output_type='pil', full_quality=p.full_quality)
+        for j in range(len(decoded)):
+            images.save_image(decoded[j], path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix=suffix)
+
+def update_sampler(p, sd_model, second_pass=False):
+    sampler_selection = p.hr_sampler_name if second_pass else p.sampler_name
+    if hasattr(sd_model, 'scheduler') and sampler_selection != 'Default':
+        sampler = sd_samplers.all_samplers_map.get(sampler_selection, None)
+        if sampler is None:
+            sampler = sd_samplers.all_samplers_map.get("UniPC")
+        if len(getattr(p, 'timesteps', [])) > 0:
+            if 'schedulers_use_karras' in shared.opts.data:
+                shared.opts.data['schedulers_use_karras'] = False
+            else:
+                shared.opts.schedulers_use_karras = False
+        sampler = sd_samplers.create_sampler(sampler.name, sd_model)
+        sampler_options = []
+        if sampler.config.get('use_karras_sigmas', False):
+            sampler_options.append('karras')
+        if sampler.config.get('rescale_betas_zero_snr', False):
+            sampler_options.append('rescale beta')
+        if sampler.config.get('thresholding', False):
+            sampler_options.append('dynamic thresholding')
+        if 'algorithm_type' in sampler.config:
+            sampler_options.append(sampler.config['algorithm_type'])
+        if shared.opts.schedulers_prediction_type != 'default':
+            sampler_options.append(shared.opts.schedulers_prediction_type)
+        if shared.opts.schedulers_beta_schedule != 'default':
+            sampler_options.append(shared.opts.schedulers_beta_schedule)
+        if 'beta_start' in sampler.config and (shared.opts.schedulers_beta_start > 0 or shared.opts.schedulers_beta_end > 0):
+            sampler_options.append(f'beta {shared.opts.schedulers_beta_start}-{shared.opts.schedulers_beta_end}')
+        if 'solver_order' in sampler.config:
+            sampler_options.append(f'order {shared.opts.schedulers_solver_order}')
+        if 'lower_order_final' in sampler.config:
+            sampler_options.append('low order')
+        p.extra_generation_params['Sampler options'] = '/'.join(sampler_options)