initial diffusers merge into dev

This commit is contained in:
Vladimir Mandic
2023-07-02 14:04:54 -04:00
parent 875d0db103
commit a2caafe4df
11 changed files with 431 additions and 52 deletions
+50 -8
View File
@@ -223,7 +223,7 @@ class StableDiffusionProcessing:
# HACK: Using introspection as the Depth2Image model doesn't appear to uniquely
# identify itself with a field common to all models. The conditioning_key is also hybrid.
if backend == Backend.DIFFUSERS: # TODO: Diffusers img2img_image_conditioning
return latent_image.new_zeros(latent_image.shape[0], 5, 1, 1)
return None
if isinstance(self.sd_model, LatentDepth2ImageDiffusion):
return self.depth2img_image_conditioning(source_image)
if self.sd_model.cond_stage_key == "edit":
@@ -520,7 +520,8 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
if k == 'sd_vae':
sd_vae.reload_vae_weights()
sd_models.apply_token_merging(p.sd_model, p.get_token_merging_ratio())
if not shared.opts.cuda_compile:
sd_models.apply_token_merging(p.sd_model, p.get_token_merging_ratio())
if cmd_opts.profile:
"""
@@ -538,7 +539,8 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
else:
res = process_images_inner(p)
finally:
sd_models.apply_token_merging(p.sd_model, 0)
if not shared.opts.cuda_compile:
sd_models.apply_token_merging(p.sd_model, 0)
if p.override_settings_restore_afterwards: # restore opts to original state
for k, v in stored_opts.items():
setattr(opts, k, v)
@@ -557,6 +559,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
assert len(p.prompt) > 0
else:
assert p.prompt is not None
seed = get_fixed_seed(p.seed)
subseed = get_fixed_seed(p.subseed)
if backend == Backend.ORIGINAL:
@@ -683,26 +686,39 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
devices.torch_gc()
if p.scripts is not None:
p.scripts.postprocess_batch(p, x_samples_ddim, batch_number=n)
else: # TODO Diffusers main processing
elif backend == Backend.DIFFUSERS:
generator = [torch.Generator(device="cpu").manual_seed(s) for s in seeds]
if shared.sd_model.scheduler.name != p.sampler_name:
sampler = sd_samplers.all_samplers_map.get(p.sampler_name, None)
if sampler is None:
sampler = sd_samplers.all_samplers_map.get("UniPC")
scheduler = sampler.constructor(shared.sd_model.sd_checkpoint_info.filename)
# TODO(Patrick): For wrapped pipelines this is currently a no-op
shared.sd_model.scheduler = scheduler.sampler
if sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE:
task_specific_kwargs = {"height": p.height, "width": p.width}
elif sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.IMAGE_2_IMAGE:
task_specific_kwargs = {"image": p.init_images[0], "strength": p.denoising_strength}
elif sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.INPAINTING:
# TODO(PVP): change out to latents once possible with `diffusers`
task_specific_kwargs = {"image": p.init_images[0], "mask_image": p.image_mask, "strength": p.denoising_strength}
output = shared.sd_model(
prompt=prompts,
negative_prompt=negative_prompts,
num_inference_steps=p.steps,
guidance_scale=p.cfg_scale,
height=p.height,
width=p.width,
generator=generator,
output_type="np",
**task_specific_kwargs
)
x_samples_ddim = output.images
else:
raise ValueError(f"Unknown backend {backend}")
for i, x_sample in enumerate(x_samples_ddim):
p.batch_index = i
if backend == Backend.ORIGINAL:
@@ -820,8 +836,12 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
self.applied_old_hires_behavior_to = None
def init(self, all_prompts, all_seeds, all_subseeds):
if backend == Backend.DIFFUSERS:
sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE)
self.width = self.width or 512
self.height = self.height or 512
if self.enable_hr:
if opts.use_old_hires_fix_width_height and self.applied_old_hires_behavior_to != (self.width, self.height):
self.hr_resize_x = self.width
@@ -873,6 +893,9 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
self.extra_generation_params["Hires upscaler"] = self.hr_upscaler
def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts): # TODO this is majority of processing time
if backend == Backend.DIFFUSERS:
sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE)
self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model)
latent_scale_mode = shared.latent_upscale_modes.get(self.hr_upscaler, None) if self.hr_upscaler is not None else shared.latent_upscale_modes.get(shared.latent_upscale_default_mode, "nearest")
if self.enable_hr and latent_scale_mode is None:
@@ -978,12 +1001,18 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
self.image_conditioning = None
def init(self, all_prompts, all_seeds, all_subseeds):
image_mask = self.image_mask
if backend == Backend.DIFFUSERS and image_mask is None:
sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE)
elif backend == Backend.DIFFUSERS and image_mask is not None:
sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.INPAINTING)
self.sd_model.dtype = self.sd_model.unet.dtype
force_latent_upscaler = shared.opts.data.get('force_latent_sampler')
if self.sampler_name in ['PLMS']:
self.sampler_name = force_latent_upscaler if force_latent_upscaler != 'None' else shared.opts.fallback_sampler # PLMS does not support img2img, use fallback instead
self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model)
crop_region = None
image_mask = self.image_mask
if image_mask is not None:
image_mask = image_mask.convert('L')
if self.inpainting_mask_invert:
@@ -1048,7 +1077,13 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
image = torch.from_numpy(batch_images)
image = 2. * image - 1.
image = image.to(shared.device)
self.init_latent = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(image))
if backend == Backend.ORIGINAL:
self.init_latent = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(image))
else:
# we don't pre-encode the latents for diffusers to allow the UI to stay general for different model types
self.init_latent = None
if self.resize_mode == 3:
self.init_latent = torch.nn.functional.interpolate(self.init_latent, size=(self.height // opt_f, self.width // opt_f), mode="bilinear")
if image_mask is not None:
@@ -1068,6 +1103,13 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
self.image_conditioning = self.img2img_image_conditioning(image, self.init_latent, image_mask)
def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts):
if backend == Backend.DIFFUSERS:
if self.init_mask is None: # pylint: disable=no-member
sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE)
else:
sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.INPAINTING)
self.sd_model.dtype = self.sd_model.unet.dtype
x = create_random_tensors([opt_C, self.height // opt_f, self.width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self)
if self.initial_noise_multiplier != 1.0:
self.extra_generation_params["Noise multiplier"] = self.initial_noise_multiplier