Remove JAX/Flax - #14169
Conversation
Deletes all Flax models, schedulers, pipelines, and loading utilities, the from_flax loading path, flax/jax packaging extras, Flax CI filters, Flax training examples, and Flax content in translated docs. Hub compatibility is preserved: Flax class names in old model_index.json files are still stripped to their PyTorch equivalents, and *.msgpack files in repos are still ignored during download.
Flax-saved checkpoint metadata is no longer recognized: class names like FlaxStableDiffusionPipeline in model_index.json now fail to resolve instead of being remapped to their PyTorch equivalents. *.msgpack files in hub repos are still ignored during download.
The example scripts are standalone and still work with diffusers<=0.39.x. A warning note is added to each Flax README section pointing out the version requirement now that the library itself no longer ships Flax.
github-actions
Bot
added
documentation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
sayakpaul
left a comment
sayakpaul
left a comment
There was a problem hiding this comment.
Thanks! I think we're quite close to merge.
|
|
||
| if hf_quantizer is not None: | ||
| hf_quantizer.validate_environment(torch_dtype=torch_dtype, from_flax=from_flax, device_map=device_map) | ||
| hf_quantizer.validate_environment(torch_dtype=torch_dtype, device_map=device_map) |
There was a problem hiding this comment.
I think we should also remove the argument from_flax from the validate_environment() method?
There was a problem hiding this comment.
The base class just accepts args and kwargs.
diffusers/src/diffusers/quantizers/base.py
Line 163 in e97c271
| def test_stable_diffusion_no_safety_checker(self): | ||
| pipe = StableDiffusionPipeline.from_pretrained( | ||
| "hf-internal-testing/tiny-stable-diffusion-lms-pipe", safety_checker=None | ||
| "hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None |
There was a problem hiding this comment.
Why is this test getting affected?
There was a problem hiding this comment.
The repo id uses Flax objects, so these tests fail since we no longer support loading/parsing these objects.
In this particular test, we are loading the torch weights even though there are Flax objects defined in the config.
https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-lms-pipe/blob/main/unet/config.json
The part of the code that allows loading in this way involves stripping out the "Flax" prefix from the model class names. This is no longer supported. So we point to a repo id with proper torch objects.
| assert final_tokenizer_size == orig_tokenizer_size | ||
| assert final_emb_size == orig_emb_size | ||
|
|
||
| def test_download_ignore_files(self): |
There was a problem hiding this comment.
This change doesn't look relevant.
There was a problem hiding this comment.
The repo id : hf-internal-testing/tiny-stable-diffusion-pipe-ignore-files uses Flax objects in the config files.
However, when looking through the code, it seems that _ignore_files is not properly applied anyway. We download based on allow_patterns and ignore_patterns, and the way these are derived can lead to situations where ignore_files is omitted entirely.
Reproducer:
import os, json
from huggingface_hub import hf_hub_download
from diffusers import DiffusionPipeline
repo = "hf-internal-testing/tiny-stable-diffusion-pipe-ignore-files"
# The repo's model_index.json explicitly asks the downloader to skip these:
ignored = json.load(open(hf_hub_download(repo, "model_index.json")))["_ignore_files"]
# -> ['vae/diffusion_pytorch_model.bin', 'text_encoder/config.json']
folder = DiffusionPipeline.download(repo)
present = {
os.path.relpath(os.path.join(r, f), folder)
for r, _, fs in os.walk(folder) for f in fs
}
for f in ignored:
assert f not in present, f"{f} was listed in _ignore_files but downloaded anyway"Also searching across all public repos, we don't see this field included anywhere.
https://huggingface.co/search/full-text?q=_ignore_files
So I think we can safely remove since it doesn't seem to be effective.
# Conflicts: # tests/pipelines/test_pipelines.py
|
Failing test is unrelated. |
3 participants
What does this PR do?
Remove support for JAX/Flax models and pipelines.
Fixes # (issue)
Before submitting
.ai/review-rules.md?documentation guidelines, and
here are tips on formatting docstrings.
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.