Add Stable Diffusion #828

fchollet · 2022-09-23T22:28:16Z

This is a working Stable Diffusion text to image model. It reuses a bunch of code from Divam's original implementation. All top-level models are full rewrites as Functional models.

All in all it's about 600 LOC. We can probably go below we further refactoring. The UNet is too declarative right now, it could be refactored to be a lot more concise.

We need to:

Refactor the UNet (DiffusionModel) to be as elegant as possible
Figure out what to do with constants.py
Upload the new weights files somewhere
Further cleanup (remove the temporary test code) and beautify the code
Add some basic unit testing

bhack · 2022-09-23T23:00:20Z

/cc @divamgupta @innat

LukeWood · 2022-09-23T23:22:02Z

So cool, can't wait to have SD ready and in the package...

bhack · 2022-09-23T23:27:53Z

As this is the first materialized example where we are handling NLP+VISION I don't know if it is interesting to re-introduce some of our related topics like:

Quoting @chenmoneygithub
keras-team/keras#16181 (comment)

And you are totally correct, we need to agree on a repo for shared components. We think it makes sense to put things into Keras main repo if both KerasCV and KerasNLP would depend on them.

Me:
#52 (comment) #30 (comment)

LukeWood

Minor changes

keras_cv/models/generative/stable_diffusion/decoder.py

keras_cv/models/generative/stable_diffusion/diffusion_model.py

keras_cv/models/generative/stable_diffusion/stable_diffusion.py

bhack · 2022-09-23T23:56:24Z

keras_cv/models/generative/stable_diffusion/__internal__/layers/group_normalization.py

@@ -0,0 +1 @@
+from tensorflow_addons.layers import GroupNormalization


We need to remove this and duplicate/porting/refactor it in keras-cv or keras repo (#74)

@ianstenbit is working on this for now, it’ll live in core Keras - temporarily as a CV internal API

tanzhenyu

Excited to see this!
A general comment -- are we getting the weights from somewhere, or we will provide a training script later?

keras_cv/models/generative/stable_diffusion/text_encoder.py

tanzhenyu · 2022-09-24T02:02:32Z

keras_cv/models/generative/stable_diffusion/diffusion_model.py

+
+
+def gelu(x):
+    tanh_res = keras.activations.tanh(x * 0.7978845608 * (1 + 0.044715 * (x**2)))


numbers seem pretty magical :-)

This is just a polynomial approximation of gelu. Generally speaking there are a number of constants harcoded in the code, because the code is non-configurable beyond the arguments of StableDiffusion and text_to_image. Its only function is to load the original weights and match the original numerics.

LukeWood · 2022-09-24T02:21:14Z

Excited to see this!
A general comment -- are we getting the weights from somewhere, or we will provide a training script later?

As of now they’re ported. If I can setup LAOIN-5B (hoping to in the next few months…) we can retrain; LAOIN-5B was the dataset this model was original trained on.

…er and the constants file.

LukeWood · 2022-09-24T06:38:00Z

Right now CLIP weight downloading is not lazy:

LukeWood

Don't forget to add an entry to keras_cv/models/init.py!

LukeWood · 2022-09-24T06:59:01Z

Managed to get it working in a notebook!

bhack · 2022-09-24T11:30:25Z

Excited to see this!
A general comment -- are we getting the weights from somewhere, or we will provide a training script later?

As of now they’re ported. If I can setup LAOIN-5B (hoping to in the next few months…) we can retrain; LAOIN-5B was the dataset this model was original trained on.

Are we going to target recent OpenCLIP releases?

https://laion.ai/blog/large-openclip/

keras_cv/models/generative/stable_diffusion/stable_diffusion.py

bhack · 2022-09-24T17:33:52Z

keras_cv/models/generative/stable_diffusion/stable_diffusion.py

@@ -0,0 +1,138 @@
+import numpy as np


Are we sure that we want to handle all these numpy objects and numpy TF experimental wrapped TF ops instead of handling directly Tensor and TF ops?

I am asking this also for coherence:

grep -r "import numpy" * --exclude="*test.py" --exclude-dir="*examples*" --exclude-dir="*test*" --exclude-dir="*benchmarks*"

keras_cv/layers/preprocessing/random_rotation.py:import numpy as np keras_cv/models/object_detection/__internal__.py:import numpy as np keras_cv/models/object_detection/retina_net/retina_net.py:import numpy as np

Then also in these few cases it was used only for np.log and np.pi.
(I don't know if there was a real cause where we could not use pi from python math and/or tf.math.log)

LukeWood · 2022-09-24T18:25:55Z

keras_cv/models/generative/stable_diffusion/clip_tokenizer.py

+
+
+class SimpleTokenizer:
+    def __init__(self, bpe_path: str = default_bpe()):


here's the bug causing your lazy loading to not work:

bpe_path=None bpe_path = bpe_path or default_bpe()

YaoYuanxin · 2022-09-24T22:43:36Z

Hi, I have a fairly basic question. I installed the updated version of keras_cv, but it's showing "no attribute".

Is it because StableDiffusion hasn't been officially released for keras_cv yet?

LukeWood · 2022-09-25T01:00:17Z

keras_cv/models/generative/stable_diffusion/stable_diffusion.py

@@ -126,13 +141,3 @@ def get_initial_parameters(self, timesteps, batch_size, seed=None):
            (batch_size, self.img_height // 8, self.img_width // 8, 4), seed=seed
        )
        return noise, alphas, alphas_prev
-


Feel free to add this to examples/models/stable_diffusion/

right now they’re basically glorified debugging demos; but we could render them someday down the line with proper polish

keras_cv/models/generative/stable_diffusion/__internal__/layers/group_normalization.py

keras_cv/models/generative/stable_diffusion/stable_diffusion.py

fchollet · 2022-09-25T02:25:45Z

I added a docstring, removed the TFA dependency, added a code example, and disabled the golden value test so that it doesn't run on CI (which would be overly expensive -- we don't want to download the weights on CI).

I believe we're ready to go, or pretty close.

LukeWood · 2022-09-25T02:33:30Z

/gcbrun

divamgupta · 2022-09-25T02:34:03Z

There will be a lot of transformer based models released in the future. Should there by separate modules for transformers?

LukeWood

I have reviewed the full PR, and this LGTM. I have also pulled locally and played with the implementation; great job on the readability and organization.

keras_cv/models/__init__.py

keras_cv/models/generative/stable_diffusion/constants.py

LukeWood · 2022-09-25T03:33:27Z

@fchollet you also need to add "regex" to install_requires in setup.py

fchollet · 2022-09-25T03:45:49Z

Somehow I'm able to run format.sh and lint.sh locally just fine, yet it fails on CI.

LukeWood · 2022-09-25T04:50:28Z

/gcbrun

LukeWood · 2022-09-25T04:57:56Z

Adding a dependency to our GCB cluster; one monent...

LukeWood · 2022-09-25T05:02:58Z

/gcbrun

LukeWood · 2022-09-25T05:25:41Z

All checks pass @fchollet ! Thank you for working so swiftly to prepare this port! Merging and rebasing #831 on top.

Excited to have such a powerful offering in our package; and can't wait to show off the performance in the keras.io guide.

* Add Stable Diffusion * Further simplification; add files * Update imports * Further minor simplications * Style fixes * Further beautification. The code is now 500 LOC excluding the tokenizer and the constants file. * Readability improvements * Improve generation loop * Simplify generation code * Fix bpe_path * Add init imports * Minor style fixes * Remove unnecessary dependencies and add file headers * Add test * Add example * Add group normalization layer * Disable test so it doesn't run on CI * Fix code style * Update docs. * Format imports * Remove unused import * Add file header * Add more copyright notices * Add last copyright notice * Add regex requirement * Hopefully last copyright notice

Add Stable Diffusion

51243ea

LukeWood suggested changes Sep 23, 2022

View reviewed changes

keras_cv/models/generative/stable_diffusion/decoder.py Outdated Show resolved Hide resolved

keras_cv/models/generative/stable_diffusion/diffusion_model.py Outdated Show resolved Hide resolved

keras_cv/models/generative/stable_diffusion/stable_diffusion.py Outdated Show resolved Hide resolved

bhack reviewed Sep 23, 2022

View reviewed changes

fchollet added 3 commits September 23, 2022 18:02

Further simplification; add files

9cac708

Update imports

ac3cbdb

Further minor simplications

bc4e148

tanzhenyu reviewed Sep 24, 2022

View reviewed changes

fchollet added 4 commits September 23, 2022 20:30

Style fixes

2124136

Further beautification. The code is now 500 LOC excluding the tokeniz…

bd52b46

…er and the constants file.

Readability improvements

ba8387b

Improve generation loop

45e3013

LukeWood suggested changes Sep 24, 2022

View reviewed changes

Simplify generation code

9c5ddeb

bhack reviewed Sep 24, 2022

View reviewed changes

LukeWood reviewed Sep 24, 2022

View reviewed changes

Fix bpe_path

f077b31

fchollet added 5 commits September 24, 2022 15:47

Add init imports

476b360

Merge branch 'keras-team:master' into master

bda6f81

Minor style fixes

6ab9b8c

Remove unnecessary dependencies and add file headers

4d99ee4

Add test

65d8357

LukeWood reviewed Sep 25, 2022

View reviewed changes

keras_cv/models/generative/stable_diffusion/__internal__/layers/group_normalization.py Outdated Show resolved Hide resolved

ianstenbit reviewed Sep 25, 2022

View reviewed changes

keras_cv/models/generative/stable_diffusion/stable_diffusion.py Show resolved Hide resolved

ianstenbit reviewed Sep 25, 2022

View reviewed changes

keras_cv/models/generative/stable_diffusion/stable_diffusion.py Show resolved Hide resolved

fchollet added 5 commits September 24, 2022 19:08

Add example

917346c

Add group normalization layer

8ef2ff3

Disable test so it doesn't run on CI

d425af0

Fix code style

889f3d0

Update docs.

979f594

Format imports

0265f5a

LukeWood approved these changes Sep 25, 2022

View reviewed changes

fchollet added 3 commits September 24, 2022 19:47

Remove unused import

75d6774

Add file header

93b2c5a

Add more copyright notices

608b27d

LukeWood reviewed Sep 25, 2022

View reviewed changes

keras_cv/models/__init__.py Show resolved Hide resolved

LukeWood reviewed Sep 25, 2022

View reviewed changes

keras_cv/models/generative/stable_diffusion/constants.py Show resolved Hide resolved

Add last copyright notice

bf0e622

Add regex requirement

d824ef3

Hopefully last copyright notice

4706c93

LukeWood merged commit e0c4053 into keras-team:master Sep 25, 2022

bhack mentioned this pull request Sep 25, 2022

Stable diffusion small fixes #836

Merged

3 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add Stable Diffusion #828

Add Stable Diffusion #828

fchollet commented Sep 23, 2022

bhack commented Sep 23, 2022

LukeWood commented Sep 23, 2022

bhack commented Sep 23, 2022

LukeWood left a comment

bhack Sep 23, 2022

LukeWood Sep 24, 2022

tanzhenyu left a comment

tanzhenyu Sep 24, 2022

fchollet Sep 24, 2022

LukeWood commented Sep 24, 2022

LukeWood commented Sep 24, 2022

LukeWood left a comment

LukeWood commented Sep 24, 2022 •

edited

Loading

bhack commented Sep 24, 2022 •

edited

Loading

bhack Sep 24, 2022 •

edited

Loading

bhack Sep 24, 2022 •

edited

Loading

LukeWood Sep 24, 2022

YaoYuanxin commented Sep 24, 2022

LukeWood Sep 25, 2022

fchollet commented Sep 25, 2022

LukeWood commented Sep 25, 2022

divamgupta commented Sep 25, 2022

LukeWood left a comment

LukeWood commented Sep 25, 2022

fchollet commented Sep 25, 2022

LukeWood commented Sep 25, 2022

LukeWood commented Sep 25, 2022

LukeWood commented Sep 25, 2022

LukeWood commented Sep 25, 2022

		@@ -0,0 +1 @@
		from tensorflow_addons.layers import GroupNormalization



		def gelu(x):
		tanh_res = keras.activations.tanh(x * 0.7978845608 * (1 + 0.044715 * (x**2)))



		class SimpleTokenizer:
		def __init__(self, bpe_path: str = default_bpe()):

Add Stable Diffusion #828

Add Stable Diffusion #828

Conversation

fchollet commented Sep 23, 2022

bhack commented Sep 23, 2022

LukeWood commented Sep 23, 2022

bhack commented Sep 23, 2022

LukeWood left a comment

Choose a reason for hiding this comment

bhack Sep 23, 2022

Choose a reason for hiding this comment

LukeWood Sep 24, 2022

Choose a reason for hiding this comment

tanzhenyu left a comment

Choose a reason for hiding this comment

tanzhenyu Sep 24, 2022

Choose a reason for hiding this comment

fchollet Sep 24, 2022

Choose a reason for hiding this comment

LukeWood commented Sep 24, 2022

LukeWood commented Sep 24, 2022

LukeWood left a comment

Choose a reason for hiding this comment

LukeWood commented Sep 24, 2022 • edited Loading

bhack commented Sep 24, 2022 • edited Loading

bhack Sep 24, 2022 • edited Loading

Choose a reason for hiding this comment

bhack Sep 24, 2022 • edited Loading

Choose a reason for hiding this comment

LukeWood Sep 24, 2022

Choose a reason for hiding this comment

YaoYuanxin commented Sep 24, 2022

LukeWood Sep 25, 2022

Choose a reason for hiding this comment

fchollet commented Sep 25, 2022

LukeWood commented Sep 25, 2022

divamgupta commented Sep 25, 2022

LukeWood left a comment

Choose a reason for hiding this comment

LukeWood commented Sep 25, 2022

fchollet commented Sep 25, 2022

LukeWood commented Sep 25, 2022

LukeWood commented Sep 25, 2022

LukeWood commented Sep 25, 2022

LukeWood commented Sep 25, 2022

LukeWood commented Sep 24, 2022 •

edited

Loading

bhack commented Sep 24, 2022 •

edited

Loading

bhack Sep 24, 2022 •

edited

Loading

bhack Sep 24, 2022 •

edited

Loading