from typing import List, Union, Tuple, Callable, Optional, Dict import random import inspect import functools import os from functools import reduce import os.path as osp from packaging.version import Version, parse import logging import importlib import sys import operator as op from requests import HTTPError from pathlib import Path if sys.version_info < (3, 8): import importlib_metadata else: import importlib.metadata as importlib_metadata import accelerate import torch.nn as nn import torch from einops import rearrange from torchvision.transforms.functional import pil_to_tensor from torch import Tensor import numpy as np from PIL import Image from huggingface_hub.utils import ( EntryNotFoundError, RepositoryNotFoundError, RevisionNotFoundError, validate_hf_hub_args, ) from huggingface_hub import hf_hub_download import safetensors import torchvision.transforms.functional as tv_functional from .io_utils import load_image, load_exec_list logger = logging.getLogger(__name__) # pylint: disable=invalid-name _is_gguf_available = importlib.util.find_spec("gguf") is not None if _is_gguf_available: try: _gguf_version = importlib_metadata.version("gguf") logger.debug(f"Successfully import gguf version {_gguf_version}") except importlib_metadata.PackageNotFoundError: _is_gguf_available = False _torch_available = importlib.util.find_spec("torch") is not None if _torch_available: try: _torch_version = importlib_metadata.version("torch") logger.info(f"PyTorch version {_torch_version} available.") except importlib_metadata.PackageNotFoundError: _torch_available = False STR_OPERATION_TO_FUNC = {">": op.gt, ">=": op.ge, "==": op.eq, "!=": op.ne, "<=": op.le, "<": op.lt} DEPRECATED_REVISION_ARGS = ["fp16", "non-ema"] HUGGINGFACE_CO_RESOLVE_ENDPOINT = os.environ.get("HF_ENDPOINT", "https://huggingface.co") WEIGHTS_NAME = "diffusion_pytorch_model.bin" SAFETENSORS_WEIGHTS_NAME = "diffusion_pytorch_model.safetensors" def is_gguf_available(): return _is_gguf_available def is_torch_available(): return _torch_available def load_gguf_checkpoint(gguf_checkpoint_path, return_tensors=False): """ Load a GGUF file and return a dictionary of parsed parameters containing tensors, the parsed tokenizer and config attributes. Args: gguf_checkpoint_path (`str`): The path the to GGUF file to load return_tensors (`bool`, defaults to `True`): Whether to read the tensors from the file and return them. Not doing so is faster and only loads the metadata in memory. """ if is_gguf_available() and is_torch_available(): import gguf from gguf import GGUFReader from ..quantizers.gguf.utils import SUPPORTED_GGUF_QUANT_TYPES, GGUFParameter else: logger.error( "Loading a GGUF checkpoint in PyTorch, requires both PyTorch and GGUF>=0.10.0 to be installed. Please see " "https://pytorch.org/ and https://github.com/ggerganov/llama.cpp/tree/master/gguf-py for installation instructions." ) raise ImportError("Please install torch and gguf>=0.10.0 to load a GGUF checkpoint in PyTorch.") reader = GGUFReader(gguf_checkpoint_path) parsed_parameters = {} for tensor in reader.tensors: name = tensor.name quant_type = tensor.tensor_type # if the tensor is a torch supported dtype do not use GGUFParameter is_gguf_quant = quant_type not in [gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16] if is_gguf_quant and quant_type not in SUPPORTED_GGUF_QUANT_TYPES: _supported_quants_str = "\n".join([str(type) for type in SUPPORTED_GGUF_QUANT_TYPES]) raise ValueError( ( f"{name} has a quantization type: {str(quant_type)} which is unsupported." "\n\nCurrently the following quantization types are supported: \n\n" f"{_supported_quants_str}" "\n\nTo request support for this quantization type please open an issue here: https://github.com/huggingface/diffusers" ) ) weights = torch.from_numpy(tensor.data.copy()) parsed_parameters[name] = GGUFParameter(weights, quant_type=quant_type) if is_gguf_quant else weights return parsed_parameters def make_grid(x_max=.9693, y_max=.9375, n_x=28, n_y=16, device='cpu', dtype=torch.float32, flatten=False, target_size=None): ''' returns: grid: shape (n_y, n_x, 3) ''' grid_y, grid_x = torch.meshgrid( torch.linspace(-y_max, y_max, n_y), torch.linspace(-x_max, x_max, n_x), indexing="ij", ) v = torch.ones_like(grid_x) grid = torch.stack([grid_x, grid_y, v], dim=-1) if target_size is not None: grid[..., 0] = grid[..., 0] * target_size[0] / 2 grid[..., 1] = grid[..., 1] * target_size[1] / 2 if flatten: grid = grid.reshape(1, -1, 3) return grid.to(device=device, dtype=dtype) def tensor2img(t: torch.Tensor, output_type='numpy', denormalize=False, mean = 0., std = 255., from_mode: str = 'RGB', convert_mode: str = None, src_dim_order: str = 'chw', dtype=np.uint8, clip=(0, 255)) -> Union[Image.Image, np.array]: def _check_denormalize_params(values, num_channels): if isinstance(values, (int, float, np.ScalarType)): return values else: if isinstance(values, list): values = np.array(values) elif isinstance(values, torch.Tensor): values = values.to(device='cpu', dtype=torch.float32).numpy() else: raise Exception(f'invalid normalizing values: {values}') if len(values) > num_channels: values = values[:num_channels] assert len(values) == num_channels values = values.reshape((1, 1, -1)) return values t = t.detach().to(device='cpu', dtype=torch.float32).squeeze().numpy() if t.ndim == 3: if src_dim_order == 'chw': t = rearrange(t, 'c h w -> h w c') c = t.shape[-1] else: assert t.ndim == 2, "t.ndim should be 2 or 3 after squeeze" c = 1 if denormalize: t = (t * _check_denormalize_params(std, c)) + _check_denormalize_params(mean, c) if clip is not None: t = np.clip(t, clip[0], clip[1]) image = t.astype(dtype) if output_type == 'pil': if len(image.shape) == 2: from_mode = 'L' image = Image.fromarray(image, mode=from_mode) if convert_mode is not None: image = image.convert(convert_mode) else: image = np.ascontiguousarray(image) assert output_type == 'numpy' return image _IMG2TENSOR_IMGTYPE = (Image.Image, np.ndarray, str) _IMG2TENSOR_DIMORDER = ('bchw', 'chw', 'hwc') def img2tensor(img: Union[Image.Image, np.ndarray, str, torch.Tensor], normalize = False, mean = 0., std = 255., dim_order: str = 'bchw', dtype=torch.float32, device: str = 'cpu', imread_mode='RGB') -> Tensor: def _check_normalize_values(values, num_channels): if isinstance(values, tuple): values = list(values) elif isinstance(values, (int, float, np.ScalarType)): values = [values] * num_channels else: assert isinstance(values, (np.ndarray, list)) if len(values) > num_channels: values = values[:num_channels] assert len(values) == num_channels return values assert isinstance(img, _IMG2TENSOR_IMGTYPE) assert dim_order in _IMG2TENSOR_DIMORDER if isinstance(img, str): img = load_image(img, mode=imread_mode) if isinstance(img, Image.Image): img = pil_to_tensor(img) if dim_order == 'bchw': img = img.unsqueeze(0) elif dim_order == 'hwc': img = img.permute((1, 2, 0)) else: if img.ndim == 2: img = img[..., None] else: assert img.ndim == 3 if dim_order == 'bchw': img = rearrange(img, 'h w c -> c h w')[None, ...] elif dim_order == 'chw': img = rearrange(img, 'h w c -> c h w') img = torch.from_numpy(np.ascontiguousarray(img)) img = img.to(device=device, dtype=dtype) if normalize: if dim_order == 'bchw': c = img.shape[1] elif dim_order == 'chw': c = img.shape[0] else: c = img.shape[2] if mean is not None and std is not None: mean = _check_normalize_values(mean, c) std = _check_normalize_values(std, c) img = tv_functional.normalize(img, mean=mean, std=std) return img def convert_tensor(t, dtype=torch.float32, device='cpu'): if isinstance(t, List) or isinstance(t, tuple) or isinstance(t, float) or isinstance(t, int): return torch.tensor(t, dtype=dtype, device=device) elif isinstance(t, np.ndarray): return torch.from_numpy(t).to(dtype=dtype, device=device) elif isinstance(t, torch.Tensor): return t.to(device=device, dtype=dtype) else: raise TORCH_DTYPE_DICT = {'fp32': torch.float32, 'fp16': torch.float16, 'bf16': torch.bfloat16} def get_torch_dtype(torch_dtype): if isinstance(torch_dtype, str): return TORCH_DTYPE_DICT[torch_dtype] return torch_dtype def seed_everything(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) def accelerate_should_save(accelerator): from accelerate import Accelerator, DistributedType # https://github.com/huggingface/diffusers/issues/2606#issuecomment-1463193509 return accelerator.distributed_type == DistributedType.DEEPSPEED \ or accelerator.is_main_process def run_mp_func(rank, ws, args_split_by_rank, target_func, *args): file_list = args[0] devices = args[1] args = args[2:] arg_list = [] if args_split_by_rank is not None: func_input_args = inspect.getfullargspec(target_func) _args_split_by_rank = set(args_split_by_rank) for arg_idx, argname in enumerate(func_input_args.args): if arg_idx < 2: # skip file_list continue arg_val = args[arg_idx - 2] if isinstance(arg_val, str) and argname in _args_split_by_rank: if ',' in arg_val: arg_val = arg_val.split(',') arg_val = arg_val[rank] arg_list.append(arg_val) else: arg_list = args torch.cuda.set_device(devices[rank]) os.environ['CUDA_VISIBLE_DEVICES'] = str(devices[rank]) file_list = load_exec_list(file_list, rank, ws, check_exist=False) # with open(f'test_{rank}.txt', 'w') as f: # f.write('\n'.join(file_list)) arg_list = [file_list, devices] + arg_list target_func(*arg_list) def torch_mp_wrapper(args_split_by_rank: list = None): def decorator(func): @functools.wraps(func) def wrapper(**kwargs): func_name = func.__name__ assert func_name.endswith('_mp') target_func = getattr(inspect.getmodule(func), func_name[:-3]) func_input_args = inspect.getfullargspec(target_func) target_func_kwargs = {} for arg_idx, argname in enumerate(func_input_args.args): target_func_kwargs[argname] = kwargs[argname] kwargs = target_func_kwargs assert 'devices' in kwargs and 'file_list' in kwargs kwargs_keys = list(kwargs.keys()) assert kwargs_keys[1] == 'devices' and kwargs_keys[0] == 'file_list' file_list = kwargs['file_list'] if 'seed' in kwargs: seed = kwargs['seed'] else: seed = 0 seed_everything(seed) file_list = load_exec_list(file_list, check_exist=False) if kwargs['devices'] == None: devices = [0] else: devices = kwargs['devices'].split(',') devices = [int(d) for d in devices] if len(devices) <= 1: kwargs['file_list'] = file_list target_func(**kwargs) else: random.shuffle(file_list) import torch.multiprocessing as mp ws = len(devices) arg_list = [v for v in kwargs.values()] arg_list[1] = devices arg_list[0] = file_list func_args = [ws, args_split_by_rank, target_func] + arg_list # func_args = tuple(func_args) mp.spawn(run_mp_func, nprocs=ws, args=func_args) return wrapper return decorator def extract_into_tensor(a, t, x_shape): b, *_ = t.shape out = a.gather(-1, t) return out.reshape(b, *((1,) * (len(x_shape) - 1))) def get_module_by_name(module: Union[torch.Tensor, nn.Module], access_string: str): """Retrieve a module nested in another by its access string. Works even when there is a Sequential in the module. https://discuss.pytorch.org/t/how-to-access-to-a-layer-by-module-name/83797/8 """ names = access_string.split(sep='.') return reduce(getattr, names, module) # This function was copied from: https://github.com/huggingface/accelerate/blob/874c4967d94badd24f893064cc3bef45f57cadf7/src/accelerate/utils/versions.py#L319 def compare_versions(library_or_version: Union[str, Version], operation: str, requirement_version: str): """ Compares a library version to some requirement using a given operation. Args: library_or_version (`str` or `packaging.version.Version`): A library name or a version to check. operation (`str`): A string representation of an operator, such as `">"` or `"<="`. requirement_version (`str`): The version to compare the library version against """ if operation not in STR_OPERATION_TO_FUNC.keys(): raise ValueError(f"`operation` must be one of {list(STR_OPERATION_TO_FUNC.keys())}, received {operation}") operation = STR_OPERATION_TO_FUNC[operation] if isinstance(library_or_version, str): library_or_version = parse(importlib_metadata.version(library_or_version)) return operation(library_or_version, parse(requirement_version)) # This function was copied from: https://github.com/huggingface/accelerate/blob/874c4967d94badd24f893064cc3bef45f57cadf7/src/accelerate/utils/versions.py#L338 def is_torch_version(operation: str, version: str): """ Compares the current PyTorch version to a given reference with an operation. Args: operation (`str`): A string representation of an operator, such as `">"` or `"<="` version (`str`): A string version of PyTorch """ return compare_versions(parse(_torch_version), operation, version) SAFETENSORS_FILE_EXTENSION = "safetensors" GGUF_FILE_EXTENSION = "gguf" def load_state_dict(checkpoint_file: Union[str, os.PathLike], variant: Optional[str] = None): """ Reads a checkpoint file, returning properly formatted errors if they arise. """ # TODO: We merge the sharded checkpoints in case we're doing quantization. We can revisit this change # when refactoring the _merge_sharded_checkpoints() method later. if isinstance(checkpoint_file, dict): return checkpoint_file try: file_extension = os.path.basename(checkpoint_file).split(".")[-1] if file_extension == SAFETENSORS_FILE_EXTENSION: return safetensors.torch.load_file(checkpoint_file, device="cpu") elif file_extension == GGUF_FILE_EXTENSION: return load_gguf_checkpoint(checkpoint_file) else: weights_only_kwarg = {"weights_only": True} if is_torch_version(">=", "1.13") else {} return torch.load( checkpoint_file, map_location="cpu", **weights_only_kwarg, ) except Exception as e: try: with open(checkpoint_file) as f: if f.read().startswith("version"): raise OSError( "You seem to have cloned a repository without having git-lfs installed. Please install " "git-lfs and run `git lfs install` followed by `git lfs pull` in the folder " "you cloned." ) else: raise ValueError( f"Unable to locate the file {checkpoint_file} which is necessary to load this pretrained " "model. Make sure you have saved the model properly." ) from e except (UnicodeDecodeError, ValueError): raise OSError( f"Unable to load weights from checkpoint file for '{checkpoint_file}' " f"at '{checkpoint_file}'. " ) @validate_hf_hub_args def _get_model_file( pretrained_model_name_or_path: Union[str, Path], *, weights_name: str, subfolder: Optional[str] = None, cache_dir: Optional[str] = None, force_download: bool = False, proxies: Optional[Dict] = None, local_files_only: bool = False, token: Optional[str] = None, user_agent: Optional[Union[Dict, str]] = None, revision: Optional[str] = None, commit_hash: Optional[str] = None, ): pretrained_model_name_or_path = str(pretrained_model_name_or_path) if os.path.isfile(pretrained_model_name_or_path): return pretrained_model_name_or_path elif os.path.isdir(pretrained_model_name_or_path): if os.path.isfile(os.path.join(pretrained_model_name_or_path, weights_name)): # Load from a PyTorch checkpoint model_file = os.path.join(pretrained_model_name_or_path, weights_name) return model_file elif subfolder is not None and os.path.isfile( os.path.join(pretrained_model_name_or_path, subfolder, weights_name) ): model_file = os.path.join(pretrained_model_name_or_path, subfolder, weights_name) return model_file else: raise EnvironmentError( f"Error no file named {weights_name} found in directory {pretrained_model_name_or_path}." ) else: try: # 2. Load model file as usual model_file = hf_hub_download( pretrained_model_name_or_path, filename=weights_name, cache_dir=cache_dir, force_download=force_download, proxies=proxies, local_files_only=local_files_only, token=token, user_agent=user_agent, subfolder=subfolder, revision=revision or commit_hash, ) return model_file except RepositoryNotFoundError as e: raise EnvironmentError( f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier " "listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a " "token having permission to this repo with `token` or log in with `huggingface-cli " "login`." ) from e except RevisionNotFoundError as e: raise EnvironmentError( f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for " "this model name. Check the model page at " f"'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions." ) from e except EntryNotFoundError as e: raise EnvironmentError( f"{pretrained_model_name_or_path} does not appear to have a file named {weights_name}." ) from e except HTTPError as e: raise EnvironmentError( f"There was a specific connection error when trying to load {pretrained_model_name_or_path}:\n{e}" ) from e except ValueError as e: raise EnvironmentError( f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it" f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a" f" directory containing a file named {weights_name} or" " \nCheckout your internet connection or see how to run the library in" " offline mode at 'https://huggingface.co/docs/diffusers/installation#offline-mode'." ) from e except EnvironmentError as e: raise EnvironmentError( f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it from " "'https://huggingface.co/models', make sure you don't have a local directory with the same name. " f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory " f"containing a file named {weights_name}" ) from e def init_model_from_pretrained( pretrained_model_name_or_path: str, module_cls, subfolder=None, model_args = None, weights_name=None, patch_state_dict_func: Callable = None, download_from_hf=True, device='cpu', pass_statedict_to_model_init=False ): ''' skip unnecessary param init for faster model creation, allow mismatch Args: module_cls (Callable): model class or model build function ''' if download_from_hf: model_file = _get_model_file(pretrained_model_name_or_path=pretrained_model_name_or_path, subfolder=subfolder, weights_name=weights_name) state_dict = load_state_dict(model_file) else: if osp.exists(pretrained_model_name_or_path): state_dict = load_state_dict(pretrained_model_name_or_path) else: state_dict = torch.hub.load_state_dict_from_url(pretrained_model_name_or_path) if model_args is None: model_args = {} with accelerate.init_empty_weights(include_buffers=False): if pass_statedict_to_model_init: model, state_dict = module_cls(**model_args, state_dict=state_dict) else: model = module_cls(**model_args) if patch_state_dict_func is not None: _state_dict = patch_state_dict_func(model, state_dict) if _state_dict is not None: state_dict = _state_dict if 'state_dict' in state_dict: state_dict = state_dict['state_dict'] incompatible_keys = model.load_state_dict(state_dict, strict=False, assign=True) missing_keys_set = set(incompatible_keys.missing_keys) for k in incompatible_keys.missing_keys: # assert k.endswith('.bias') or k.endswith('.weight') if k.endswith('.bias'): if k.replace('.bias', '.weight') in missing_keys_set: continue module = get_module_by_name(model, k.replace('.weight', '').replace('.bias', '')) if isinstance(module, torch.nn.Parameter): # print(module.data) module.data = torch.randn(module.data.size(), device='cpu') else: module.to_empty(device="cpu") module.reset_parameters() if device != 'cpu': model = model.to(device=device) return model def image2np(image: Union[torch.Tensor, Image.Image, str, ], denormalize=True): if isinstance(image, torch.Tensor): image = tensor2img(image, mean=127.5, std=127.5, normalize=denormalize) elif isinstance(image, Image.Image): image = np.array(image) elif isinstance(image, np.ndarray): pass elif isinstance(image, str): image = load_image(image, output_type='numpy') else: raise Exception(f'invalid image type: {type(image)}') return image def fix_params(model): for name, param in model.named_parameters(): param.requires_grad = False def zero_module(module): """ Zero out the parameters of a module and return it. """ for p in module.parameters(): p.detach().zero_() return module