import torch
from DeepCache import DeepCacheSDHelper
from diffusers import ControlNetModel
from diffusers.models.attention_processor import AttnProcessor2_0, IPAdapterAttnProcessor2_0
from .config import Config
from .logger import Logger
from .upscaler import RealESRGAN
from .utils import timer
class Loader:
"""
A lazy-loading resource manager for Stable Diffusion pipelines. Lifecycles are managed by
comparing the current state with desired. Can be used as a singleton when created by the
`get_loader()` helper.
Usage:
loader = get_loader(singleton=True)
loader.load(
pipeline_id="controlnet_txt2img",
ip_adapter_model="full-face",
model="XpucT/Reliberate",
scheduler="UniPC",
controlnet_annotator="canny",
deepcache_interval=2,
scale=2,
use_karras=True
)
"""
def __init__(self):
self.model = ""
self.pipeline = None
self.upscaler = None
self.controlnet = None
self.annotator = "" # controlnet annotator (canny)
self.ip_adapter = "" # ip-adapter kind (full-face or plus)
self.log = Logger("Loader")
def should_unload_upscaler(self, scale=1):
return self.upscaler is not None and self.upscaler.scale != scale
def should_unload_deepcache(self, cache_interval=1):
has_deepcache = hasattr(self.pipeline, "deepcache")
if has_deepcache and cache_interval == 1:
return True
if has_deepcache and self.pipeline.deepcache.params["cache_interval"] != cache_interval:
# Unload if interval is different so it can be reloaded
return True
return False
def should_unload_ip_adapter(self, ip_adapter_model=""):
if not self.ip_adapter:
return False
if not ip_adapter_model:
return True
if self.ip_adapter != ip_adapter_model:
# Unload if model is different so it can be reloaded
return True
return False
def should_unload_controlnet(self, pipeline_id="", annotator=""):
if self.controlnet is None:
return False
if self.annotator != annotator:
return True
if not pipeline_id.startswith("controlnet_"):
return True
return False
def should_unload_pipeline(self, model=""):
if self.pipeline is None:
return False
if self.model != model:
return True
return False
# Copied from https://github.com/huggingface/diffusers/blob/v0.28.0/src/diffusers/loaders/ip_adapter.py#L300
def unload_ip_adapter(self):
# Remove the image encoder if text-to-image
if isinstance(self.pipeline, Config.PIPELINES["txt2img"]):
self.pipeline.image_encoder = None
self.pipeline.register_to_config(image_encoder=[None, None])
# Remove hidden projection layer added by IP-Adapter
self.pipeline.unet.encoder_hid_proj = None
self.pipeline.unet.config.encoder_hid_dim_type = None
# Remove the feature extractor
self.pipeline.feature_extractor = None
self.pipeline.register_to_config(feature_extractor=[None, None])
# Replace the custom attention processors with defaults
attn_procs = {}
for name, value in self.pipeline.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.pipeline.unet.set_attn_processor(attn_procs)
self.ip_adapter = ""
def unload_all(
self,
pipeline_id="",
ip_adapter_model="",
model="",
controlnet_annotator="",
deepcache_interval=1,
scale=1,
):
if self.should_unload_deepcache(deepcache_interval): # remove deepcache first
self.log.info("Disabling DeepCache")
self.pipeline.deepcache.disable()
delattr(self.pipeline, "deepcache")
if self.should_unload_ip_adapter(ip_adapter_model):
self.log.info("Unloading IP-Adapter")
self.unload_ip_adapter()
if self.should_unload_controlnet(pipeline_id, controlnet_annotator):
self.log.info("Unloading ControlNet")
self.controlnet = None
self.annotator = ""
if self.should_unload_upscaler(scale):
self.log.info("Unloading upscaler")
self.upscaler = None
if self.should_unload_pipeline(model):
self.log.info("Unloading pipeline")
self.pipeline = None
self.model = ""
def should_load_upscaler(self, scale=1):
return self.upscaler is None and scale > 1
def should_load_deepcache(self, cache_interval=1):
has_deepcache = hasattr(self.pipeline, "deepcache")
if not has_deepcache and cache_interval > 1:
return True
return False
def should_load_controlnet(self, pipeline_id=""):
return self.controlnet is None and pipeline_id.startswith("controlnet_")
def should_load_ip_adapter(self, ip_adapter_model=""):
has_ip_adapter = (
hasattr(self.pipeline.unet, "encoder_hid_proj")
and self.pipeline.unet.config.encoder_hid_dim_type == "ip_image_proj"
)
return not has_ip_adapter and ip_adapter_model != ""
def should_load_scheduler(self, cls, use_karras=False):
has_karras = hasattr(self.pipeline.scheduler.config, "use_karras_sigmas")
if not isinstance(self.pipeline.scheduler, cls):
return True
if has_karras and self.pipeline.scheduler.config.use_karras_sigmas != use_karras:
return True
return False
def should_load_pipeline(self, pipeline_id=""):
if self.pipeline is None:
return True
if not isinstance(self.pipeline, Config.PIPELINES[pipeline_id]):
return True
return False
def load_upscaler(self, scale=1):
with timer(f"Loading {scale}x upscaler", logger=self.log.info):
self.upscaler = RealESRGAN(scale, device=self.pipeline.device)
self.upscaler.load_weights()
def load_deepcache(self, cache_interval=1):
self.log.info(f"Enabling DeepCache interval {cache_interval}")
self.pipeline.deepcache = DeepCacheSDHelper(self.pipeline)
self.pipeline.deepcache.set_params(cache_interval=cache_interval)
self.pipeline.deepcache.enable()
def load_controlnet(self, controlnet_annotator):
with timer("Loading ControlNet", logger=self.log.info):
self.controlnet = ControlNetModel.from_pretrained(
Config.ANNOTATORS[controlnet_annotator],
variant="fp16",
torch_dtype=torch.float16,
)
self.annotator = controlnet_annotator
def load_ip_adapter(self, ip_adapter_model=""):
with timer("Loading IP-Adapter", logger=self.log.info):
self.pipeline.load_ip_adapter(
"h94/IP-Adapter",
subfolder="models",
weight_name=f"ip-adapter-{ip_adapter_model}_sd15.safetensors",
)
self.pipeline.set_ip_adapter_scale(0.5) # 50% works the best
self.ip_adapter = ip_adapter_model
def load_scheduler(self, cls, use_karras=False, **kwargs):
self.log.info(f"Loading {cls.__name__}{' with Karras' if use_karras else ''}")
self.pipeline.scheduler = cls(**kwargs)
def load_pipeline(
self,
pipeline_id,
model,
**kwargs,
):
Pipeline = Config.PIPELINES[pipeline_id]
# Load from scratch
if self.pipeline is None:
with timer(f"Loading {model} ({pipeline_id})", logger=self.log.info):
if self.controlnet is not None:
kwargs["controlnet"] = self.controlnet
if model in Config.SINGLE_FILE_MODELS:
checkpoint = Config.HF_REPOS[model][0]
self.pipeline = Pipeline.from_single_file(
f"https://huggingface.co/{model}/{checkpoint}",
**kwargs,
).to("cuda")
else:
self.pipeline = Pipeline.from_pretrained(model, **kwargs).to("cuda")
# Change to a different one
else:
with timer(f"Changing pipeline to {pipeline_id}", logger=self.log.info):
kwargs = {}
if self.controlnet is not None:
kwargs["controlnet"] = self.controlnet
self.pipeline = Pipeline.from_pipe(
self.pipeline,
**kwargs,
).to("cuda")
# Update model and disable terminal progress bars
self.model = model
self.pipeline.set_progress_bar_config(disable=True)
def load(
self,
pipeline_id,
ip_adapter_model,
model,
scheduler,
controlnet_annotator,
deepcache_interval,
scale,
use_karras,
):
Scheduler = Config.SCHEDULERS[scheduler]
scheduler_kwargs = {
"beta_start": 0.00085,
"beta_end": 0.012,
"beta_schedule": "scaled_linear",
"timestep_spacing": "leading",
"steps_offset": 1,
}
if scheduler not in ["Euler a"]:
scheduler_kwargs["use_karras_sigmas"] = use_karras
pipeline_kwargs = {
"torch_dtype": torch.float16, # defaults to fp32
"safety_checker": None,
"requires_safety_checker": False,
"scheduler": Scheduler(**scheduler_kwargs),
}
# Single-file models don't need a variant
if model not in Config.SINGLE_FILE_MODELS:
pipeline_kwargs["variant"] = "fp16"
else:
pipeline_kwargs["variant"] = None
# Prepare state for loading checks
self.unload_all(
pipeline_id,
ip_adapter_model,
model,
controlnet_annotator,
deepcache_interval,
scale,
)
# Load controlnet model before pipeline
if self.should_load_controlnet(pipeline_id):
self.load_controlnet(controlnet_annotator)
if self.should_load_pipeline(pipeline_id):
self.load_pipeline(pipeline_id, model, **pipeline_kwargs)
if self.should_load_scheduler(Scheduler, use_karras):
self.load_scheduler(Scheduler, use_karras, **scheduler_kwargs)
if self.should_load_deepcache(deepcache_interval):
self.load_deepcache(deepcache_interval)
if self.should_load_ip_adapter(ip_adapter_model):
self.load_ip_adapter(ip_adapter_model)
if self.should_load_upscaler(scale):
self.load_upscaler(scale)
# Get a singleton or a new instance of the Loader
def get_loader(singleton=False):
if not singleton:
return Loader()
else:
if not hasattr(get_loader, "_instance"):
get_loader._instance = Loader()
assert isinstance(get_loader._instance, Loader)
return get_loader._instance