Add VAEBackbone and use it for SD3
#1892
Conversation
| return input_shape | ||
|
|
||
|
|
||
| class ResNetBlock(keras.layers.Layer): |
There was a problem hiding this comment.
can the ResNet implementation we already have be reused here?
There was a problem hiding this comment.
The normalization layer and activation function are different.
| return output_shape | ||
|
|
||
|
|
||
| class AutoEncoder(Backbone): |
There was a problem hiding this comment.
You should add this in our model folder.
Create a model folder named autoencoder
add auto_encoder_layers.py - add all the layers and tests there
add autoencoder_backbone.py and add this backbone
add basic backbone tests.
It is fine to not expose it and only use it in stable diffusion
There was a problem hiding this comment.
Got it. This is better!
I have added a folder called vae and I believe it can be also used for FLUX models.
sampathweb
added
kokoro:force-run
james77777778
force-pushed
the
add-sd3-autoencoder
branch
from
e4dd23c
to
132b886
Compare
james77777778
changed the title
VAEBackbone and use it for SD3
|
@divyashreepathihalli Colab for demo: https://colab.research.google.com/drive/1ucx2lEck1ZO3HVTjZ4qgVsQrBWrHWfVQ I can add image2image and inpaint after this PR is merged. |
keras_hub/src/models/stable_diffusion_3/flow_match_euler_discrete_scheduler.py
Show resolved
Hide resolved
james77777778
force-pushed
the
add-sd3-autoencoder
branch
from
132b886
to
22db582
Compare
keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_backbone_test.py
Show resolved
Hide resolved
| info = "Audio shape: " | ||
| info += highlight_shape(audio_converter.audio_shape()) | ||
| add_layer(audio_converter, info) | ||
| if preprocessor and isinstance(preprocessor, keras.Layer): |
There was a problem hiding this comment.
can you please explain why you have made these changes?
There was a problem hiding this comment.
This is needed for a nested preprocessor. I have added a comment for this change:
# Since the preprocessor might be nested with multiple `Tokenizer`,
# `ImageConverter`, `AudioConverter` and even other `Preprocessor`
# instances, we should recursively iterate through them.
This PR adds the VAE encoder to SD3 and also fixes the summary of
Taskto recursively parse the info, rather than relying on the layer attribute.The preset on kaggle has also been updated.
Demo colab:
https://colab.research.google.com/drive/1ucx2lEck1ZO3HVTjZ4qgVsQrBWrHWfVQ?usp=sharing