From a2caafe4df486afc51ae141c559ac3539ba56d4a Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 2 Jul 2023 14:04:54 -0400 Subject: [PATCH] initial diffusers merge into dev --- DIFFUSERS.md | 49 +++++++++ modules/modelloader.py | 53 ++++++--- modules/processing.py | 58 ++++++++-- modules/sd_hijack.py | 6 +- modules/sd_models.py | 179 +++++++++++++++++++++++++++---- modules/sd_samplers.py | 4 +- modules/sd_samplers_diffusers.py | 37 +++++++ modules/shared.py | 82 +++++++++++++- modules/ui.py | 2 +- modules/ui_models.py | 12 +++ requirements.txt | 1 + 11 files changed, 431 insertions(+), 52 deletions(-) create mode 100644 DIFFUSERS.md create mode 100644 modules/sd_samplers_diffusers.py diff --git a/DIFFUSERS.md b/DIFFUSERS.md new file mode 100644 index 000000000..199623c4d --- /dev/null +++ b/DIFFUSERS.md @@ -0,0 +1,49 @@ +# Diffusers WiP + +initial support merged into `dev` branch + + git clone https://github.com/vladmandic/automatic -b dev diffusers + cd diffusers + webui --debug --backend diffusers + +default sd 1.5 model will be downloaded automatically to `models/Diffusers` + +on first startup, disable **controlnet** and **multi-diffusion** extensions as right now they are not compatible with diffusers + +to update repo, do not use `--upgrade` flag, use manual `git pull` instead + +## Test + +### Standard + +- run with `webui --debug --backend original` +- goal is to test standard workflows (so not diffusers) to ensure there are no regressions + so diffusers code can be merged into `master` and we can continue with development there + +### Diffusers + +- sd 1.5 and sd 2.1 model + models can be downloaded from huggingface hub + but focus on default model for now and i'll add downloader soon +- lora, textual inversion + only loras/textual-inversions downloaded from huggingface hub are supported for now + i'll add standard safetensors soon +- txt2img, img2img, inpaint, outpaint, process + +### Experimental + +- cuda model compile using `reduce overhead` model with and without `fullgraph` +- kandinsky model + +## Todo + +- enable loading of safetensors models +- cleanup logging +- search&download models from hfhub +- controlnet extension +- multidiffusion extension +- sdxl model + +## Issues + +- TBD diff --git a/modules/modelloader.py b/modules/modelloader.py index c49cdeb5d..799537240 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -1,6 +1,8 @@ import os import shutil import importlib +import json +from typing import Dict from urllib.parse import urlparse from modules import shared @@ -9,29 +11,54 @@ from modules.paths import script_path, models_path diffuser_repos = [] -def load_diffusers(model_path: str, hub_url: str = None, command_path: str = None): - import huggingface_hub as hf +def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config: Dict[str, str] = None): from diffusers import DiffusionPipeline + import huggingface_hub as hf + if download_config is None: + download_config = { + "force_download": False, + "resume_download": True, + "cache_dir": shared.opts.diffusers_dir, + } + + if cache_dir is not None: + download_config["cache_dir"] = cache_dir + + pipeline_dir = DiffusionPipeline.download(hub_id, **download_config) + model_info_dict = hf.model_info(hub_id).cardData # TODO hfhub card-data? + + # some checkpoints need to be downloaded as "hidden" as they just serve as pre- or post-pipelines of other pipelines + if model_info_dict is not None and "prior" in model_info_dict: + download_dir = DiffusionPipeline.download(model_info_dict["prior"], **download_config) + model_info_dict["prior"] = download_dir + # mark prior as hidden + with open(os.path.join(download_dir, "hidden"), "w", encoding="utf-8") as f: + f.write("True") + + with open(os.path.join(pipeline_dir, "model_info.json"), "w", encoding="utf-8") as json_file: + json.dump(model_info_dict, json_file) + + return pipeline_dir + +def load_diffusers_models(model_path: str, command_path: str = None): + import huggingface_hub as hf places = [] - - # download repo - if hub_url is not None: - DiffusionPipeline.download(hub_url, cache_dir=model_path) - places.append(model_path) if command_path is not None and command_path != model_path and os.path.isdir(command_path): places.append(command_path) diffuser_repos.clear() output = [] - try: - for place in places: + for place in places: + try: res = hf.scan_cache_dir(cache_dir=place) for r in list(res.repos): - diffuser_repos.append({ 'name': r.repo_id, 'filename': r.repo_id, 'path': str(r.repo_path), 'size': r.size_on_disk, 'mtime': r.last_modified, 'hash': list(r.revisions)[-1].commit_hash }) - output.append(str(r.repo_id)) - except Exception as e: - shared.log.error(f"Error listing diffusers: {place} {e}") + cache_path = os.path.join(r.repo_path, "snapshots", list(r.revisions)[-1].commit_hash) + diffuser_repos.append({ 'name': r.repo_id, 'filename': r.repo_id, 'path': cache_path, 'size': r.size_on_disk, 'mtime': r.last_modified, 'hash': list(r.revisions)[-1].commit_hash, 'model_info': str(os.path.join(cache_path, "model_info.json")) }) + if not os.path.isfile(os.path.join(cache_path, "hidden")): + output.append(str(r.repo_id)) + except Exception as e: + shared.log.error(f"Error listing diffusers: {place} {e}") shared.log.debug(f'Scanning diffusers cache: {len(output)} {model_path} {command_path}') return output diff --git a/modules/processing.py b/modules/processing.py index 6d05beed4..e8808a910 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -223,7 +223,7 @@ class StableDiffusionProcessing: # HACK: Using introspection as the Depth2Image model doesn't appear to uniquely # identify itself with a field common to all models. The conditioning_key is also hybrid. if backend == Backend.DIFFUSERS: # TODO: Diffusers img2img_image_conditioning - return latent_image.new_zeros(latent_image.shape[0], 5, 1, 1) + return None if isinstance(self.sd_model, LatentDepth2ImageDiffusion): return self.depth2img_image_conditioning(source_image) if self.sd_model.cond_stage_key == "edit": @@ -520,7 +520,8 @@ def process_images(p: StableDiffusionProcessing) -> Processed: if k == 'sd_vae': sd_vae.reload_vae_weights() - sd_models.apply_token_merging(p.sd_model, p.get_token_merging_ratio()) + if not shared.opts.cuda_compile: + sd_models.apply_token_merging(p.sd_model, p.get_token_merging_ratio()) if cmd_opts.profile: """ @@ -538,7 +539,8 @@ def process_images(p: StableDiffusionProcessing) -> Processed: else: res = process_images_inner(p) finally: - sd_models.apply_token_merging(p.sd_model, 0) + if not shared.opts.cuda_compile: + sd_models.apply_token_merging(p.sd_model, 0) if p.override_settings_restore_afterwards: # restore opts to original state for k, v in stored_opts.items(): setattr(opts, k, v) @@ -557,6 +559,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: assert len(p.prompt) > 0 else: assert p.prompt is not None + seed = get_fixed_seed(p.seed) subseed = get_fixed_seed(p.subseed) if backend == Backend.ORIGINAL: @@ -683,26 +686,39 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: devices.torch_gc() if p.scripts is not None: p.scripts.postprocess_batch(p, x_samples_ddim, batch_number=n) - else: # TODO Diffusers main processing + + elif backend == Backend.DIFFUSERS: generator = [torch.Generator(device="cpu").manual_seed(s) for s in seeds] if shared.sd_model.scheduler.name != p.sampler_name: sampler = sd_samplers.all_samplers_map.get(p.sampler_name, None) if sampler is None: sampler = sd_samplers.all_samplers_map.get("UniPC") scheduler = sampler.constructor(shared.sd_model.sd_checkpoint_info.filename) + # TODO(Patrick): For wrapped pipelines this is currently a no-op shared.sd_model.scheduler = scheduler.sampler + + if sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE: + task_specific_kwargs = {"height": p.height, "width": p.width} + elif sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.IMAGE_2_IMAGE: + task_specific_kwargs = {"image": p.init_images[0], "strength": p.denoising_strength} + elif sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.INPAINTING: + # 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} + output = shared.sd_model( prompt=prompts, negative_prompt=negative_prompts, num_inference_steps=p.steps, guidance_scale=p.cfg_scale, - height=p.height, - width=p.width, generator=generator, output_type="np", + **task_specific_kwargs ) x_samples_ddim = output.images + else: + raise ValueError(f"Unknown backend {backend}") + for i, x_sample in enumerate(x_samples_ddim): p.batch_index = i if backend == Backend.ORIGINAL: @@ -820,8 +836,12 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing): self.applied_old_hires_behavior_to = None def init(self, all_prompts, all_seeds, all_subseeds): + if backend == Backend.DIFFUSERS: + sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE) + self.width = self.width or 512 self.height = self.height or 512 + if self.enable_hr: if opts.use_old_hires_fix_width_height and self.applied_old_hires_behavior_to != (self.width, self.height): self.hr_resize_x = self.width @@ -873,6 +893,9 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing): self.extra_generation_params["Hires upscaler"] = self.hr_upscaler def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts): # TODO this is majority of processing time + if backend == Backend.DIFFUSERS: + sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE) + self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model) latent_scale_mode = shared.latent_upscale_modes.get(self.hr_upscaler, None) if self.hr_upscaler is not None else shared.latent_upscale_modes.get(shared.latent_upscale_default_mode, "nearest") if self.enable_hr and latent_scale_mode is None: @@ -978,12 +1001,18 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing): self.image_conditioning = None def init(self, all_prompts, all_seeds, all_subseeds): + image_mask = self.image_mask + if backend == Backend.DIFFUSERS and image_mask is None: + sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE) + elif backend == Backend.DIFFUSERS and image_mask is not None: + sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.INPAINTING) + self.sd_model.dtype = self.sd_model.unet.dtype + force_latent_upscaler = shared.opts.data.get('force_latent_sampler') if self.sampler_name in ['PLMS']: self.sampler_name = force_latent_upscaler if force_latent_upscaler != 'None' else shared.opts.fallback_sampler # PLMS does not support img2img, use fallback instead self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model) crop_region = None - image_mask = self.image_mask if image_mask is not None: image_mask = image_mask.convert('L') if self.inpainting_mask_invert: @@ -1048,7 +1077,13 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing): image = torch.from_numpy(batch_images) image = 2. * image - 1. image = image.to(shared.device) - self.init_latent = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(image)) + + if backend == Backend.ORIGINAL: + self.init_latent = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(image)) + else: + # we don't pre-encode the latents for diffusers to allow the UI to stay general for different model types + self.init_latent = None + if self.resize_mode == 3: self.init_latent = torch.nn.functional.interpolate(self.init_latent, size=(self.height // opt_f, self.width // opt_f), mode="bilinear") if image_mask is not None: @@ -1068,6 +1103,13 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing): self.image_conditioning = self.img2img_image_conditioning(image, self.init_latent, image_mask) def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts): + if backend == Backend.DIFFUSERS: + if self.init_mask is None: # pylint: disable=no-member + sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE) + else: + sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.INPAINTING) + self.sd_model.dtype = self.sd_model.unet.dtype + x = create_random_tensors([opt_C, self.height // opt_f, self.width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self) if self.initial_noise_multiplier != 1.0: self.extra_generation_params["Noise multiplier"] = self.initial_noise_multiplier diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index 420fbcf18..5e397ebb9 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -179,11 +179,11 @@ class StableDiffusionModelHijack: shared.log.info("Model compile enabled: IPEX Optimize Graph Mode") else: shared.log.warning("Model compile skipped: IPEX Method is for Intel GPU's with OneAPI") - elif opts.cuda_compile and opts.cuda_compile_mode != 'none': + elif opts.cuda_compile and opts.cuda_compile_mode != 'none' and shared.backend == shared.Backend.ORIGINAL: try: import logging import torch._dynamo as dynamo # pylint: disable=unused-import - torch._dynamo.config.log_level = logging.WARNING if opts.cuda_compile_verbose else logging.CRITICAL # pylint: disable=protected-access + # torch._dynamo.config.log_level = logging.WARNING if opts.cuda_compile_verbose else logging.CRITICAL # pylint: disable=protected-access torch._dynamo.config.verbose = opts.cuda_compile_verbose # pylint: disable=protected-access torch._dynamo.config.suppress_errors = opts.cuda_compile_errors # pylint: disable=protected-access torch.backends.cudnn.benchmark = True @@ -191,7 +191,7 @@ class StableDiffusionModelHijack: import hidet hidet.torch.dynamo_config.use_tensor_core(True) hidet.torch.dynamo_config.search_space(2) - m.model = torch.compile(m.model, mode="default", backend=opts.cuda_compile_mode, fullgraph=False, dynamic=False) + m.model = torch.compile(m.model, mode="default", backend=opts.cuda_compile_mode, fullgraph=opts.cuda_compile_fullgraph, dynamic=False) shared.log.info(f"Model compile enabled: {opts.cuda_compile_mode}") except Exception as err: shared.log.warning(f"Model compile not supported: {err}") diff --git a/modules/sd_models.py b/modules/sd_models.py index 30d5dfdb8..d237f8728 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -6,6 +6,7 @@ import json import threading from os import mkdir from urllib import request +from enum import Enum import filelock from rich import progress # pylint: disable=redefined-builtin import torch @@ -13,6 +14,7 @@ import safetensors.torch from omegaconf import OmegaConf import tomesd from transformers import logging as transformers_logging +import diffusers import ldm.modules.midas as midas from ldm.util import instantiate_from_config from modules import paths, shared, modelloader, devices, script_callbacks, sd_vae, sd_disable_initialization, errors, hashes, sd_models_config @@ -52,7 +54,7 @@ class CheckpointInfo: self.name = name self.hash = model_hash(self.filename) self.sha256 = hashes.sha256_from_cache(self.filename, f"checkpoint/{name}") - else: # TODO Diffusers + elif shared.backend == shared.Backend.DIFFUSERS: repo = [r for r in modelloader.diffuser_repos if filename == r['filename']] if len(repo) == 0: error_message = f'Cannot find diffuser model: {filename}' @@ -61,6 +63,17 @@ class CheckpointInfo: self.name = repo[0]['name'] self.hash = repo[0]['hash'][:8] self.sha256 = repo[0]['hash'] + self.path = repo[0]['path'] + + if os.path.isfile(repo[0]['model_info']): + file_path = repo[0]['model_info'] + with open(file_path, "r", encoding="utf-8") as json_file: + self.model_info = json.load(json_file) + else: + self.model_info = None + else: + raise ValueError(f'Unknown backend: {shared.backend}') + self.name_for_extra = os.path.splitext(os.path.basename(filename))[0] self.model_name = os.path.splitext(name.replace("/", "_").replace("\\", "_"))[0] self.shorthash = self.sha256[0:10] if self.sha256 else None @@ -116,7 +129,8 @@ def list_models(): else: global model_path # pylint: disable=global-statement model_path = os.path.join(models_path, 'Diffusers') - model_list = modelloader.load_diffusers(model_path=model_path, command_path=shared.opts.diffusers_dir) + model_list = modelloader.load_diffusers_models(model_path=model_path, command_path=shared.opts.diffusers_dir) + for filename in sorted(model_list, key=str.lower): checkpoint_info = CheckpointInfo(filename) if checkpoint_info.name is not None: @@ -143,15 +157,15 @@ def list_models(): shared.opts.data['sd_model_checkpoint'] = "v1-5-pruned-emaonly.safetensors" model_list = modelloader.load_models(model_path=model_path, model_url=model_url, command_path=shared.opts.ckpt_dir, ext_filter=[".ckpt", ".safetensors"], download_name="v1-5-pruned-emaonly.safetensors", ext_blacklist=[".vae.ckpt", ".vae.safetensors"]) else: - hub_url = "runwayml/stable-diffusion-v1-5" - model_list = modelloader.load_diffusers(model_path=model_path, hub_url=hub_url, command_path=shared.opts.diffusers_dir) + default_model_id = "runwayml/stable-diffusion-v1-5" + modelloader.download_diffusers_model(default_model_id, model_path) + model_list = modelloader.load_diffusers_models(model_path=model_path, command_path=shared.opts.diffusers_dir) for filename in sorted(model_list, key=str.lower): checkpoint_info = CheckpointInfo(filename) if checkpoint_info.name is not None: checkpoint_info.register() - def update_model_hashes(): txt = [] lst = [ckpt for ckpt in checkpoints_list.values() if ckpt.hash is None] @@ -474,49 +488,115 @@ class SdModelData: model_data = SdModelData() +class PriorPipeline: + def __init__(self, prior, main): + self.prior = prior + self.main = main + self.scheduler = main.scheduler + self.tokenizer = self.prior.tokenizer + + def to(self, *args, **kwargs): + # only the prior is moved to CUDA in a first step + self.prior.to(*args, **kwargs) + + def enable_model_cpu_offload(self, *args, **kwargs): + self.prior.enable_model_cpu_offload(*args, **kwargs) + self.main.enable_model_cpu_offload(*args, **kwargs) + + def enable_sequential_cpu_offload(self, *args, **kwargs): + self.prior.enable_sequential_cpu_offload(*args, **kwargs) + self.main.enable_sequential_cpu_offload(*args, **kwargs) + + def enable_xformers_memory_efficient_attention(self, *args, **kwargs): + self.prior.enable_xformers_memory_efficient_attention(*args, **kwargs) + self.main.enable_xformers_memory_efficient_attention(*args, **kwargs) + + def __call__(self, *args, **kwargs): + unclip_outputs = self.prior(prompt=kwargs.get("prompt"), negative_prompt=kwargs.get("negative_prompt")) + + if self.prior.device.type == "cuda": + prior_device = self.prior.device + self.prior.to("cpu") + self.main.to(prior_device) + + kwargs = {**kwargs, **unclip_outputs} + result = self.main(*args, **kwargs) + + if self.main.device.type == "cuda": + main_device = self.main.device + self.main.to("cpu") + self.prior.to(main_device) + + return result + + def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=None): # pylint: disable=unused-argument if timer is None: timer = Timer() - import diffusers import logging logging.getLogger("diffusers").setLevel(logging.ERROR) timer.record("diffusers") - diffusor_config = { - "force_download": False, - "safety_checker": None, - "resume_download": True, + diffusers_load_config = { "low_cpu_mem_usage": True, - "use_safetensors": True, - "cache_dir": shared.opts.diffusers_dir, "torch_dtype": devices.dtype, + "safety_checker": None, + # "use_safetensors": True, # TODO(PVP) - we can't enable this for all checkpoints just yet } + if shared.opts.data['sd_model_checkpoint'] == 'model.ckpt': shared.opts.data['sd_model_checkpoint'] = "runwayml/stable-diffusion-v1-5" sd_model = None try: if shared.cmd_opts.ckpt is not None and model_data.initial: # initial load model_name = modelloader.find_diffuser(shared.cmd_opts.ckpt) + if model_name is not None: shared.log.info(f'Loading diffuser model: {model_name}') - scheduler = diffusers.UniPCMultistepScheduler.from_pretrained(model_name, subfolder="scheduler") - sd_model = diffusers.DiffusionPipeline.from_pretrained(model_name, scheduler=scheduler, **diffusor_config) + model_file = modelloader.download_diffusers_model(hub_id=model_name) + sd_model = diffusers.DiffusionPipeline.from_pretrained(model_file, **diffusers_load_config) + list_models() # rescan for downloaded model checkpoint_info = CheckpointInfo(model_name) + if sd_model is None: checkpoint_info = checkpoint_info or select_checkpoint() shared.log.info(f'Loading diffuser model: {checkpoint_info.filename}') - scheduler = diffusers.UniPCMultistepScheduler.from_pretrained(checkpoint_info.filename, subfolder="scheduler") - sd_model = diffusers.DiffusionPipeline.from_pretrained(checkpoint_info.filename, scheduler=scheduler, **diffusor_config) + sd_model = diffusers.DiffusionPipeline.from_pretrained(checkpoint_info.path, **diffusers_load_config) + + if "StableDiffusion" in sd_model.__class__.__name__: + sd_model.scheduler = diffusers.UniPCMultistepScheduler.from_config(sd_model.scheduler.config) + sd_model.scheduler.name = 'UniPC' + elif "Kandinsky" in sd_model.__class__.__name__: + sd_model.scheduler.name = 'DDIM' + + # Prior pipelines + if checkpoint_info.model_info is not None and "prior" in checkpoint_info.model_info: + prior_id = checkpoint_info.model_info["prior"] + shared.log.info(f"Loading prior {prior_id} for {checkpoint_info.filename}") + prior = diffusers.DiffusionPipeline.from_pretrained(prior_id, **diffusers_load_config) + sd_model = PriorPipeline(prior=prior, main=sd_model) # wrap sd_model + if shared.cmd_opts.medvram: sd_model.enable_model_cpu_offload() if shared.cmd_opts.lowvram: sd_model.enable_sequential_cpu_offload() if shared.opts.cross_attention_optimization == "xFormers": sd_model.enable_xformers_memory_efficient_attention() - sd_model.sd_checkpoint_info = checkpoint_info - sd_model.sd_model_checkpoint = checkpoint_info.filename - sd_model.sd_model_hash = checkpoint_info.hash - scheduler.name = 'UniPC' + if shared.opts.cuda_compile and torch.cuda.is_available(): + sd_model.to(devices.device) + sd_model.unet.to(memory_format=torch.channels_last) + import torch._dynamo as dynamo # pylint: disable=unused-import + # torch._dynamo.config.log_level = logging.WARNING if shared.opts.cuda_compile_verbose else logging.CRITICAL # pylint: disable=protected-access + torch._dynamo.config.verbose = shared.opts.cuda_compile_verbose # pylint: disable=protected-access + torch._dynamo.config.suppress_errors = shared.opts.cuda_compile_errors # pylint: disable=protected-access + sd_model.unet = torch.compile(sd_model.unet, mode=shared.opts.cuda_compile_mode, fullgraph=shared.opts.cuda_compile_fullgraph) # pylint: disable=attribute-defined-outside-init + shared.log.info(f"Compiling pipeline={sd_model.__class__.__name__} shape={8 * sd_model.unet.config.sample_size} mode={shared.opts.cuda_compile_mode}") + sd_model("dummy prompt") + shared.log.info("Complilation done.") + + sd_model.sd_checkpoint_info = checkpoint_info # pylint: disable=attribute-defined-outside-init + sd_model.sd_model_checkpoint = checkpoint_info.filename # pylint: disable=attribute-defined-outside-init + sd_model.sd_model_hash = checkpoint_info.hash # pylint: disable=attribute-defined-outside-init sd_model.to(devices.device) except Exception as e: shared.log.error("Failed to load diffusers model") @@ -528,6 +608,60 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No shared.log.info(f'Model load finished: {memory_stats()}') +class DiffusersTaskType(Enum): + TEXT_2_IMAGE = 1 + IMAGE_2_IMAGE = 2 + INPAINTING = 3 + +def set_diffuser_pipe(pipe, new_pipe_type): + wrapper_pipe = None + + sd_checkpoint_info = pipe.sd_checkpoint_info + sd_model_checkpoint = pipe.sd_model_checkpoint + sd_model_hash = pipe.sd_model_hash + + if pipe.__class__ == PriorPipeline: + wrapper_pipe = pipe + pipe = pipe.main + + pipe_name = pipe.__class__.__name__ + pipe_name = pipe_name.replace("Img2Img", "").replace("Inpaint", "") + if new_pipe_type == DiffusersTaskType.TEXT_2_IMAGE: + new_pipe_cls_str = pipe_name + elif new_pipe_type == DiffusersTaskType.IMAGE_2_IMAGE: + new_pipe_cls_str = pipe_name.replace("Pipeline", "Img2ImgPipeline") + elif new_pipe_type == DiffusersTaskType.INPAINTING: + new_pipe_cls_str = pipe_name.replace("Pipeline", "InpaintPipeline") + + new_pipe_cls = getattr(diffusers, new_pipe_cls_str) + + if pipe.__class__ == new_pipe_cls: + return + + new_pipe = new_pipe_cls(**pipe.components) + + if wrapper_pipe is not None: + wrapper_pipe.main = new_pipe + new_pipe = wrapper_pipe + + new_pipe.sd_checkpoint_info = sd_checkpoint_info + new_pipe.sd_model_checkpoint = sd_model_checkpoint + new_pipe.sd_model_hash = sd_model_hash + + shared.sd_model = new_pipe + shared.log.info(f"Pipeline class changed from {pipe.__class__.__name__} to {new_pipe_cls.__name__}") + + +def get_diffusers_task(pipe: diffusers.DiffusionPipeline) -> DiffusersTaskType: + if pipe.__class__ == PriorPipeline: + pipe = pipe.main + + if "Img2Img" in pipe.__class__.__name__: + return DiffusersTaskType.IMAGE_2_IMAGE + elif "Inpaint" in pipe.__class__.__name__: + return DiffusersTaskType.INPAINTING + return DiffusersTaskType.TEXT_2_IMAGE + def load_model(checkpoint_info=None, already_loaded_state_dict=None, timer=None): from modules import lowvram, sd_hijack @@ -684,6 +818,11 @@ def apply_token_merging(sd_model, token_merging_ratio): return if current_token_merging_ratio > 0: tomesd.remove_patch(sd_model) + + if sd_model.__class__ == PriorPipeline: + # token merging is not supported for PriorPipelines currently + return + if token_merging_ratio > 0: tomesd.apply_patch( sd_model, diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index eeda69d0b..6033f1bf4 100644 --- a/modules/sd_samplers.py +++ b/modules/sd_samplers.py @@ -1,4 +1,4 @@ -from modules import sd_samplers_compvis, sd_samplers_kdiffusion, sd_samplers_diffusors, shared +from modules import sd_samplers_compvis, sd_samplers_kdiffusion, sd_samplers_diffusers, shared from modules.sd_samplers_common import samples_to_image_grid, sample_to_image # pylint: disable=unused-import from modules.shared import backend, Backend @@ -9,7 +9,7 @@ if backend == Backend.ORIGINAL: ] else: all_samplers = [ - *sd_samplers_diffusors.samplers_data_diffusors, + *sd_samplers_diffusers.samplers_data_diffusers, ] all_samplers_map = {x.name: x for x in all_samplers} samplers = all_samplers diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py new file mode 100644 index 000000000..ebb988109 --- /dev/null +++ b/modules/sd_samplers_diffusers.py @@ -0,0 +1,37 @@ +from diffusers import ( + DDIMScheduler, + DDPMScheduler, + DEISMultistepScheduler, + DPMSolverMultistepScheduler, + DPMSolverSinglestepScheduler, + EulerAncestralDiscreteScheduler, + EulerDiscreteScheduler, + HeunDiscreteScheduler, + KDPM2DiscreteScheduler, + PNDMScheduler, + UniPCMultistepScheduler, +) +from modules import sd_samplers_common + +samplers_data_diffusers = [ + sd_samplers_common.SamplerData('UniPC', lambda model: DiffusionSampler('UniPC', UniPCMultistepScheduler, model), [], {}), + sd_samplers_common.SamplerData('DDIM', lambda model: DiffusionSampler('DDIM', DDIMScheduler, model), [], {}), + sd_samplers_common.SamplerData('DDPM', lambda model: DiffusionSampler('DDPM', DDPMScheduler, model), [], {}), + sd_samplers_common.SamplerData('DEIS', lambda model: DiffusionSampler('DEIS', DEISMultistepScheduler, model), [], {}), + sd_samplers_common.SamplerData('DPM++ 2M', lambda model: DiffusionSampler('DPM++ 2M', DPMSolverMultistepScheduler, model), [], {}), + sd_samplers_common.SamplerData('DPM++ 1S', lambda model: DiffusionSampler('DPM++ 1S', DPMSolverSinglestepScheduler, model), [], {}), + sd_samplers_common.SamplerData('DPM++ 2M SDE', lambda model: DiffusionSampler('DPM++ 2M SDE', DPMSolverMultistepScheduler, model, algorithm_type="sde-dpmsolver++"), [], {}), + sd_samplers_common.SamplerData('DPM++ 2M Karras', lambda model: DiffusionSampler('DPM++ 2M Karras', DPMSolverMultistepScheduler, model, use_karras_sigmas=True), [], {}), + sd_samplers_common.SamplerData('DPM++ 1S Karras', lambda model: DiffusionSampler('DPM++ 1S Karras', DPMSolverSinglestepScheduler, model, use_karras_sigmas=True), [], {}), + sd_samplers_common.SamplerData('DPM++ 2M SDE Karras', lambda model: DiffusionSampler('DPM++ 2M SDE Karras', DPMSolverMultistepScheduler, model, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++"), [], {}), + sd_samplers_common.SamplerData('Euler', lambda model: DiffusionSampler('Euler', EulerDiscreteScheduler, model), [], {}), + sd_samplers_common.SamplerData('Euler a', lambda model: DiffusionSampler('Euler a', EulerAncestralDiscreteScheduler, model), [], {}), + sd_samplers_common.SamplerData('Heun', lambda model: DiffusionSampler('Heun', HeunDiscreteScheduler, model), [], {}), + sd_samplers_common.SamplerData('DPM2++ 2M', lambda model: DiffusionSampler('KDPM2', KDPM2DiscreteScheduler, model), [], {}), + sd_samplers_common.SamplerData('PNDM', lambda model: DiffusionSampler('PNDM', PNDMScheduler, model), [], {}), +] + +class DiffusionSampler: + def __init__(self, name, constructor, sd_model, **kwargs): + self.sampler = constructor.from_pretrained(sd_model, subfolder="scheduler", **kwargs) + self.sampler.name = name diff --git a/modules/shared.py b/modules/shared.py index 7d5b939d3..f9ce0a981 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -4,11 +4,14 @@ import time import json import datetime import urllib.request +from urllib.parse import urlparse from enum import Enum +import tempfile import gradio as gr import tqdm import requests -from modules import errors, ui_components, shared_items, cmd_args +import diffusers +from modules import errors, ui_components, shared_items, cmd_args, modelloader from modules.paths_internal import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir # pylint: disable=W0611 import modules.interrogate import modules.memmon @@ -72,6 +75,11 @@ ui_reorder_categories = [ ] +def is_url(string): + parsed_url = urlparse(string) + return all([parsed_url.scheme, parsed_url.netloc]) + + class Backend(Enum): ORIGINAL = 1 DIFFUSERS = 2 @@ -185,7 +193,7 @@ state.server_start = time.time() class OptionInfo: - def __init__(self, default=None, label="", component=None, component_args=None, onchange=None, section=None, refresh=None, comment_before='', comment_after=''): + def __init__(self, default=None, label="", component=None, component_args=None, onchange=None, section=None, refresh=None, submit=None, comment_before='', comment_after=''): self.default = default self.label = label self.component = component @@ -195,6 +203,7 @@ class OptionInfo: self.refresh = refresh self.comment_before = comment_before # HTML text that will be added after label in UI self.comment_after = comment_after # HTML text that will be added before label in UI + self.submit = submit def link(self, label, uri): self.comment_before += f"[{label}]" @@ -223,9 +232,70 @@ def list_checkpoint_tiles(): import modules.sd_models # pylint: disable=W0621 return modules.sd_models.checkpoint_tiles() - default_checkpoint = list_checkpoint_tiles()[0] if len(list_checkpoint_tiles()) > 0 else "model.ckpt" +def load_diffusers_ckpt(model_repo: str): + cached_dir = modelloader.download_diffusers_model(model_repo) + print(f"Downloaded {cached_dir}") + return "" + +def load_diffusers_lora(lora_repo: str): + pipe = sys.modules[__name__].sd_model + + if lora_repo == "": + pipe._remove_text_encoder_monkey_patch() # pylint: disable=W0212 + proc_cls_name = next(iter(pipe.unet.attn_processors.values())).__class__.__name__ + non_lora_proc_cls = getattr(diffusers.models.attention_processor, proc_cls_name[len("LORA"):]) + pipe.unet.set_attn_processor(non_lora_proc_cls()) + print("Removed LoRA.") + return "" + elif is_url(lora_repo): + with tempfile.TemporaryDirectory() as temp_dir: + os.system(f"wget -P {temp_dir} {lora_repo}") + temp_file_path = os.path.join(temp_dir, lora_repo.split('/')[-1]) + pipe.load_lora_weights(temp_file_path) + + lora_repo = '/'.join(lora_repo.split('/')[-2:]) + + print(f"Loaded Civit.ai LoRA: {lora_repo}") + return f"{lora_repo} is loaded. Pass empty text field to remove LoRA or pass new LoRA id." + elif len(lora_repo.split('/')) == 2: + lora_dir = os.path.dirname(opts.data["diffusers_dir"]) + cache_dir = os.path.join(lora_dir, "Diffusers_LoRA") + pipe.load_lora_weights(lora_repo, cache_dir=cache_dir) + print(f"Loaded {lora_repo}") + return f"{lora_repo} is loaded. Pass empty text field to remove LoRA or pass new LoRA id." + else: + print(f"{lora_repo} is not a valid LoRA identifier.") + return "" + +def load_diffusers_text_inv(text_inv_repo: str): + pipe = sys.modules[__name__].sd_model + + if text_inv_repo == "": + pipe.tokenizer = pipe.tokenizer.__class__.from_pretrained(pipe.tokenizer.name_or_path) + pipe.text_encoder.resize_token_embeddings(len(pipe.tokenizer)) + print("Removed all textual inversions.") + return "" + elif is_url(text_inv_repo): + with tempfile.TemporaryDirectory() as temp_dir: + os.system(f"wget -P {temp_dir} {text_inv_repo}") + temp_file_path = os.path.join(temp_dir, text_inv_repo.split('/')[-1]) + pipe.load_textual_inversion(temp_file_path) + + text_inv_repo = '/'.join(text_inv_repo.split('/')[-2:]) + + print(f"Loaded Civit.ai Textual Inv: {text_inv_repo}") + elif len(text_inv_repo.split('/')) == 2: + text_inv_dir = os.path.dirname(opts.data["diffusers_dir"]) + cache_dir = os.path.join(text_inv_dir, "Diffusers_Text_Inv") + pipe.load_textual_inversion(text_inv_repo, cache_dir=cache_dir) + print(f"Loaded {text_inv_repo}") + + text_inv_tokens = pipe.tokenizer.added_tokens_encoder.keys() + text_inv_tokens = [t for t in text_inv_tokens if not (len(t.split("_")) > 1 and t.split("_")[-1].isdigit())] + + return f"{', '.join(text_inv_tokens)} loaded. Pass empty text field to remove all or add new textual inversion id." def refresh_checkpoints(): import modules.sd_models # pylint: disable=W0621 @@ -329,7 +399,8 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { "cuda_allow_tf32": OptionInfo(True, "Allow TF32 math ops"), "cuda_allow_tf16_reduced": OptionInfo(True, "Allow TF16 reduced precision math ops"), "cuda_compile": OptionInfo(False, "Enable model compile (experimental)"), - "cuda_compile_mode": OptionInfo("none", "Model compile mode (experimental)", gr.Radio, lambda: {"choices": ['none', 'inductor', 'cudagraphs', 'aot_ts_nvfuser', 'hidet', 'ipex']}), + "cuda_compile_mode": OptionInfo("none", "Model compile mode (experimental)", gr.Radio, lambda: {"choices": ['none', 'inductor', 'reduce-overhead', 'cudagraphs', 'aot_ts_nvfuser', 'hidet', 'ipex']}), + "cuda_compile_fullgraph": OptionInfo(False, "Model compile fullgraph"), "cuda_compile_verbose": OptionInfo(False, "Model compile verbose mode"), "cuda_compile_errors": OptionInfo(True, "Model compile suppress errors"), "disable_gc": OptionInfo(False, "Disable Torch memory garbage collection"), @@ -449,8 +520,9 @@ options_templates.update(options_section(('live-preview', "Live Previews"), { "logmonitor_refresh_period": OptionInfo(5000, "Log view update period, in milliseconds", gr.Slider, {"minimum": 0, "maximum": 30000, "step": 25}), })) + options_templates.update(options_section(('sampler-params', "Sampler Settings"), { - "show_samplers": OptionInfo(["Euler a", "UniPC", "DDIM", "DPM++ 2M SDE", "DPM++ 2M SDE Karras", "DPM2 Karras", "DPM++ 2M Karras"], "Show samplers in user interface", gr.CheckboxGroup, lambda: {"choices": [x.name for x in list_samplers() if x.name != "PLMS"]}), + "show_samplers": OptionInfo(["Euler a", "UniPC", "DDIM", "DPM++ 2M SDE", "DPM++ 2M SDE Karras", "DPM2 Karras", "DPM++ 2M Karras", "DEIS"], "Show samplers in user interface", gr.CheckboxGroup, lambda: {"choices": [x.name for x in list_samplers() if x.name != "PLMS"]}), "fallback_sampler": OptionInfo("Euler a", "Secondary sampler", gr.Dropdown, lambda: {"choices": ["None"] + [x.name for x in list_samplers()]}), "force_latent_sampler": OptionInfo("None", "Force latent upscaler sampler", gr.Dropdown, lambda: {"choices": ["None"] + [x.name for x in list_samplers()]}), "always_batch_cond_uncond": OptionInfo(False, "Disable conditional batching enabled on low memory systems"), diff --git a/modules/ui.py b/modules/ui.py index 3edab8f2b..2a27df235 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -971,6 +971,7 @@ def create_ui(): quicksettings_names = opts.quicksettings_list quicksettings_names = {x: i for i, x in enumerate(quicksettings_names) if x != 'quicksettings'} quicksettings_list = [] + previous_section = [] tab_item_keys = [] current_tab = None @@ -1146,7 +1147,6 @@ def webpath(fn): web_path = os.path.relpath(fn, script_path).replace('\\', '/') else: web_path = os.path.abspath(fn) - return f'file={web_path}?{os.path.getmtime(fn)}' diff --git a/modules/ui_models.py b/modules/ui_models.py index aa1c082f6..afa6dc702 100644 --- a/modules/ui_models.py +++ b/modules/ui_models.py @@ -161,3 +161,15 @@ def create_ui(): return model_data, txt model_list_btn.click(fn=list_models, inputs=[], outputs=[model_table, models_outcome]) + + with gr.Tab(label="HF Hub"): + """" + options_templates.update(options_section(('diffusers', "Diffusers"), { + "diffusers_ckpt_download": OptionInfo("", "HFHub Checkpoint download", gr.Textbox, {"placeholder": "e.g. runwayml/stable-diffusion-v1-5"}, submit=load_diffusers_ckpt), + "diffusers_lora_download": OptionInfo("", "HFHub LoRA download", gr.Textbox, {"placeholder": "e.g. pcuenq/pokemon-lora"}, submit=load_diffusers_lora), + "diffusers_text_inv_download": OptionInfo("", "HFHub Textual Inversion download", gr.Textbox, {"placeholder": "e.g. sd-concepts-library/midjourney-style"}, submit=load_diffusers_text_inv), + })) + """ + + with gr.Tab(label="CivitAI"): + pass diff --git a/requirements.txt b/requirements.txt index 0aeebdd1a..fced77fd2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -43,6 +43,7 @@ yapf scikit-image basicsr compel +antlr4-python3-runtime==4.9.3 typing-extensions==4.6.3 pydantic==1.10.9 requests==2.31.0