lib/loader.py

import torch
from DeepCache import DeepCacheSDHelper
from diffusers import (
    DDIMScheduler,
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionImg2ImgPipeline,
    StableDiffusionPipeline,
)
from diffusers.models import AutoencoderKL, AutoencoderTiny
from diffusers.models.attention_processor import AttnProcessor2_0, IPAdapterAttnProcessor2_0
from torch._dynamo import OptimizedModule

from .upscaler import RealESRGAN

__import__("warnings").filterwarnings("ignore", category=FutureWarning, module="diffusers")


# inspired by ComfyUI
# https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/model_management.py
class Loader:
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(Loader, cls).__new__(cls)
            cls._instance.pipe = None
            cls._instance.upscaler = None
            cls._instance.ip_adapter = None
        return cls._instance

    def _load_upscaler(self, device=None, scale=4):
        same_scale = self.upscaler is not None and self.upscaler.scale == scale
        if scale == 1:
            self.upscaler = None
        if scale > 1 and not same_scale:
            self.upscaler = RealESRGAN(device=device, scale=scale)
            self.upscaler.load_weights()

    def _load_deepcache(self, interval=1):
        has_deepcache = hasattr(self.pipe, "deepcache")
        if has_deepcache and self.pipe.deepcache.params["cache_interval"] == interval:
            return
        if has_deepcache:
            self.pipe.deepcache.disable()
        else:
            self.pipe.deepcache = DeepCacheSDHelper(pipe=self.pipe)
        self.pipe.deepcache.set_params(cache_interval=interval)
        self.pipe.deepcache.enable()

    def _load_freeu(self, freeu=False):
        # https://github.com/huggingface/diffusers/blob/v0.30.0/src/diffusers/models/unets/unet_2d_condition.py
        block = self.pipe.unet.up_blocks[0]
        attrs = ["b1", "b2", "s1", "s2"]
        has_freeu = all(getattr(block, attr, None) is not None for attr in attrs)
        if has_freeu and not freeu:
            print("Disabling FreeU...")
            self.pipe.disable_freeu()
        elif not has_freeu and freeu:
            # https://github.com/ChenyangSi/FreeU
            print("Enabling FreeU...")
            self.pipe.enable_freeu(b1=1.5, b2=1.6, s1=0.9, s2=0.2)

    def _load_ip_adapter(self, ip_adapter=None):
        if self.ip_adapter is None and self.ip_adapter != ip_adapter:
            print(f"Loading IP Adapter: {ip_adapter}...")
            self.pipe.load_ip_adapter(
                "h94/IP-Adapter",
                subfolder="models",
                weight_name=f"ip-adapter-{ip_adapter}_sd15.safetensors",
            )
            # TODO: slider for ip_scale
            self.pipe.set_ip_adapter_scale(0.5)
            self.ip_adapter = ip_adapter

        if self.ip_adapter is not None and ip_adapter is None:
            print("Unloading IP Adapter...")
            if not isinstance(self.pipe, StableDiffusionImg2ImgPipeline):
                self.pipe.image_encoder = None
                self.pipe.register_to_config(image_encoder=[None, None])

            self.pipe.feature_extractor = None
            self.pipe.unet.encoder_hid_proj = None
            self.pipe.unet.config.encoder_hid_dim_type = None
            self.pipe.register_to_config(feature_extractor=[None, None])

            attn_procs = {}
            for name, value in self.pipe.unet.attn_processors.items():
                attn_processor_class = AttnProcessor2_0()  # raises if not torch 2
                attn_procs[name] = (
                    attn_processor_class
                    if isinstance(value, IPAdapterAttnProcessor2_0)
                    else value.__class__()
                )
            self.pipe.unet.set_attn_processor(attn_procs)
            self.pipe.ip_adapter = None

    def _load_vae(self, taesd=False, model_name=None, variant=None):
        vae_type = type(self.pipe.vae)
        is_kl = issubclass(vae_type, (AutoencoderKL, OptimizedModule))
        is_tiny = issubclass(vae_type, AutoencoderTiny)

        # by default all models use KL
        if is_kl and taesd:
            # can't compile tiny VAE
            print("Switching to Tiny VAE...")
            self.pipe.vae = AutoencoderTiny.from_pretrained(
                pretrained_model_name_or_path="madebyollin/taesd",
                torch_dtype=self.pipe.dtype,
            ).to(self.pipe.device)
            return

        if is_tiny and not taesd:
            print("Switching to KL VAE...")
            model = AutoencoderKL.from_pretrained(
                pretrained_model_name_or_path=model_name,
                torch_dtype=self.pipe.dtype,
                subfolder="vae",
                variant=variant,
            ).to(self.pipe.device)
            self.pipe.vae = torch.compile(
                mode="reduce-overhead",
                fullgraph=True,
                model=model,
            )

    def _load_pipeline(self, kind, model, device, **kwargs):
        pipelines = {
            "txt2img": StableDiffusionPipeline,
            "img2img": StableDiffusionImg2ImgPipeline,
        }
        if self.pipe is None:
            self.pipe = pipelines[kind].from_pretrained(model, **kwargs).to(device)
        if not isinstance(self.pipe, pipelines[kind]):
            self.pipe = pipelines[kind].from_pipe(self.pipe).to(device)
            self.ip_adapter = None

    def load(
        self,
        kind,
        ip_adapter,
        model,
        scheduler,
        karras,
        taesd,
        freeu,
        deepcache,
        scale,
        device,
        dtype,
    ):
        model_lower = model.lower()

        schedulers = {
            "DDIM": DDIMScheduler,
            "DEIS 2M": DEISMultistepScheduler,
            "DPM++ 2M": DPMSolverMultistepScheduler,
            "Euler": EulerDiscreteScheduler,
            "Euler a": EulerAncestralDiscreteScheduler,
            "PNDM": PNDMScheduler,
        }

        scheduler_kwargs = {
            "beta_schedule": "scaled_linear",
            "timestep_spacing": "leading",
            "beta_start": 0.00085,
            "beta_end": 0.012,
            "steps_offset": 1,
        }

        if scheduler not in ["DDIM", "Euler a", "PNDM"]:
            scheduler_kwargs["use_karras_sigmas"] = karras

        # https://github.com/huggingface/diffusers/blob/8a3f0c1/scripts/convert_original_stable_diffusion_to_diffusers.py#L939
        if scheduler == "DDIM":
            scheduler_kwargs["clip_sample"] = False
            scheduler_kwargs["set_alpha_to_one"] = False

        # no fp16 variant
        if model_lower not in [
            "sg161222/realistic_vision_v5.1_novae",
            "prompthero/openjourney-v4",
            "linaqruf/anything-v3-1",
        ]:
            variant = "fp16"
        else:
            variant = None

        pipe_kwargs = {
            "scheduler": schedulers[scheduler](**scheduler_kwargs),
            "requires_safety_checker": False,
            "safety_checker": None,
            "torch_dtype": dtype,
            "variant": variant,
        }

        if self.pipe is None:
            print(f"Loading {model_lower} with {'Tiny' if taesd else 'KL'} VAE...")

        self._load_pipeline(kind, model_lower, device, **pipe_kwargs)
        model_name = self.pipe.config._name_or_path
        same_model = model_name.lower() == model_lower
        same_scheduler = isinstance(self.pipe.scheduler, schedulers[scheduler])
        same_karras = (
            not hasattr(self.pipe.scheduler.config, "use_karras_sigmas")
            or self.pipe.scheduler.config.use_karras_sigmas == karras
        )

        if same_model:
            if not same_scheduler:
                print(f"Switching to {scheduler}...")
            if not same_karras:
                print(f"{'Enabling' if karras else 'Disabling'} Karras sigmas...")
            if not same_scheduler or not same_karras:
                self.pipe.scheduler = schedulers[scheduler](**scheduler_kwargs)
        else:
            self.pipe = None
            self._load_pipeline(kind, model_lower, device, **pipe_kwargs)

        self._load_ip_adapter(ip_adapter)
        self._load_vae(taesd, model_lower, variant)
        self._load_freeu(freeu)
        self._load_deepcache(deepcache)
        self._load_upscaler(device, scale)
        torch.cuda.empty_cache()
        return self.pipe, self.upscaler