|
20 | 20 | from torch.nn import functional as F |
21 | 21 |
|
22 | 22 | from diffusers.configuration_utils import ConfigMixin, register_to_config |
23 | | -from diffusers.loaders import FromOriginalControlnetMixin |
| 23 | +# from diffusers.loaders import FromOriginalControlnetMixin |
24 | 24 | from diffusers.utils import BaseOutput, logging |
25 | 25 | from diffusers.models.attention_processor import ( |
26 | 26 | ADDED_KV_ATTENTION_PROCESSORS, |
|
35 | 35 | from .unet_3d_blocks import UNetMidBlockSpatioTemporal, get_down_block, get_up_block |
36 | 36 | from diffusers.models import UNetSpatioTemporalConditionModel |
37 | 37 |
|
| 38 | +from huggingface_hub import hf_hub_download |
| 39 | +from io import BytesIO |
| 40 | +import os |
| 41 | +import requests |
| 42 | +from pathlib import Path |
| 43 | + |
| 44 | +from huggingface_hub.constants import HF_HUB_OFFLINE, HF_HOME |
| 45 | +DIFFUSERS_CACHE = os.path.join(HF_HOME, "diffusers") |
38 | 46 |
|
39 | 47 | logger = logging.get_logger(__name__) # pylint: disable=invalid-name |
40 | 48 |
|
| 49 | +class FromOriginalControlnetMixin: |
| 50 | + @classmethod |
| 51 | + def from_single_file(cls, pretrained_model_link_or_path, **kwargs): |
| 52 | + r""" |
| 53 | + Instantiate a [`ControlNetModel`] from pretrained controlnet weights saved in the original `.ckpt` or |
| 54 | + `.safetensors` format. The pipeline is set in evaluation mode (`model.eval()`) by default. |
| 55 | +
|
| 56 | + Parameters: |
| 57 | + pretrained_model_link_or_path (`str` or `os.PathLike`, *optional*): |
| 58 | + Can be either: |
| 59 | + - A link to the `.ckpt` file (for example |
| 60 | + `"https://huggingface.co/<repo_id>/blob/main/<path_to_file>.ckpt"`) on the Hub. |
| 61 | + - A path to a *file* containing all pipeline weights. |
| 62 | + torch_dtype (`str` or `torch.dtype`, *optional*): |
| 63 | + Override the default `torch.dtype` and load the model with another dtype. If `"auto"` is passed, the |
| 64 | + dtype is automatically derived from the model's weights. |
| 65 | + force_download (`bool`, *optional*, defaults to `False`): |
| 66 | + Whether or not to force the (re-)download of the model weights and configuration files, overriding the |
| 67 | + cached versions if they exist. |
| 68 | + cache_dir (`Union[str, os.PathLike]`, *optional*): |
| 69 | + Path to a directory where a downloaded pretrained model configuration is cached if the standard cache |
| 70 | + is not used. |
| 71 | + resume_download (`bool`, *optional*, defaults to `False`): |
| 72 | + Whether or not to resume downloading the model weights and configuration files. If set to `False`, any |
| 73 | + incompletely downloaded files are deleted. |
| 74 | + proxies (`Dict[str, str]`, *optional*): |
| 75 | + A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', |
| 76 | + 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. |
| 77 | + local_files_only (`bool`, *optional*, defaults to `False`): |
| 78 | + Whether to only load local model weights and configuration files or not. If set to True, the model |
| 79 | + won't be downloaded from the Hub. |
| 80 | + use_auth_token (`str` or *bool*, *optional*): |
| 81 | + The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from |
| 82 | + `diffusers-cli login` (stored in `~/.huggingface`) is used. |
| 83 | + revision (`str`, *optional*, defaults to `"main"`): |
| 84 | + The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier |
| 85 | + allowed by Git. |
| 86 | + use_safetensors (`bool`, *optional*, defaults to `None`): |
| 87 | + If set to `None`, the safetensors weights are downloaded if they're available **and** if the |
| 88 | + safetensors library is installed. If set to `True`, the model is forcibly loaded from safetensors |
| 89 | + weights. If set to `False`, safetensors weights are not loaded. |
| 90 | + image_size (`int`, *optional*, defaults to 512): |
| 91 | + The image size the model was trained on. Use 512 for all Stable Diffusion v1 models and the Stable |
| 92 | + Diffusion v2 base model. Use 768 for Stable Diffusion v2. |
| 93 | + upcast_attention (`bool`, *optional*, defaults to `None`): |
| 94 | + Whether the attention computation should always be upcasted. |
| 95 | + kwargs (remaining dictionary of keyword arguments, *optional*): |
| 96 | + Can be used to overwrite load and saveable variables (for example the pipeline components of the |
| 97 | + specific pipeline class). The overwritten components are directly passed to the pipelines `__init__` |
| 98 | + method. See example below for more information. |
| 99 | +
|
| 100 | + Examples: |
| 101 | +
|
| 102 | + ```py |
| 103 | + from diffusers import StableDiffusionControlNetPipeline, ControlNetModel |
| 104 | +
|
| 105 | + url = "https://huggingface.co/lllyasviel/ControlNet-v1-1/blob/main/control_v11p_sd15_canny.pth" # can also be a local path |
| 106 | + model = ControlNetModel.from_single_file(url) |
| 107 | +
|
| 108 | + url = "https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/v1-5-pruned.safetensors" # can also be a local path |
| 109 | + pipe = StableDiffusionControlNetPipeline.from_single_file(url, controlnet=controlnet) |
| 110 | + ``` |
| 111 | + """ |
| 112 | + # import here to avoid circular dependency |
| 113 | + from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_controlnet_from_original_ckpt |
| 114 | + |
| 115 | + config_file = kwargs.pop("config_file", None) |
| 116 | + cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE) |
| 117 | + resume_download = kwargs.pop("resume_download", False) |
| 118 | + force_download = kwargs.pop("force_download", False) |
| 119 | + proxies = kwargs.pop("proxies", None) |
| 120 | + local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE) |
| 121 | + use_auth_token = kwargs.pop("use_auth_token", None) |
| 122 | + num_in_channels = kwargs.pop("num_in_channels", None) |
| 123 | + use_linear_projection = kwargs.pop("use_linear_projection", None) |
| 124 | + revision = kwargs.pop("revision", None) |
| 125 | + extract_ema = kwargs.pop("extract_ema", False) |
| 126 | + image_size = kwargs.pop("image_size", None) |
| 127 | + upcast_attention = kwargs.pop("upcast_attention", None) |
| 128 | + |
| 129 | + torch_dtype = kwargs.pop("torch_dtype", None) |
| 130 | + |
| 131 | + use_safetensors = kwargs.pop("use_safetensors", None) |
| 132 | + |
| 133 | + file_extension = pretrained_model_link_or_path.rsplit(".", 1)[-1] |
| 134 | + from_safetensors = file_extension == "safetensors" |
| 135 | + |
| 136 | + if from_safetensors and use_safetensors is False: |
| 137 | + raise ValueError("Make sure to install `safetensors` with `pip install safetensors`.") |
| 138 | + |
| 139 | + # remove huggingface url |
| 140 | + for prefix in ["https://huggingface.co/", "huggingface.co/", "hf.co/", "https://hf.co/"]: |
| 141 | + if pretrained_model_link_or_path.startswith(prefix): |
| 142 | + pretrained_model_link_or_path = pretrained_model_link_or_path[len(prefix) :] |
| 143 | + |
| 144 | + # Code based on diffusers.pipelines.pipeline_utils.DiffusionPipeline.from_pretrained |
| 145 | + ckpt_path = Path(pretrained_model_link_or_path) |
| 146 | + if not ckpt_path.is_file(): |
| 147 | + # get repo_id and (potentially nested) file path of ckpt in repo |
| 148 | + repo_id = "/".join(ckpt_path.parts[:2]) |
| 149 | + file_path = "/".join(ckpt_path.parts[2:]) |
| 150 | + |
| 151 | + if file_path.startswith("blob/"): |
| 152 | + file_path = file_path[len("blob/") :] |
| 153 | + |
| 154 | + if file_path.startswith("main/"): |
| 155 | + file_path = file_path[len("main/") :] |
| 156 | + |
| 157 | + pretrained_model_link_or_path = hf_hub_download( |
| 158 | + repo_id, |
| 159 | + filename=file_path, |
| 160 | + cache_dir=cache_dir, |
| 161 | + resume_download=resume_download, |
| 162 | + proxies=proxies, |
| 163 | + local_files_only=local_files_only, |
| 164 | + use_auth_token=use_auth_token, |
| 165 | + revision=revision, |
| 166 | + force_download=force_download, |
| 167 | + ) |
| 168 | + |
| 169 | + if config_file is None: |
| 170 | + config_url = "https://raw.githubusercontent.com/lllyasviel/ControlNet/main/models/cldm_v15.yaml" |
| 171 | + config_file = BytesIO(requests.get(config_url).content) |
| 172 | + |
| 173 | + image_size = image_size or 512 |
| 174 | + |
| 175 | + controlnet = download_controlnet_from_original_ckpt( |
| 176 | + pretrained_model_link_or_path, |
| 177 | + original_config_file=config_file, |
| 178 | + image_size=image_size, |
| 179 | + extract_ema=extract_ema, |
| 180 | + num_in_channels=num_in_channels, |
| 181 | + upcast_attention=upcast_attention, |
| 182 | + from_safetensors=from_safetensors, |
| 183 | + use_linear_projection=use_linear_projection, |
| 184 | + ) |
| 185 | + |
| 186 | + if torch_dtype is not None: |
| 187 | + controlnet.to(dtype=torch_dtype) |
| 188 | + |
| 189 | + return controlnet |
41 | 190 |
|
42 | 191 | @dataclass |
43 | 192 | class ControlNetOutput(BaseOutput): |
|
0 commit comments