Lazy Import for Diffusers #4829

DN6 · 2023-08-29T12:51:36Z

What does this PR do?

Adds lazy import functionality to Diffusers, similar to what exists in Transformers

Benchmark:

Initially tested import speed up by running time python -c "import diffusers" with all backends (torch, transformers, flax, onnxruntime) installed.

With Lazy Import

real    0m0.417s
user    0m0.714s
sys     0m0.499s

Without Lazy Import

real    0m5.391s
user    0m5.299s
sys     0m1.273s

This PR:

Adds lazy import to the following modules

diffusers
models
pipelines
schedulers

Moves objects defined in the __init__.py files of a module to their own dedicated files. e.g. StableDiffusionPipelineOutput is now in a pipeline_output.py file within the stable_diffusion module. This follows the same format as transformers and also keeps __init__.py reserved for import related code.

TODO:

utils/check_dummies behaviour and the corresponding test is broken because of the new import structure. This is the last thing that needs to be addressed before merging.

Fixes # (issue)
#4260

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline?
Did you read our philosophy doc (important for complex PRs)?
Was this discussed/approved via a Github issue or the forum? Please add a link to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

patrickvonplaten · 2023-08-29T16:56:36Z

I'd actually do one big PR here and we can try to merge it quickly to not run into merge conflicts :-)

We should measure the speed-up of importing things once this is merged and also try to clean up many of these imports inside functions in the loaders.py file, e.g.:

diffusers/src/diffusers/loaders.py

Line 274 in a76f2ad

from .models.attention_processor import (

DN6 · 2023-09-06T06:51:51Z

@patrickvonplaten
I added type checking to the diffusers init and moved the accelerate import out of the utils init and changed the relevant files.

The import speed should be improved. Results on my machine with all backends + accelerate installed.

time python -c "import diffusers"

real    0m0.417s
user    0m0.714s
sys     0m0.499s

Could you verify you're seeing is similar numbers after the changes?

patrickvonplaten · 2023-09-06T08:54:41Z

Ok I think I found one last package that slows down the import which is xformers: https://github.com/facebookresearch/xformers . I'd say most of our users have xformers installed - can you try to also make sure xformers is not important by default?

DN6 · 2023-09-06T10:16:04Z

@patrickvonplaten changed it so that we use _torch_version = importlib_metadata.version("torch") instead of import torch; torch.__version__ when checking for xformers availability.

patrickvonplaten · 2023-09-06T14:16:30Z

@DN6 could you take a look at the final failing tests or do you need help here?

patrickvonplaten · 2023-09-06T14:17:14Z

Think we also need to integrate the changes from the just merged Wuerstchen PR

DN6 · 2023-09-07T17:31:18Z

@patrickvonplaten
Fixed conflicts and added Wuerstchen.

Seeing some strange behaviour in the Fast tests for PRs / Fast PyTorch Models & Schedulers CPU tests.

When running

python -m pytest -n 2 --max-worker-restart=0 --dist=loadfile     -s -v -k "not Flax and not Onnx"  tests/models tests/schedulers

All tests pass

The failures occur when tests/others is also included. I isolated it to this test tests/others/test_check_copies.py

If you run this test with the model/scheduler test, you get a few failures in the scheduler tests.

FAILED tests/schedulers/test_schedulers.py::SchedulerBaseTests::test_save_load_compatible_schedulers - AssertionError: assert 'The config a...ation file.\n' == 'The config a...ation file.\n'
FAILED tests/schedulers/test_schedulers.py::SchedulerBaseTests::test_save_load_from_different_config - AssertionError: assert 'The config a...ation file.\n' == ''
FAILED tests/schedulers/test_schedulers.py::SchedulerBaseTests::test_save_load_from_different_config_comp_schedulers - assert 'The config a...ult values.\n' == "{'f'} was no...ult values.\n"

Any thoughts on what's happening? I can't seem to figure it out.

DN6 · 2023-09-08T10:59:46Z

utils/check_copies.py

-# This is to make sure the diffusers module imported is the one in the repo.
-spec = importlib.util.spec_from_file_location(
-    "diffusers",
-    os.path.join(DIFFUSERS_PATH, "__init__.py"),
-    submodule_search_locations=[DIFFUSERS_PATH],
-)
-diffusers_module = spec.loader.load_module()
-
-


@patrickvonplaten This section is causing the failing fast tests for models and schedulers. diffusers_module doesn't appear to be used anywhere. So I removed the snippet.

LMK if there's something I've missed here.

That works!

patrickvonplaten · 2023-09-11T07:56:15Z

Alright, great job - let's get this one in to avoid any more merge conflicts.
Over the next days, let's be aware of test failures that might be related to this cc @DN6

patrickvonplaten · 2023-09-11T08:13:59Z

src/diffusers/pipelines/stable_diffusion/__init__.py

+for name, value in _dummy_objects.items():
+    setattr(sys.modules[__name__], name, value)
+for name, value in _additional_imports.items():
+    setattr(sys.modules[__name__], name, value)


We also need the "If TYPE_CHECKING" branch here

Abhinay1997 · 2023-09-12T03:55:28Z

@DN6, Looks like StableDiffusionGLIGENPipeline is added twice (See Line 238-240) in src/diffusers/__init__.py. I don't think it'll affect anything but just letting you know.

* initial commit * move modules to import struct * add dummy objects and _LazyModule * add lazy import to schedulers * clean up unused imports * lazy import on models module * lazy import for schedulers module * add lazy import to pipelines module * lazy import altdiffusion * lazy import audio diffusion * lazy import audioldm * lazy import consistency model * lazy import controlnet * lazy import dance diffusion ddim ddpm * lazy import deepfloyd * lazy import kandinksy * lazy imports * lazy import semantic diffusion * lazy imports * lazy import stable diffusion * move sd output to its own module * clean up * lazy import t2iadapter * lazy import unclip * lazy import versatile and vq diffsuion * lazy import vq diffusion * helper to fetch objects from modules * lazy import sdxl * lazy import txt2vid * lazy import stochastic karras * fix model imports * fix bug * lazy import * clean up * clean up * fixes for tests * fixes for tests * clean up * remove import of torch_utils from utils module * clean up * clean up * fix mistake import statement * dedicated modules for exporting and loading * remove testing utils from utils module * fixes from merge conflicts * Update src/diffusers/pipelines/kandinsky2_2/__init__.py * fix docs * fix alt diffusion copied from * fix check dummies * fix more docs * remove accelerate import from utils module * add type checking * make style * fix check dummies * remove torch import from xformers check * clean up error message * fixes after upstream merges * dummy objects fix * fix tests * remove unused module import --------- Co-authored-by: Patrick von Platen <[email protected]>

DN6 added 3 commits August 25, 2023 11:19

initial commit

ace9bf4

move modules to import struct

8b0b0cf

add dummy objects and _LazyModule

e0b971a

DN6 requested a review from patrickvonplaten August 29, 2023 13:01

DN6 added 25 commits August 30, 2023 07:14

merge upstream

7ed3edb

add lazy import to schedulers

b39ab6d

clean up unused imports

25b52de

lazy import on models module

0f2cbb2

lazy import for schedulers module

be58b3c

add lazy import to pipelines module

9f536bf

lazy import altdiffusion

cd1aa81

lazy import audio diffusion

c37a7fa

lazy import audioldm

52b9ea7

lazy import consistency model

3ecae3f

lazy import controlnet

7de8650

lazy import dance diffusion ddim ddpm

784b861

lazy import deepfloyd

5672ea4

lazy import kandinksy

d2a9cfe

lazy imports

3f23b3b

lazy import semantic diffusion

d28da56

lazy imports

7b4e7c3

lazy import stable diffusion

b25e77e

move sd output to its own module

298670b

clean up

598d87f

lazy import t2iadapter

13c8787

lazy import unclip

2208e27

lazy import versatile and vq diffsuion

4483aca

lazy import vq diffusion

5009fe3

helper to fetch objects from modules

7debdd7

DN6 added 2 commits September 6, 2023 06:38

add type checking

8b8547a

make style

775f3ca

patrickvonplaten and others added 3 commits September 6, 2023 11:02

fix check dummies

eedf9e1

remove torch import from xformers check

2b82164

clean up error message

e1cb44c

Merge branch 'main' into lazy-import

d89f569

DN6 added 4 commits September 7, 2023 10:09

Merge branch 'main' into lazy-import

caab88c

fixes after upstream merges

287cfa7

Merge branch 'main' into lazy-import

5fb4ce0

dummy objects fix

a5ce781

DN6 added 2 commits September 8, 2023 09:13

fix tests

3a0a45c

remove unused module import

52a76ee

DN6 commented Sep 8, 2023

View reviewed changes

patrickvonplaten merged commit b6e0b01 into main Sep 11, 2023

patrickvonplaten reviewed Sep 11, 2023

View reviewed changes

younesbelkada mentioned this pull request Sep 11, 2023

[core] Fix import of randn_tensor huggingface/trl#751

Merged

kashif deleted the lazy-import branch September 11, 2023 19:06

Abhinay1997 added a commit to Abhinay1997/diffusers that referenced this pull request Sep 12, 2023

Compatability with changes from huggingface#4829 lazy-loading

8ad40e1

DN6 mentioned this pull request Sep 12, 2023

remove extra gligen in import #4987

Merged

6 tasks

fmmoret mentioned this pull request Feb 21, 2025

[Feature]: Improve import speed BerriAI/litellm#7605

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Lazy Import for Diffusers #4829

Lazy Import for Diffusers #4829

Uh oh!

DN6 commented Aug 29, 2023 •

edited by patrickvonplaten

Loading

Uh oh!

patrickvonplaten commented Aug 29, 2023

Uh oh!

DN6 commented Sep 6, 2023

Uh oh!

patrickvonplaten commented Sep 6, 2023

Uh oh!

DN6 commented Sep 6, 2023

Uh oh!

patrickvonplaten commented Sep 6, 2023

Uh oh!

patrickvonplaten commented Sep 6, 2023

Uh oh!

DN6 commented Sep 7, 2023

Uh oh!

DN6 Sep 8, 2023

Uh oh!

patrickvonplaten Sep 11, 2023

Uh oh!

patrickvonplaten commented Sep 11, 2023

Uh oh!

patrickvonplaten Sep 11, 2023

Uh oh!

Abhinay1997 commented Sep 12, 2023 •

edited

Loading

Uh oh!

Uh oh!

Lazy Import for Diffusers #4829

Lazy Import for Diffusers #4829

Uh oh!

Conversation

DN6 commented Aug 29, 2023 • edited by patrickvonplaten Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Benchmark:

With Lazy Import

Without Lazy Import

Before submitting

Who can review?

Uh oh!

patrickvonplaten commented Aug 29, 2023

Uh oh!

DN6 commented Sep 6, 2023

Uh oh!

patrickvonplaten commented Sep 6, 2023

Uh oh!

DN6 commented Sep 6, 2023

Uh oh!

patrickvonplaten commented Sep 6, 2023

Uh oh!

patrickvonplaten commented Sep 6, 2023

Uh oh!

DN6 commented Sep 7, 2023

Uh oh!

DN6 Sep 8, 2023

Choose a reason for hiding this comment

Uh oh!

patrickvonplaten Sep 11, 2023

Choose a reason for hiding this comment

Uh oh!

patrickvonplaten commented Sep 11, 2023

Uh oh!

patrickvonplaten Sep 11, 2023

Choose a reason for hiding this comment

Uh oh!

Abhinay1997 commented Sep 12, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

DN6 commented Aug 29, 2023 •

edited by patrickvonplaten

Loading

Abhinay1997 commented Sep 12, 2023 •

edited

Loading