Skip to content

Commit ce56662

Browse files
patil-surajpatil-suraj
committed
make from hub import work
1 parent 1a6196e commit ce56662

3 files changed

Lines changed: 347 additions & 7 deletions

File tree

‎models/vision/ddpm/modeling_ddpm.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ class DDPM(DiffusionPipeline):
2323

2424
modeling_file = "modeling_ddpm.py"
2525

26-
def __init__(self, unet, noise_scheduler, vqvae):
26+
def __init__(self, unet, noise_scheduler):
2727
super().__init__()
2828
self.register_modules(unet=unet, noise_scheduler=noise_scheduler)
2929

‎src/diffusers/dynamic_modules_utils.py‎

Lines changed: 339 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,339 @@
1+
# coding=utf-8
2+
# Copyright 2021 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 at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# 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 and
14+
# limitations under the License.
15+
"""Utilities to dynamically load objects from the Hub."""
16+
17+
import importlib
18+
import os
19+
import re
20+
import shutil
21+
import sys
22+
from pathlib import Path
23+
from typing import Dict, Optional, Union
24+
25+
from huggingface_hub import HfFolder, model_info
26+
27+
from transformers.utils import (
28+
HF_MODULES_CACHE,
29+
TRANSFORMERS_DYNAMIC_MODULE_NAME,
30+
cached_path,
31+
hf_bucket_url,
32+
is_offline_mode,
33+
logging,
34+
)
35+
36+
37+
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
38+
39+
40+
def init_hf_modules():
41+
"""
42+
Creates the cache directory for modules with an init, and adds it to the Python path.
43+
"""
44+
# This function has already been executed if HF_MODULES_CACHE already is in the Python path.
45+
if HF_MODULES_CACHE in sys.path:
46+
return
47+
48+
sys.path.append(HF_MODULES_CACHE)
49+
os.makedirs(HF_MODULES_CACHE, exist_ok=True)
50+
init_path = Path(HF_MODULES_CACHE) / "__init__.py"
51+
if not init_path.exists():
52+
init_path.touch()
53+
54+
55+
def create_dynamic_module(name: Union[str, os.PathLike]):
56+
"""
57+
Creates a dynamic module in the cache directory for modules.
58+
"""
59+
init_hf_modules()
60+
dynamic_module_path = Path(HF_MODULES_CACHE) / name
61+
# If the parent module does not exist yet, recursively create it.
62+
if not dynamic_module_path.parent.exists():
63+
create_dynamic_module(dynamic_module_path.parent)
64+
os.makedirs(dynamic_module_path, exist_ok=True)
65+
init_path = dynamic_module_path / "__init__.py"
66+
if not init_path.exists():
67+
init_path.touch()
68+
69+
70+
def get_relative_imports(module_file):
71+
"""
72+
Get the list of modules that are relatively imported in a module file.
73+
74+
Args:
75+
module_file (`str` or `os.PathLike`): The module file to inspect.
76+
"""
77+
with open(module_file, "r", encoding="utf-8") as f:
78+
content = f.read()
79+
80+
# Imports of the form `import .xxx`
81+
relative_imports = re.findall("^\s*import\s+\.(\S+)\s*$", content, flags=re.MULTILINE)
82+
# Imports of the form `from .xxx import yyy`
83+
relative_imports += re.findall("^\s*from\s+\.(\S+)\s+import", content, flags=re.MULTILINE)
84+
# Unique-ify
85+
return list(set(relative_imports))
86+
87+
88+
def get_relative_import_files(module_file):
89+
"""
90+
Get the list of all files that are needed for a given module. Note that this function recurses through the relative
91+
imports (if a imports b and b imports c, it will return module files for b and c).
92+
93+
Args:
94+
module_file (`str` or `os.PathLike`): The module file to inspect.
95+
"""
96+
no_change = False
97+
files_to_check = [module_file]
98+
all_relative_imports = []
99+
100+
# Let's recurse through all relative imports
101+
while not no_change:
102+
new_imports = []
103+
for f in files_to_check:
104+
new_imports.extend(get_relative_imports(f))
105+
106+
module_path = Path(module_file).parent
107+
new_import_files = [str(module_path / m) for m in new_imports]
108+
new_import_files = [f for f in new_import_files if f not in all_relative_imports]
109+
files_to_check = [f"{f}.py" for f in new_import_files]
110+
111+
no_change = len(new_import_files) == 0
112+
all_relative_imports.extend(files_to_check)
113+
114+
return all_relative_imports
115+
116+
117+
def check_imports(filename):
118+
"""
119+
Check if the current Python environment contains all the libraries that are imported in a file.
120+
"""
121+
with open(filename, "r", encoding="utf-8") as f:
122+
content = f.read()
123+
124+
# Imports of the form `import xxx`
125+
imports = re.findall("^\s*import\s+(\S+)\s*$", content, flags=re.MULTILINE)
126+
# Imports of the form `from xxx import yyy`
127+
imports += re.findall("^\s*from\s+(\S+)\s+import", content, flags=re.MULTILINE)
128+
# Only keep the top-level module
129+
imports = [imp.split(".")[0] for imp in imports if not imp.startswith(".")]
130+
131+
# Unique-ify and test we got them all
132+
imports = list(set(imports))
133+
missing_packages = []
134+
for imp in imports:
135+
try:
136+
importlib.import_module(imp)
137+
except ImportError:
138+
missing_packages.append(imp)
139+
140+
if len(missing_packages) > 0:
141+
raise ImportError(
142+
"This modeling file requires the following packages that were not found in your environment: "
143+
f"{', '.join(missing_packages)}. Run `pip install {' '.join(missing_packages)}`"
144+
)
145+
146+
return get_relative_imports(filename)
147+
148+
149+
def get_class_in_module(class_name, module_path):
150+
"""
151+
Import a module on the cache directory for modules and extract a class from it.
152+
"""
153+
module_path = module_path.replace(os.path.sep, ".")
154+
module = importlib.import_module(module_path)
155+
return getattr(module, class_name)
156+
157+
158+
def get_cached_module_file(
159+
pretrained_model_name_or_path: Union[str, os.PathLike],
160+
module_file: str,
161+
cache_dir: Optional[Union[str, os.PathLike]] = None,
162+
force_download: bool = False,
163+
resume_download: bool = False,
164+
proxies: Optional[Dict[str, str]] = None,
165+
use_auth_token: Optional[Union[bool, str]] = None,
166+
revision: Optional[str] = None,
167+
local_files_only: bool = False,
168+
):
169+
"""
170+
Prepares Downloads a module from a local folder or a distant repo and returns its path inside the cached
171+
Transformers module.
172+
173+
Args:
174+
pretrained_model_name_or_path (`str` or `os.PathLike`):
175+
This can be either:
176+
177+
- a string, the *model id* of a pretrained model configuration hosted inside a model repo on
178+
huggingface.co. Valid model ids can be located at the root-level, like `bert-base-uncased`, or namespaced
179+
under a user or organization name, like `dbmdz/bert-base-german-cased`.
180+
- a path to a *directory* containing a configuration file saved using the
181+
[`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
182+
183+
module_file (`str`):
184+
The name of the module file containing the class to look for.
185+
cache_dir (`str` or `os.PathLike`, *optional*):
186+
Path to a directory in which a downloaded pretrained model configuration should be cached if the standard
187+
cache should not be used.
188+
force_download (`bool`, *optional*, defaults to `False`):
189+
Whether or not to force to (re-)download the configuration files and override the cached versions if they
190+
exist.
191+
resume_download (`bool`, *optional*, defaults to `False`):
192+
Whether or not to delete incompletely received file. Attempts to resume the download if such a file exists.
193+
proxies (`Dict[str, str]`, *optional*):
194+
A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
195+
'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
196+
use_auth_token (`str` or *bool*, *optional*):
197+
The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
198+
when running `transformers-cli login` (stored in `~/.huggingface`).
199+
revision (`str`, *optional*, defaults to `"main"`):
200+
The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
201+
git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
202+
identifier allowed by git.
203+
local_files_only (`bool`, *optional*, defaults to `False`):
204+
If `True`, will only try to load the tokenizer configuration from local files.
205+
206+
<Tip>
207+
208+
Passing `use_auth_token=True` is required when you want to use a private model.
209+
210+
</Tip>
211+
212+
Returns:
213+
`str`: The path to the module inside the cache.
214+
"""
215+
# Download and cache module_file from the repo `pretrained_model_name_or_path` of grab it if it's a local file.
216+
pretrained_model_name_or_path = str(pretrained_model_name_or_path)
217+
module_file_or_url = os.path.join(pretrained_model_name_or_path, module_file)
218+
submodule = "local"
219+
220+
try:
221+
# Load from URL or cache if already cached
222+
resolved_module_file = cached_path(
223+
module_file_or_url,
224+
cache_dir=cache_dir,
225+
force_download=force_download,
226+
proxies=proxies,
227+
resume_download=resume_download,
228+
local_files_only=local_files_only,
229+
use_auth_token=use_auth_token,
230+
)
231+
232+
except EnvironmentError:
233+
logger.error(f"Could not locate the {module_file} inside {pretrained_model_name_or_path}.")
234+
raise
235+
236+
# Check we have all the requirements in our environment
237+
modules_needed = check_imports(resolved_module_file)
238+
239+
# Now we move the module inside our cached dynamic modules.
240+
full_submodule = TRANSFORMERS_DYNAMIC_MODULE_NAME + os.path.sep + submodule
241+
create_dynamic_module(full_submodule)
242+
submodule_path = Path(HF_MODULES_CACHE) / full_submodule
243+
# We always copy local files (we could hash the file to see if there was a change, and give them the name of
244+
# that hash, to only copy when there is a modification but it seems overkill for now).
245+
# The only reason we do the copy is to avoid putting too many folders in sys.path.
246+
shutil.copy(resolved_module_file, submodule_path / module_file)
247+
for module_needed in modules_needed:
248+
module_needed = f"{module_needed}.py"
249+
shutil.copy(os.path.join(pretrained_model_name_or_path, module_needed), submodule_path / module_needed)
250+
return os.path.join(full_submodule, module_file)
251+
252+
253+
def get_class_from_dynamic_module(
254+
pretrained_model_name_or_path: Union[str, os.PathLike],
255+
module_file: str,
256+
class_name: str,
257+
cache_dir: Optional[Union[str, os.PathLike]] = None,
258+
force_download: bool = False,
259+
resume_download: bool = False,
260+
proxies: Optional[Dict[str, str]] = None,
261+
use_auth_token: Optional[Union[bool, str]] = None,
262+
revision: Optional[str] = None,
263+
local_files_only: bool = False,
264+
**kwargs,
265+
):
266+
"""
267+
Extracts a class from a module file, present in the local folder or repository of a model.
268+
269+
<Tip warning={true}>
270+
271+
Calling this function will execute the code in the module file found locally or downloaded from the Hub. It should
272+
therefore only be called on trusted repos.
273+
274+
</Tip>
275+
276+
Args:
277+
pretrained_model_name_or_path (`str` or `os.PathLike`):
278+
This can be either:
279+
280+
- a string, the *model id* of a pretrained model configuration hosted inside a model repo on
281+
huggingface.co. Valid model ids can be located at the root-level, like `bert-base-uncased`, or namespaced
282+
under a user or organization name, like `dbmdz/bert-base-german-cased`.
283+
- a path to a *directory* containing a configuration file saved using the
284+
[`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
285+
286+
module_file (`str`):
287+
The name of the module file containing the class to look for.
288+
class_name (`str`):
289+
The name of the class to import in the module.
290+
cache_dir (`str` or `os.PathLike`, *optional*):
291+
Path to a directory in which a downloaded pretrained model configuration should be cached if the standard
292+
cache should not be used.
293+
force_download (`bool`, *optional*, defaults to `False`):
294+
Whether or not to force to (re-)download the configuration files and override the cached versions if they
295+
exist.
296+
resume_download (`bool`, *optional*, defaults to `False`):
297+
Whether or not to delete incompletely received file. Attempts to resume the download if such a file exists.
298+
proxies (`Dict[str, str]`, *optional*):
299+
A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
300+
'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
301+
use_auth_token (`str` or `bool`, *optional*):
302+
The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
303+
when running `transformers-cli login` (stored in `~/.huggingface`).
304+
revision (`str`, *optional*, defaults to `"main"`):
305+
The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
306+
git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
307+
identifier allowed by git.
308+
local_files_only (`bool`, *optional*, defaults to `False`):
309+
If `True`, will only try to load the tokenizer configuration from local files.
310+
311+
<Tip>
312+
313+
Passing `use_auth_token=True` is required when you want to use a private model.
314+
315+
</Tip>
316+
317+
Returns:
318+
`type`: The class, dynamically imported from the module.
319+
320+
Examples:
321+
322+
```python
323+
# Download module `modeling.py` from huggingface.co and cache then extract the class `MyBertModel` from this
324+
# module.
325+
cls = get_class_from_dynamic_module("sgugger/my-bert-model", "modeling.py", "MyBertModel")
326+
```"""
327+
# And lastly we get the class inside our newly created module
328+
final_module = get_cached_module_file(
329+
pretrained_model_name_or_path,
330+
module_file,
331+
cache_dir=cache_dir,
332+
force_download=force_download,
333+
resume_download=resume_download,
334+
proxies=proxies,
335+
use_auth_token=use_auth_token,
336+
revision=revision,
337+
local_files_only=local_files_only,
338+
)
339+
return get_class_in_module(class_name, final_module.replace(".py", ""))

‎src/diffusers/pipeline_utils.py‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,15 @@
1616

1717
import importlib
1818
import os
19+
from pathlib import Path
1920
from typing import Optional, Union
2021
from huggingface_hub import snapshot_download
2122

2223
# CHANGE to diffusers.utils
2324
from transformers.utils import logging
2425

2526
from .configuration_utils import ConfigMixin
27+
from .dynamic_modules_utils import get_class_from_dynamic_module
2628

2729

2830
INDEX_FILE = "diffusion_model.pt"
@@ -91,12 +93,10 @@ def save_pretrained(self, save_directory: Union[str, os.PathLike]):
9193
def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
9294
# use snapshot download here to get it working from from_pretrained
9395
cached_folder = snapshot_download(pretrained_model_name_or_path)
94-
config_dict, pipeline_kwargs = cls.get_config_dict(cached_folder)
96+
_, config_dict = cls.get_config_dict(cached_folder)
9597

96-
module = pipeline_kwargs["_module"]
97-
# TODO(Suraj) - make from hub import work
98-
# Make `ddpm = DiffusionPipeline.from_pretrained("fusing/ddpm-lsun-bedroom-pipe")` work
99-
# Add Sylvains code from transformers
98+
module = config_dict.pop("_module", None)
99+
class_name_ = config_dict.pop("_class_name")
100100

101101
init_kwargs = {}
102102

@@ -122,5 +122,6 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
122122

123123
init_kwargs[name] = loaded_sub_model # UNet(...), # DiffusionSchedule(...)
124124

125-
model = cls(**init_kwargs)
125+
class_obj = get_class_from_dynamic_module(cached_folder, module, class_name_, cached_folder)
126+
model = class_obj(**init_kwargs)
126127
return model

0 commit comments

Comments
 (0)