implement hires for diffusers

This commit is contained in:
Vladimir Mandic
2023-08-19 12:25:41 +00:00
parent 5eac99d3f5
commit 87bb354f4c
7 changed files with 124 additions and 73 deletions
+14 -2
View File
@@ -1,12 +1,24 @@
# Change Log for SD.Next
## Update for 2023-08-18
## Update for 2023-08-19
Another larger release thats been baking in dev branch for a while...
- general:
- caching of extra network information to enable much faster create/refresh operations
thanks @midcoastal
- diffusers:
- redo "move model to cpu" logic to be more reliable
- add **hires** support (*experimental*)
applies to all model types that support img2img, including **sd** and **sd-xl**
also supports all hires upscaler types as well as standard params like steps and denoising strength
when used with **sd-xl**, it can be used with or without refiner loaded
how to enable - there are no explicit checkboxes other than second pass itself:
- hires: upscaler is set and target resolution is not at default
- refiner: if refiner model is loaded
- images save options: *before hires*, *before refiner*
- redo `move model to cpu` logic in settings -> diffusers to be more reliable
note that system defaults have also changed, so you may need to tweak to your liking
- update dependencies
## Update for 2023-08-17
+1 -1
View File
@@ -306,7 +306,7 @@ infotext_to_setting_name_mapping = [
('Noise multiplier', 'initial_noise_multiplier'),
('Eta', 'eta_ancestral'),
('Eta DDIM', 'eta_ddim'),
('Lora method', 'diffusers_lora_loader'),
('LoRA method', 'diffusers_lora_loader'),
('Discard penultimate sigma', 'always_discard_next_to_last_sigma'),
('UniPC variant', 'uni_pc_variant'),
('UniPC skip type', 'uni_pc_skip_type'),
+4 -3
View File
@@ -209,9 +209,10 @@ def resize_image(resize_mode, im, width, height, upscaler_name=None):
Resizes an image with the specified resize_mode, width, and height.
Args:
resize_mode: The mode to use when resizing the image.
0: Resize the image to the specified width and height.
1: Resize the image to fill the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, cropping the excess.
2: Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, filling empty with data from image.
0: No resie
1: Resize the image to the specified width and height.
2: Resize the image to fill the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, cropping the excess.
3: Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, filling empty with data from image.
im: The image to resize.
width: The width to resize the image to.
height: The height to resize the image to.
+1 -1
View File
@@ -711,7 +711,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
else:
raise ValueError(f"Unknown backend {shared.backend}")
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram and shared.backend == shared.Backend.ORIGINAL:
lowvram.send_everything_to_cpu()
devices.torch_gc()
if p.scripts is not None:
+90 -53
View File
@@ -1,8 +1,6 @@
import inspect
import typing
import torch
# import numpy as np
# from PIL import Image
import modules.devices as devices
import modules.shared as shared
import modules.sd_samplers as sd_samplers
@@ -23,27 +21,38 @@ except Exception as ex:
def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_prompts):
results = []
if p.enable_hr and p.hr_upscaler != 'None' and p.denoising_strength > 0 and len(getattr(p, 'init_images', [])) == 0:
p.is_hr_pass = True
is_refiner_enabled = p.enable_hr and shared.sd_refiner is not None
def diffusers_callback(step: int, _timestep: int, latents: torch.FloatTensor):
shared.state.sampling_step = step
def hires_resize(latents): # input=latents output=pil
latent_upscaler = shared.latent_upscale_modes.get(p.hr_upscaler, None)
shared.log.info(f'Diffusers 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')
p.init_images = []
for first_pass_image in first_pass_images:
init_image = images.resize_image(1, first_pass_image, p.hr_upscale_to_x, p.hr_upscale_to_y, upscaler_name=p.hr_upscaler) if latent_upscaler is None else first_pass_image
p.init_images.append(init_image)
p.width = p.hr_upscale_to_x
p.height = p.hr_upscale_to_y
def save_intermediate(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 = vae_decode(latents=latents, model=shared.sd_model, output_type='pil', full_quality=p.full_quality)
for i in range(len(decoded)):
images.save_image(decoded[i], path=p.outpath_samples, basename="", seed=seeds[i], prompt=prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix=suffix)
def diffusers_callback(_step: int, _timestep: int, latents: torch.FloatTensor):
shared.state.sampling_step += 1
shared.state.sampling_steps = p.steps
if p.is_hr_pass:
shared.state.sampling_steps += p.hr_second_pass_steps
shared.state.current_latent = latents
def hires_resize(latents):
return latents # TODO finish hires
if p.hr_upscaler == 'None':
return latents
scale = shared.latent_upscale_modes.get(p.hr_upscaler, None)
if scale is not None:
p.init_hr()
p.ops.append('hires')
shared.log.info(f'Diffusers Hires: upscaler={p.hr_upscaler} mode={scale["mode"]} antialias={scale["antialias"]} width={p.hr_upscale_to_x} height={p.hr_upscale_to_y} images={latents.shape[0]}')
hires_image = torch.nn.functional.interpolate(latents, size=(p.hr_upscale_to_y // 8, p.hr_upscale_to_x // 8), mode=scale["mode"], antialias=scale["antialias"])
else:
shared.log.warning(f'Diffusers hires unsupported: upscaler={p.hr_upscaler} supported=latent modes')
hires_image = latents
return hires_image
def full_vae_decode(latents, model):
shared.log.debug(f'Diffusers 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]}')
if shared.opts.diffusers_move_unet and not model.has_accelerate:
@@ -66,19 +75,18 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
return decoded
def vae_decode(latents, model, output_type='np', full_quality=True):
if not torch.is_tensor(latents): # already decoded
return latents
if latents.shape[0] == 0:
shared.log.error(f'VAE nothing to decode: {latents.shape}')
return []
if shared.state.interrupted or shared.state.skipped:
return []
if not hasattr(model, 'vae'):
shared.log.error('VAE not found in model')
return []
if not torch.is_tensor(latents):
shared.log.error(f'VAE input is not latents: {type(latents)}')
return []
if latents.shape[0] == 0:
shared.log.error(f'VAE nothing to decode: {latents.shape}')
return []
if p.enable_hr:
latents = hires_resize(latents=latents)
if len(latents.shape) == 3: # lost a batch dim in hires
latents = latents.unsqueeze(0)
if full_quality:
decoded = full_vae_decode(latents=latents, model=shared.sd_model)
else:
@@ -105,7 +113,9 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
negative_prompts_2.append(negative_prompts_2[-1])
return prompts, negative_prompts, prompts_2, negative_prompts_2
def set_pipeline_args(model, prompts: list, negative_prompts: list, prompts_2: typing.Optional[list]=None, negative_prompts_2: typing.Optional[list]=None, is_refiner: bool=False, **kwargs):
def set_pipeline_args(model, prompts: list, negative_prompts: list, prompts_2: typing.Optional[list]=None, negative_prompts_2: typing.Optional[list]=None, is_refiner: bool=False, desc:str='', **kwargs):
if hasattr(model, "set_progress_bar_config"):
model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} '+desc, ncols=80, colour='#327fba')
args = {}
pipeline = model
signature = inspect.signature(type(pipeline).__call__)
@@ -138,7 +148,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
else:
args['negative_prompt'] = negative_prompts
if 'num_inference_steps' in possible:
args['num_inference_steps'] = p.steps
args['num_inference_steps'] = p.steps if not p.is_hr_pass else p.hr_second_pass_steps
if 'guidance_scale' in possible:
args['guidance_scale'] = p.cfg_scale
if 'generator' in possible:
@@ -182,8 +192,9 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
return args
is_karras_compatible = shared.sd_model.__class__.__init__.__annotations__.get("scheduler", None) == diffusers.schedulers.scheduling_utils.KarrasDiffusionSchedulers
if (not hasattr(shared.sd_model.scheduler, 'name')) or (shared.sd_model.scheduler.name != p.sampler_name) and (p.sampler_name != 'Default') and is_karras_compatible:
sampler = sd_samplers.all_samplers_map.get(p.sampler_name, None)
use_sampler = p.sampler_name if not p.is_hr_pass else p.latent_sampler
if (not hasattr(shared.sd_model.scheduler, 'name')) or (shared.sd_model.scheduler.name != use_sampler) and (use_sampler != 'Default') and is_karras_compatible:
sampler = sd_samplers.all_samplers_map.get(use_sampler, None)
if sampler is None:
sampler = sd_samplers.all_samplers_map.get("UniPC")
sd_samplers.create_sampler(sampler.name, shared.sd_model) # TODO(Patrick): For wrapped pipelines this is currently a no-op
@@ -220,8 +231,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
if shared.opts.diffusers_move_base and not shared.sd_model.has_accelerate:
shared.sd_model.to(devices.device)
refiner_enabled = shared.sd_refiner is not None and p.enable_hr
pipe_args = set_pipeline_args(
base_args = set_pipeline_args(
model=shared.sd_model,
prompts=prompts,
negative_prompts=negative_prompts,
@@ -229,35 +239,56 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
negative_prompts_2=[p.refiner_negative] if len(p.refiner_negative) > 0 else negative_prompts,
eta=shared.opts.eta_ddim,
guidance_rescale=p.diffusers_guidance_rescale,
denoising_start=0 if refiner_enabled and p.refiner_start > 0 and p.refiner_start < 1 else None,
denoising_end=p.refiner_start if refiner_enabled and p.refiner_start > 0 and p.refiner_start < 1 else None,
denoising_start=0 if is_refiner_enabled and p.refiner_start > 0 and p.refiner_start < 1 else None,
denoising_end=p.refiner_start if is_refiner_enabled and p.refiner_start > 0 and p.refiner_start < 1 else None,
output_type='latent' if hasattr(shared.sd_model, 'vae') else 'np',
is_refiner=False,
clip_skip=p.clip_skip,
desc='Base',
**task_specific_kwargs
)
p.extra_generation_params['CFG rescale'] = p.diffusers_guidance_rescale
p.extra_generation_params["Eta DDIM"] = shared.opts.eta_ddim if shared.opts.eta_ddim is not None and shared.opts.eta_ddim > 0 else None
output = shared.sd_model(**pipe_args) # pylint: disable=not-callable
if shared.state.interrupted or shared.state.skipped:
unload_diffusers_lora()
return results
output = shared.sd_model(**base_args) # pylint: disable=not-callable
if lora_state['active']:
p.extra_generation_params['Lora method'] = shared.opts.diffusers_lora_loader
p.extra_generation_params['LoRA method'] = shared.opts.diffusers_lora_loader
unload_diffusers_lora()
if not refiner_enabled:
results = vae_decode(latents=output.images, model=shared.sd_model, full_quality=p.full_quality)
else:
for i in range(len(output.images)): # save images before refiner
if shared.opts.save and not p.do_not_save_samples and shared.opts.save_images_before_refiner and hasattr(shared.sd_model, 'vae'):
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 = vae_decode(latents=output.images, model=shared.sd_model, output_type='pil', full_quality=p.full_quality)
for i in range(len(decoded)):
images.save_image(decoded[i], path=p.outpath_samples, basename="", seed=seeds[i], prompt=prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-refiner")
if shared.state.interrupted or shared.state.skipped:
return results
# optional hires pass
if p.is_hr_pass:
p.init_hr()
if p.width != p.hr_upscale_to_x or p.height != p.hr_upscale_to_y:
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)
print('HERE', p.init_images)
sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE)
p.ops.append('hires')
hires_args = set_pipeline_args(
model=shared.sd_model,
prompts=prompts,
negative_prompts=negative_prompts,
prompts_2=[p.refiner_prompt] if len(p.refiner_prompt) > 0 else prompts,
negative_prompts_2=[p.refiner_negative] if len(p.refiner_negative) > 0 else negative_prompts,
eta=shared.opts.eta_ddim,
guidance_rescale=p.diffusers_guidance_rescale,
output_type='latent' if hasattr(shared.sd_model, 'vae') else 'np',
is_refiner=False,
clip_skip=p.clip_skip,
image=p.init_images,
strength=p.denoising_strength,
desc='Hires',
)
output = shared.sd_model(**hires_args) # pylint: disable=not-callable
# optional refiner pass or decode
if is_refiner_enabled:
if shared.opts.save and not p.do_not_save_samples and shared.opts.save_images_before_refiner and hasattr(shared.sd_model, 'vae'):
save_intermediate(latents=output.images, suffix="-before-refiner")
if shared.opts.diffusers_move_base and not shared.sd_model.has_accelerate:
shared.log.debug('Diffusers: Moving base model to CPU')
shared.sd_model.to(devices.cpu)
@@ -276,7 +307,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
shared.sd_refiner.to(devices.device)
p.ops.append('refine')
for i in range(len(output.images)):
pipe_args = set_pipeline_args(
refiner_args = set_pipeline_args(
model=shared.sd_refiner,
prompts=[p.refiner_prompt] if len(p.refiner_prompt) > 0 else prompts[i],
negative_prompts=[p.refiner_negative] if len(p.refiner_negative) > 0 else negative_prompts[i],
@@ -291,19 +322,25 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
output_type='latent' if hasattr(shared.sd_refiner, 'vae') else 'np',
is_refiner=True,
clip_skip=p.clip_skip,
desc='Refiner',
)
refiner_output = shared.sd_refiner(**pipe_args) # pylint: disable=not-callable
refiner_output = shared.sd_refiner(**refiner_args) # pylint: disable=not-callable
p.extra_generation_params['Image CFG scale'] = p.image_cfg_scale if p.image_cfg_scale is not None else None
p.extra_generation_params['Refiner start'] = p.refiner_start
p.extra_generation_params["Hires steps"] = p.hr_second_pass_steps
if not shared.state.interrupted and not shared.state.skipped:
refiner_images = vae_decode(latents=refiner_output.images, model=shared.sd_refiner, full_quality=True)
results.append(refiner_images[0])
for refiner_image in refiner_images:
results.append(refiner_image)
if shared.opts.diffusers_move_refiner and not shared.sd_refiner.has_accelerate:
shared.log.debug('Diffusers: Moving refiner model to CPU')
shared.sd_refiner.to(devices.cpu)
devices.torch_gc()
# final decode since there is no refiner
if not is_refiner_enabled:
results = vae_decode(latents=output.images, model=shared.sd_model, full_quality=p.full_quality)
return results
+10 -9
View File
@@ -268,7 +268,7 @@ def list_themes():
def disable_extensions():
if opts.lyco_patch_lora:
if opts.lyco_patch_lora and backend != Backend.DIFFUSERS:
if 'Lora' not in opts.disabled_extensions:
opts.data['disabled_extensions'].append('Lora')
else:
@@ -398,8 +398,8 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
options_templates.update(options_section(('diffusers', "Diffusers Settings"), {
"diffusers_pipeline": OptionInfo(pipelines[0], 'Diffusers pipeline', gr.Dropdown, lambda: {"choices": pipelines}),
"diffusers_move_base": OptionInfo(False, "Move base model to CPU when using refiner"),
"diffusers_move_unet": OptionInfo(False, "Move base model to CPU when using VAE"),
"diffusers_move_base": OptionInfo(True, "Move base model to CPU when using refiner"),
"diffusers_move_unet": OptionInfo(True, "Move base model to CPU when using VAE"),
"diffusers_move_refiner": OptionInfo(True, "Move refiner model to CPU when not in use"),
"diffusers_extract_ema": OptionInfo(True, "Use model EMA weights when possible"),
"diffusers_generator_device": OptionInfo("default", "Generator device", gr.Radio, lambda: {"choices": ["default", "cpu"]}),
@@ -407,7 +407,7 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), {
"diffusers_seq_cpu_offload": OptionInfo(False, "Enable sequential CPU offload (--lowvram)"),
"diffusers_vae_upcast": OptionInfo("default", "VAE upcasting", gr.Radio, lambda: {"choices": ['default', 'true', 'false']}),
"diffusers_vae_slicing": OptionInfo(True, "Enable VAE slicing"),
"diffusers_vae_tiling": OptionInfo(False, "Enable VAE tiling"),
"diffusers_vae_tiling": OptionInfo(True, "Enable VAE tiling"),
"diffusers_attention_slicing": OptionInfo(False, "Enable attention slicing"),
"diffusers_model_load_variant": OptionInfo("default", "Diffusers model loading variant", gr.Radio, lambda: {"choices": ['default', 'fp32', 'fp16']}),
"diffusers_vae_load_variant": OptionInfo("default", "Diffusers VAE loading variant", gr.Radio, lambda: {"choices": ['default', 'fp32', 'fp16']}),
@@ -422,7 +422,7 @@ options_templates.update(options_section(('system-paths', "System Paths"), {
"ckpt_dir": OptionInfo(os.path.join(paths.models_path, 'Stable-diffusion'), "Path to directory with stable diffusion checkpoints"),
"diffusers_dir": OptionInfo(os.path.join(paths.models_path, 'Diffusers'), "Path to directory with stable diffusion diffusers"),
"vae_dir": OptionInfo(os.path.join(paths.models_path, 'VAE'), "Path to directory with VAE files"),
"lora_dir": OptionInfo(os.path.join(paths.models_path, 'Lora'), "Path to directory with Lora network(s)"),
"lora_dir": OptionInfo(os.path.join(paths.models_path, 'Lora'), "Path to directory with LoRA network(s)"),
"lyco_dir": OptionInfo(os.path.join(paths.models_path, 'LyCORIS'), "Path to directory with LyCORIS network(s)"),
"styles_dir": OptionInfo(os.path.join(paths.data_path, 'styles.csv'), "Path to user-defined styles file"),
"embeddings_dir": OptionInfo(os.path.join(paths.models_path, 'embeddings'), "Embeddings directory for textual inversion"),
@@ -626,9 +626,9 @@ options_templates.update(options_section(('extra_networks', "Extra Networks"), {
"extra_networks_card_square": OptionInfo(True, "UI disable variable aspect ratio"),
"extra_networks_card_fit": OptionInfo("cover", "UI image contain method", gr.Radio, lambda: {"choices": ["contain", "cover", "fill"]}),
"extra_network_skip_indexing": OptionInfo(False, "Do not automatically build extra network pages", gr.Checkbox),
"lyco_patch_lora": OptionInfo(False, "Use LyCoris handler for all Lora types", gr.Checkbox),
"lyco_patch_lora": OptionInfo(False, "Use LyCoris handler for all LoRA types", gr.Checkbox),
# "lora_disable": OptionInfo(False, "Disable built-in Lora handler", gr.Checkbox, { "visible": True }, onchange=disable_extensions),
"lora_functional": OptionInfo(False, "Use Kohya method for handling multiple Loras", gr.Checkbox),
"lora_functional": OptionInfo(False, "Use Kohya method for handling multiple LoRA", gr.Checkbox),
"extra_networks_add_text_separator": OptionInfo(" ", "Extra text to add before <...> when adding extra network to prompt", gr.Text, { "visible": False }),
"extra_networks_default_multiplier": OptionInfo(1.0, "Multiplier for extra networks", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
"sd_hypernetwork": OptionInfo("None", "Add hypernetwork to prompt", gr.Dropdown, lambda: {"choices": ["None"] + list(hypernetworks.keys())}, refresh=reload_hypernetworks),
@@ -708,10 +708,11 @@ class Options:
diff = {}
for k, v in self.data.items():
if k in self.data_labels:
if type(v) is list:
diff[k] = v
if self.data_labels[k].default != v:
diff[k] = v
output = json.dumps(diff, indent=2)
writefile(output, filename)
writefile(diff, filename)
except Exception as e:
log.error(f'Saving settings failed: {filename} {e}')
+4 -4
View File
@@ -92,8 +92,8 @@ def calc_resolution_hires(enable, width, height, hr_scale, hr_resize_x, hr_resiz
from modules import processing, devices
if not enable:
return ""
if modules.shared.backend == modules.shared.Backend.DIFFUSERS:
return "Hires resize: disabled"
# if modules.shared.backend == modules.shared.Backend.DIFFUSERS:
# return "Hires resize: disabled"
p = processing.StableDiffusionProcessingTxt2Img(width=width, height=height, enable_hr=True, hr_scale=hr_scale, hr_resize_x=hr_resize_x, hr_resize_y=hr_resize_y)
p.init_hr()
with devices.autocast():
@@ -106,8 +106,8 @@ def resize_from_to_html(width, height, scale_by):
target_height = int(height * scale_by)
if not target_width or not target_height:
return "no image selected"
if modules.shared.backend == modules.shared.Backend.DIFFUSERS:
return "Hires resize: disabled"
# if modules.shared.backend == modules.shared.Backend.DIFFUSERS:
# return "Hires resize: disabled"
return f"Hires resize: from <span class='resolution'>{width}x{height}</span> to <span class='resolution'>{target_width}x{target_height}</span>"