Skip to content

Commit 25feac9

Browse files
patrickvonplatenpatrickvonplaten
committed
add pipeline
1 parent e83c536 commit 25feac9

5 files changed

Lines changed: 170 additions & 2 deletions

File tree

‎models/vision/ddpm/example.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
#!/usr/bin/env python3
2+
from diffusers import UNetModel, GaussianDiffusion
3+
from modeling_ddpm import DDPM
4+
import tempfile
5+
6+
unet = UNetModel.from_pretrained("fusing/ddpm_dummy")
7+
sampler = GaussianDiffusion.from_config("fusing/ddpm_dummy")
8+
9+
# compose Diffusion Pipeline
10+
ddpm = DDPM(unet, sampler)
11+
# generate / sample
12+
image = ddpm()
13+
print(image)
14+
15+
16+
# save and load with 0 extra code (handled by general `DiffusionPipeline` class)
17+
with tempfile.TemporaryDirectory() as tmpdirname:
18+
ddpm.save_pretrained(tmpdirname)
19+
print("Model saved")
20+
ddpm_new = DDPM.from_pretrained(tmpdirname)
21+
print("Model loaded")
22+
print(ddpm_new)

‎models/vision/ddpm/modeling_ddpm.py‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
# Copyright 2022 The HuggingFace Team. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
14+
# limitations under the License.
15+
16+
17+
from diffusers import DiffusionPipeline
18+
19+
20+
class DDPM(DiffusionPipeline):
21+
22+
def __init__(self, unet, gaussian_sampler):
23+
super().__init__(unet=unet, gaussian_sampler=gaussian_sampler)
24+
25+
def __call__(self, batch_size=1):
26+
image = self.gaussian_sampler.sample(self.unet, batch_size=batch_size)
27+
return image

‎src/diffusers/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,3 +6,6 @@
66

77
from .models.unet import UNetModel
88
from .samplers.gaussian import GaussianDiffusion
9+
10+
from .pipeline_utils import DiffusionPipeline
11+
from .modeling_utils import PreTrainedModel

‎src/diffusers/configuration_utils.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -91,8 +91,8 @@ def save_config(self, save_directory: Union[str, os.PathLike], push_to_hub: bool
9191
logger.info(f"Configuration saved in {output_config_file}")
9292

9393
@classmethod
94-
def from_config(
95-
cls, pretrained_model_name_or_path: Union[str, os.PathLike], return_unused_kwargs=False, **kwargs
94+
def get_config_dict(
95+
cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs
9696
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
9797
cache_dir = kwargs.pop("cache_dir", None)
9898
force_download = kwargs.pop("force_download", False)
@@ -198,6 +198,14 @@ def from_config(
198198
f"Values will be initialized to default values."
199199
)
200200

201+
return config_dict, unused_kwargs
202+
203+
@classmethod
204+
def from_config(
205+
cls, pretrained_model_name_or_path: Union[str, os.PathLike], return_unused_kwargs=False, **kwargs
206+
):
207+
config_dict, unused_kwargs = cls.get_config_dict(pretrained_model_name_or_path=pretrained_model_name_or_path, **kwargs)
208+
201209
model = cls(**config_dict)
202210

203211
if return_unused_kwargs:

‎src/diffusers/pipeline_utils.py‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
# coding=utf-8
2+
# Copyright 2022 The HuggingFace Inc. team.
3+
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
4+
#
5+
# Licensed under the Apache License, Version 2.0 (the "License");
6+
# you may not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
17+
import os
18+
from typing import Optional, Union
19+
import importlib
20+
21+
from .configuration_utils import Config
22+
23+
# CHANGE to diffusers.utils
24+
from transformers.utils import logging
25+
26+
27+
INDEX_FILE = "diffusion_model.pt"
28+
29+
30+
logger = logging.get_logger(__name__)
31+
32+
33+
LOADABLE_CLASSES = {
34+
"diffusers": {
35+
"PreTrainedModel": ["save_pretrained", "from_pretrained"],
36+
"GaussianDiffusion": ["save_config", "from_config"],
37+
},
38+
"transformers": {
39+
"PreTrainedModel": ["save_pretrained", "from_pretrained"],
40+
},
41+
}
42+
43+
44+
class DiffusionPipeline(Config):
45+
46+
config_name = "model_index.json"
47+
48+
def __init__(self, **kwargs):
49+
for name, module in kwargs.items():
50+
# retrive library
51+
library = module.__module__.split(".")[0]
52+
# retrive class_name
53+
class_name = module.__class__.__name__
54+
55+
# save model index config
56+
self.register(**{name: (library, class_name)})
57+
58+
# set models
59+
setattr(self, name, module)
60+
61+
def save_pretrained(self, save_directory: Union[str, os.PathLike]):
62+
self.save_config(save_directory)
63+
64+
model_index_dict = self._dict_to_save
65+
model_index_dict.pop("_class_name")
66+
67+
for name, (library_name, class_name) in self._dict_to_save.items():
68+
importable_classes = LOADABLE_CLASSES[library_name]
69+
70+
library = importlib.import_module(library_name)
71+
class_obj = getattr(library, class_name)
72+
class_candidates = {c: getattr(library, c) for c in importable_classes.keys()}
73+
74+
save_method_name = None
75+
for class_name, class_candidate in class_candidates.items():
76+
if issubclass(class_obj, class_candidate):
77+
save_method_name = importable_classes[class_name][0]
78+
79+
save_method = getattr(getattr(self, name), save_method_name)
80+
save_method(os.path.join(save_directory, name))
81+
82+
@classmethod
83+
def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
84+
# use snapshot download here to get it working from from_pretrained
85+
config_dict, _ = cls.get_config_dict(pretrained_model_name_or_path)
86+
87+
init_kwargs = {}
88+
89+
for name, (library_name, class_name) in config_dict.items():
90+
importable_classes = LOADABLE_CLASSES[library_name]
91+
92+
library = importlib.import_module(library_name)
93+
class_obj = getattr(library, class_name)
94+
class_candidates = {c: getattr(library, c) for c in importable_classes.keys()}
95+
96+
load_method_name = None
97+
for class_name, class_candidate in class_candidates.items():
98+
if issubclass(class_obj, class_candidate):
99+
load_method_name = importable_classes[class_name][1]
100+
101+
load_method = getattr(class_obj, load_method_name)
102+
103+
loaded_sub_model = load_method(os.path.join(pretrained_model_name_or_path, name))
104+
105+
init_kwargs[name] = loaded_sub_model
106+
107+
model = cls(**init_kwargs)
108+
return model

0 commit comments

Comments
 (0)