declare-lab/tango2

1# coding=utf-82# Copyright 2023 The HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""Utilities to dynamically load objects from the Hub."""16 17import importlib18import inspect19import json20import os21import re22import shutil23import sys24from distutils.version import StrictVersion25from pathlib import Path26from typing import Dict, Optional, Union27from urllib import request28 29from huggingface_hub import HfFolder, cached_download, hf_hub_download, model_info30 31from .. import __version__32from . import DIFFUSERS_DYNAMIC_MODULE_NAME, HF_MODULES_CACHE, logging33 34 35COMMUNITY_PIPELINES_URL = (36    "https://raw.githubusercontent.com/huggingface/diffusers/{revision}/examples/community/{pipeline}.py"37)38 39 40logger = logging.get_logger(__name__)  # pylint: disable=invalid-name41 42 43def get_diffusers_versions():44    url = "https://pypi.org/pypi/diffusers/json"45    releases = json.loads(request.urlopen(url).read())["releases"].keys()46    return sorted(releases, key=StrictVersion)47 48 49def init_hf_modules():50    """51    Creates the cache directory for modules with an init, and adds it to the Python path.52    """53    # This function has already been executed if HF_MODULES_CACHE already is in the Python path.54    if HF_MODULES_CACHE in sys.path:55        return56 57    sys.path.append(HF_MODULES_CACHE)58    os.makedirs(HF_MODULES_CACHE, exist_ok=True)59    init_path = Path(HF_MODULES_CACHE) / "__init__.py"60    if not init_path.exists():61        init_path.touch()62 63 64def create_dynamic_module(name: Union[str, os.PathLike]):65    """66    Creates a dynamic module in the cache directory for modules.67    """68    init_hf_modules()69    dynamic_module_path = Path(HF_MODULES_CACHE) / name70    # If the parent module does not exist yet, recursively create it.71    if not dynamic_module_path.parent.exists():72        create_dynamic_module(dynamic_module_path.parent)73    os.makedirs(dynamic_module_path, exist_ok=True)74    init_path = dynamic_module_path / "__init__.py"75    if not init_path.exists():76        init_path.touch()77 78 79def get_relative_imports(module_file):80    """81    Get the list of modules that are relatively imported in a module file.82 83    Args:84        module_file (`str` or `os.PathLike`): The module file to inspect.85    """86    with open(module_file, "r", encoding="utf-8") as f:87        content = f.read()88 89    # Imports of the form `import .xxx`90    relative_imports = re.findall("^\s*import\s+\.(\S+)\s*$", content, flags=re.MULTILINE)91    # Imports of the form `from .xxx import yyy`92    relative_imports += re.findall("^\s*from\s+\.(\S+)\s+import", content, flags=re.MULTILINE)93    # Unique-ify94    return list(set(relative_imports))95 96 97def get_relative_import_files(module_file):98    """99    Get the list of all files that are needed for a given module. Note that this function recurses through the relative100    imports (if a imports b and b imports c, it will return module files for b and c).101 102    Args:103        module_file (`str` or `os.PathLike`): The module file to inspect.104    """105    no_change = False106    files_to_check = [module_file]107    all_relative_imports = []108 109    # Let's recurse through all relative imports110    while not no_change:111        new_imports = []112        for f in files_to_check:113            new_imports.extend(get_relative_imports(f))114 115        module_path = Path(module_file).parent116        new_import_files = [str(module_path / m) for m in new_imports]117        new_import_files = [f for f in new_import_files if f not in all_relative_imports]118        files_to_check = [f"{f}.py" for f in new_import_files]119 120        no_change = len(new_import_files) == 0121        all_relative_imports.extend(files_to_check)122 123    return all_relative_imports124 125 126def check_imports(filename):127    """128    Check if the current Python environment contains all the libraries that are imported in a file.129    """130    with open(filename, "r", encoding="utf-8") as f:131        content = f.read()132 133    # Imports of the form `import xxx`134    imports = re.findall("^\s*import\s+(\S+)\s*$", content, flags=re.MULTILINE)135    # Imports of the form `from xxx import yyy`136    imports += re.findall("^\s*from\s+(\S+)\s+import", content, flags=re.MULTILINE)137    # Only keep the top-level module138    imports = [imp.split(".")[0] for imp in imports if not imp.startswith(".")]139 140    # Unique-ify and test we got them all141    imports = list(set(imports))142    missing_packages = []143    for imp in imports:144        try:145            importlib.import_module(imp)146        except ImportError:147            missing_packages.append(imp)148 149    if len(missing_packages) > 0:150        raise ImportError(151            "This modeling file requires the following packages that were not found in your environment: "152            f"{', '.join(missing_packages)}. Run `pip install {' '.join(missing_packages)}`"153        )154 155    return get_relative_imports(filename)156 157 158def get_class_in_module(class_name, module_path):159    """160    Import a module on the cache directory for modules and extract a class from it.161    """162    module_path = module_path.replace(os.path.sep, ".")163    module = importlib.import_module(module_path)164 165    if class_name is None:166        return find_pipeline_class(module)167    return getattr(module, class_name)168 169 170def find_pipeline_class(loaded_module):171    """172    Retrieve pipeline class that inherits from `DiffusionPipeline`. Note that there has to be exactly one class173    inheriting from `DiffusionPipeline`.174    """175    from ..pipelines import DiffusionPipeline176 177    cls_members = dict(inspect.getmembers(loaded_module, inspect.isclass))178 179    pipeline_class = None180    for cls_name, cls in cls_members.items():181        if (182            cls_name != DiffusionPipeline.__name__183            and issubclass(cls, DiffusionPipeline)184            and cls.__module__.split(".")[0] != "diffusers"185        ):186            if pipeline_class is not None:187                raise ValueError(188                    f"Multiple classes that inherit from {DiffusionPipeline.__name__} have been found:"189                    f" {pipeline_class.__name__}, and {cls_name}. Please make sure to define only one in"190                    f" {loaded_module}."191                )192            pipeline_class = cls193 194    return pipeline_class195 196 197def get_cached_module_file(198    pretrained_model_name_or_path: Union[str, os.PathLike],199    module_file: str,200    cache_dir: Optional[Union[str, os.PathLike]] = None,201    force_download: bool = False,202    resume_download: bool = False,203    proxies: Optional[Dict[str, str]] = None,204    use_auth_token: Optional[Union[bool, str]] = None,205    revision: Optional[str] = None,206    local_files_only: bool = False,207):208    """209    Prepares Downloads a module from a local folder or a distant repo and returns its path inside the cached210    Transformers module.211 212    Args:213        pretrained_model_name_or_path (`str` or `os.PathLike`):214            This can be either:215 216            - a string, the *model id* of a pretrained model configuration hosted inside a model repo on217              huggingface.co. Valid model ids can be located at the root-level, like `bert-base-uncased`, or namespaced218              under a user or organization name, like `dbmdz/bert-base-german-cased`.219            - a path to a *directory* containing a configuration file saved using the220              [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.221 222        module_file (`str`):223            The name of the module file containing the class to look for.224        cache_dir (`str` or `os.PathLike`, *optional*):225            Path to a directory in which a downloaded pretrained model configuration should be cached if the standard226            cache should not be used.227        force_download (`bool`, *optional*, defaults to `False`):228            Whether or not to force to (re-)download the configuration files and override the cached versions if they229            exist.230        resume_download (`bool`, *optional*, defaults to `False`):231            Whether or not to delete incompletely received file. Attempts to resume the download if such a file exists.232        proxies (`Dict[str, str]`, *optional*):233            A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',234            'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.235        use_auth_token (`str` or *bool*, *optional*):236            The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated237            when running `transformers-cli login` (stored in `~/.huggingface`).238        revision (`str`, *optional*, defaults to `"main"`):239            The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a240            git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any241            identifier allowed by git.242        local_files_only (`bool`, *optional*, defaults to `False`):243            If `True`, will only try to load the tokenizer configuration from local files.244 245    <Tip>246 247    You may pass a token in `use_auth_token` if you are not logged in (`huggingface-cli long`) and want to use private248    or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models).249 250    </Tip>251 252    Returns:253        `str`: The path to the module inside the cache.254    """255    # Download and cache module_file from the repo `pretrained_model_name_or_path` of grab it if it's a local file.256    pretrained_model_name_or_path = str(pretrained_model_name_or_path)257 258    module_file_or_url = os.path.join(pretrained_model_name_or_path, module_file)259 260    if os.path.isfile(module_file_or_url):261        resolved_module_file = module_file_or_url262        submodule = "local"263    elif pretrained_model_name_or_path.count("/") == 0:264        available_versions = get_diffusers_versions()265        # cut ".dev0"266        latest_version = "v" + ".".join(__version__.split(".")[:3])267 268        # retrieve github version that matches269        if revision is None:270            revision = latest_version if latest_version in available_versions else "main"271            logger.info(f"Defaulting to latest_version: {revision}.")272        elif revision in available_versions:273            revision = f"v{revision}"274        elif revision == "main":275            revision = revision276        else:277            raise ValueError(278                f"`custom_revision`: {revision} does not exist. Please make sure to choose one of"279                f" {', '.join(available_versions + ['main'])}."280            )281 282        # community pipeline on GitHub283        github_url = COMMUNITY_PIPELINES_URL.format(revision=revision, pipeline=pretrained_model_name_or_path)284        try:285            resolved_module_file = cached_download(286                github_url,287                cache_dir=cache_dir,288                force_download=force_download,289                proxies=proxies,290                resume_download=resume_download,291                local_files_only=local_files_only,292                use_auth_token=False,293            )294            submodule = "git"295            module_file = pretrained_model_name_or_path + ".py"296        except EnvironmentError:297            logger.error(f"Could not locate the {module_file} inside {pretrained_model_name_or_path}.")298            raise299    else:300        try:301            # Load from URL or cache if already cached302            resolved_module_file = hf_hub_download(303                pretrained_model_name_or_path,304                module_file,305                cache_dir=cache_dir,306                force_download=force_download,307                proxies=proxies,308                resume_download=resume_download,309                local_files_only=local_files_only,310                use_auth_token=use_auth_token,311            )312            submodule = os.path.join("local", "--".join(pretrained_model_name_or_path.split("/")))313        except EnvironmentError:314            logger.error(f"Could not locate the {module_file} inside {pretrained_model_name_or_path}.")315            raise316 317    # Check we have all the requirements in our environment318    modules_needed = check_imports(resolved_module_file)319 320    # Now we move the module inside our cached dynamic modules.321    full_submodule = DIFFUSERS_DYNAMIC_MODULE_NAME + os.path.sep + submodule322    create_dynamic_module(full_submodule)323    submodule_path = Path(HF_MODULES_CACHE) / full_submodule324    if submodule == "local" or submodule == "git":325        # We always copy local files (we could hash the file to see if there was a change, and give them the name of326        # that hash, to only copy when there is a modification but it seems overkill for now).327        # The only reason we do the copy is to avoid putting too many folders in sys.path.328        shutil.copy(resolved_module_file, submodule_path / module_file)329        for module_needed in modules_needed:330            module_needed = f"{module_needed}.py"331            shutil.copy(os.path.join(pretrained_model_name_or_path, module_needed), submodule_path / module_needed)332    else:333        # Get the commit hash334        # TODO: we will get this info in the etag soon, so retrieve it from there and not here.335        if isinstance(use_auth_token, str):336            token = use_auth_token337        elif use_auth_token is True:338            token = HfFolder.get_token()339        else:340            token = None341 342        commit_hash = model_info(pretrained_model_name_or_path, revision=revision, token=token).sha343 344        # The module file will end up being placed in a subfolder with the git hash of the repo. This way we get the345        # benefit of versioning.346        submodule_path = submodule_path / commit_hash347        full_submodule = full_submodule + os.path.sep + commit_hash348        create_dynamic_module(full_submodule)349 350        if not (submodule_path / module_file).exists():351            shutil.copy(resolved_module_file, submodule_path / module_file)352        # Make sure we also have every file with relative353        for module_needed in modules_needed:354            if not (submodule_path / module_needed).exists():355                get_cached_module_file(356                    pretrained_model_name_or_path,357                    f"{module_needed}.py",358                    cache_dir=cache_dir,359                    force_download=force_download,360                    resume_download=resume_download,361                    proxies=proxies,362                    use_auth_token=use_auth_token,363                    revision=revision,364                    local_files_only=local_files_only,365                )366    return os.path.join(full_submodule, module_file)367 368 369def get_class_from_dynamic_module(370    pretrained_model_name_or_path: Union[str, os.PathLike],371    module_file: str,372    class_name: Optional[str] = None,373    cache_dir: Optional[Union[str, os.PathLike]] = None,374    force_download: bool = False,375    resume_download: bool = False,376    proxies: Optional[Dict[str, str]] = None,377    use_auth_token: Optional[Union[bool, str]] = None,378    revision: Optional[str] = None,379    local_files_only: bool = False,380    **kwargs,381):382    """383    Extracts a class from a module file, present in the local folder or repository of a model.384 385    <Tip warning={true}>386 387    Calling this function will execute the code in the module file found locally or downloaded from the Hub. It should388    therefore only be called on trusted repos.389 390    </Tip>391 392    Args:393        pretrained_model_name_or_path (`str` or `os.PathLike`):394            This can be either:395 396            - a string, the *model id* of a pretrained model configuration hosted inside a model repo on397              huggingface.co. Valid model ids can be located at the root-level, like `bert-base-uncased`, or namespaced398              under a user or organization name, like `dbmdz/bert-base-german-cased`.399            - a path to a *directory* containing a configuration file saved using the400              [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.401 402        module_file (`str`):403            The name of the module file containing the class to look for.404        class_name (`str`):405            The name of the class to import in the module.406        cache_dir (`str` or `os.PathLike`, *optional*):407            Path to a directory in which a downloaded pretrained model configuration should be cached if the standard408            cache should not be used.409        force_download (`bool`, *optional*, defaults to `False`):410            Whether or not to force to (re-)download the configuration files and override the cached versions if they411            exist.412        resume_download (`bool`, *optional*, defaults to `False`):413            Whether or not to delete incompletely received file. Attempts to resume the download if such a file exists.414        proxies (`Dict[str, str]`, *optional*):415            A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',416            'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.417        use_auth_token (`str` or `bool`, *optional*):418            The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated419            when running `transformers-cli login` (stored in `~/.huggingface`).420        revision (`str`, *optional*, defaults to `"main"`):421            The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a422            git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any423            identifier allowed by git.424        local_files_only (`bool`, *optional*, defaults to `False`):425            If `True`, will only try to load the tokenizer configuration from local files.426 427    <Tip>428 429    You may pass a token in `use_auth_token` if you are not logged in (`huggingface-cli long`) and want to use private430    or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models).431 432    </Tip>433 434    Returns:435        `type`: The class, dynamically imported from the module.436 437    Examples:438 439    ```python440    # Download module `modeling.py` from huggingface.co and cache then extract the class `MyBertModel` from this441    # module.442    cls = get_class_from_dynamic_module("sgugger/my-bert-model", "modeling.py", "MyBertModel")443    ```"""444    # And lastly we get the class inside our newly created module445    final_module = get_cached_module_file(446        pretrained_model_name_or_path,447        module_file,448        cache_dir=cache_dir,449        force_download=force_download,450        resume_download=resume_download,451        proxies=proxies,452        use_auth_token=use_auth_token,453        revision=revision,454        local_files_only=local_files_only,455    )456    return get_class_in_module(class_name, final_module.replace(".py", ""))457