From 79e6c51c068745d5f69bc74fd7c1c084632760e9 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sat, 23 Dec 2023 10:02:34 -0500 Subject: [PATCH] fix control img2img --- modules/control/run.py | 7 +- modules/processing.py | 39 +++++---- modules/processing_diffusers.py | 113 +----------------------- modules/processing_vae.py | 147 ++++++++++++++++++++++++++++++++ modules/ui.py | 19 +++-- wiki | 2 +- 6 files changed, 189 insertions(+), 138 deletions(-) create mode 100644 modules/processing_vae.py diff --git a/modules/control/run.py b/modules/control/run.py index 40868d9a4..e2dad2ea6 100644 --- a/modules/control/run.py +++ b/modules/control/run.py @@ -118,6 +118,8 @@ def control_run(units: List[unit.Unit], inputs, inits, unit_type: str, is_genera n_iter = batch_count, batch_size = batch_size, ) + processing.process_init(p) + if resize_mode != 0 or inputs is None or inputs == [None]: p.width = width # pylint: disable=attribute-defined-outside-init p.height = height # pylint: disable=attribute-defined-outside-init @@ -128,7 +130,6 @@ def control_run(units: List[unit.Unit], inputs, inits, unit_type: str, is_genera del p.width del p.height - t0 = time.time() for u in units: if not u.enabled or u.type != unit_type: @@ -364,9 +365,9 @@ def control_run(units: List[unit.Unit], inputs, inits, unit_type: str, is_genera processed_image = Image.fromarray(processed_image) if unit_type == 'controlnet' and input_type == 1: # Init image same as control + p.task_args['image'] = input_image p.task_args['control_image'] = p.image p.task_args['strength'] = p.denoising_strength - p.task_args['image'] = input_image elif unit_type == 'controlnet' and input_type == 2: # Separate init image p.task_args['control_image'] = p.image p.task_args['strength'] = p.denoising_strength @@ -394,8 +395,10 @@ def control_run(units: List[unit.Unit], inputs, inits, unit_type: str, is_genera else: pipe = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE) elif unit_type == 'reference': + p.is_control = True pipe = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE) else: # actual control + p.is_control = True if 'control_image' in p.task_args: pipe = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE) # only controlnet supports img2img else: diff --git a/modules/processing.py b/modules/processing.py index 7710117fe..b2ea8c175 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -179,6 +179,7 @@ class StableDiffusionProcessing: self.all_subseeds = None self.clip_skip = clip_skip self.iteration = 0 + self.is_control = False self.is_hr_pass = False self.is_refiner_pass = False self.hr_force = False @@ -797,20 +798,9 @@ def validate_sample(tensor): return cast -def process_images_inner(p: StableDiffusionProcessing) -> Processed: - """this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch""" - - if type(p.prompt) == list: - assert len(p.prompt) > 0 - else: - assert p.prompt is not None - +def process_init(p: StableDiffusionProcessing): seed = get_fixed_seed(p.seed) subseed = get_fixed_seed(p.subseed) - if shared.backend == shared.Backend.ORIGINAL: - modules.sd_hijack.model_hijack.apply_circular(p.tiling) - modules.sd_hijack.model_hijack.clear_comments() - comments = {} if type(p.prompt) == list: p.all_prompts = [shared.prompt_styles.apply_styles_to_prompt(x, p.styles) for x in p.prompt] else: @@ -827,15 +817,32 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: p.all_subseeds = subseed else: p.all_subseeds = [int(subseed) + x for x in range(len(p.all_prompts))] - if os.path.exists(shared.opts.embeddings_dir) and not p.do_not_reload_embeddings and shared.backend == shared.Backend.ORIGINAL: - modules.sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings(force_reload=False) - if p.scripts is not None and isinstance(p.scripts, modules.scripts.ScriptRunner): - p.scripts.process(p) + + +def process_images_inner(p: StableDiffusionProcessing) -> Processed: + """this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch""" + + if type(p.prompt) == list: + assert len(p.prompt) > 0 + else: + assert p.prompt is not None + + if shared.backend == shared.Backend.ORIGINAL: + modules.sd_hijack.model_hijack.apply_circular(p.tiling) + modules.sd_hijack.model_hijack.clear_comments() + comments = {} infotexts = [] output_images = [] cached_uc = [None, None] cached_c = [None, None] + process_init(p) + if os.path.exists(shared.opts.embeddings_dir) and not p.do_not_reload_embeddings and shared.backend == shared.Backend.ORIGINAL: + modules.sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings(force_reload=False) + if p.scripts is not None and isinstance(p.scripts, modules.scripts.ScriptRunner): + p.scripts.process(p) + + def get_conds_with_caching(function, required_prompts, steps, cache): if cache[0] is not None and (required_prompts, steps) == cache[0]: return cache[1] diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 2415d2cde..c97c81c8b 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -9,14 +9,13 @@ import modules.devices as devices import modules.shared as shared import modules.sd_samplers as sd_samplers import modules.sd_models as sd_models -import modules.sd_vae as sd_vae -import modules.taesd.sd_vae_taesd as sd_vae_taesd import modules.images as images import modules.errors as errors from modules.processing import StableDiffusionProcessing, create_random_tensors import modules.prompt_parser_diffusers as prompt_parser_diffusers from modules.sd_hijack_hypertile import hypertile_set from modules.processing_correction import correction_callback +from modules.processing_vae import vae_encode, vae_decode debug = shared.log.trace if os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -115,113 +114,6 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro shared.profiler.step() return kwargs - def full_vae_decode(latents, model): - t0 = time.time() - if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): - 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 and hasattr(model, 'vae'): - model.vae.to(devices.device) - latents.to(model.vae.device) - - upcast = (model.vae.dtype == torch.float16) and getattr(model.vae.config, 'force_upcast', False) and hasattr(model, 'upcast_vae') - if upcast: # this is done by diffusers automatically if output_type != 'latent' - model.upcast_vae() - latents = latents.to(next(iter(model.vae.post_quant_conv.parameters())).dtype) - - decoded = model.vae.decode(latents / model.vae.config.scaling_factor, return_dict=False)[0] - # Downcast VAE after OpenVINO compile - if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx" and shared.compiled_model_state.first_pass_vae: - shared.compiled_model_state.first_pass_vae = False - if hasattr(shared.sd_model, "vae"): - shared.sd_model.vae.to(dtype=torch.float8_e4m3fn) - devices.torch_gc(force=True) - if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): - model.unet.to(unet_device) - t1 = time.time() - 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={upcast} images={latents.shape[0]} latents={latents.shape} time={round(t1-t0, 3)}') - 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 getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): - 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 and hasattr(model, 'vae'): - model.vae.to(devices.device) - encoded = model.vae.encode(image.to(model.vae.device, model.vae.dtype)).latent_dist.sample() - if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): - model.unet.to(unet_device) - return encoded - - def taesd_vae_decode(latents): - 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(latents.shape[0]): - decoded[i] = sd_vae_taesd.decode(latents[i]) - 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): - t0 = time.time() - prev_job = shared.state.job - shared.state.job = 'vae' - 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 latents.shape[0] == 4 and latents.shape[1] != 4: # likely animatediff latent - latents = latents.permute(1, 0, 2, 3) - 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: - decoded = taesd_vae_decode(latents=latents) - # TODO validate decoded sample diffusers - # decoded = validate_sample(decoded) - if hasattr(model, 'image_processor'): - imgs = model.image_processor.postprocess(decoded, output_type=output_type) - else: - import diffusers - image_processor = diffusers.image_processor.VaeImageProcessor() - imgs = image_processor.postprocess(decoded, output_type=output_type) - shared.state.job = prev_job - if shared.cmd_opts.profile: - t1 = time.time() - shared.log.debug(f'Profile: VAE decode: {t1-t0:.2f}') - 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: - tensor = tensor * 2 - 1 - 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] @@ -480,7 +372,8 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro shared.sd_model.to(devices.device) # pipeline type is set earlier in processing, but check for sanity - if sd_models.get_diffusers_task(shared.sd_model) != sd_models.DiffusersTaskType.TEXT_2_IMAGE and len(getattr(p, 'init_images' ,[])) == 0: + has_images = len(getattr(p, 'init_images' ,[])) > 0 or getattr(p, 'is_control', False) is True + if sd_models.get_diffusers_task(shared.sd_model) != sd_models.DiffusersTaskType.TEXT_2_IMAGE and not has_images: shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE) # reset pipeline if hasattr(shared.sd_model, 'unet') and hasattr(shared.sd_model.unet, 'config') and hasattr(shared.sd_model.unet.config, 'in_channels') and shared.sd_model.unet.config.in_channels == 9: shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.INPAINTING) # force pipeline diff --git a/modules/processing_vae.py b/modules/processing_vae.py new file mode 100644 index 000000000..5da449469 --- /dev/null +++ b/modules/processing_vae.py @@ -0,0 +1,147 @@ +import os +import time +import torch +import torchvision.transforms.functional as TF +from modules import shared, devices, sd_vae +import modules.taesd.sd_vae_taesd as sd_vae_taesd + + +debug = shared.log.trace if os.environ.get('SD_VAE_DEBUG', None) is not None else lambda *args, **kwargs: None +debug('Trace: VAE') + + +def create_latents(image, p, dtype=None, device=None): + from modules.processing import create_random_tensors + from PIL import Image + if image is None: + return image + elif isinstance(image, Image.Image): + latents = vae_encode(image, model=shared.sd_model, full_quality=p.full_quality) + elif isinstance(image, list): + latents = [vae_encode(i, model=shared.sd_model, full_quality=p.full_quality).squeeze(dim=0) for i in image] + latents = torch.stack(latents, dim=0).to(shared.device) + else: + shared.log.warning(f'Latents: input type: {type(image)} {image}') + return image + noise = p.denoising_strength * create_random_tensors(latents.shape[1:], seeds=p.all_seeds, subseeds=p.all_subseeds, subseed_strength=p.subseed_strength, p=p) + latents = (1 - p.denoising_strength) * latents + noise + if dtype is not None: + latents = latents.to(dtype=dtype) + if device is not None: + latents = latents.to(device=device) + return latents + + +def full_vae_decode(latents, model): + t0 = time.time() + if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): + 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 and hasattr(model, 'vae'): + model.vae.to(devices.device) + latents.to(model.vae.device) + + upcast = (model.vae.dtype == torch.float16) and getattr(model.vae.config, 'force_upcast', False) and hasattr(model, 'upcast_vae') + if upcast: # this is done by diffusers automatically if output_type != 'latent' + model.upcast_vae() + latents = latents.to(next(iter(model.vae.post_quant_conv.parameters())).dtype) + + decoded = model.vae.decode(latents / model.vae.config.scaling_factor, return_dict=False)[0] + + # Downcast VAE after OpenVINO compile + if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx" and shared.compiled_model_state.first_pass_vae: + shared.compiled_model_state.first_pass_vae = False + if hasattr(shared.sd_model, "vae"): + shared.sd_model.vae.to(dtype=torch.float8_e4m3fn) + devices.torch_gc(force=True) + + if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): + model.unet.to(unet_device) + t1 = time.time() + 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={upcast} images={latents.shape[0]} latents={latents.shape} time={round(t1-t0, 3)}') + return decoded + + +def full_vae_encode(image, model): + 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 getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): + 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 and hasattr(model, 'vae'): + model.vae.to(devices.device) + encoded = model.vae.encode(image.to(model.vae.device, model.vae.dtype)).latent_dist.sample() + if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): + model.unet.to(unet_device) + return encoded + + +def taesd_vae_decode(latents): + 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(latents.shape[0]): + decoded[i] = sd_vae_taesd.decode(latents[i]) + return decoded + + +def taesd_vae_encode(image): + 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): + t0 = time.time() + prev_job = shared.state.job + shared.state.job = 'vae' + 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 latents.shape[0] == 4 and latents.shape[1] != 4: # likely animatediff latent + latents = latents.permute(1, 0, 2, 3) + 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: + decoded = taesd_vae_decode(latents=latents) + # TODO validate decoded sample diffusers + # decoded = validate_sample(decoded) + if hasattr(model, 'image_processor'): + imgs = model.image_processor.postprocess(decoded, output_type=output_type) + else: + import diffusers + image_processor = diffusers.image_processor.VaeImageProcessor() + imgs = image_processor.postprocess(decoded, output_type=output_type) + shared.state.job = prev_job + if shared.cmd_opts.profile: + t1 = time.time() + shared.log.debug(f'Profile: VAE decode: {t1-t0:.2f}') + 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: + tensor = tensor * 2 - 1 + latents = full_vae_encode(image=tensor, model=shared.sd_model) + else: + latents = taesd_vae_encode(image=tensor) + return latents diff --git a/modules/ui.py b/modules/ui.py index a47556115..87b1e0553 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -908,6 +908,8 @@ def create_ui(startup_timer = None): from modules import ui_control ui_control.create_ui() timer.startup.record("ui-control") + else: + control_interface = None with gr.Blocks(analytics_enabled=False) as extras_interface: from modules import ui_postprocessing @@ -1129,15 +1131,14 @@ def create_ui(startup_timer = None): timer.startup.record("ui-settings") - interfaces = [ - (txt2img_interface, "Text", "txt2img"), - (img2img_interface, "Image", "img2img"), - (control_interface, "Control", "control"), - (extras_interface, "Process", "process"), - (train_interface, "Train", "train"), - (models_interface, "Models", "models"), - (interrogate_interface, "Interrogate", "interrogate"), - ] + interfaces = [] + interfaces += [(txt2img_interface, "Text", "txt2img")] + interfaces += [(img2img_interface, "Image", "img2img")] + interfaces += [(control_interface, "Control", "control")] if control_interface is not None else [] + interfaces += [(extras_interface, "Process", "process")] + interfaces += [(train_interface, "Train", "train")] + interfaces += [(models_interface, "Models", "models")] + interfaces += [(interrogate_interface, "Interrogate", "interrogate")] interfaces += script_callbacks.ui_tabs_callback() interfaces += [(settings_interface, "System", "system")] diff --git a/wiki b/wiki index 9994450f3..62c496b6c 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 9994450f36413a1a72500c397ff39e024aab13a1 +Subproject commit 62c496b6cf936eec933d58d912c47372212c1281