[2.1.822] New Inpaint System

See related documents for more details.
This commit is contained in:
lllyasviel
2023-11-19 17:37:22 -08:00
committed by GitHub
parent 8f9f020e8f
commit dececbd060
19 changed files with 400 additions and 191 deletions
+81 -67
View File
@@ -194,8 +194,10 @@ def prepare_text_encoder(async_call=True):
@torch.no_grad()
@torch.inference_mode()
def refresh_everything(refiner_model_name, base_model_name, loras, base_model_additional_loras=None):
global final_unet, final_clip, final_vae, final_refiner_unet, final_refiner_vae, final_expansion
def refresh_everything(refiner_model_name, base_model_name, loras,
base_model_additional_loras=None, use_synthetic_refiner=False):
global final_unet, final_clip, final_vae, final_refiner_unet, final_refiner_vae, \
final_expansion, model_refiner, model_base
final_unet = None
final_clip = None
@@ -203,8 +205,23 @@ def refresh_everything(refiner_model_name, base_model_name, loras, base_model_ad
final_refiner_unet = None
final_refiner_vae = None
refresh_refiner_model(refiner_model_name)
refresh_base_model(base_model_name)
if use_synthetic_refiner and refiner_model_name == 'None':
print('Synthetic Refiner Activated')
refresh_base_model(base_model_name)
model_refiner = core.StableDiffusionModel(
unet=model_base.unet,
vae=model_base.vae,
clip=model_base.clip,
clip_vision=model_base.clip_vision,
filename=model_base.filename
)
model_refiner.vae = None
model_refiner.clip = None
model_refiner.clip_vision = None
else:
refresh_refiner_model(refiner_model_name)
refresh_base_model(base_model_name)
refresh_loras(loras, base_model_additional_loras=base_model_additional_loras)
assert_model_integrity()
@@ -212,14 +229,9 @@ def refresh_everything(refiner_model_name, base_model_name, loras, base_model_ad
final_clip = model_base.clip_with_lora
final_vae = model_base.vae
final_unet.model.diffusion_model.in_inpaint = False
final_refiner_unet = model_refiner.unet_with_lora
final_refiner_vae = model_refiner.vae
if final_refiner_unet is not None:
final_refiner_unet.model.diffusion_model.in_inpaint = False
if final_expansion is None:
final_expansion = FooocusExpansion()
@@ -276,32 +288,52 @@ def calculate_sigmas(sampler, model, scheduler, steps, denoise):
@torch.no_grad()
@torch.inference_mode()
def process_diffusion(positive_cond, negative_cond, steps, switch, width, height, image_seed, callback, sampler_name, scheduler_name, latent=None, denoise=1.0, tiled=False, cfg_scale=7.0, refiner_swap_method='joint'):
global final_unet, final_refiner_unet, final_vae, final_refiner_vae
def get_candidate_vae(steps, switch, denoise=1.0, refiner_swap_method='joint'):
assert refiner_swap_method in ['joint', 'separate', 'vae']
assert refiner_swap_method in ['joint', 'separate', 'vae', 'upscale']
refiner_use_different_vae = final_refiner_vae is not None and final_refiner_unet is not None
if refiner_swap_method == 'upscale':
if not refiner_use_different_vae:
refiner_swap_method = 'joint'
else:
if refiner_use_different_vae:
if denoise > 0.95:
refiner_swap_method = 'vae'
if final_refiner_vae is not None and final_refiner_unet is not None:
if denoise > 0.9:
return final_vae, final_refiner_vae
else:
if denoise > (float(steps - switch) / float(steps)) ** 0.834: # karras 0.834
return final_vae, None
else:
# VAE swap only support full denoise
# Disable refiner to avoid SD15 in joint/separate swap
final_refiner_unet = None
final_refiner_vae = None
return final_refiner_vae, None
return final_vae, final_refiner_vae
@torch.no_grad()
@torch.inference_mode()
def process_diffusion(positive_cond, negative_cond, steps, switch, width, height, image_seed, callback, sampler_name, scheduler_name, latent=None, denoise=1.0, tiled=False, cfg_scale=7.0, refiner_swap_method='joint'):
target_unet, target_vae, target_refiner_unet, target_refiner_vae, target_clip \
= final_unet, final_vae, final_refiner_unet, final_refiner_vae, final_clip
assert refiner_swap_method in ['joint', 'separate', 'vae']
if final_refiner_vae is not None and final_refiner_unet is not None:
# Refiner Use Different VAE (then it is SD15)
if denoise > 0.9:
refiner_swap_method = 'vae'
else:
refiner_swap_method = 'joint'
if denoise > (float(steps - switch) / float(steps)) ** 0.834: # karras 0.834
target_unet, target_vae, target_refiner_unet, target_refiner_vae \
= final_unet, final_vae, None, None
print(f'[Sampler] only use Base because of partial denoise.')
else:
positive_cond = clip_separate(positive_cond, target_model=final_refiner_unet.model, target_clip=final_clip)
negative_cond = clip_separate(negative_cond, target_model=final_refiner_unet.model, target_clip=final_clip)
target_unet, target_vae, target_refiner_unet, target_refiner_vae \
= final_refiner_unet, final_refiner_vae, None, None
print(f'[Sampler] only use Refiner because of partial denoise.')
print(f'[Sampler] refiner_swap_method = {refiner_swap_method}')
if latent is None:
empty_latent = core.generate_empty_latent(width=width, height=height, batch_size=1)
initial_latent = core.generate_empty_latent(width=width, height=height, batch_size=1)
else:
empty_latent = latent
initial_latent = latent
minmax_sigmas = calculate_sigmas(sampler=sampler_name, scheduler=scheduler_name, model=final_unet.model, steps=steps, denoise=denoise)
sigma_min, sigma_max = minmax_sigmas[minmax_sigmas > 0].min(), minmax_sigmas.max()
@@ -310,18 +342,18 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
print(f'[Sampler] sigma_min = {sigma_min}, sigma_max = {sigma_max}')
modules.patch.BrownianTreeNoiseSamplerPatched.global_init(
empty_latent['samples'].to(fcbh.model_management.get_torch_device()),
initial_latent['samples'].to(fcbh.model_management.get_torch_device()),
sigma_min, sigma_max, seed=image_seed, cpu=False)
decoded_latent = None
if refiner_swap_method == 'joint':
sampled_latent = core.ksampler(
model=final_unet,
refiner=final_refiner_unet,
model=target_unet,
refiner=target_refiner_unet,
positive=positive_cond,
negative=negative_cond,
latent=empty_latent,
latent=initial_latent,
steps=steps, start_step=0, last_step=steps, disable_noise=False, force_full_denoise=True,
seed=image_seed,
denoise=denoise,
@@ -333,32 +365,14 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
previewer_start=0,
previewer_end=steps,
)
decoded_latent = core.decode_vae(vae=final_vae, latent_image=sampled_latent, tiled=tiled)
if refiner_swap_method == 'upscale':
sampled_latent = core.ksampler(
model=final_refiner_unet,
positive=clip_separate(positive_cond, target_model=final_refiner_unet.model, target_clip=final_clip),
negative=clip_separate(negative_cond, target_model=final_refiner_unet.model, target_clip=final_clip),
latent=empty_latent,
steps=steps, start_step=0, last_step=steps, disable_noise=False, force_full_denoise=True,
seed=image_seed,
denoise=denoise,
callback_function=callback,
cfg=cfg_scale,
sampler_name=sampler_name,
scheduler=scheduler_name,
previewer_start=0,
previewer_end=steps,
)
decoded_latent = core.decode_vae(vae=final_refiner_vae, latent_image=sampled_latent, tiled=tiled)
decoded_latent = core.decode_vae(vae=target_vae, latent_image=sampled_latent, tiled=tiled)
if refiner_swap_method == 'separate':
sampled_latent = core.ksampler(
model=final_unet,
model=target_unet,
positive=positive_cond,
negative=negative_cond,
latent=empty_latent,
latent=initial_latent,
steps=steps, start_step=0, last_step=switch, disable_noise=False, force_full_denoise=False,
seed=image_seed,
denoise=denoise,
@@ -371,15 +385,15 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
)
print('Refiner swapped by changing ksampler. Noise preserved.')
target_model = final_refiner_unet
target_model = target_refiner_unet
if target_model is None:
target_model = final_unet
target_model = target_unet
print('Use base model to refine itself - this may because of developer mode.')
sampled_latent = core.ksampler(
model=target_model,
positive=clip_separate(positive_cond, target_model=target_model.model, target_clip=final_clip),
negative=clip_separate(negative_cond, target_model=target_model.model, target_clip=final_clip),
positive=clip_separate(positive_cond, target_model=target_model.model, target_clip=target_clip),
negative=clip_separate(negative_cond, target_model=target_model.model, target_clip=target_clip),
latent=sampled_latent,
steps=steps, start_step=switch, last_step=steps, disable_noise=True, force_full_denoise=True,
seed=image_seed,
@@ -392,9 +406,9 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
previewer_end=steps,
)
target_model = final_refiner_vae
target_model = target_refiner_vae
if target_model is None:
target_model = final_vae
target_model = target_vae
decoded_latent = core.decode_vae(vae=target_model, latent_image=sampled_latent, tiled=tiled)
if refiner_swap_method == 'vae':
@@ -404,10 +418,10 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
modules.inpaint_worker.current_task.unswap()
sampled_latent = core.ksampler(
model=final_unet,
model=target_unet,
positive=positive_cond,
negative=negative_cond,
latent=empty_latent,
latent=initial_latent,
steps=steps, start_step=0, last_step=switch, disable_noise=False, force_full_denoise=True,
seed=image_seed,
denoise=denoise,
@@ -420,9 +434,9 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
)
print('Fooocus VAE-based swap.')
target_model = final_refiner_unet
target_model = target_refiner_unet
if target_model is None:
target_model = final_unet
target_model = target_unet
print('Use base model to refine itself - this may because of developer mode.')
sampled_latent = vae_parse(sampled_latent)
@@ -442,8 +456,8 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
sampled_latent = core.ksampler(
model=target_model,
positive=clip_separate(positive_cond, target_model=target_model.model, target_clip=final_clip),
negative=clip_separate(negative_cond, target_model=target_model.model, target_clip=final_clip),
positive=clip_separate(positive_cond, target_model=target_model.model, target_clip=target_clip),
negative=clip_separate(negative_cond, target_model=target_model.model, target_clip=target_clip),
latent=sampled_latent,
steps=len_sigmas, start_step=0, last_step=len_sigmas, disable_noise=False, force_full_denoise=True,
seed=image_seed+1,
@@ -458,9 +472,9 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
noise_mean=noise_mean
)
target_model = final_refiner_vae
target_model = target_refiner_vae
if target_model is None:
target_model = final_vae
target_model = target_vae
decoded_latent = core.decode_vae(vae=target_model, latent_image=sampled_latent, tiled=tiled)
images = core.pytorch_to_numpy(decoded_latent)