initial diffusers merge into dev

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

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
 

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

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}")

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,

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

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

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"[<a href='{uri}' target='_blank'>{label}</a>]"
@@ -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"),

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)}'
 
 

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

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