Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-20 01:31:13 +02:00
Code Issues Packages Projects Releases Wiki Activity

add new hires with refiner and non-latent modes

Browse Source
This commit is contained in:
Vladimir Mandic
2023-09-12 11:54:07 -04:00
parent 69b36532ad
commit 9cf7fc4a75
12 changed files with 113 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/images.py
+3 -1
View File
@@ -201,7 +201,7 @@ def draw_prompt_matrix(im, width, height, all_prompts, margin=0):
return draw_grid_annotations(im, width, height, hor_texts, ver_texts, margin)
def resize_image(resize_mode, im, width, height, upscaler_name=None):
def resize_image(resize_mode, im, width, height, upscaler_name=None, output_type='image'):
"""
Resizes an image with the specified resize_mode, width, and height.
Args:
@@ -261,6 +261,8 @@ def resize_image(resize_mode, im, width, height, upscaler_name=None):
fill_width = width // 2 - src_w // 2
res.paste(resized.resize((fill_width, height), box=(0, 0, 0, height)), box=(0, 0))
res.paste(resized.resize((fill_width, height), box=(resized.width, 0, resized.width, height)), box=(fill_width + src_w, 0))
if output_type == 'np':
return np.array(res)
return res
modules/processing.py
+9 -4
View File
@@ -488,7 +488,7 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts, all_seeds, all_su
"Backend": 'Diffusers' if shared.backend == shared.Backend.DIFFUSERS else 'Original',
"Version": git_commit,
"Comment": comment,
"Operations": ', '.join(list(set(p.ops))).replace('"', '') if len(p.ops) > 0 else None,
"Operations": '; '.join(p.ops).replace('"', '') if len(p.ops) > 0 else 'none',
}
if 'txt2img' in p.ops:
pass
@@ -800,8 +800,13 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
for i, x_sample in enumerate(x_samples_ddim):
p.batch_index = i
x_sample = 255. * (np.moveaxis(x_sample.cpu().numpy(), 0, 2) if shared.backend == shared.Backend.ORIGINAL else x_sample)
x_sample = validate_sample(x_sample)
if type(x_sample) == Image.Image:
image = x_sample
x_sample = np.array(x_sample)
else:
x_sample = 255. * (np.moveaxis(x_sample.cpu().numpy(), 0, 2) if shared.backend == shared.Backend.ORIGINAL else x_sample)
x_sample = validate_sample(x_sample)
image = Image.fromarray(x_sample)
if p.restore_faces:
if shared.opts.save and not p.do_not_save_samples and shared.opts.save_images_before_face_restoration:
orig = p.restore_faces
@@ -811,7 +816,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
images.save_image(Image.fromarray(x_sample), path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-face-restoration")
p.ops.append('face')
x_sample = modules.face_restoration.restore_faces(x_sample)
image = Image.fromarray(x_sample)
image = Image.fromarray(x_sample)
if p.scripts is not None:
pp = modules.scripts.PostprocessImageArgs(image)
p.scripts.postprocess_image(p, pp)
modules/processing_diffusers.py
+48 -9
View File
@@ -2,6 +2,7 @@ import time
import inspect
import typing
import torch
import torchvision.transforms.functional as TF
import modules.devices as devices
import modules.shared as shared
import modules.sd_samplers as sd_samplers
@@ -31,14 +32,17 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
shared.log.info(f'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')
first_pass_images = vae_decode(latents=latents, model=shared.sd_model, full_quality=p.full_quality, output_type='pil')
p.init_images = []
for first_pass_image in first_pass_images:
if latent_upscaler is None:
init_image = images.resize_image(1, first_pass_image, p.hr_upscale_to_x, p.hr_upscale_to_y, upscaler_name=p.hr_upscaler)
else:
init_image = first_pass_image
# if is_refiner_enabled:
# init_image = vae_encode(init_image, model=shared.sd_model, full_quality=p.full_quality)
p.init_images.append(init_image)
return p.init_images
def save_intermediate(latents, suffix):
for i in range(len(latents)):
@@ -64,7 +68,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
time.sleep(0.1)
def full_vae_decode(latents, model):
shared.log.debug(f'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]}')
shared.log.debug(f'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]} latents={latents.shape}')
if shared.opts.diffusers_move_unet and not model.has_accelerate:
shared.log.debug('Moving to CPU: model=UNet')
unet_device = model.unet.device
@@ -78,13 +82,34 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
model.unet.to(unet_device)
return decoded
def full_vae_encode(image, model):
shared.log.debug(f'VAE encode: 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)}')
if shared.opts.diffusers_move_unet and not model.has_accelerate:
shared.log.debug('Moving to CPU: model=UNet')
unet_device = model.unet.device
model.unet.to(devices.cpu)
devices.torch_gc()
if not shared.cmd_opts.lowvram and not shared.opts.diffusers_seq_cpu_offload:
model.vae.to(devices.device)
encoded = model.vae.encode(image.to(model.vae.device, model.vae.dtype))
if shared.opts.diffusers_move_unet and not model.has_accelerate:
model.unet.to(unet_device)
return encoded
def taesd_vae_decode(latents):
shared.log.debug(f'VAE decode: name=TAESD images={latents.shape[0]}')
decoded = torch.zeros((len(latents), 3, p.height, p.width), dtype=devices.dtype_vae, device=devices.device)
shared.log.debug(f'VAE decode: name=TAESD images={len(latents)} latents={latents.shape}')
if len(latents) == 0:
return []
decoded = torch.zeros((len(latents), 3, latents.shape[2] * 8, latents.shape[3] * 8), dtype=devices.dtype_vae, device=devices.device)
for i in range(len(output.images)):
decoded[i] = (sd_vae_taesd.decode(latents[i]) * 2.0) - 1.0
return decoded
def taesd_vae_encode(image):
shared.log.debug(f'VAE encode: name=TAESD image={image.shape}')
encoded = sd_vae_taesd.encode(image)
return encoded
def vae_decode(latents, model, output_type='np', full_quality=True):
if not torch.is_tensor(latents): # already decoded
return latents
@@ -105,6 +130,19 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
imgs = model.image_processor.postprocess(decoded, output_type=output_type)
return imgs
def vae_encode(image, model, full_quality=True): # pylint: disable=unused-variable
if shared.state.interrupted or shared.state.skipped:
return []
if not hasattr(model, 'vae'):
shared.log.error('VAE not found in model')
return []
tensor = TF.to_tensor(image.convert("RGB")).unsqueeze(0).to(devices.device, devices.dtype_vae)
if full_quality:
latents = full_vae_encode(image=tensor, model=shared.sd_model)
else:
latents = taesd_vae_encode(image=tensor)
return latents
def fix_prompts(prompts, negative_prompts, prompts_2, negative_prompts_2):
if type(prompts) is str:
prompts = [prompts]
@@ -323,8 +361,9 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
p.ops.append('upscale')
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)
output.images = hires_resize(latents=output.images)
if latent_scale_mode is not None or p.hr_force:
p.ops.append('hires')
recompile_model(hires=True)
sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE)
hires_args = set_pipeline_args(
@@ -372,10 +411,10 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
p.ops.append('refine')
for i in range(len(output.images)):
image = output.images[i]
if (image.shape[2] == 3) and (image.shape[0] % 8 != 0 or image.shape[1] % 8 != 0):
shared.log.warning(f'Refiner requires image size to be divisible by 8: {image.shape}')
results.append(image)
return results
# if (image.shape[2] == 3) and (image.shape[0] % 8 != 0 or image.shape[1] % 8 != 0):
# shared.log.warning(f'Refiner requires image size to be divisible by 8: {image.shape}')
# results.append(image)
# return results
refiner_args = set_pipeline_args(
model=shared.sd_refiner,
prompts=[p.refiner_prompt] if len(p.refiner_prompt) > 0 else prompts[i],
modules/taesd/sd_vae_taesd.py
+19
View File
@@ -61,3 +61,22 @@ def decode(latents):
enc = latents.unsqueeze(0).to(devices.device, devices.dtype_vae)
image = vae.decoder(enc).clamp(0, 1).detach()
return image[0]
def encode(image):
from modules import shared
model_class = shared.sd_model_type
if model_class == 'ldm':
model_class = 'sd'
if 'sd' not in model_class:
shared.log.warning(f'TAESD unsupported model type: {model_class}')
return Image.new('RGB', (8, 8), color = (0, 0, 0))
vae = taesd_models[f'{model_class}-encoder']
if vae is None:
model_path = os.path.join(paths_internal.models_path, "TAESD", f"tae{model_class}_encoder.pth")
download_model(model_path)
if os.path.exists(model_path):
taesd_models[f'{model_class}-encoder'] = TAESD(encoder_path=model_path, decoder_path=None)
vae = taesd_models[f'{model_class}-encoder']
vae.to(devices.device, devices.dtype_vae)
latents = vae.encoder(image).detach()
return latents
modules/ui.py
+4 -3
View File
@@ -466,13 +466,14 @@ def create_ui(startup_timer = None):
txt2img_prompt.submit(**txt2img_args)
submit.click(**txt2img_args)
def enable_hr_change(visible: bool):
return {"visible": visible, "__type__": "update"}, f'Refiner: {"disabled" if modules.shared.opts.sd_model_refiner == "None" else "enabled"}'
def enable_hr_change(visible: bool, refiner_start):
enabled = modules.shared.opts.sd_model_refiner != "None" and refiner_start > 0 and refiner_start < 1
return {"visible": visible, "__type__": "update"}, f'Refiner: {"enabled" if enabled else "disabled"}'
res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height], show_progress=False)
batch_switch_btn.click(lambda w, h: (h, w), inputs=[batch_count, batch_size], outputs=[batch_count, batch_size], show_progress=False)
txt_prompt_img.change(fn=modules.images.image_data, inputs=[txt_prompt_img], outputs=[txt2img_prompt, txt_prompt_img])
show_second_pass.change(enable_hr_change, inputs=[show_second_pass], outputs=[second_pass_group, hr_refiner], show_progress = False)
show_second_pass.change(enable_hr_change, inputs=[show_second_pass, refiner_start], outputs=[second_pass_group, hr_refiner], show_progress = False)
show_seed.change(gr_show, inputs=[show_seed], outputs=[seed_group], show_progress = False)
show_batch.change(gr_show, inputs=[show_batch], outputs=[batch_group], show_progress = False)
show_advanced.change(gr_show, inputs=[show_advanced], outputs=[advanced_group], show_progress = False)