allow loading pipe from normal repo

0a1d4c58 · Patrick von Platen · 7f6a36c3 · 0a1d4c58
Commit 0a1d4c58 authored Jun 07, 2022 by Patrick von Platen
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

src/diffusers/pipeline_utils.py src/diffusers/pipeline_utils.py +4 -1

No files found.
--- a/src/diffusers/pipeline_utils.py
+++ b/src/diffusers/pipeline_utils.py
@@ -118,7 +118,10 @@ class DiffusionPipeline(ConfigMixin):

            load_method = getattr(class_obj, load_method_name)

-            loaded_sub_model = load_method(os.path.join(cached_folder, name))
+            if os.path.dir(os.path.join(cached_folder, name)):
+                loaded_sub_model = load_method(os.path.join(cached_folder, name))
+            else:
+                loaded_sub_model = load_method(cached_folder)

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