Rewrite loading and inference

lib/inference.py

@@ -5,153 +5,112 @@ from datetime import datetime
 import torch
 from compel import Compel, DiffusersTextualInversionManager, ReturnedEmbeddingsType
 from compel.prompt_parser import PromptParser
-from huggingface_hub.utils import HFValidationError, RepositoryNotFoundError
-from spaces import GPU
+from gradio import Error, Info, Progress
+from spaces import GPU, config
 
-from .config import Config
-from .loader import Loader
+from .loader import get_loader
 from .logger import Logger
-from .utils import (
-    annotate_image,
-    clear_cuda_cache,
-    resize_image,
-    safe_progress,
-    timer,
-)
-
-
-# Dynamic signature for the GPU duration function
-def gpu_duration(**kwargs):
-    loading = 20
-    duration = 10
-    width = kwargs.get("width", 512)
-    height = kwargs.get("height", 512)
-    scale = kwargs.get("scale", 1)
-    num_images = kwargs.get("num_images", 1)
-    size = width * height
-    if size > 500_000:
-        duration += 5
-    if scale == 4:
-        duration += 5
-    return loading + (duration * num_images)
-
-
-# Request GPU when deployed to Hugging Face
-@GPU(duration=gpu_duration)
+from .utils import annotate_image, cuda_collect, resize_image, timer
+
+
+@GPU
 def generate(
-    positive_prompt,
+    positive_prompt="",
     negative_prompt="",
-    image_prompt=None,
-    control_image_prompt=None,
-    ip_image_prompt=None,
+    image_input=None,
+    controlnet_input=None,
+    ip_adapter_input=None,
     seed=None,
-    model="Lykon/dreamshaper-8",
-    scheduler="DDIM",
-    annotator="canny",
+    model="XpucT/Reliberate",
+    scheduler="UniPC",
+    controlnet_annotator="canny",
     width=512,
     height=512,
     guidance_scale=6.0,
     inference_steps=40,
     denoising_strength=0.8,
-    deepcache=1,
+    deepcache_interval=1,
     scale=1,
     num_images=1,
-    karras=False,
-    ip_face=False,
-    Error=Exception,
-    Info=None,
-    progress=None,
+    use_karras=False,
+    use_ip_adapter_face=False,
+    _=Progress(track_tqdm=True),
 ):
-    start = time.perf_counter()
-    log = Logger("generate")
-    log.info(f"Generating {num_images} image{'s' if num_images > 1 else ''}...")
-
-    if Config.ZERO_GPU:
-        safe_progress(progress, 100, 100, "ZeroGPU init")
-
     if not torch.cuda.is_available():
         raise Error("CUDA not available")
 
-    # https://pytorch.org/docs/stable/generated/torch.manual_seed.html
-    if seed is None or seed < 0:
-        seed = int(datetime.now().timestamp() * 1_000_000) % (2**64)
+    if positive_prompt.strip() == "":
+        raise Error("You must enter a prompt")
 
-    CURRENT_STEP = 0
-    CURRENT_IMAGE = 1
+    start = time.perf_counter()
+    log = Logger("generate")
+    log.info(f"Generating {num_images} image{'s' if num_images > 1 else ''}...")
 
-    KIND = "img2img" if image_prompt is not None else "txt2img"
-    KIND = f"controlnet_{KIND}" if control_image_prompt is not None else KIND
+    KIND = "img2img" if image_input is not None else "txt2img"
+    KIND = f"controlnet_{KIND}" if controlnet_input is not None else KIND
 
     EMBEDDINGS_TYPE = ReturnedEmbeddingsType.LAST_HIDDEN_STATES_NORMALIZED
 
     FAST_NEGATIVE = "<fast_negative>" in negative_prompt
 
-    if ip_image_prompt:
-        IP_ADAPTER = "full-face" if ip_face else "plus"
+    if ip_adapter_input:
+        IP_KIND = "full-face" if use_ip_adapter_face else "plus"
     else:
-        IP_ADAPTER = ""
-
-    # Custom progress bar for multiple images
-    def callback_on_step_end(pipeline, step, timestep, latents):
-        nonlocal CURRENT_STEP, CURRENT_IMAGE
-        if progress is not None:
-            # calculate total steps for img2img based on denoising strength
-            strength = denoising_strength if KIND == "img2img" else 1
-            total_steps = min(int(inference_steps * strength), inference_steps)
-            CURRENT_STEP = step + 1
-            progress(
-                (CURRENT_STEP, total_steps),
-                desc=f"Generating image {CURRENT_IMAGE}/{num_images}",
-            )
-        return latents
+        IP_KIND = ""
 
-    loader = Loader()
+    # ZeroGPU is serverless so you want ephemeral instances
+    # You want a singleton on localhost so the pipeline stays in memory
+    loader = get_loader(singleton=not config.Config.zero_gpu)
     loader.load(
         KIND,
-        IP_ADAPTER,
+        IP_KIND,
         model,
         scheduler,
-        annotator,
-        deepcache,
+        controlnet_annotator,
+        deepcache_interval,
         scale,
-        karras,
-        progress,
+        use_karras,
     )
 
-    if loader.pipe is None:
-        raise Error(f"Error loading {model}")
-
-    pipe = loader.pipe
+    pipeline = loader.pipeline
     upscaler = loader.upscaler
 
+    # Probably a typo in the config
+    if pipeline is None:
+        raise Error(f"Error loading {model}")
+
     # Load fast negative embedding
     if FAST_NEGATIVE:
         embeddings_dir = os.path.abspath(
             os.path.join(os.path.dirname(__file__), "..", "embeddings")
         )
-        pipe.load_textual_inversion(
+        pipeline.load_textual_inversion(
             pretrained_model_name_or_path=f"{embeddings_dir}/fast_negative.pt",
             token="<fast_negative>",
         )
 
     # Embed prompts with weights
     compel = Compel(
-        device=pipe.device,
-        tokenizer=pipe.tokenizer,
+        device=pipeline.device,
+        tokenizer=pipeline.tokenizer,
         truncate_long_prompts=False,
-        text_encoder=pipe.text_encoder,
+        text_encoder=pipeline.text_encoder,
         returned_embeddings_type=EMBEDDINGS_TYPE,
-        dtype_for_device_getter=lambda _: pipe.dtype,
-        textual_inversion_manager=DiffusersTextualInversionManager(pipe),
+        dtype_for_device_getter=lambda _: pipeline.dtype,
+        textual_inversion_manager=DiffusersTextualInversionManager(pipeline),
     )
 
+    # https://pytorch.org/docs/stable/generated/torch.manual_seed.html
+    if seed is None or seed < 0:
+        seed = int(datetime.now().timestamp() * 1_000_000) % (2**64)
+
+    # Increment the seed after each iteration
     images = []
     current_seed = seed
-    safe_progress(progress, 0, num_images, f"Generating image 0/{num_images}")
 
     for i in range(num_images):
         try:
-            generator = torch.Generator(device=pipe.device).manual_seed(current_seed)
+            generator = torch.Generator(device=pipeline.device).manual_seed(current_seed)
             positive_embeds, negative_embeds = compel.pad_conditioning_tensors_to_same_length(
                 [compel(positive_prompt), compel(negative_prompt)]
             )
@@ -169,53 +128,44 @@ def generate(
             "output_type": "np" if scale > 1 else "pil",
         }
 
-        if progress is not None:
-            kwargs["callback_on_step_end"] = callback_on_step_end
-
-        # Resizing so the initial latents are the same size as the generated image
-        if KIND == "img2img":
+        if KIND == "img2img" or KIND == "controlnet_img2img":
             kwargs["strength"] = denoising_strength
-            kwargs["image"] = resize_image(image_prompt, (width, height))
+            kwargs["image"] = resize_image(image_input, (width, height))
 
         if KIND == "controlnet_txt2img":
-            kwargs["image"] = annotate_image(control_image_prompt, annotator)
+            kwargs["image"] = annotate_image(controlnet_input, controlnet_annotator)
 
         if KIND == "controlnet_img2img":
-            kwargs["control_image"] = annotate_image(control_image_prompt, annotator)
+            kwargs["control_image"] = annotate_image(controlnet_input, controlnet_annotator)
 
-        if IP_ADAPTER:
-            kwargs["ip_adapter_image"] = resize_image(ip_image_prompt)
+        if IP_KIND:
+            # No size means preserve aspect ratio
+            kwargs["ip_adapter_image"] = resize_image(ip_adapter_input)
 
         try:
-            image = pipe(**kwargs).images[0]
-            images.append((image, str(current_seed)))
+            image = pipeline(**kwargs).images[0]
+            images.append((image, str(current_seed)))  # tuple with seed for gallery caption
             current_seed += 1
         finally:
             if FAST_NEGATIVE:
-                pipe.unload_textual_inversion()
-
-            CURRENT_STEP = 0
-            CURRENT_IMAGE += 1
+                pipeline.unload_textual_inversion()
 
     # Upscale
     if scale > 1:
-        msg = f"Upscaling {scale}x"
-        with timer(msg, logger=log.info):
-            safe_progress(progress, 0, num_images, desc=msg)
+        with timer(f"Upscaling {num_images} images {scale}x", logger=log.info):
             for i, image in enumerate(images):
                 image = upscaler.predict(image[0])
-                images[i] = image
-                safe_progress(progress, i + 1, num_images, desc=msg)
-
-    # Flush memory after generating
-    clear_cuda_cache()
+                seed = images[i][1]
+                images[i] = (image, seed)  # tuple again
 
     end = time.perf_counter()
     msg = f"Generating {len(images)} image{'s' if len(images) > 1 else ''} took {end - start:.2f}s"
     log.info(msg)
 
-    # Alert if notifier provided
     if Info:
         Info(msg)
 
+    # Flush cache before returning
+    cuda_collect()
+
     return images

lib/loader.py

@@ -1,6 +1,3 @@
-import gc
-from threading import Lock
-
 import torch
 from DeepCache import DeepCacheSDHelper
 from diffusers import ControlNetModel
@@ -9,328 +6,313 @@ from diffusers.models.attention_processor import AttnProcessor2_0, IPAdapterAttn
 from .config import Config
 from .logger import Logger
 from .upscaler import RealESRGAN
-from .utils import clear_cuda_cache, safe_progress, timer
+from .utils import timer
 
 
 class Loader:
-    _instance = None
-    _lock = Lock()
-
-    def __new__(cls):
-        with cls._lock:
-            if cls._instance is None:
-                cls._instance = super().__new__(cls)
-                cls._instance.pipe = None
-                cls._instance.model = None
-                cls._instance.upscaler = None
-                cls._instance.controlnet = None
-                cls._instance.ip_adapter = None
-                cls._instance.log = Logger("Loader")
-        return cls._instance
-
-    def _should_unload_upscaler(self, scale=1):
-        if self.upscaler is not None and self.upscaler.scale != scale:
-            return True
-        return False
-
-    def _should_unload_deepcache(self, interval=1):
-        has_deepcache = hasattr(self.pipe, "deepcache")
-        if has_deepcache and interval == 1:
+    """
+    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.pipe.deepcache.params["cache_interval"] != interval:
+        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, model="", ip_adapter=""):
-        # unload if model changed
-        if self.model and self.model.lower() != model.lower():
+    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 and not ip_adapter:
+        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, kind="", controlnet=""):
+    def should_unload_controlnet(self, pipeline_id="", annotator=""):
         if self.controlnet is None:
             return False
-        if self.controlnet.lower() != controlnet.lower():
+        if self.annotator != annotator:
             return True
-        if not kind.startswith("controlnet_"):
+        if not pipeline_id.startswith("controlnet_"):
             return True
         return False
 
-    def _should_unload_pipeline(self, kind="", model="", controlnet=""):
-        if self.pipe is None:
+    def should_unload_pipeline(self, model=""):
+        if self.pipeline is None:
             return False
-        if self.model.lower() != model.lower():
-            return True
-        if kind == "txt2img" and not isinstance(self.pipe, Config.PIPELINES["txt2img"]):
-            return True
-        if kind == "img2img" and not isinstance(self.pipe, Config.PIPELINES["img2img"]):
-            return True
-        if kind == "controlnet_txt2img" and not isinstance(
-            self.pipe,
-            Config.PIPELINES["controlnet_txt2img"],
-        ):
-            return True
-        if kind == "controlnet_img2img" and not isinstance(
-            self.pipe,
-            Config.PIPELINES["controlnet_img2img"],
-        ):
-            return True
-        if self._should_unload_controlnet(kind, controlnet):
+        if self.model != model:
             return True
         return False
 
-    def _unload_upscaler(self):
-        if self.upscaler is not None:
-            with timer(f"Unloading {self.upscaler.scale}x upscaler", logger=self.log.info):
-                self.upscaler.to("cpu")
-
-    def _unload_deepcache(self):
-        if self.pipe.deepcache is not None:
-            self.log.info("Disabling DeepCache")
-            self.pipe.deepcache.disable()
-            delattr(self.pipe, "deepcache")
-
     # Copied from https://github.com/huggingface/diffusers/blob/v0.28.0/src/diffusers/loaders/ip_adapter.py#L300
-    def _unload_ip_adapter(self):
-        if self.ip_adapter is not None:
-            with timer("Unloading IP-Adapter", logger=self.log.info):
-                if not isinstance(self.pipe, Config.PIPELINES["img2img"]):
-                    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)
-
-    def _unload_pipeline(self):
-        if self.pipe is not None:
-            with timer(f"Unloading {self.model}", logger=self.log.info):
-                self.pipe.to("cpu")
-
-    def _unload(
+    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,
-        kind="",
+        pipeline_id="",
+        ip_adapter_model="",
         model="",
-        controlnet="",
-        ip_adapter="",
-        deepcache=1,
+        controlnet_annotator="",
+        deepcache_interval=1,
         scale=1,
     ):
-        to_unload = []
-        if self._should_unload_deepcache(deepcache):  # remove deepcache first
-            self._unload_deepcache()
-
-        if self._should_unload_upscaler(scale):
-            self._unload_upscaler()
-            to_unload.append("upscaler")
-
-        if self._should_unload_ip_adapter(model, ip_adapter):
-            self._unload_ip_adapter()
-            to_unload.append("ip_adapter")
-
-        if self._should_unload_controlnet(kind, controlnet):
-            to_unload.append("controlnet")
-
-        if self._should_unload_pipeline(kind, model, controlnet):
-            self._unload_pipeline()
-            to_unload.append("model")
-            to_unload.append("pipe")
-
-        # Flush cache and run garbage collector
-        clear_cuda_cache()
-        for component in to_unload:
-            setattr(self, component, None)
-        gc.collect()
-
-    def _should_load_upscaler(self, scale=1):
-        if self.upscaler is None and 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_deepcache(self, interval=1):
-        has_deepcache = hasattr(self.pipe, "deepcache")
-        if not has_deepcache and interval != 1:
+    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_deepcache and self.pipe.deepcache.params["cache_interval"] != interval:
+        if has_karras and self.pipeline.scheduler.config.use_karras_sigmas != use_karras:
             return True
         return False
 
-    def _should_load_ip_adapter(self, ip_adapter=""):
-        if not self.ip_adapter and ip_adapter:
+    def should_load_pipeline(self, pipeline_id=""):
+        if self.pipeline is None:
             return True
-        return False
-
-    def _should_load_pipeline(self):
-        if self.pipe is None:
+        if not isinstance(self.pipeline, Config.PIPELINES[pipeline_id]):
             return True
         return False
 
-    def _load_upscaler(self, scale=1):
-        if self._should_load_upscaler(scale):
-            try:
-                msg = f"Loading {scale}x upscaler"
-                with timer(msg, logger=self.log.info):
-                    self.upscaler = RealESRGAN(scale, device=self.pipe.device)
-                    self.upscaler.load_weights()
-            except Exception as e:
-                self.log.error(f"Error loading {scale}x upscaler: {e}")
-                self.upscaler = None
-
-    def _load_deepcache(self, interval=1):
-        if self._should_load_deepcache(interval):
-            self.log.info("Enabling DeepCache")
-            self.pipe.deepcache = DeepCacheSDHelper(self.pipe)
-            self.pipe.deepcache.set_params(cache_interval=interval)
-            self.pipe.deepcache.enable()
-
-    def _load_ip_adapter(self, ip_adapter=""):
-        if self._should_load_ip_adapter(ip_adapter):
-            msg = "Loading IP-Adapter"
-            with timer(msg, logger=self.log.info):
-                self.pipe.load_ip_adapter(
-                    "h94/IP-Adapter",
-                    subfolder="models",
-                    weight_name=f"ip-adapter-{ip_adapter}_sd15.safetensors",
-                )
-                # 50% works the best
-                self.pipe.set_ip_adapter_scale(0.5)
-                self.ip_adapter = ip_adapter
-
-    def _load_pipeline(
+    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,
-        kind,
+        pipeline_id,
         model,
-        progress,
         **kwargs,
     ):
-        pipeline = Config.PIPELINES[kind]
-        if self._should_load_pipeline():
-            try:
-                with timer(f"Loading {model} ({kind})", logger=self.log.info):
-                    self.model = model
-                    if model.lower() in Config.MODEL_CHECKPOINTS.keys():
-                        self.pipe = pipeline.from_single_file(
-                            f"https://huggingface.co/{model}/{Config.MODEL_CHECKPOINTS[model.lower()]}",
-                            progress,
-                            **kwargs,
-                        ).to("cuda")
-                    else:
-                        self.pipe = pipeline.from_pretrained(model, progress, **kwargs).to("cuda")
-            except Exception as e:
-                self.log.error(f"Error loading {model}: {e}")
-                self.model = None
-                self.pipe = None
-                return
-        if not isinstance(self.pipe, pipeline):
-            self.pipe = pipeline.from_pipe(self.pipe).to("cuda")
-        if self.pipe is not None:
-            self.pipe.set_progress_bar_config(disable=progress is not None)
+        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,
-        kind,
-        ip_adapter,
+        pipeline_id,
+        ip_adapter_model,
         model,
         scheduler,
-        annotator,
-        deepcache,
+        controlnet_annotator,
+        deepcache_interval,
         scale,
-        karras,
-        progress,
+        use_karras,
     ):
+        Scheduler = Config.SCHEDULERS[scheduler]
+
         scheduler_kwargs = {
-            "beta_schedule": "scaled_linear",
-            "timestep_spacing": "leading",
             "beta_start": 0.00085,
             "beta_end": 0.012,
+            "beta_schedule": "scaled_linear",
+            "timestep_spacing": "leading",
             "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
+        if scheduler not in ["Euler a"]:
+            scheduler_kwargs["use_karras_sigmas"] = use_karras
 
-        pipe_kwargs = {
+        pipeline_kwargs = {
+            "torch_dtype": torch.float16,  # defaults to fp32
             "safety_checker": None,
             "requires_safety_checker": False,
-            "scheduler": Config.SCHEDULERS[scheduler](**scheduler_kwargs),
+            "scheduler": Scheduler(**scheduler_kwargs),
         }
 
-        # diffusers fp16 variant
-        if model.lower() not in Config.MODEL_CHECKPOINTS.keys():
-            pipe_kwargs["variant"] = "fp16"
+        # Single-file models don't need a variant
+        if model not in Config.SINGLE_FILE_MODELS:
+            pipeline_kwargs["variant"] = "fp16"
         else:
-            pipe_kwargs["variant"] = None
-
-        # converts to fp32 by default
-        pipe_kwargs["torch_dtype"] = torch.float16
+            pipeline_kwargs["variant"] = None
+
+        # Prepare state for loading checks
+        self.unload_all(
+            pipeline_id,
+            ip_adapter_model,
+            model,
+            controlnet_annotator,
+            deepcache_interval,
+            scale,
+        )
 
-        # config maps the repo to the ID: canny -> lllyasviel/control_sd15_canny
-        if kind.startswith("controlnet_"):
-            pipe_kwargs["controlnet"] = ControlNetModel.from_pretrained(
-                Config.ANNOTATORS[annotator],
-                torch_dtype=torch.float16,
-                variant="fp16",
-            )
-            self.controlnet = annotator
+        # Load controlnet model before pipeline
+        if self.should_load_controlnet(pipeline_id):
+            self.load_controlnet(controlnet_annotator)
 
-        self._unload(kind, model, annotator, ip_adapter, deepcache, scale)
-        self._load_pipeline(kind, model, progress, **pipe_kwargs)
+        if self.should_load_pipeline(pipeline_id):
+            self.load_pipeline(pipeline_id, model, **pipeline_kwargs)
 
-        # error loading model
-        if self.pipe is None:
-            return
+        if self.should_load_scheduler(Scheduler, use_karras):
+            self.load_scheduler(Scheduler, use_karras, **scheduler_kwargs)
 
-        same_scheduler = isinstance(self.pipe.scheduler, Config.SCHEDULERS[scheduler])
-        same_karras = (
-            not hasattr(self.pipe.scheduler.config, "use_karras_sigmas")
-            or self.pipe.scheduler.config.use_karras_sigmas == karras
-        )
+        if self.should_load_deepcache(deepcache_interval):
+            self.load_deepcache(deepcache_interval)
 
-        # same model, different scheduler
-        if self.model.lower() == model.lower():
-            if not same_scheduler:
-                self.log.info(f"Enabling {scheduler} scheduler")
-            if not same_karras:
-                self.log.info(f"{'Enabling' if karras else 'Disabling'} Karras sigmas")
-            if not same_scheduler or not same_karras:
-                self.pipe.scheduler = Config.SCHEDULERS[scheduler](**scheduler_kwargs)
-
-        CURRENT_STEP = 1
-        TOTAL_STEPS = sum(
-            [
-                self._should_load_deepcache(deepcache),
-                self._should_load_ip_adapter(ip_adapter),
-                self._should_load_upscaler(scale),
-            ]
-        )
+        if self.should_load_ip_adapter(ip_adapter_model):
+            self.load_ip_adapter(ip_adapter_model)
 
-        desc = "Configuring pipeline"
-        if self._should_load_deepcache(deepcache):
-            self._load_deepcache(deepcache)
-            safe_progress(progress, CURRENT_STEP, TOTAL_STEPS, desc)
-            CURRENT_STEP += 1
+        if self.should_load_upscaler(scale):
+            self.load_upscaler(scale)
 
-        if self._should_load_ip_adapter(ip_adapter):
-            self._load_ip_adapter(ip_adapter)
-            safe_progress(progress, CURRENT_STEP, TOTAL_STEPS, desc)
-            CURRENT_STEP += 1
 
-        if self._should_load_upscaler(scale):
-            self._load_upscaler(scale)
-            safe_progress(progress, CURRENT_STEP, TOTAL_STEPS, desc)
+# 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