GitLab

Signed-off-by: Ryan Russell <git@ryanrussell.org>
Signed-off-by: Ryan Russell <git@ryanrussell.org>

* add grad ckpt to downsample blocks

* make it work

* don't pass gradient_checkpointing to upsample block

* add tests for UNet2DConditionModel

* add test_gradient_checkpointing

* add gradient_checkpointing for up and down blocks

* add functions to enable and disable grad ckpt

* remove the forward argument

* better naming

* make supports_gradient_checkpointing private

* Handle the PIL.Image.Resampling deprecation

* style

Signed-off-by: Ryan Russell <git@ryanrussell.org>
Signed-off-by: Ryan Russell <git@ryanrussell.org>

* Optionally return state in from_config.

Useful for Flax schedulers.

* has_state is now a property, make check more strict.

I don't check the class is `SchedulerMixin` to prevent circular
dependencies. It should be enough that the class name starts with "Flax"
the object declares it "has_state" and the "create_state" exists too.

* Use state in pipeline from_pretrained.

* Make style

* Add smoke tests for the training examples

* upd

* use a dummy dataset

* mark as slow

* cleanup

* Update test cases

* naming

replace `dropout_prob` by `dropout` in `vae`

* Fix typo in docstring.

* Allow dtype to be overridden on model load.

This may be a temporary solution until #567 is addressed.

* Create latents in float32

The denoising loop always computes the next step in float32, so this
would fail when using `bfloat16`.

Fix params replication when sing the dummy checker.

* [Flax] Fix unet and ddim scheduler

* correct

* finish

* WIP: flax FlaxDiffusionPipeline & FlaxStableDiffusionPipeline

* todo comment

* Fix imports

* Fix imports

* add dummies

* Fix empty init

* make pipeline work

* up

* Use Flax schedulers (typing, docstring)

* Wrap model imports inside availability checks.

* more updates

* make sure flax is not broken

* make style

* more fixes

* up
Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
Co-authored-by: Pedro Cuenca <pedro@latenitesoft.com>

* use FlaxPreTrainedModel for flax safety module

* fix name

* fix one more

* Apply suggestions from code review

* Add the K-LMS scheduler to the inpainting pipeline + tests

* Remove redundant casts

* Fix BaseOutput initialization from dict

* style

* Simplify post-init, add tests

* remove debug

* rename weights to align with PT

* DiagonalGaussianDistribution => FlaxDiagonalGaussianDistribution

* fix name

* first commit:

- add `from_pt` argument in `from_pretrained` function
- add `modeling_flax_pytorch_utils.py` file

* small nit

- fix a small nit - to not enter in the second if condition

* major changes

- modify FlaxUnet modules
- first conversion script
- more keys to be matched

* keys match

- now all keys match
- change module names for correct matching
- upsample module name changed

* working v1

- test pass with atol and rtol= `4e-02`

* replace unsued arg

* make quality

* add small docstring

* add more comments

- add TODO for embedding layers

* small change

- use `jnp.expand_dims` for converting `timesteps` in case it is a 0-dimensional array

* add more conditions on conversion

- add better test to check for keys conversion

* make shapes consistent

- output `img_w x img_h x n_channels` from the VAE

* Revert "make shapes consistent"

This reverts commit 4cad1aeb4aeb224402dad13c018a5d42e96267f6.

* fix unet shape

- channels first!

* Fix a setting bug

* Fix typos

* Reverted params to parms

* Fix CrossAttention._sliced_attention
Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>

* add CLIPGuidedStableDiffusion

* add credits

* add readme

* style

* add clip prompt

* fnfix cond_n

* fix cond fn

* fix cond fn for lms

* [Tests] Mark the ncsnpp model tests as slow

* style

* make_reports

* add test utils

* style

* style

* [Flax] Add Vae

* correct

* Apply suggestions from code review
Co-authored-by: Suraj Patil <surajp815@gmail.com>

* Finish
Co-authored-by: Suraj Patil <surajp815@gmail.com>

* Fix _upsample_2d
Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>

Flax: ignore dtype for configuration.

This makes it possible to save models and configuration files.

* Starting to integrate safety checker.

* Fix initialization of CLIPVisionConfig

* Remove commented lines.

* make style

* Remove unused import

* Pass dtype to modules
Co-authored-by: Suraj Patil <surajp815@gmail.com>

* Pass dtype to modules
Co-authored-by: Suraj Patil <surajp815@gmail.com>
Co-authored-by: Suraj Patil <surajp815@gmail.com>

* remove match_shape

* ported fixes from #479 to flax

* remove unused argument

* typo

* remove warnings

* add `get_modified_files.py`

- file copied from https://github.com/huggingface/transformers/blob/main/utils/get_modified_files.py

* make fixup

* [Config] improve logging

* finish