Skip to content

GitLab

Switch branch/tag
  1. 08 Feb, 2024 1 commit
  3. 06 Dec, 2023 1 commit
  5. 20 Nov, 2023 2 commits
    • Roy Hvaara's avatar
      [JAX] Replace uses of jax.devices("cpu") with jax.local_devices(backend="cpu") (#5864) · 2695ba8e
      Roy Hvaara authored 
      

An upcoming change to JAX will include non-local (addressable) CPU devices in jax.devices() when JAX is used multicontroller-style, where there are multiple Python processes.

This change preserves the current behavior by replacing uses of jax.devices("cpu"), which previously only returned local devices, with jax.local_devices("cpu"), which will return local devices both now and in the future.

This change is always safe (i.e., it should always preserve the previous behavior), but it may sometimes be unnecessary if code is never used in a multicontroller setting.
Co-authored-by: default avatarPeter Hawkins <phawkins@google.com>
      2695ba8e
    • Kashif Rasul's avatar
      [Styling] stylify using ruff (#5841) · 6b04d61c
      Kashif Rasul authored 
      * ruff format

* not need to use doc-builder's black styling as the doc is styled in ruff

* make fix-copies

* comment

* use run_ruff
      6b04d61c
  7. 09 Oct, 2023 1 commit
  9. 22 Sep, 2023 1 commit
    • Pedro Cuenca's avatar
      SDXL flax (#4254) · 3651b14c
      Pedro Cuenca authored 
      

* support transformer_layers_per block in flax UNet

* add support for text_time additional embeddings to Flax UNet

* rename attention layers for VAE

* add shape asserts when renaming attention layers

* transpose VAE attention layers

* add pipeline flax SDXL code [WIP]

* continue add pipeline flax SDXL code [WIP]

* cleanup

* Working on JIT support

Fixed prompt embedding shapes so they work in parallel mode. Assuming we
always have both text encoders for now, for simplicity.

* Fixing embeddings (untested)

* Remove spurious line

* Shard guidance_scale when jitting.

* Decode images

* Fix sharding

* style

* Refiner UNet can be loaded.

* Refiner / img2img pipeline

* Allow latent outputs from base and latent inputs in refiner

This makes it possible to chain base + refiner without having to use the
vae decoder in the base model, the vae encoder in the refiner, skipping
conversions to/from PIL, and avoiding TPU <-> CPU memory copies.

* Adapt to FlaxCLIPTextModelOutput

* Update Flax XL pipeline to FlaxCLIPTextModelOutput

* make fix-copies

* make style

* add euler scheduler

* Fix import

* Fix copies, comment unused code.

* Fix SDXL Flax imports

* Fix euler discrete begin

* improve init import

* finish

* put discrete euler in init

* fix flax euler

* Fix more

* make style

* correct init

* correct init

* Temporarily remove FlaxStableDiffusionXLImg2ImgPipeline

* correct pipelines

* finish

---------
Co-authored-by: default avatarMartin Müller <martin.muller.me@gmail.com>
Co-authored-by: default avatarpatil-suraj <surajp815@gmail.com>
Co-authored-by: default avatarPatrick von Platen <patrick.v.platen@gmail.com>
      3651b14c
  11. 15 Aug, 2023 1 commit
  13. 30 Jun, 2023 1 commit
    • Steven Liu's avatar
      [docs] Model API (#3562) · 174dcd69
      Steven Liu authored 
      * add modelmixin and unets

* remove old model page

* minor fixes

* fix unet2dcondition

* add vqmodel and autoencoderkl

* add rest of models

* fix autoencoderkl path

* fix toctree

* fix toctree again

* apply feedback

* apply feedback

* fix copies

* fix controlnet copy

* fix copies
      174dcd69
  15. 01 Mar, 2023 1 commit
  17. 30 Dec, 2022 1 commit
  19. 29 Dec, 2022 1 commit
    • Simon Kirsten's avatar
      Flax: Fix img2img and align with other pipeline (#1824) · ab0e92fd
      Simon Kirsten authored 
      

* Flax: Add components function

* Flax: Fix img2img and align with other pipeline

* Flax: Fix PRNGKey type

* Refactor strength to start_timestep

* Fix preprocess images

* Fix processed_images dimen

* latents.shape -> latents_shape

* Fix typo

* Remove "static" comment

* Remove unnecessary optional types in _generate

* Apply doc-builder code style.
Co-authored-by: default avatarPedro Cuenca <pedro@huggingface.co>
      ab0e92fd
  21. 14 Dec, 2022 1 commit
  23. 07 Dec, 2022 1 commit
  25. 28 Nov, 2022 1 commit
  27. 24 Oct, 2022 1 commit
  29. 21 Oct, 2022 1 commit
  31. 12 Oct, 2022 3 commits
  33. 05 Oct, 2022 1 commit
  35. 03 Oct, 2022 1 commit
    • Pedro Cuenca's avatar
      Fix import with Flax but without PyTorch (#688) · 688031c5
      Pedro Cuenca authored 
      * Don't use `load_state_dict` if torch is not installed.

* Define `SchedulerOutput` to use torch or flax arrays.

* Don't import LMSDiscreteScheduler without torch.

* Create distinct FlaxSchedulerOutput.

* Additional changes required for FlaxSchedulerMixin

* Do not import torch pipelines in Flax.

* Revert "Define `SchedulerOutput` to use torch or flax arrays."

This reverts commit f653140134b74d9ffec46d970eb46925fe3a409d.

* Prefix Flax scheduler outputs for consistency.

* make style

* FlaxSchedulerOutput is now a dataclass.

* Don't use f-string without placeholders.

* Add blank line.

* Style (docstrings)
      688031c5
  37. 29 Sep, 2022 1 commit
  39. 24 Sep, 2022 1 commit
  41. 22 Sep, 2022 1 commit
  43. 21 Sep, 2022 2 commits
  45. 20 Sep, 2022 2 commits
    • Mishig Davaadorj's avatar
      FlaxDiffusionPipeline & FlaxStableDiffusionPipeline (#559) · d934d3d7
      Mishig Davaadorj authored 
      

* 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: default avatarPatrick von Platen <patrick.v.platen@gmail.com>
Co-authored-by: default avatarPedro Cuenca <pedro@latenitesoft.com>
      d934d3d7
    • Younes Belkada's avatar
      Add `from_pt` argument in `.from_pretrained` (#527) · 0902449e
      Younes Belkada authored 
      * 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!
      0902449e
  47. 19 Sep, 2022 2 commits
  49. 15 Sep, 2022 2 commits
    • Mishig Davaadorj's avatar
      Add `init_weights` method to `FlaxMixin` (#513) · fb5468a6
      Mishig Davaadorj authored 
      

* Add `init_weights` method to `FlaxMixin`

* Rn `random_state` -> `shape_state`

* `PRNGKey(0)` for `jax.eval_shape`

* No allow mismatched sizes

* Update src/diffusers/modeling_flax_utils.py
Co-authored-by: default avatarSuraj Patil <surajp815@gmail.com>

* Update src/diffusers/modeling_flax_utils.py
Co-authored-by: default avatarSuraj Patil <surajp815@gmail.com>

* docstring diffusers
Co-authored-by: default avatarSuraj Patil <surajp815@gmail.com>
      fb5468a6
    • Kashif Rasul's avatar
      Karras VE, DDIM and DDPM flax schedulers (#508) · b34be039
      Kashif Rasul authored 
      * beta never changes removed from state

* fix typos in docs

* removed unused var

* initial ddim flax scheduler

* import

* added dummy objects

* fix style

* fix typo

* docs

* fix typo in comment

* set return type

* added flax ddom

* fix style

* remake

* pass PRNG key as argument and split before use

* fix doc string

* use config

* added flax Karras VE scheduler

* make style

* fix dummy

* fix ndarray type annotation

* replace returns a new state

* added lms_discrete scheduler

* use self.config

* add_noise needs state

* use config

* use config

* docstring

* added flax score sde ve

* fix imports

* fix typos
      b34be039
  51. 14 Sep, 2022 1 commit