declare-lab/tango2
92
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 16 17import os18import re19import sys20import traceback21import warnings22from pathlib import Path23from typing import Dict, Optional, Union24from uuid import uuid425 26from huggingface_hub import HfFolder, ModelCard, ModelCardData, hf_hub_download, whoami27from huggingface_hub.file_download import REGEX_COMMIT_HASH28from huggingface_hub.utils import (29 EntryNotFoundError,30 RepositoryNotFoundError,31 RevisionNotFoundError,32 is_jinja_available,33)34from packaging import version35from requests import HTTPError36 37from .. import __version__38from .constants import (39 DEPRECATED_REVISION_ARGS,40 DIFFUSERS_CACHE,41 HUGGINGFACE_CO_RESOLVE_ENDPOINT,42 SAFETENSORS_WEIGHTS_NAME,43 WEIGHTS_NAME,44)45from .import_utils import (46 ENV_VARS_TRUE_VALUES,47 _flax_version,48 _jax_version,49 _onnxruntime_version,50 _torch_version,51 is_flax_available,52 is_onnx_available,53 is_torch_available,54)55from .logging import get_logger56 57 58logger = get_logger(__name__)59 60 61MODEL_CARD_TEMPLATE_PATH = Path(__file__).parent / "model_card_template.md"62SESSION_ID = uuid4().hex63HF_HUB_OFFLINE = os.getenv("HF_HUB_OFFLINE", "").upper() in ENV_VARS_TRUE_VALUES64DISABLE_TELEMETRY = os.getenv("DISABLE_TELEMETRY", "").upper() in ENV_VARS_TRUE_VALUES65HUGGINGFACE_CO_TELEMETRY = HUGGINGFACE_CO_RESOLVE_ENDPOINT + "/api/telemetry/"66 67 68def http_user_agent(user_agent: Union[Dict, str, None] = None) -> str:69 """70 Formats a user-agent string with basic info about a request.71 """72 ua = f"diffusers/{__version__}; python/{sys.version.split()[0]}; session_id/{SESSION_ID}"73 if DISABLE_TELEMETRY or HF_HUB_OFFLINE:74 return ua + "; telemetry/off"75 if is_torch_available():76 ua += f"; torch/{_torch_version}"77 if is_flax_available():78 ua += f"; jax/{_jax_version}"79 ua += f"; flax/{_flax_version}"80 if is_onnx_available():81 ua += f"; onnxruntime/{_onnxruntime_version}"82 # CI will set this value to True83 if os.environ.get("DIFFUSERS_IS_CI", "").upper() in ENV_VARS_TRUE_VALUES:84 ua += "; is_ci/true"85 if isinstance(user_agent, dict):86 ua += "; " + "; ".join(f"{k}/{v}" for k, v in user_agent.items())87 elif isinstance(user_agent, str):88 ua += "; " + user_agent89 return ua90 91 92def get_full_repo_name(model_id: str, organization: Optional[str] = None, token: Optional[str] = None):93 if token is None:94 token = HfFolder.get_token()95 if organization is None:96 username = whoami(token)["name"]97 return f"{username}/{model_id}"98 else:99 return f"{organization}/{model_id}"100 101 102def create_model_card(args, model_name):103 if not is_jinja_available():104 raise ValueError(105 "Modelcard rendering is based on Jinja templates."106 " Please make sure to have `jinja` installed before using `create_model_card`."107 " To install it, please run `pip install Jinja2`."108 )109 110 if hasattr(args, "local_rank") and args.local_rank not in [-1, 0]:111 return112 113 hub_token = args.hub_token if hasattr(args, "hub_token") else None114 repo_name = get_full_repo_name(model_name, token=hub_token)115 116 model_card = ModelCard.from_template(117 card_data=ModelCardData( # Card metadata object that will be converted to YAML block118 language="en",119 license="apache-2.0",120 library_name="diffusers",121 tags=[],122 datasets=args.dataset_name,123 metrics=[],124 ),125 template_path=MODEL_CARD_TEMPLATE_PATH,126 model_name=model_name,127 repo_name=repo_name,128 dataset_name=args.dataset_name if hasattr(args, "dataset_name") else None,129 learning_rate=args.learning_rate,130 train_batch_size=args.train_batch_size,131 eval_batch_size=args.eval_batch_size,132 gradient_accumulation_steps=(133 args.gradient_accumulation_steps if hasattr(args, "gradient_accumulation_steps") else None134 ),135 adam_beta1=args.adam_beta1 if hasattr(args, "adam_beta1") else None,136 adam_beta2=args.adam_beta2 if hasattr(args, "adam_beta2") else None,137 adam_weight_decay=args.adam_weight_decay if hasattr(args, "adam_weight_decay") else None,138 adam_epsilon=args.adam_epsilon if hasattr(args, "adam_epsilon") else None,139 lr_scheduler=args.lr_scheduler if hasattr(args, "lr_scheduler") else None,140 lr_warmup_steps=args.lr_warmup_steps if hasattr(args, "lr_warmup_steps") else None,141 ema_inv_gamma=args.ema_inv_gamma if hasattr(args, "ema_inv_gamma") else None,142 ema_power=args.ema_power if hasattr(args, "ema_power") else None,143 ema_max_decay=args.ema_max_decay if hasattr(args, "ema_max_decay") else None,144 mixed_precision=args.mixed_precision,145 )146 147 card_path = os.path.join(args.output_dir, "README.md")148 model_card.save(card_path)149 150 151def extract_commit_hash(resolved_file: Optional[str], commit_hash: Optional[str] = None):152 """153 Extracts the commit hash from a resolved filename toward a cache file.154 """155 if resolved_file is None or commit_hash is not None:156 return commit_hash157 resolved_file = str(Path(resolved_file).as_posix())158 search = re.search(r"snapshots/([^/]+)/", resolved_file)159 if search is None:160 return None161 commit_hash = search.groups()[0]162 return commit_hash if REGEX_COMMIT_HASH.match(commit_hash) else None163 164 165# Old default cache path, potentially to be migrated.166# This logic was more or less taken from `transformers`, with the following differences:167# - Diffusers doesn't use custom environment variables to specify the cache path.168# - There is no need to migrate the cache format, just move the files to the new location.169hf_cache_home = os.path.expanduser(170 os.getenv("HF_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "huggingface"))171)172old_diffusers_cache = os.path.join(hf_cache_home, "diffusers")173 174 175def move_cache(old_cache_dir: Optional[str] = None, new_cache_dir: Optional[str] = None) -> None:176 if new_cache_dir is None:177 new_cache_dir = DIFFUSERS_CACHE178 if old_cache_dir is None:179 old_cache_dir = old_diffusers_cache180 181 old_cache_dir = Path(old_cache_dir).expanduser()182 new_cache_dir = Path(new_cache_dir).expanduser()183 for old_blob_path in old_cache_dir.glob("**/blobs/*"):184 if old_blob_path.is_file() and not old_blob_path.is_symlink():185 new_blob_path = new_cache_dir / old_blob_path.relative_to(old_cache_dir)186 new_blob_path.parent.mkdir(parents=True, exist_ok=True)187 os.replace(old_blob_path, new_blob_path)188 try:189 os.symlink(new_blob_path, old_blob_path)190 except OSError:191 logger.warning(192 "Could not create symlink between old cache and new cache. If you use an older version of diffusers again, files will be re-downloaded."193 )194 # At this point, old_cache_dir contains symlinks to the new cache (it can still be used).195 196 197cache_version_file = os.path.join(DIFFUSERS_CACHE, "version_diffusers_cache.txt")198if not os.path.isfile(cache_version_file):199 cache_version = 0200else:201 with open(cache_version_file) as f:202 cache_version = int(f.read())203 204if cache_version < 1:205 old_cache_is_not_empty = os.path.isdir(old_diffusers_cache) and len(os.listdir(old_diffusers_cache)) > 0206 if old_cache_is_not_empty:207 logger.warning(208 "The cache for model files in Diffusers v0.14.0 has moved to a new location. Moving your "209 "existing cached models. This is a one-time operation, you can interrupt it or run it "210 "later by calling `diffusers.utils.hub_utils.move_cache()`."211 )212 try:213 move_cache()214 except Exception as e:215 trace = "\n".join(traceback.format_tb(e.__traceback__))216 logger.error(217 f"There was a problem when trying to move your cache:\n\n{trace}\n{e.__class__.__name__}: {e}\n\nPlease "218 "file an issue at https://github.com/huggingface/diffusers/issues/new/choose, copy paste this whole "219 "message and we will do our best to help."220 )221 222if cache_version < 1:223 try:224 os.makedirs(DIFFUSERS_CACHE, exist_ok=True)225 with open(cache_version_file, "w") as f:226 f.write("1")227 except Exception:228 logger.warning(229 f"There was a problem when trying to write in your cache folder ({DIFFUSERS_CACHE}). Please, ensure "230 "the directory exists and can be written to."231 )232 233 234def _add_variant(weights_name: str, variant: Optional[str] = None) -> str:235 if variant is not None:236 splits = weights_name.split(".")237 splits = splits[:-1] + [variant] + splits[-1:]238 weights_name = ".".join(splits)239 240 return weights_name241 242 243def _get_model_file(244 pretrained_model_name_or_path,245 *,246 weights_name,247 subfolder,248 cache_dir,249 force_download,250 proxies,251 resume_download,252 local_files_only,253 use_auth_token,254 user_agent,255 revision,256 commit_hash=None,257):258 pretrained_model_name_or_path = str(pretrained_model_name_or_path)259 if os.path.isfile(pretrained_model_name_or_path):260 return pretrained_model_name_or_path261 elif os.path.isdir(pretrained_model_name_or_path):262 if os.path.isfile(os.path.join(pretrained_model_name_or_path, weights_name)):263 # Load from a PyTorch checkpoint264 model_file = os.path.join(pretrained_model_name_or_path, weights_name)265 return model_file266 elif subfolder is not None and os.path.isfile(267 os.path.join(pretrained_model_name_or_path, subfolder, weights_name)268 ):269 model_file = os.path.join(pretrained_model_name_or_path, subfolder, weights_name)270 return model_file271 else:272 raise EnvironmentError(273 f"Error no file named {weights_name} found in directory {pretrained_model_name_or_path}."274 )275 else:276 # 1. First check if deprecated way of loading from branches is used277 if (278 revision in DEPRECATED_REVISION_ARGS279 and (weights_name == WEIGHTS_NAME or weights_name == SAFETENSORS_WEIGHTS_NAME)280 and version.parse(version.parse(__version__).base_version) >= version.parse("0.17.0")281 ):282 try:283 model_file = hf_hub_download(284 pretrained_model_name_or_path,285 filename=_add_variant(weights_name, revision),286 cache_dir=cache_dir,287 force_download=force_download,288 proxies=proxies,289 resume_download=resume_download,290 local_files_only=local_files_only,291 use_auth_token=use_auth_token,292 user_agent=user_agent,293 subfolder=subfolder,294 revision=revision or commit_hash,295 )296 warnings.warn(297 f"Loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'` is deprecated. Loading instead from `revision='main'` with `variant={revision}`. Loading model variants via `revision='{revision}'` will be removed in diffusers v1. Please use `variant='{revision}'` instead.",298 FutureWarning,299 )300 return model_file301 except: # noqa: E722302 warnings.warn(303 f"You are loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'`. This behavior is deprecated and will be removed in diffusers v1. One should use `variant='{revision}'` instead. However, it appears that {pretrained_model_name_or_path} currently does not have a {_add_variant(weights_name, revision)} file in the 'main' branch of {pretrained_model_name_or_path}. \n The Diffusers team and community would be very grateful if you could open an issue: https://github.com/huggingface/diffusers/issues/new with the title '{pretrained_model_name_or_path} is missing {_add_variant(weights_name, revision)}' so that the correct variant file can be added.",304 FutureWarning,305 )306 try:307 # 2. Load model file as usual308 model_file = hf_hub_download(309 pretrained_model_name_or_path,310 filename=weights_name,311 cache_dir=cache_dir,312 force_download=force_download,313 proxies=proxies,314 resume_download=resume_download,315 local_files_only=local_files_only,316 use_auth_token=use_auth_token,317 user_agent=user_agent,318 subfolder=subfolder,319 revision=revision or commit_hash,320 )321 return model_file322 323 except RepositoryNotFoundError:324 raise EnvironmentError(325 f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier "326 "listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a "327 "token having permission to this repo with `use_auth_token` or log in with `huggingface-cli "328 "login`."329 )330 except RevisionNotFoundError:331 raise EnvironmentError(332 f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for "333 "this model name. Check the model page at "334 f"'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions."335 )336 except EntryNotFoundError:337 raise EnvironmentError(338 f"{pretrained_model_name_or_path} does not appear to have a file named {weights_name}."339 )340 except HTTPError as err:341 raise EnvironmentError(342 f"There was a specific connection error when trying to load {pretrained_model_name_or_path}:\n{err}"343 )344 except ValueError:345 raise EnvironmentError(346 f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it"347 f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a"348 f" directory containing a file named {weights_name} or"349 " \nCheckout your internet connection or see how to run the library in"350 " offline mode at 'https://huggingface.co/docs/diffusers/installation#offline-mode'."351 )352 except EnvironmentError:353 raise EnvironmentError(354 f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it from "355 "'https://huggingface.co/models', make sure you don't have a local directory with the same name. "356 f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory "357 f"containing a file named {weights_name}"358 )359