From 3e1a6a96d07933c1397dec6946e834ab35720da1 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Fri, 7 Jul 2023 09:38:16 -0400 Subject: [PATCH] add additional pipelines --- CHANGELOG.md | 4 +++- DIFFUSERS.md | 2 +- modules/processing.py | 12 +++++++++++ modules/sd_models.py | 35 +++++++++++++++++++++++++++----- modules/sd_vae.py | 39 ++++++++++++++++++++++++++++++------ modules/shared.py | 8 ++++++-- modules/ui_extra_networks.py | 14 +++++++++---- 7 files changed, 95 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c4447625..18d85957b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,8 +6,10 @@ - add settings -> extra networks -> do not automatically build extra network pages speeds up app start if you have a lot of extra networks and you want to build them manually when needed - extra network ui tweaks +- cache extra networks between tabs + this should result in neat 2x speedup on building extra networks - merge experimental diffusers support - this will be covered in details in separate post + covered in details in a separate post ## Update for 07/01/2023 diff --git a/DIFFUSERS.md b/DIFFUSERS.md index 47b86af47..df2c462fc 100644 --- a/DIFFUSERS.md +++ b/DIFFUSERS.md @@ -7,7 +7,7 @@ initial support merged into `dev` branch - download from branch and start as normal: > git clone https://github.com/vladmandic/automatic -b dev diffusers > cd diffusers - > webui --debug --backend original + > webui --debug --backend diffusers - to go back to standard execution pipeline, start with > webui --debug --backend original diff --git a/modules/processing.py b/modules/processing.py index 6bf8b5eaf..04058d519 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -705,12 +705,22 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: # 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} + def diffusers_callback(step: int, _timestep: int, latents: torch.FloatTensor): # TODO simplified callback for now + shared.state.sampling_step = step + shared.state.sampling_steps = p.steps + shared.state.current_latent = latents + shared.state.set_current_image() + if p.scripts is not None: + p.scripts.process(p) + output = shared.sd_model( # pylint: disable=not-callable prompt=prompts, negative_prompt=negative_prompts, num_inference_steps=p.steps, guidance_scale=p.cfg_scale, generator=generator, + callback_steps = 1, + callback = diffusers_callback, output_type='np' if shared.sd_refiner is None else 'latent', cross_attention_kwargs=cross_attention_kwargs, **task_specific_kwargs @@ -724,6 +734,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: num_inference_steps=p.steps, guidance_scale=p.cfg_scale, generator=generator, + callback_steps = 1, + callback = diffusers_callback, output_type='np', cross_attention_kwargs=cross_attention_kwargs, image=init_image diff --git a/modules/sd_models.py b/modules/sd_models.py index 83661836c..688aa3638 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -129,7 +129,7 @@ def list_models(): checkpoint_aliases.clear() ext_filter=[".safetensors"] if shared.opts.sd_disable_ckpt else [".ckpt", ".safetensors"] model_list = [] - if shared.backend == shared.Backend.ORIGINAL or shared.opts.diffusers_pipeline == shared.pipelines[0]: + if shared.backend == shared.Backend.ORIGINAL or shared.opts.diffusers_allow_safetensors: model_list += modelloader.load_models(model_path=model_path, model_url=None, command_path=shared.opts.ckpt_dir, ext_filter=ext_filter, download_name=None, ext_blacklist=[".vae.ckpt", ".vae.safetensors"]) if shared.backend == shared.Backend.DIFFUSERS: model_list += modelloader.load_diffusers_models(model_path=os.path.join(models_path, 'Diffusers'), command_path=shared.opts.diffusers_dir) @@ -577,7 +577,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No # "use_safetensors": True, # TODO(PVP) - we can't enable this for all checkpoints just yet } - if shared.opts.data['sd_model_checkpoint'] == 'model.ckpt': + if shared.opts.data.get('sd_model_checkpoint', '') == 'model.ckpt' or shared.opts.data.get('sd_model_checkpoint', '') == '': shared.opts.data['sd_model_checkpoint'] = "runwayml/stable-diffusion-v1-5" if op == 'model' or op == 'dict': @@ -608,6 +608,12 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No unload_model_weights(op=op) return shared.log.info(f'Loading diffuser {op}: {checkpoint_info.filename}') + + vae_file, vae_source = sd_vae.resolve_vae(checkpoint_info.filename) + vae = sd_vae.load_vae_diffusers(None, vae_file, vae_source) + if vae is not None: + diffusers_load_config["vae"] = vae + if not os.path.isfile(checkpoint_info.path): try: sd_model = diffusers.DiffusionPipeline.from_pretrained(checkpoint_info.path, **diffusers_load_config) @@ -617,7 +623,6 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No diffusers_load_config["local_files_only "] = True diffusers_load_config["extract_ema"] = shared.opts.diffusers_extract_ema try: - # pipelines = ['Stable Diffusion', 'Stable Diffusion XL', 'Kandinsky V1', 'Kandinsky V2', 'DeepFloyd IF', 'Shap-E'] if shared.opts.diffusers_pipeline == shared.pipelines[0]: pipeline = diffusers.StableDiffusionPipeline elif shared.opts.diffusers_pipeline == shared.pipelines[1]: @@ -630,13 +635,32 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No pipeline = diffusers.IFPipeline elif shared.opts.diffusers_pipeline == shared.pipelines[5]: pipeline = diffusers.ShapEPipeline + elif shared.opts.diffusers_pipeline == shared.pipelines[6]: + pipeline = diffusers.StableDiffusionImg2ImgPipeline + elif shared.opts.diffusers_pipeline == shared.pipelines[7]: + pipeline = diffusers.StableDiffusionXLImg2ImgPipeline + elif shared.opts.diffusers_pipeline == shared.pipelines[8]: + pipeline = diffusers.KandinskyImg2ImgPipeline + elif shared.opts.diffusers_pipeline == shared.pipelines[9]: + pipeline = diffusers.KandinskyV22Img2ImgPipeline + elif shared.opts.diffusers_pipeline == shared.pipelines[10]: + pipeline = diffusers.IFImg2ImgPipeline + elif shared.opts.diffusers_pipeline == shared.pipelines[11]: + pipeline = diffusers.ShapEImg2ImgPipeline else: shared.log.error(f'Diffusers unknown pipeline: {shared.opts.diffusers_pipeline}') except Exception as e: shared.log.error(f'Diffusers failed initializing pipeline: {shared.opts.diffusers_pipeline} {e}') return try: - sd_model = pipeline.from_ckpt(checkpoint_info.path, **diffusers_load_config) + if hasattr(pipeline, 'from_single_file'): + diffusers_load_config['use_safetensors'] = True + sd_model = pipeline.from_single_file(checkpoint_info.path, **diffusers_load_config) + elif hasattr(pipeline, 'from_ckpt'): + sd_model = pipeline.from_ckpt(checkpoint_info.path, **diffusers_load_config) + else: + shared.log.error(f'Diffusers cannot load safetensor model: {checkpoint_info.path} {shared.opts.diffusers_pipeline}') + return except Exception as e: shared.log.error(f'Diffusers failed loading model using pipeline: {checkpoint_info.path} {shared.opts.diffusers_pipeline} {e}') return @@ -938,13 +962,14 @@ def unload_model_weights(op='model'): if shared.backend == shared.Backend.ORIGINAL: sd_hijack.model_hijack.undo_hijack(model_data.sd_model) model_data.sd_model = None + shared.log.debug(f'Weights unloaded {op}: {memory_stats()}') else: if model_data.sd_refiner: model_data.sd_refiner.to(devices.cpu) if shared.backend == shared.Backend.ORIGINAL: sd_hijack.model_hijack.undo_hijack(model_data.sd_refiner) model_data.sd_refiner = None - shared.log.debug(f'Weights unloaded {op}: {memory_stats()}') + shared.log.debug(f'Weights unloaded {op}: {memory_stats()}') devices.torch_gc(force=True) diff --git a/modules/sd_vae.py b/modules/sd_vae.py index 1943ec5f9..8ad295ccf 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -5,6 +5,7 @@ from copy import deepcopy import torch from modules import shared, paths, devices, script_callbacks, sd_models + vae_ignore_keys = {"model_ema.decay", "model_ema.num_updates"} vae_dict = {} base_vae = None @@ -13,6 +14,7 @@ checkpoint_info = None vae_path = os.path.abspath(os.path.join(paths.models_path, 'VAE')) checkpoints_loaded = collections.OrderedDict() + def get_base_vae(model): if base_vae is not None and checkpoint_info == model.sd_checkpoint_info and model: return base_vae @@ -147,6 +149,26 @@ def load_vae(model, vae_file=None, vae_source="from unknown source"): loaded_vae_file = vae_file +def load_vae_diffusers(_model, vae_file=None, vae_source="from unknown source"): + global loaded_vae_file # pylint: disable=global-statement + if loaded_vae_file == vae_file: + return + loaded_vae_file = None + if vae_file is None: + return + if not os.path.isfile(vae_file): + shared.log.error('VAE not found: {vae_file}') + return + shared.log.info(f"Loading diffusers VAE: {vae_source}: {vae_file}") + try: + import diffusers + diffusers_vae = diffusers.AutoencoderKL.from_pretrained(vae_file) + except Exception as e: + shared.log.error(f"Loading diffusers VAE failed: {vae_file} {e}") + diffusers_vae = None + return diffusers_vae + + # don't call this from outside def _load_vae_dict(model, vae_dict_1): model.first_stage_model.load_state_dict(vae_dict_1) @@ -178,12 +200,17 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified): lowvram.send_everything_to_cpu() else: sd_model.to(devices.cpu) - sd_hijack.model_hijack.undo_hijack(sd_model) - if shared.cmd_opts.rollback_vae and devices.dtype_vae == torch.bfloat16: - devices.dtype_vae = torch.float16 - load_vae(sd_model, vae_file, vae_source) - sd_hijack.model_hijack.hijack(sd_model) - script_callbacks.model_loaded_callback(sd_model) + + if shared.backend == shared.Backend.ORIGINAL: + sd_hijack.model_hijack.undo_hijack(sd_model) + if shared.cmd_opts.rollback_vae and devices.dtype_vae == torch.bfloat16: + devices.dtype_vae = torch.float16 + load_vae(sd_model, vae_file, vae_source) + sd_hijack.model_hijack.hijack(sd_model) + script_callbacks.model_loaded_callback(sd_model) + elif shared.backend == shared.Backend.DIFFUSERS: + load_vae_diffusers(sd_model, vae_file, vae_source) + if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram: sd_model.to(devices.device) shared.log.info(f"VAE weights loaded: {vae_file}") diff --git a/modules/shared.py b/modules/shared.py index a0aa8be8d..f654c69bc 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -38,7 +38,10 @@ hypernetworks = {} loaded_hypernetworks = [] gradio_theme = gr.themes.Base() settings_components = None -pipelines = ['Stable Diffusion', 'Stable Diffusion XL', 'Kandinsky V1', 'Kandinsky V2', 'DeepFloyd IF', 'Shap-E'] +pipelines = [ + 'Stable Diffusion', 'Stable Diffusion XL', 'Kandinsky V1', 'Kandinsky V2', 'DeepFloyd IF', 'Shap-E', + 'Stable Diffusion Img2Img', 'Stable Diffusion XL Img2Img', 'Kandinsky V1 Img2Img', 'Kandinsky V2 Img2Img', 'DeepFloyd IF Img2Img', 'Shap-E Img2Img' +] latent_upscale_default_mode = "Latent" latent_upscale_modes = { "Latent": {"mode": "bilinear", "antialias": False}, @@ -356,7 +359,8 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { })) options_templates.update(options_section(('diffusers', "Diffusers Settings"), { - "diffusers_pipeline": OptionInfo(pipelines[0], 'Diffuser Pipeline', gr.Dropdown, lambda: {"choices": pipelines}), + "diffusers_allow_safetensors": OptionInfo(False, 'Diffuser Pipeline when loading from safetensors'), + "diffusers_pipeline": OptionInfo(pipelines[0], 'Diffuser Pipeline when loading from safetensors', gr.Dropdown, lambda: {"choices": pipelines}), "diffusers_extract_ema": OptionInfo(True, "Use model EMA weights when possible"), "diffusers_generator_device": OptionInfo("default", "Generator device", gr.Radio, lambda: {"choices": ["default", "cpu"]}), "diffusers_seq_cpu_offload": OptionInfo(False, "Enable sequential CPU offload"), diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py index cfb4e7525..c893b6de0 100644 --- a/modules/ui_extra_networks.py +++ b/modules/ui_extra_networks.py @@ -66,6 +66,7 @@ class ExtraNetworksPage: self.allow_negative_prompt = False self.metadata = {} self.info = {} + self.html = '' self.items = [] self.missing_thumbs = [] self.card = ''' @@ -150,7 +151,6 @@ class ExtraNetworksPage: self_name_id = self.name.replace(" ", "_") if skip: return f"
Extra network page not ready
Click refresh to try again
" - items_html = '' subdirs = {} allowed_folders = [os.path.abspath(x) for x in self.allowed_directories_for_previews()] for parentdir in [*set(allowed_folders)]: @@ -174,16 +174,21 @@ class ExtraNetworksPage: {html.escape(subdir) if subdir!="" else "all"}
""" for subdir in subdirs]) try: + if len(self.html) > 0: + res = f"
{subdirs_html}
{self.html}
" + return res + self.html = '' self.items = list(self.list_items()) self.create_xyz_grid() for item in self.items: self.metadata[item["name"]] = item.get("metadata", {}) self.info[item["name"]] = self.find_info(item['filename']) - items_html += self.create_html_for_item(item, tabname) - if len(subdirs_html) > 0 or len(items_html) > 0: - res = f"
{subdirs_html}
{items_html}
" + self.html += self.create_html_for_item(item, tabname) + if len(subdirs_html) > 0 or len(self.html) > 0: + res = f"
{subdirs_html}
{self.html}
" else: return '' + shared.log.debug(f'Extra networks: {self.name} items={len(self.items)} subdirs={len(subdirs)}') threading.Thread(target=self.create_thumb).start() return res except Exception as e: @@ -327,6 +332,7 @@ def create_ui(container, button, tabname, skip_indexing = False): def refresh(): res = [] for pg in ui.stored_extra_pages: + pg.html = '' pg.refresh() res.append(pg.create_html(ui.tabname)) ui.search.update(value = ui.search.value)