Skip to content

fix auto_pipeline: pass kwargs to load_config #4793

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Aug 28, 2023
Merged

fix auto_pipeline: pass kwargs to load_config #4793

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
9e6a029
fix
Aug 26, 2023
908a01f
Update src/diffusers/pipelines/auto_pipeline.py
yiyixuxu Aug 26, 2023
af22d43
apply to other auto pipelines
Aug 26, 2023
8b89b66
add a test
Aug 27, 2023
1acb4d4
add test
Aug 28, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 73 additions & 3 deletions src/diffusers/pipelines/auto_pipeline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from collections import OrderedDict

from ..configuration_utils import ConfigMixin
from ..utils import DIFFUSERS_CACHE
from .controlnet import (
StableDiffusionControlNetImg2ImgPipeline,
StableDiffusionControlNetInpaintPipeline,
Expand Down Expand Up @@ -295,14 +296,37 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs):
>>> image = pipeline(prompt).images[0]
```
"""
config = cls.load_config(pretrained_model_or_path)
cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)
force_download = kwargs.pop("force_download", False)
resume_download = kwargs.pop("resume_download", False)
proxies = kwargs.pop("proxies", None)
use_auth_token = kwargs.pop("use_auth_token", None)
local_files_only = kwargs.pop("local_files_only", False)
revision = kwargs.pop("revision", None)
subfolder = kwargs.pop("subfolder", None)
user_agent = kwargs.pop("user_agent", {})

load_config_kwargs = {
"cache_dir": cache_dir,
"force_download": force_download,
"resume_download": resume_download,
"proxies": proxies,
"use_auth_token": use_auth_token,
"local_files_only": local_files_only,
"revision": revision,
"subfolder": subfolder,
"user_agent": user_agent,
}

config = cls.load_config(pretrained_model_or_path, **load_config_kwargs)
orig_class_name = config["_class_name"]

if "controlnet" in kwargs:
orig_class_name = config["_class_name"].replace("Pipeline", "ControlNetPipeline")

text_2_image_cls = _get_task_class(AUTO_TEXT2IMAGE_PIPELINES_MAPPING, orig_class_name)

kwargs = {**load_config_kwargs, **kwargs}
return text_2_image_cls.from_pretrained(pretrained_model_or_path, **kwargs)

@classmethod
Expand Down Expand Up @@ -535,14 +559,37 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs):
>>> image = pipeline(prompt, image).images[0]
```
"""
config = cls.load_config(pretrained_model_or_path)
cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)
force_download = kwargs.pop("force_download", False)
resume_download = kwargs.pop("resume_download", False)
proxies = kwargs.pop("proxies", None)
use_auth_token = kwargs.pop("use_auth_token", None)
local_files_only = kwargs.pop("local_files_only", False)
revision = kwargs.pop("revision", None)
subfolder = kwargs.pop("subfolder", None)
user_agent = kwargs.pop("user_agent", {})

load_config_kwargs = {
"cache_dir": cache_dir,
"force_download": force_download,
"resume_download": resume_download,
"proxies": proxies,
"use_auth_token": use_auth_token,
"local_files_only": local_files_only,
"revision": revision,
"subfolder": subfolder,
"user_agent": user_agent,
}

config = cls.load_config(pretrained_model_or_path, **load_config_kwargs)
orig_class_name = config["_class_name"]

if "controlnet" in kwargs:
orig_class_name = config["_class_name"].replace("Pipeline", "ControlNetPipeline")

image_2_image_cls = _get_task_class(AUTO_IMAGE2IMAGE_PIPELINES_MAPPING, orig_class_name)

kwargs = {**load_config_kwargs, **kwargs}
return image_2_image_cls.from_pretrained(pretrained_model_or_path, **kwargs)

@classmethod
Expand Down Expand Up @@ -776,14 +823,37 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs):
>>> image = pipeline(prompt, image=init_image, mask_image=mask_image).images[0]
```
"""
config = cls.load_config(pretrained_model_or_path)
cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)
force_download = kwargs.pop("force_download", False)
resume_download = kwargs.pop("resume_download", False)
proxies = kwargs.pop("proxies", None)
use_auth_token = kwargs.pop("use_auth_token", None)
local_files_only = kwargs.pop("local_files_only", False)
revision = kwargs.pop("revision", None)
subfolder = kwargs.pop("subfolder", None)
user_agent = kwargs.pop("user_agent", {})

load_config_kwargs = {
"cache_dir": cache_dir,
"force_download": force_download,
"resume_download": resume_download,
"proxies": proxies,
"use_auth_token": use_auth_token,
"local_files_only": local_files_only,
"revision": revision,
"subfolder": subfolder,
"user_agent": user_agent,
}

config = cls.load_config(pretrained_model_or_path, **load_config_kwargs)
orig_class_name = config["_class_name"]

if "controlnet" in kwargs:
orig_class_name = config["_class_name"].replace("Pipeline", "ControlNetPipeline")

inpainting_cls = _get_task_class(AUTO_INPAINT_PIPELINES_MAPPING, orig_class_name)

kwargs = {**load_config_kwargs, **kwargs}
return inpainting_cls.from_pretrained(pretrained_model_or_path, **kwargs)

@classmethod
Expand Down
27 changes: 27 additions & 0 deletions tests/pipelines/test_pipelines_auto.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,11 @@
# limitations under the License.

import gc
import os
import shutil
import unittest
from collections import OrderedDict
from pathlib import Path

import torch

Expand All @@ -24,6 +27,7 @@
AutoPipelineForInpainting,
AutoPipelineForText2Image,
ControlNetModel,
DiffusionPipeline,
)
from diffusers.pipelines.auto_pipeline import (
AUTO_IMAGE2IMAGE_PIPELINES_MAPPING,
Expand Down Expand Up @@ -81,6 +85,29 @@ def test_from_pipe_consistent_sdxl(self):

assert dict(pipe.config) == original_config

def test_kwargs_local_files_only(self):
repo = "hf-internal-testing/tiny-stable-diffusion-torch"
tmpdirname = DiffusionPipeline.download(repo)
tmpdirname = Path(tmpdirname)

# edit commit_id to so that it's not the latest commit
commit_id = tmpdirname.name
new_commit_id = commit_id + "hug"

ref_dir = tmpdirname.parent.parent / "refs/main"
with open(ref_dir, "w") as f:
f.write(new_commit_id)

new_tmpdirname = tmpdirname.parent / new_commit_id
os.rename(tmpdirname, new_tmpdirname)

try:
AutoPipelineForText2Image.from_pretrained(repo, local_files_only=True)
except OSError:
assert False, "not able to load local files"

shutil.rmtree(tmpdirname.parent.parent)


@slow
class AutoPipelineIntegrationTest(unittest.TestCase):
Expand Down