Remove custom pipelines

lib/config.py

@@ -1,35 +1,20 @@
-import os
-from importlib import import_module
-from importlib.util import find_spec
 from types import SimpleNamespace
 from warnings import filterwarnings
 
 from diffusers import (
-    DDIMScheduler,
     DEISMultistepScheduler,
     DPMSolverMultistepScheduler,
     EulerAncestralDiscreteScheduler,
     EulerDiscreteScheduler,
-    PNDMScheduler,
+    StableDiffusionControlNetImg2ImgPipeline,
+    StableDiffusionControlNetPipeline,
+    StableDiffusionImg2ImgPipeline,
+    StableDiffusionPipeline,
     UniPCMultistepScheduler,
 )
 from diffusers.utils import logging as diffusers_logging
 from transformers import logging as transformers_logging
 
-from .pipelines import (
-    CustomStableDiffusionControlNetImg2ImgPipeline,
-    CustomStableDiffusionControlNetPipeline,
-    CustomStableDiffusionImg2ImgPipeline,
-    CustomStableDiffusionPipeline,
-)
-
-# Improved GPU handling and progress bars; set before importing spaces
-os.environ["ZEROGPU_V2"] = "1"
-
-# Errors if enabled and not installed
-if find_spec("hf_transfer"):
-    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
-
 filterwarnings("ignore", category=FutureWarning, module="diffusers")
 filterwarnings("ignore", category=FutureWarning, module="transformers")
 
@@ -37,7 +22,7 @@ diffusers_logging.set_verbosity_error()
 transformers_logging.set_verbosity_error()
 
 # Standard Stable Diffusion 1.5 file structure
-sd_files = [
+_sd_files = [
     "feature_extractor/preprocessor_config.json",
     "safety_checker/config.json",
     "scheduler/scheduler_config.json",
@@ -56,36 +41,29 @@ sd_files = [
 
 # Using namespace instead of dataclass for simplicity
 Config = SimpleNamespace(
-    HF_TOKEN=os.environ.get("HF_TOKEN", None),
-    ZERO_GPU=import_module("spaces").config.Config.zero_gpu,
-    # TODO: fix model config redundancy
-    HF_MODELS={
-        # downloaded on startup
+    PIPELINES={
+        "txt2img": StableDiffusionPipeline,
+        "img2img": StableDiffusionImg2ImgPipeline,
+        "controlnet_txt2img": StableDiffusionControlNetPipeline,
+        "controlnet_img2img": StableDiffusionControlNetImg2ImgPipeline,
+    },
+    HF_REPOS={
         "ai-forever/Real-ESRGAN": ["RealESRGAN_x2.pth", "RealESRGAN_x4.pth"],
         "cyberdelia/CyberRealistic": ["CyberRealistic_V5_FP16.safetensors"],
         "fluently/Fluently-v4": ["Fluently-v4.safetensors"],
+        "h94/IP-Adapter": [
+            "models/ip-adapter-full-face_sd15.safetensors",
+            "models/ip-adapter-plus_sd15.safetensors",
+            "models/image_encoder/model.safetensors",
+        ],
         "lllyasviel/control_v11p_sd15_canny": ["diffusion_pytorch_model.fp16.safetensors"],
-        "Lykon/dreamshaper-8": [*sd_files],
+        "Lykon/dreamshaper-8": _sd_files,
         "s6yx/ReV_Animated": ["rev_1.2.2/rev_1.2.2-fp16.safetensors"],
         "SG161222/Realistic_Vision_V5.1_noVAE": ["Realistic_Vision_V5.1_fp16-no-ema.safetensors"],
-        "stable-diffusion-v1-5/stable-diffusion-v1-5": [*sd_files],
+        "stable-diffusion-v1-5/stable-diffusion-v1-5": _sd_files,
         "XpucT/Deliberate": ["Deliberate_v6.safetensors"],
         "XpucT/Reliberate": ["Reliberate_v3.safetensors"],
     },
-    MONO_FONTS=["monospace"],
-    SANS_FONTS=[
-        "sans-serif",
-        "Apple Color Emoji",
-        "Segoe UI Emoji",
-        "Segoe UI Symbol",
-        "Noto Color Emoji",
-    ],
-    PIPELINES={
-        "txt2img": CustomStableDiffusionPipeline,
-        "img2img": CustomStableDiffusionImg2ImgPipeline,
-        "controlnet_txt2img": CustomStableDiffusionControlNetPipeline,
-        "controlnet_img2img": CustomStableDiffusionControlNetImg2ImgPipeline,
-    },
     MODEL="XpucT/Reliberate",
     MODELS=[
         "cyberdelia/CyberRealistic",
@@ -97,25 +75,21 @@ Config = SimpleNamespace(
         "XpucT/Deliberate",
         "XpucT/Reliberate",
     ],
-    # Single-file model weights
-    MODEL_CHECKPOINTS={
-        # keep keys lowercase for case-insensitive matching in the loader
-        "cyberdelia/cyberrealistic": "CyberRealistic_V5_FP16.safetensors",
-        "fluently/fluently-v4": "Fluently-v4.safetensors",
-        "s6yx/rev_animated": "rev_1.2.2/rev_1.2.2-fp16.safetensors",
-        "sg161222/realistic_vision_v5.1_novae": "Realistic_Vision_V5.1_fp16-no-ema.safetensors",
-        "xpuct/deliberate": "Deliberate_v6.safetensors",
-        "xpuct/reliberate": "Reliberate_v3.safetensors",
-    },
-    SCHEDULER="UniPC 2M",
+    SINGLE_FILE_MODELS=[
+        "cyberdelia/CyberRealistic",
+        "fluently/Fluently-v4",
+        "s6yx/ReV_Animated",
+        "SG161222/Realistic_Vision_V5.1_noVAE",
+        "XpucT/Deliberate",
+        "XpucT/Reliberate",
+    ],
+    SCHEDULER="UniPC",
     SCHEDULERS={
-        "DDIM": DDIMScheduler,
-        "DEIS 2M": DEISMultistepScheduler,
+        "DEIS": DEISMultistepScheduler,
         "DPM++ 2M": DPMSolverMultistepScheduler,
         "Euler": EulerDiscreteScheduler,
         "Euler a": EulerAncestralDiscreteScheduler,
-        "PNDM": PNDMScheduler,
-        "UniPC 2M": UniPCMultistepScheduler,
+        "UniPC": UniPCMultistepScheduler,
     },
     ANNOTATOR="canny",
     ANNOTATORS={
@@ -124,7 +98,6 @@ Config = SimpleNamespace(
     WIDTH=512,
     HEIGHT=512,
     NUM_IMAGES=1,
-    SEED=-1,
     GUIDANCE_SCALE=6,
     INFERENCE_STEPS=40,
     DENOISING_STRENGTH=0.8,

lib/pipelines.py

@@ -1,241 +0,0 @@
-import os
-from importlib import import_module
-
-from diffusers import (
-    StableDiffusionControlNetImg2ImgPipeline,
-    StableDiffusionControlNetPipeline,
-    StableDiffusionImg2ImgPipeline,
-    StableDiffusionPipeline,
-)
-from diffusers.loaders.single_file import (
-    SINGLE_FILE_OPTIONAL_COMPONENTS,
-    load_single_file_sub_model,
-)
-from diffusers.loaders.single_file_utils import fetch_diffusers_config, load_single_file_checkpoint
-from diffusers.models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT
-from diffusers.pipelines.pipeline_loading_utils import (
-    ALL_IMPORTABLE_CLASSES,
-    _get_pipeline_class,
-    load_sub_model,
-)
-from diffusers.utils import logging
-from huggingface_hub import snapshot_download
-from huggingface_hub.utils import validate_hf_hub_args
-
-
-class CustomDiffusionMixin:
-    r"""
-    Overrides DiffusionPipeline `from_pretrained` and `from_single_file` methods to allow passing a progress function.
-    """
-
-    # Copied from https://github.com/huggingface/diffusers/blob/v0.30.3/src/diffusers/pipelines/pipeline_utils.py#L480
-    @classmethod
-    @validate_hf_hub_args
-    def from_pretrained(cls, pretrained_model_name_or_path, progress=None, **kwargs):
-        torch_dtype = kwargs.pop("torch_dtype", None)
-        variant = kwargs.pop("variant", None)
-        token = kwargs.pop("token", None)
-
-        # download the checkpoints and configs
-        cached_folder = cls.download(
-            pretrained_model_name_or_path,
-            variant=variant,
-            token=token,
-            **kwargs,
-        )
-
-        # pop out "_ignore_files" as it is only needed for download
-        config_dict = cls.load_config(cached_folder)
-        config_dict.pop("_ignore_files", None)
-
-        # Define which model components should load variants.
-        # We retrieve the information by matching whether variant model checkpoints exist in the subfolders.
-        # Example: `diffusion_pytorch_model.safetensors` -> `diffusion_pytorch_model.fp16.safetensors` with variant being `"fp16"`.
-        model_variants = {}
-        if variant is not None:
-            for folder in os.listdir(cached_folder):
-                folder_path = os.path.join(cached_folder, folder)
-                is_folder = os.path.isdir(folder_path) and folder in config_dict
-                variant_exists = is_folder and any(
-                    p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)
-                )
-                if variant_exists:
-                    model_variants[folder] = variant
-
-        # load the pipeline class
-        pipeline_class = _get_pipeline_class(cls, config=config_dict)
-
-        # define expected modules given pipeline signature and define non-None initialized modules (=`init_kwargs`)
-        expected_modules, optional_kwargs = cls._get_signature_keys(pipeline_class)
-        passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs}
-        passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs}
-
-        def load_module(name, value):
-            if value[0] is None:
-                return False
-            if name in passed_class_obj and passed_class_obj[name] is None:
-                return False
-            return True
-
-        init_dict, _, _ = pipeline_class.extract_init_dict(config_dict, **kwargs)
-        init_kwargs = {
-            k: init_dict.pop(k)
-            for k in optional_kwargs
-            if k in init_dict and k not in pipeline_class._optional_components
-        }
-        init_kwargs = {**init_kwargs, **passed_pipe_kwargs}
-        init_dict = {k: v for k, v in init_dict.items() if load_module(k, v)}
-
-        # load each module in the pipeline
-        pipelines = import_module("diffusers.pipelines")
-        tqdm = logging.tqdm if progress is None else progress.tqdm
-        for name, (library_name, class_name) in tqdm(
-            sorted(init_dict.items()),
-            desc="Loading pipeline components",
-        ):
-            # use passed sub model or load class_name from library_name
-            loaded_sub_model = None
-            if name in passed_class_obj:
-                # passed as an argument like "scheduler"
-                loaded_sub_model = passed_class_obj[name]
-            else:
-                loaded_sub_model = load_sub_model(
-                    library_name=library_name,
-                    class_name=class_name,
-                    importable_classes=ALL_IMPORTABLE_CLASSES,
-                    pipelines=pipelines,
-                    is_pipeline_module=hasattr(pipelines, library_name),
-                    pipeline_class=pipeline_class,
-                    torch_dtype=torch_dtype,
-                    provider=None,
-                    sess_options=None,
-                    device_map=None,
-                    max_memory=None,
-                    offload_folder=None,
-                    offload_state_dict=False,
-                    model_variants=model_variants,
-                    name=name,
-                    from_flax=False,
-                    variant=variant,
-                    low_cpu_mem_usage=_LOW_CPU_MEM_USAGE_DEFAULT,
-                    cached_folder=cached_folder,
-                )
-            init_kwargs[name] = loaded_sub_model
-
-        # potentially add passed objects if expected
-        missing_modules = set(expected_modules) - set(init_kwargs.keys())
-        if len(missing_modules) > 0:
-            for module in missing_modules:
-                init_kwargs[module] = passed_class_obj.get(module, None)
-
-        # instantiate the pipeline
-        model = pipeline_class(**init_kwargs)
-
-        # save where the model was instantiated from
-        model.register_to_config(_name_or_path=pretrained_model_name_or_path)
-        return model
-
-    # Copied from https://github.com/huggingface/diffusers/blob/v0.30.3/src/diffusers/loaders/single_file.py#L270
-    @classmethod
-    @validate_hf_hub_args
-    def from_single_file(cls, pretrained_model_link_or_path, progress=None, **kwargs):
-        token = kwargs.pop("token", None)
-        torch_dtype = kwargs.pop("torch_dtype", None)
-
-        # load the pipeline class
-        pipeline_class = _get_pipeline_class(cls, config=None)
-        checkpoint = load_single_file_checkpoint(pretrained_model_link_or_path, token=token)
-
-        config = fetch_diffusers_config(checkpoint)
-        default_pretrained_model_config_name = config["pretrained_model_name_or_path"]
-
-        # attempt to download the config files for the pipeline
-        cached_model_config_path = snapshot_download(
-            default_pretrained_model_config_name,
-            token=token,
-            allow_patterns=["**/*.json", "*.json", "*.txt", "**/*.txt", "**/*.model"],
-        )
-
-        # pop out "_ignore_files" as it is only needed for download
-        config_dict = pipeline_class.load_config(cached_model_config_path)
-        config_dict.pop("_ignore_files", None)
-
-        # define expected modules given pipeline signature and define non-None initialized modules (=`init_kwargs`)
-        expected_modules, optional_kwargs = pipeline_class._get_signature_keys(cls)
-        passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs}
-        passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs}
-
-        def load_module(name, value):
-            if value[0] is None:
-                return False
-            if name in passed_class_obj and passed_class_obj[name] is None:
-                return False
-            if name in SINGLE_FILE_OPTIONAL_COMPONENTS:
-                return False
-            return True
-
-        init_dict, _, _ = pipeline_class.extract_init_dict(config_dict, **kwargs)
-        init_kwargs = {k: init_dict.pop(k) for k in optional_kwargs if k in init_dict}
-        init_kwargs = {**init_kwargs, **passed_pipe_kwargs}
-        init_dict = {k: v for k, v in init_dict.items() if load_module(k, v)}
-
-        # load each module in the pipeline
-        pipelines = import_module("diffusers.pipelines")
-        tqdm = logging.tqdm if progress is None else progress.tqdm
-        for name, (library_name, class_name) in tqdm(
-            sorted(init_dict.items()),
-            desc="Loading pipeline components",
-        ):
-            # use passed sub model or load class_name from library_name
-            loaded_sub_model = None
-            if name in passed_class_obj:
-                # passed as an argument like "scheduler"
-                loaded_sub_model = passed_class_obj[name]
-            else:
-                loaded_sub_model = load_single_file_sub_model(
-                    library_name=library_name,
-                    class_name=class_name,
-                    name=name,
-                    checkpoint=checkpoint,
-                    is_pipeline_module=hasattr(pipelines, library_name),
-                    cached_model_config_path=cached_model_config_path,
-                    pipelines=pipelines,
-                    torch_dtype=torch_dtype,
-                    **kwargs,
-                )
-            init_kwargs[name] = loaded_sub_model
-
-        # potentially add passed objects if expected
-        missing_modules = set(expected_modules) - set(init_kwargs.keys())
-        if len(missing_modules) > 0:
-            for module in missing_modules:
-                init_kwargs[module] = passed_class_obj.get(module, None)
-
-        # instantiate the pipeline
-        pipe = pipeline_class(**init_kwargs)
-
-        # save where the model was instantiated from
-        pipe.register_to_config(_name_or_path=pretrained_model_link_or_path)
-        return pipe
-
-
-class CustomStableDiffusionPipeline(CustomDiffusionMixin, StableDiffusionPipeline):
-    pass
-
-
-class CustomStableDiffusionImg2ImgPipeline(CustomDiffusionMixin, StableDiffusionImg2ImgPipeline):
-    pass
-
-
-class CustomStableDiffusionControlNetPipeline(
-    CustomDiffusionMixin,
-    StableDiffusionControlNetPipeline,
-):
-    pass
-
-
-class CustomStableDiffusionControlNetImg2ImgPipeline(
-    CustomDiffusionMixin,
-    StableDiffusionControlNetImg2ImgPipeline,
-):
-    pass