diff --git a/.github/ISSUE_TEMPLATE/bug-report.yml b/.github/ISSUE_TEMPLATE/bug-report.yml index a0517725284e..7f741180fa3d 100644 --- a/.github/ISSUE_TEMPLATE/bug-report.yml +++ b/.github/ISSUE_TEMPLATE/bug-report.yml @@ -101,7 +101,7 @@ body: Questions on Documentation: @stevhliu - Questions on JAX- and MPS-related things: @pcuenca + Questions on MPS-related things: @pcuenca Questions on audio pipelines: @sanchit-gandhi diff --git a/.github/workflows/nightly_tests.yml b/.github/workflows/nightly_tests.yml index daa4361a417f..678d106a5fc7 100644 --- a/.github/workflows/nightly_tests.yml +++ b/.github/workflows/nightly_tests.yml @@ -95,7 +95,7 @@ jobs: CUBLAS_WORKSPACE_CONFIG: :16:8 run: | pytest -n 1 --max-worker-restart=0 --dist=loadfile \ - -k "not Flax and not Onnx" \ + -k "not Onnx" \ --make-reports=tests_pipeline_${{ matrix.module }}_cuda \ --report-log=tests_pipeline_${{ matrix.module }}_cuda.log \ tests/pipelines/${{ matrix.module }} @@ -151,7 +151,7 @@ jobs: CUBLAS_WORKSPACE_CONFIG: :16:8 run: | pytest -n 1 --max-worker-restart=0 --dist=loadfile \ - -k "not Flax and not Onnx" \ + -k "not Onnx" \ --make-reports=tests_torch_${{ matrix.module }}_cuda \ --report-log=tests_torch_${{ matrix.module }}_cuda.log \ tests/${{ matrix.module }} @@ -312,7 +312,7 @@ jobs: CUBLAS_WORKSPACE_CONFIG: :16:8 run: | pytest -n 1 --max-worker-restart=0 --dist=loadfile \ - -k "not Flax and not Onnx" \ + -k "not Onnx" \ --make-reports=tests_torch_minimum_version_cuda \ tests/models/test_modeling_common.py \ tests/pipelines/test_pipelines_common.py \ @@ -466,7 +466,6 @@ jobs: run_nightly_pipeline_level_quantization_tests, # run_nightly_onnx_tests, torch_minimum_version_cuda_tests, - # run_flax_tpu_tests ] if: always() runs-on: diff --git a/.github/workflows/pr_modular_tests.yml b/.github/workflows/pr_modular_tests.yml index 91d710c119a0..6e62e64ca18f 100644 --- a/.github/workflows/pr_modular_tests.yml +++ b/.github/workflows/pr_modular_tests.yml @@ -139,7 +139,7 @@ jobs: - name: Run fast PyTorch Pipeline CPU tests run: | pytest -n 8 --max-worker-restart=0 --dist=loadfile \ - -k "not Flax and not Onnx" \ + -k "not Onnx" \ --make-reports=tests_torch_cpu_modular_pipelines \ tests/modular_pipelines diff --git a/.github/workflows/pr_tests.yml b/.github/workflows/pr_tests.yml index a7c244481117..bdf869943a28 100644 --- a/.github/workflows/pr_tests.yml +++ b/.github/workflows/pr_tests.yml @@ -136,7 +136,7 @@ jobs: if: ${{ matrix.config.framework == 'pytorch_pipelines' }} run: | pytest -n 8 --max-worker-restart=0 --dist=loadfile \ - -k "not Flax and not Onnx" \ + -k "not Onnx" \ --make-reports=tests_${{ matrix.config.report }} \ tests/pipelines @@ -144,7 +144,7 @@ jobs: if: ${{ matrix.config.framework == 'pytorch_models' }} run: | pytest -n 4 --max-worker-restart=0 --dist=loadfile \ - -k "not Flax and not Onnx and not Dependency" \ + -k "not Onnx and not Dependency" \ --make-reports=tests_${{ matrix.config.report }} \ tests/models tests/schedulers tests/hooks tests/others diff --git a/.github/workflows/pr_tests_gpu.yml b/.github/workflows/pr_tests_gpu.yml index 540d0966c0be..e972c80da1d6 100644 --- a/.github/workflows/pr_tests_gpu.yml +++ b/.github/workflows/pr_tests_gpu.yml @@ -163,13 +163,13 @@ jobs: run: | if [ "${{ matrix.module }}" = "ip_adapters" ]; then pytest -n 1 --max-worker-restart=0 --dist=loadfile \ - -k "not Flax and not Onnx" \ + -k "not Onnx" \ --make-reports=tests_pipeline_${{ matrix.module }}_cuda \ tests/pipelines/${{ matrix.module }} else pattern=$(cat ${{ steps.extract_tests.outputs.pattern_file }}) pytest -n 1 --max-worker-restart=0 --dist=loadfile \ - -k "not Flax and not Onnx and $pattern" \ + -k "not Onnx and $pattern" \ --make-reports=tests_pipeline_${{ matrix.module }}_cuda \ tests/pipelines/${{ matrix.module }} fi @@ -235,10 +235,10 @@ jobs: run: | pattern=$(cat ${{ steps.extract_tests.outputs.pattern_file }}) if [ -z "$pattern" ]; then - pytest -n 1 --max-worker-restart=0 --dist=loadfile -k "not Flax and not Onnx" tests/${{ matrix.module }} \ + pytest -n 1 --max-worker-restart=0 --dist=loadfile -k "not Onnx" tests/${{ matrix.module }} \ --make-reports=tests_torch_cuda_${{ matrix.module }} else - pytest -n 1 --max-worker-restart=0 --dist=loadfile -k "not Flax and not Onnx and $pattern" tests/${{ matrix.module }} \ + pytest -n 1 --max-worker-restart=0 --dist=loadfile -k "not Onnx and $pattern" tests/${{ matrix.module }} \ --make-reports=tests_torch_cuda_${{ matrix.module }} fi diff --git a/.github/workflows/push_tests.yml b/.github/workflows/push_tests.yml index 17696ca17efe..33e632e849e6 100644 --- a/.github/workflows/push_tests.yml +++ b/.github/workflows/push_tests.yml @@ -98,7 +98,7 @@ jobs: CUBLAS_WORKSPACE_CONFIG: :16:8 run: | pytest -n 1 --max-worker-restart=0 --dist=loadfile \ - -k "not Flax and not Onnx" \ + -k "not Onnx" \ --make-reports=tests_pipeline_${{ matrix.module }}_cuda \ tests/pipelines/${{ matrix.module }} - name: Failure short reports @@ -153,7 +153,7 @@ jobs: CUBLAS_WORKSPACE_CONFIG: :16:8 run: | pytest -n 1 --max-worker-restart=0 --dist=loadfile \ - -k "not Flax and not Onnx" \ + -k "not Onnx" \ --make-reports=tests_torch_cuda_${{ matrix.module }} \ tests/${{ matrix.module }} diff --git a/.github/workflows/push_tests_fast.yml b/.github/workflows/push_tests_fast.yml index 44677ab72c0d..10e89b38c487 100644 --- a/.github/workflows/push_tests_fast.yml +++ b/.github/workflows/push_tests_fast.yml @@ -73,7 +73,7 @@ jobs: if: ${{ matrix.config.framework == 'pytorch' }} run: | pytest -n 4 --max-worker-restart=0 --dist=loadfile \ - -k "not Flax and not Onnx" \ + -k "not Onnx" \ --make-reports=tests_${{ matrix.config.report }} \ tests/ diff --git a/.github/workflows/release_tests_fast.yml b/.github/workflows/release_tests_fast.yml index d8e73f7e7a73..6a532496f688 100644 --- a/.github/workflows/release_tests_fast.yml +++ b/.github/workflows/release_tests_fast.yml @@ -98,7 +98,7 @@ jobs: CUBLAS_WORKSPACE_CONFIG: :16:8 run: | pytest -n 1 --max-worker-restart=0 --dist=loadfile \ - -k "not Flax and not Onnx" \ + -k "not Onnx" \ --make-reports=tests_pipeline_${{ matrix.module }}_cuda \ tests/pipelines/${{ matrix.module }} - name: Failure short reports @@ -153,7 +153,7 @@ jobs: CUBLAS_WORKSPACE_CONFIG: :16:8 run: | pytest -n 1 --max-worker-restart=0 --dist=loadfile \ - -k "not Flax and not Onnx" \ + -k "not Onnx" \ --make-reports=tests_torch_${{ matrix.module }}_cuda \ tests/${{ matrix.module }} @@ -205,7 +205,7 @@ jobs: CUBLAS_WORKSPACE_CONFIG: :16:8 run: | pytest -n 1 --max-worker-restart=0 --dist=loadfile \ - -k "not Flax and not Onnx" \ + -k "not Onnx" \ --make-reports=tests_torch_minimum_cuda \ tests/models/test_modeling_common.py \ tests/pipelines/test_pipelines_common.py \ diff --git a/docs/source/ja/installation.md b/docs/source/ja/installation.md index fd6f4eda0fca..b05aa4b2047f 100644 --- a/docs/source/ja/installation.md +++ b/docs/source/ja/installation.md @@ -14,10 +14,9 @@ specific language governing permissions and limitations under the License. お使いのディープラーニングライブラリに合わせてDiffusersをインストールできます。 -🤗 DiffusersはPython 3.8+、PyTorch 1.7.0+、Flaxでテストされています。使用するディープラーニングライブラリの以下のインストール手順に従ってください： +🤗 DiffusersはPython 3.8+、PyTorch 1.7.0+でテストされています。使用するディープラーニングライブラリの以下のインストール手順に従ってください： - [PyTorch](https://pytorch.org/get-started/locally/)のインストール手順。 -- [Flax](https://flax.readthedocs.io/en/latest/)のインストール手順。 ## pip でインストール @@ -45,11 +44,6 @@ source .env/bin/activate pip install diffusers["torch"] transformers ``` - -```bash -pip install diffusers["flax"] transformers -``` - ## ソースからのインストール @@ -97,11 +91,6 @@ cd diffusers pip install -e ".[torch]" ``` - -```bash -pip install -e ".[flax]" -``` - これらのコマンドは、リポジトリをクローンしたフォルダと Python のライブラリパスをリンクします。 @@ -123,7 +112,7 @@ Python環境は次の実行時に `main` バージョンの🤗 Diffusersを見 ## テレメトリー・ロギングに関するお知らせ このライブラリは `from_pretrained()` リクエスト中にデータを収集します。 -このデータには Diffusers と PyTorch/Flax のバージョン、要求されたモデルやパイプラインクラスが含まれます。 +このデータには Diffusers と PyTorch のバージョン、要求されたモデルやパイプラインクラスが含まれます。 また、Hubでホストされている場合は、事前に学習されたチェックポイントへのパスが含まれます。 この使用データは問題のデバッグや新機能の優先順位付けに役立ちます。 テレメトリーはHuggingFace Hubからモデルやパイプラインをロードするときのみ送信されます。ローカルでの使用中は収集されません。 diff --git a/docs/source/ko/_toctree.yml b/docs/source/ko/_toctree.yml index 9bd5e8e9e240..c649e4d155ff 100644 --- a/docs/source/ko/_toctree.yml +++ b/docs/source/ko/_toctree.yml @@ -162,8 +162,6 @@ - local: in_translation # optimization/tgate title: (번역중) TGATE - sections: - - local: using-diffusers/stable_diffusion_jax_how_to - title: JAX/Flax - local: optimization/onnx title: ONNX - local: optimization/open_vino diff --git a/docs/source/ko/installation.md b/docs/source/ko/installation.md index 198ca4b7c760..1061484df816 100644 --- a/docs/source/ko/installation.md +++ b/docs/source/ko/installation.md @@ -14,10 +14,9 @@ specific language governing permissions and limitations under the License. 사용하시는 라이브러리에 맞는 🤗 Diffusers를 설치하세요. -🤗 Diffusers는 Python 3.8+, PyTorch 1.7.0+ 및 flax에서 테스트되었습니다. 사용중인 딥러닝 라이브러리에 대한 아래의 설치 안내를 따르세요. +🤗 Diffusers는 Python 3.8+ 및 PyTorch 1.7.0+에서 테스트되었습니다. 사용중인 딥러닝 라이브러리에 대한 아래의 설치 안내를 따르세요. - [PyTorch 설치 안내](https://pytorch.org/get-started/locally/) -- [Flax 설치 안내](https://flax.readthedocs.io/en/latest/) ## pip를 이용한 설치 @@ -45,12 +44,6 @@ source .env/bin/activate pip install diffusers["torch"] ``` -**Flax의 경우** - -```bash -pip install diffusers["flax"] -``` - ## 소스로부터 설치 소스에서 `diffusers`를 설치하기 전에, `torch` 및 `accelerate`이 설치되어 있는지 확인하세요. @@ -97,12 +90,6 @@ cd diffusers pip install -e ".[torch]" ``` -**Flax의 경우** - -```sh -pip install -e ".[flax]" -``` - 이러한 명령어들은 저장소를 복제한 폴더와 Python 라이브러리 경로를 연결합니다. Python은 이제 일반 라이브러리 경로에 더하여 복제한 폴더 내부를 살펴봅니다. 예를들어 Python 패키지가 `~/anaconda3/envs/main/lib/python3.10/site-packages/`에 설치되어 있는 경우 Python은 복제한 폴더인 `~/diffusers/`도 검색합니다. @@ -122,7 +109,7 @@ git pull ## 텔레메트리 로깅에 대한 알림 우리 라이브러리는 `from_pretrained()` 요청 중에 텔레메트리 정보를 원격으로 수집합니다. -이 데이터에는 Diffusers 및 PyTorch/Flax의 버전, 요청된 모델 또는 파이프라인 클래스, 그리고 허브에서 호스팅되는 경우 사전학습된 체크포인트에 대한 경로를 포함합니다. +이 데이터에는 Diffusers 및 PyTorch의 버전, 요청된 모델 또는 파이프라인 클래스, 그리고 허브에서 호스팅되는 경우 사전학습된 체크포인트에 대한 경로를 포함합니다. 이 사용 데이터는 문제를 디버깅하고 새로운 기능의 우선순위를 지정하는데 도움이 됩니다. 텔레메트리는 HuggingFace 허브에서 모델과 파이프라인을 불러올 때만 전송되며, 로컬 사용 중에는 수집되지 않습니다. diff --git a/docs/source/ko/training/dreambooth.md b/docs/source/ko/training/dreambooth.md index 3e5a17d5f67c..35b5a0c4d2f1 100644 --- a/docs/source/ko/training/dreambooth.md +++ b/docs/source/ko/training/dreambooth.md @@ -18,7 +18,7 @@ specific language governing permissions and limitations under the License. 에서의 Dreambooth 예시 project's blog. -이 가이드는 다양한 GPU, Flax 사양에 대해 [`CompVis/stable-diffusion-v1-4`](https://huggingface.co/CompVis/stable-diffusion-v1-4) 모델로 DreamBooth를 파인튜닝하는 방법을 보여줍니다. 더 깊이 파고들어 작동 방식을 확인하는 데 관심이 있는 경우, 이 가이드에 사용된 DreamBooth의 모든 학습 스크립트를 [여기](https://github.com/huggingface/diffusers/tree/main/examples/dreambooth)에서 찾을 수 있습니다. +이 가이드는 다양한 GPU 사양에 대해 [`CompVis/stable-diffusion-v1-4`](https://huggingface.co/CompVis/stable-diffusion-v1-4) 모델로 DreamBooth를 파인튜닝하는 방법을 보여줍니다. 더 깊이 파고들어 작동 방식을 확인하는 데 관심이 있는 경우, 이 가이드에 사용된 DreamBooth의 모든 학습 스크립트를 [여기](https://github.com/huggingface/diffusers/tree/main/examples/dreambooth)에서 찾을 수 있습니다. 스크립트를 실행하기 전에 라이브러리의 학습에 필요한 dependencies를 설치해야 합니다. 또한 `main` GitHub 브랜치에서 🧨 Diffusers를 설치하는 것이 좋습니다. @@ -83,34 +83,6 @@ accelerate launch train_dreambooth.py \ --max_train_steps=400 ``` - - -TPU에 액세스할 수 있거나 더 빠르게 훈련하고 싶다면 [Flax 학습 스크립트](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_flax.py)를 사용해 볼 수 있습니다. Flax 학습 스크립트는 gradient checkpointing 또는 gradient accumulation을 지원하지 않으므로, 메모리가 30GB 이상인 GPU가 필요합니다. - -스크립트를 실행하기 전에 요구 사항이 설치되어 있는지 확인하십시오. - -```bash -pip install -U -r requirements.txt -``` - -그러면 다음 명령어로 학습 스크립트를 실행시킬 수 있습니다: - -```bash -export MODEL_NAME="duongna/stable-diffusion-v1-4-flax" -export INSTANCE_DIR="path-to-instance-images" -export OUTPUT_DIR="path-to-save-model" - -python train_dreambooth_flax.py \ - --pretrained_model_name_or_path=$MODEL_NAME \ - --instance_data_dir=$INSTANCE_DIR \ - --output_dir=$OUTPUT_DIR \ - --instance_prompt="a photo of sks dog" \ - --resolution=512 \ - --train_batch_size=1 \ - --learning_rate=5e-6 \ - --max_train_steps=400 -``` - ### Prior-preserving(사전 보존) loss를 사용한 파인튜닝 @@ -145,28 +117,6 @@ accelerate launch train_dreambooth.py \ --max_train_steps=800 ``` - -```bash -export MODEL_NAME="duongna/stable-diffusion-v1-4-flax" -export INSTANCE_DIR="path-to-instance-images" -export CLASS_DIR="path-to-class-images" -export OUTPUT_DIR="path-to-save-model" - -python train_dreambooth_flax.py \ - --pretrained_model_name_or_path=$MODEL_NAME \ - --instance_data_dir=$INSTANCE_DIR \ - --class_data_dir=$CLASS_DIR \ - --output_dir=$OUTPUT_DIR \ - --with_prior_preservation --prior_loss_weight=1.0 \ - --instance_prompt="a photo of sks dog" \ - --class_prompt="a photo of dog" \ - --resolution=512 \ - --train_batch_size=1 \ - --learning_rate=5e-6 \ - --num_class_images=200 \ - --max_train_steps=800 -``` - ## 텍스트 인코더와 and UNet로 파인튜닝하기 @@ -206,29 +156,6 @@ accelerate launch train_dreambooth.py \ --max_train_steps=800 ``` - -```bash -export MODEL_NAME="duongna/stable-diffusion-v1-4-flax" -export INSTANCE_DIR="path-to-instance-images" -export CLASS_DIR="path-to-class-images" -export OUTPUT_DIR="path-to-save-model" - -python train_dreambooth_flax.py \ - --pretrained_model_name_or_path=$MODEL_NAME \ - --train_text_encoder \ - --instance_data_dir=$INSTANCE_DIR \ - --class_data_dir=$CLASS_DIR \ - --output_dir=$OUTPUT_DIR \ - --with_prior_preservation --prior_loss_weight=1.0 \ - --instance_prompt="a photo of sks dog" \ - --class_prompt="a photo of dog" \ - --resolution=512 \ - --train_batch_size=1 \ - --learning_rate=2e-6 \ - --num_class_images=200 \ - --max_train_steps=800 -``` - ## LoRA로 파인튜닝하기 @@ -322,8 +249,6 @@ pipeline.save_pretrained("dreambooth-pipeline") --enable_xformers_memory_efficient_attention ``` -xFormers는 Flax에서 사용할 수 없습니다. - ### 그래디언트 없음으로 설정 메모리 사용량을 줄일 수 있는 또 다른 방법은 [기울기 설정](https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html)을 0 대신 `None`으로 하는 것입니다. 그러나 이로 인해 특정 동작이 변경될 수 있으므로 문제가 발생하면 이 인수를 제거해 보십시오. 학습 스크립트에 다음 인수를 추가하여 그래디언트를 `None`으로 설정합니다. diff --git a/docs/source/ko/training/text2image.md b/docs/source/ko/training/text2image.md index b26603bf1b34..f7cc57543357 100644 --- a/docs/source/ko/training/text2image.md +++ b/docs/source/ko/training/text2image.md @@ -16,7 +16,7 @@ specific language governing permissions and limitations under the License. > [!WARNING] > text-to-image 파인튜닝 스크립트는 experimental 상태입니다. 과적합하기 쉽고 치명적인 망각과 같은 문제에 부딪히기 쉽습니다. 자체 데이터셋에서 최상의 결과를 얻으려면 다양한 하이퍼파라미터를 탐색하는 것이 좋습니다. -Stable Diffusion과 같은 text-to-image 모델은 텍스트 프롬프트에서 이미지를 생성합니다. 이 가이드는 PyTorch 및 Flax를 사용하여 자체 데이터셋에서 [`CompVis/stable-diffusion-v1-4`](https://huggingface.co/CompVis/stable-diffusion-v1-4) 모델로 파인튜닝하는 방법을 보여줍니다. 이 가이드에 사용된 text-to-image 파인튜닝을 위한 모든 학습 스크립트에 관심이 있는 경우 이 [리포지토리](https://github.com/huggingface/diffusers/tree/main/examples/text_to_image)에서 자세히 찾을 수 있습니다. +Stable Diffusion과 같은 text-to-image 모델은 텍스트 프롬프트에서 이미지를 생성합니다. 이 가이드는 PyTorch를 사용하여 자체 데이터셋에서 [`CompVis/stable-diffusion-v1-4`](https://huggingface.co/CompVis/stable-diffusion-v1-4) 모델로 파인튜닝하는 방법을 보여줍니다. 이 가이드에 사용된 text-to-image 파인튜닝을 위한 모든 학습 스크립트에 관심이 있는 경우 이 [리포지토리](https://github.com/huggingface/diffusers/tree/main/examples/text_to_image)에서 자세히 찾을 수 있습니다. 스크립트를 실행하기 전에, 라이브러리의 학습 dependency들을 설치해야 합니다: @@ -35,12 +35,10 @@ accelerate config ### 하드웨어 요구 사항 -`gradient_checkpointing` 및 `mixed_precision`을 사용하면 단일 24GB GPU에서 모델을 파인튜닝할 수 있습니다. 더 높은 `batch_size`와 더 빠른 훈련을 위해서는 GPU 메모리가 30GB 이상인 GPU를 사용하는 것이 좋습니다. TPU 또는 GPU에서 파인튜닝을 위해 JAX나 Flax를 사용할 수도 있습니다. 자세한 내용은 [아래](#flax-jax-finetuning)를 참조하세요. +`gradient_checkpointing` 및 `mixed_precision`을 사용하면 단일 24GB GPU에서 모델을 파인튜닝할 수 있습니다. 더 높은 `batch_size`와 더 빠른 훈련을 위해서는 GPU 메모리가 30GB 이상인 GPU를 사용하는 것이 좋습니다. xFormers로 memory efficient attention을 활성화하여 메모리 사용량 훨씬 더 줄일 수 있습니다. [xFormers가 설치](./optimization/xformers)되어 있는지 확인하고 `--enable_xformers_memory_efficient_attention`를 학습 스크립트에 명시합니다. -xFormers는 Flax에 사용할 수 없습니다. - ## Hub에 모델 업로드하기 학습 스크립트에 다음 인수를 추가하여 모델을 허브에 저장합니다: @@ -120,52 +118,6 @@ accelerate launch train_text_to_image.py \ ``` - -[@duongna211](https://github.com/duongna21)의 기여로, Flax를 사용해 TPU 및 GPU에서 Stable Diffusion 모델을 더 빠르게 학습할 수 있습니다. 이는 TPU 하드웨어에서 매우 효율적이지만 GPU에서도 훌륭하게 작동합니다. Flax 학습 스크립트는 gradient checkpointing나 gradient accumulation과 같은 기능을 아직 지원하지 않으므로 메모리가 30GB 이상인 GPU 또는 TPU v3가 필요합니다. - -스크립트를 실행하기 전에 요구 사항이 설치되어 있는지 확인하십시오: - -```bash -pip install -U -r requirements_flax.txt -``` - -그러면 다음과 같이 [Flax 학습 스크립트](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_flax.py)를 실행할 수 있습니다. - -```bash -export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5" -export dataset_name="lambdalabs/naruto-blip-captions" - -python train_text_to_image_flax.py \ - --pretrained_model_name_or_path=$MODEL_NAME \ - --dataset_name=$dataset_name \ - --resolution=512 --center_crop --random_flip \ - --train_batch_size=1 \ - --max_train_steps=15000 \ - --learning_rate=1e-05 \ - --max_grad_norm=1 \ - --output_dir="sd-naruto-model" -``` - -자체 데이터셋으로 파인튜닝하려면 🤗 [Datasets](https://huggingface.co/docs/datasets/index)에서 요구하는 형식에 따라 데이터셋을 준비하세요. [데이터셋을 허브에 업로드](https://huggingface.co/docs/datasets/image_dataset#upload-dataset-to-the-hub)하거나 [파일들이 있는 로컬 폴더를 준비](https ://huggingface.co/docs/datasets/image_dataset#imagefolder)할 수 있습니다. - -사용자 커스텀 loading logic을 사용하려면 스크립트를 수정하십시오. 도움이 되도록 코드의 적절한 위치에 포인터를 남겼습니다. 🤗 아래 예제 스크립트는 `TRAIN_DIR`의 로컬 데이터셋으로를 파인튜닝하는 방법을 보여줍니다: - -```bash -export MODEL_NAME="duongna/stable-diffusion-v1-4-flax" -export TRAIN_DIR="path_to_your_dataset" - -python train_text_to_image_flax.py \ - --pretrained_model_name_or_path=$MODEL_NAME \ - --train_data_dir=$TRAIN_DIR \ - --resolution=512 --center_crop --random_flip \ - --train_batch_size=1 \ - --mixed_precision="fp16" \ - --max_train_steps=15000 \ - --learning_rate=1e-05 \ - --max_grad_norm=1 \ - --output_dir="sd-naruto-model" -``` - ## LoRA @@ -189,33 +141,4 @@ image = pipe(prompt="yoda").images[0] image.save("yoda-naruto.png") ``` - -```python -import jax -import numpy as np -from flax.jax_utils import replicate -from flax.training.common_utils import shard -from diffusers import FlaxStableDiffusionPipeline - -model_path = "path_to_saved_model" -pipe, params = FlaxStableDiffusionPipeline.from_pretrained(model_path, dtype=jax.numpy.bfloat16) - -prompt = "yoda naruto" -prng_seed = jax.random.PRNGKey(0) -num_inference_steps = 50 - -num_samples = jax.device_count() -prompt = num_samples * [prompt] -prompt_ids = pipeline.prepare_inputs(prompt) - -# shard inputs and rng -params = replicate(params) -prng_seed = jax.random.split(prng_seed, jax.device_count()) -prompt_ids = shard(prompt_ids) - -images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images -images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]))) -image.save("yoda-naruto.png") -``` - \ No newline at end of file diff --git a/docs/source/ko/training/text_inversion.md b/docs/source/ko/training/text_inversion.md index d8b44930e3fd..4c8659280d5a 100644 --- a/docs/source/ko/training/text_inversion.md +++ b/docs/source/ko/training/text_inversion.md @@ -52,7 +52,7 @@ from accelerate.utils import write_basic_config write_basic_config() ``` -마지막으로, Memory-Efficient Attention을 통해 메모리 사용량을 줄이기 위해 [xFormers](https://huggingface.co/docs/diffusers/main/en/training/optimization/xformers)를 설치합니다. xFormers를 설치한 후, 학습 스크립트에 `--enable_xformers_memory_efficient_attention` 인자를 추가합니다. xFormers는 Flax에서 지원되지 않습니다. +마지막으로, Memory-Efficient Attention을 통해 메모리 사용량을 줄이기 위해 [xFormers](https://huggingface.co/docs/diffusers/main/en/training/optimization/xformers)를 설치합니다. xFormers를 설치한 후, 학습 스크립트에 `--enable_xformers_memory_efficient_attention` 인자를 추가합니다. ## 허브에 모델 업로드하기 @@ -129,37 +129,6 @@ accelerate launch textual_inversion.py \ > --num_vectors=5 > ``` - - -TPU에 액세스할 수 있는 경우, [Flax 학습 스크립트](https://github.com/huggingface/diffusers/blob/main/examples/textual_inversion/textual_inversion_flax.py)를 사용하여 더 빠르게 모델을 학습시켜보세요. (물론 GPU에서도 작동합니다.) 동일한 설정에서 Flax 학습 스크립트는 PyTorch 학습 스크립트보다 최소 70% 더 빨라야 합니다! ⚡️ - -시작하기 앞서 Flax에 대한 의존성 라이브러리들을 설치해야 합니다. - -```bash -pip install -U -r requirements_flax.txt -``` - -모델의 리포지토리 ID(또는 모델 가중치가 포함된 디렉터리 경로)를 `MODEL_NAME` 환경 변수에 할당하고, 해당 값을 [`pretrained_model_name_or_path`](https://huggingface.co/docs/diffusers/en/api/diffusion_pipeline#diffusers.DiffusionPipeline.from_pretrained.pretrained_model_name_or_path) 인자에 전달합니다. - -그런 다음 [학습 스크립트](https://github.com/huggingface/diffusers/blob/main/examples/textual_inversion/textual_inversion_flax.py)를 시작할 수 있습니다. - -```bash -export MODEL_NAME="duongna/stable-diffusion-v1-4-flax" -export DATA_DIR="./cat" - -python textual_inversion_flax.py \ - --pretrained_model_name_or_path=$MODEL_NAME \ - --train_data_dir=$DATA_DIR \ - --learnable_property="object" \ - --placeholder_token="" --initializer_token="toy" \ - --resolution=512 \ - --train_batch_size=1 \ - --max_train_steps=3000 \ - --learning_rate=5.0e-04 --scale_lr \ - --output_dir="textual_inversion_cat" \ - --push_to_hub -``` - ### 중간 로깅 @@ -218,38 +187,6 @@ image.save("cat-backpack.png") pipe.load_textual_inversion("./charturnerv2.pt") ``` - - -현재 Flax에 대한 `load_textual_inversion` 함수는 없습니다. 따라서 학습 후 textual-inversion 임베딩 벡터가 모델의 일부로서 저장되었는지를 확인해야 합니다. 그런 다음은 다른 Flax 모델과 마찬가지로 실행할 수 있습니다. - -```python -import jax -import numpy as np -from flax.jax_utils import replicate -from flax.training.common_utils import shard -from diffusers import FlaxStableDiffusionPipeline - -model_path = "path-to-your-trained-model" -pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(model_path, dtype=jax.numpy.bfloat16) - -prompt = "A backpack" -prng_seed = jax.random.PRNGKey(0) -num_inference_steps = 50 - -num_samples = jax.device_count() -prompt = num_samples * [prompt] -prompt_ids = pipeline.prepare_inputs(prompt) - -# shard inputs and rng -params = replicate(params) -prng_seed = jax.random.split(prng_seed, jax.device_count()) -prompt_ids = shard(prompt_ids) - -images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images -images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]))) -image.save("cat-backpack.png") -``` - ## 작동 방식 diff --git a/docs/source/ko/using-diffusers/schedulers.md b/docs/source/ko/using-diffusers/schedulers.md index b12c08b8c869..02fed1d9399d 100644 --- a/docs/source/ko/using-diffusers/schedulers.md +++ b/docs/source/ko/using-diffusers/schedulers.md @@ -274,53 +274,3 @@ image 보시다시피 생성된 이미지들은 매우 비슷하고, 비슷한 퀄리티를 보이는 것 같습니다. 실제로 어떤 스케줄러를 선택할 것인가는 종종 특정 이용 사례에 기반해서 결정되곤 합니다. 결국 여러 종류의 스케줄러를 직접 실행시켜보고 눈으로 직접 비교해서 판단하는 게 좋은 선택일 것 같습니다. - - -## Flax에서 스케줄러 교체하기 - -JAX/Flax 사용자인 경우 기본 파이프라인 스케줄러를 변경할 수도 있습니다. 다음은 Flax Stable Diffusion 파이프라인과 초고속 [DDPM-Solver++ 스케줄러를](../api/schedulers/multistep_dpm_solver) 사용하여 추론을 실행하는 방법에 대한 예시입니다 . - -```Python -import jax -import numpy as np -from flax.jax_utils import replicate -from flax.training.common_utils import shard - -from diffusers import FlaxStableDiffusionPipeline, FlaxDPMSolverMultistepScheduler - -model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5" -scheduler, scheduler_state = FlaxDPMSolverMultistepScheduler.from_pretrained( - model_id, - subfolder="scheduler" -) -pipeline, params = FlaxStableDiffusionPipeline.from_pretrained( - model_id, - scheduler=scheduler, - variant="bf16", - dtype=jax.numpy.bfloat16, -) -params["scheduler"] = scheduler_state - -# Generate 1 image per parallel device (8 on TPUv2-8 or TPUv3-8) -prompt = "a photo of an astronaut riding a horse on mars" -num_samples = jax.device_count() -prompt_ids = pipeline.prepare_inputs([prompt] * num_samples) - -prng_seed = jax.random.PRNGKey(0) -num_inference_steps = 25 - -# shard inputs and rng -params = replicate(params) -prng_seed = jax.random.split(prng_seed, jax.device_count()) -prompt_ids = shard(prompt_ids) - -images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images -images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]))) -``` - -> [!WARNING] -> 다음 Flax 스케줄러는 *아직* Flax Stable Diffusion 파이프라인과 호환되지 않습니다. -> -> - `FlaxLMSDiscreteScheduler` -> - `FlaxDDPMScheduler` - diff --git a/docs/source/ko/using-diffusers/stable_diffusion_jax_how_to.md b/docs/source/ko/using-diffusers/stable_diffusion_jax_how_to.md deleted file mode 100644 index 9830ccb5b036..000000000000 --- a/docs/source/ko/using-diffusers/stable_diffusion_jax_how_to.md +++ /dev/null @@ -1,264 +0,0 @@ - - -# JAX / Flax에서의 🧨 Stable Diffusion! - -[[open-in-colab]] - -🤗 Hugging Face [Diffusers] (https://github.com/huggingface/diffusers) 는 버전 0.5.1부터 Flax를 지원합니다! 이를 통해 Colab, Kaggle, Google Cloud Platform에서 사용할 수 있는 것처럼 Google TPU에서 초고속 추론이 가능합니다. - -이 노트북은 JAX / Flax를 사용해 추론을 실행하는 방법을 보여줍니다. Stable Diffusion의 작동 방식에 대한 자세한 내용을 원하거나 GPU에서 실행하려면 이 [노트북] ](https://huggingface.co/docs/diffusers/stable_diffusion)을 참조하세요. - -먼저, TPU 백엔드를 사용하고 있는지 확인합니다. Colab에서 이 노트북을 실행하는 경우, 메뉴에서 런타임을 선택한 다음 "런타임 유형 변경" 옵션을 선택한 다음 하드웨어 가속기 설정에서 TPU를 선택합니다. - -JAX는 TPU 전용은 아니지만 각 TPU 서버에는 8개의 TPU 가속기가 병렬로 작동하기 때문에 해당 하드웨어에서 더 빛을 발한다는 점은 알아두세요. - - -## Setup - -먼저 diffusers가 설치되어 있는지 확인합니다. - -```bash -!pip install jax==0.3.25 jaxlib==0.3.25 flax transformers ftfy -!pip install diffusers -``` - -```python -import jax.tools.colab_tpu - -jax.tools.colab_tpu.setup_tpu() -import jax -``` - -```python -num_devices = jax.device_count() -device_type = jax.devices()[0].device_kind - -print(f"Found {num_devices} JAX devices of type {device_type}.") -assert ( - "TPU" in device_type -), "Available device is not a TPU, please select TPU from Edit > Notebook settings > Hardware accelerator" -``` - -```python out -Found 8 JAX devices of type Cloud TPU. -``` - -그런 다음 모든 dependencies를 가져옵니다. - -```python -import numpy as np -import jax -import jax.numpy as jnp - -from pathlib import Path -from jax import pmap -from flax.jax_utils import replicate -from flax.training.common_utils import shard -from PIL import Image - -from huggingface_hub import notebook_login -from diffusers import FlaxStableDiffusionPipeline -``` - -## 모델 불러오기 - -TPU 장치는 효율적인 half-float 유형인 bfloat16을 지원합니다. 테스트에는 이 유형을 사용하지만 대신 float32를 사용하여 전체 정밀도(full precision)를 사용할 수도 있습니다. - -```python -dtype = jnp.bfloat16 -``` - -Flax는 함수형 프레임워크이므로 모델은 무상태(stateless)형이며 매개변수는 모델 외부에 저장됩니다. 사전학습된 Flax 파이프라인을 불러오면 파이프라인 자체와 모델 가중치(또는 매개변수)가 모두 반환됩니다. 저희는 bf16 버전의 가중치를 사용하고 있으므로 유형 경고가 표시되지만 무시해도 됩니다. - -```python -pipeline, params = FlaxStableDiffusionPipeline.from_pretrained( - "CompVis/stable-diffusion-v1-4", - variant="bf16", - dtype=dtype, -) -``` - -## 추론 - -TPU에는 일반적으로 8개의 디바이스가 병렬로 작동하므로 보유한 디바이스 수만큼 프롬프트를 복제합니다. 그런 다음 각각 하나의 이미지 생성을 담당하는 8개의 디바이스에서 한 번에 추론을 수행합니다. 따라서 하나의 칩이 하나의 이미지를 생성하는 데 걸리는 시간과 동일한 시간에 8개의 이미지를 얻을 수 있습니다. - -프롬프트를 복제하고 나면 파이프라인의 `prepare_inputs` 함수를 호출하여 토큰화된 텍스트 ID를 얻습니다. 토큰화된 텍스트의 길이는 기본 CLIP 텍스트 모델의 구성에 따라 77토큰으로 설정됩니다. - -```python -prompt = "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of field, close up, split lighting, cinematic" -prompt = [prompt] * jax.device_count() -prompt_ids = pipeline.prepare_inputs(prompt) -prompt_ids.shape -``` - -```python out -(8, 77) -``` - -### 복사(Replication) 및 정렬화 - -모델 매개변수와 입력값은 우리가 보유한 8개의 병렬 장치에 복사(Replication)되어야 합니다. 매개변수 딕셔너리는 `flax.jax_utils.replicate`(딕셔너리를 순회하며 가중치의 모양을 변경하여 8번 반복하는 함수)를 사용하여 복사됩니다. 배열은 `shard`를 사용하여 복제됩니다. - -```python -p_params = replicate(params) -``` - -```python -prompt_ids = shard(prompt_ids) -prompt_ids.shape -``` - -```python out -(8, 1, 77) -``` - -이 shape은 8개의 디바이스 각각이 shape `(1, 77)`의 jnp 배열을 입력값으로 받는다는 의미입니다. 즉 1은 디바이스당 batch(배치) 크기입니다. 메모리가 충분한 TPU에서는 한 번에 여러 이미지(칩당)를 생성하려는 경우 1보다 클 수 있습니다. - -이미지를 생성할 준비가 거의 완료되었습니다! 이제 생성 함수에 전달할 난수 생성기만 만들면 됩니다. 이것은 난수를 다루는 모든 함수에 난수 생성기가 있어야 한다는, 난수에 대해 매우 진지하고 독단적인 Flax의 표준 절차입니다. 이렇게 하면 여러 분산된 기기에서 훈련할 때에도 재현성이 보장됩니다. - -아래 헬퍼 함수는 시드를 사용하여 난수 생성기를 초기화합니다. 동일한 시드를 사용하는 한 정확히 동일한 결과를 얻을 수 있습니다. 나중에 노트북에서 결과를 탐색할 때엔 다른 시드를 자유롭게 사용하세요. - -```python -def create_key(seed=0): - return jax.random.PRNGKey(seed) -``` - -rng를 얻은 다음 8번 '분할'하여 각 디바이스가 다른 제너레이터를 수신하도록 합니다. 따라서 각 디바이스마다 다른 이미지가 생성되며 전체 프로세스를 재현할 수 있습니다. - -```python -rng = create_key(0) -rng = jax.random.split(rng, jax.device_count()) -``` - -JAX 코드는 매우 빠르게 실행되는 효율적인 표현으로 컴파일할 수 있습니다. 하지만 후속 호출에서 모든 입력이 동일한 모양을 갖도록 해야 하며, 그렇지 않으면 JAX가 코드를 다시 컴파일해야 하므로 최적화된 속도를 활용할 수 없습니다. - -`jit = True`를 인수로 전달하면 Flax 파이프라인이 코드를 컴파일할 수 있습니다. 또한 모델이 사용 가능한 8개의 디바이스에서 병렬로 실행되도록 보장합니다. - -다음 셀을 처음 실행하면 컴파일하는 데 시간이 오래 걸리지만 이후 호출(입력이 다른 경우에도)은 훨씬 빨라집니다. 예를 들어, 테스트했을 때 TPU v2-8에서 컴파일하는 데 1분 이상 걸리지만 이후 추론 실행에는 약 7초가 걸립니다. - -``` -%%time -images = pipeline(prompt_ids, p_params, rng, jit=True)[0] -``` - -```python out -CPU times: user 56.2 s, sys: 42.5 s, total: 1min 38s -Wall time: 1min 29s -``` - -반환된 배열의 shape은 `(8, 1, 512, 512, 3)`입니다. 이를 재구성하여 두 번째 차원을 제거하고 512 × 512 × 3의 이미지 8개를 얻은 다음 PIL로 변환합니다. - -```python -images = images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]) -images = pipeline.numpy_to_pil(images) -``` - -### 시각화 - -이미지를 그리드에 표시하는 도우미 함수를 만들어 보겠습니다. - -```python -def image_grid(imgs, rows, cols): - w, h = imgs[0].size - grid = Image.new("RGB", size=(cols * w, rows * h)) - for i, img in enumerate(imgs): - grid.paste(img, box=(i % cols * w, i // cols * h)) - return grid -``` - -```python -image_grid(images, 2, 4) -``` - -![img](https://huggingface.co/datasets/YiYiXu/test-doc-assets/resolve/main/stable_diffusion_jax_how_to_cell_38_output_0.jpeg) - - -## 다른 프롬프트 사용 - -모든 디바이스에서 동일한 프롬프트를 복제할 필요는 없습니다. 프롬프트 2개를 각각 4번씩 생성하거나 한 번에 8개의 서로 다른 프롬프트를 생성하는 등 원하는 것은 무엇이든 할 수 있습니다. 한번 해보세요! - -먼저 입력 준비 코드를 편리한 함수로 리팩터링하겠습니다: - -```python -prompts = [ - "Labrador in the style of Hokusai", - "Painting of a squirrel skating in New York", - "HAL-9000 in the style of Van Gogh", - "Times Square under water, with fish and a dolphin swimming around", - "Ancient Roman fresco showing a man working on his laptop", - "Close-up photograph of young black woman against urban background, high quality, bokeh", - "Armchair in the shape of an avocado", - "Clown astronaut in space, with Earth in the background", -] -``` - -```python -prompt_ids = pipeline.prepare_inputs(prompts) -prompt_ids = shard(prompt_ids) - -images = pipeline(prompt_ids, p_params, rng, jit=True).images -images = images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]) -images = pipeline.numpy_to_pil(images) - -image_grid(images, 2, 4) -``` - -![img](https://huggingface.co/datasets/YiYiXu/test-doc-assets/resolve/main/stable_diffusion_jax_how_to_cell_43_output_0.jpeg) - - -## 병렬화(parallelization)는 어떻게 작동하는가? - -앞서 `diffusers` Flax 파이프라인이 모델을 자동으로 컴파일하고 사용 가능한 모든 기기에서 병렬로 실행한다고 말씀드렸습니다. 이제 그 프로세스를 간략하게 살펴보고 작동 방식을 보여드리겠습니다. - -JAX 병렬화는 여러 가지 방법으로 수행할 수 있습니다. 가장 쉬운 방법은 jax.pmap 함수를 사용하여 단일 프로그램, 다중 데이터(SPMD) 병렬화를 달성하는 것입니다. 즉, 동일한 코드의 복사본을 각각 다른 데이터 입력에 대해 여러 개 실행하는 것입니다. 더 정교한 접근 방식도 가능하므로 관심이 있으시다면 [JAX 문서](https://jax.readthedocs.io/en/latest/index.html)와 [`pjit` 페이지](https://jax.readthedocs.io/en/latest/jax-101/08-pjit.html?highlight=pjit)에서 이 주제를 살펴보시기 바랍니다! - -`jax.pmap`은 두 가지 기능을 수행합니다: - -- `jax.jit()`를 호출한 것처럼 코드를 컴파일(또는 `jit`)합니다. 이 작업은 `pmap`을 호출할 때가 아니라 pmapped 함수가 처음 호출될 때 수행됩니다. -- 컴파일된 코드가 사용 가능한 모든 기기에서 병렬로 실행되도록 합니다. - -작동 방식을 보여드리기 위해 이미지 생성을 실행하는 비공개 메서드인 파이프라인의 `_generate` 메서드를 `pmap`합니다. 이 메서드는 향후 `Diffusers` 릴리스에서 이름이 변경되거나 제거될 수 있다는 점에 유의하세요. - -```python -p_generate = pmap(pipeline._generate) -``` - -`pmap`을 사용한 후 준비된 함수 `p_generate`는 개념적으로 다음을 수행합니다: -* 각 장치에서 기본 함수 `pipeline._generate`의 복사본을 호출합니다. -* 각 장치에 입력 인수의 다른 부분을 보냅니다. 이것이 바로 샤딩이 사용되는 이유입니다. 이 경우 `prompt_ids`의 shape은 `(8, 1, 77, 768)`입니다. 이 배열은 8개로 분할되고 `_generate`의 각 복사본은 `(1, 77, 768)`의 shape을 가진 입력을 받게 됩니다. - -병렬로 호출된다는 사실을 완전히 무시하고 `_generate`를 코딩할 수 있습니다. batch(배치) 크기(이 예제에서는 `1`)와 코드에 적합한 차원만 신경 쓰면 되며, 병렬로 작동하기 위해 아무것도 변경할 필요가 없습니다. - -파이프라인 호출을 사용할 때와 마찬가지로, 다음 셀을 처음 실행할 때는 시간이 걸리지만 그 이후에는 훨씬 빨라집니다. - -``` -%%time -images = p_generate(prompt_ids, p_params, rng) -images = images.block_until_ready() -images.shape -``` - -```python out -CPU times: user 1min 15s, sys: 18.2 s, total: 1min 34s -Wall time: 1min 15s -``` - -```python -images.shape -``` - -```python out -(8, 1, 512, 512, 3) -``` - -JAX는 비동기 디스패치를 사용하고 가능한 한 빨리 제어권을 Python 루프에 반환하기 때문에 추론 시간을 정확하게 측정하기 위해 `block_until_ready()`를 사용합니다. 아직 구체화되지 않은 계산 결과를 사용하려는 경우 자동으로 차단이 수행되므로 코드에서 이 함수를 사용할 필요가 없습니다. \ No newline at end of file diff --git a/docs/source/pt/installation.md b/docs/source/pt/installation.md index 24c7e4bb7385..3e6960c9f6b2 100644 --- a/docs/source/pt/installation.md +++ b/docs/source/pt/installation.md @@ -12,10 +12,9 @@ specific language governing permissions and limitations under the License. # Instalação -🤗 Diffusers é testado no Python 3.8+, PyTorch 1.7.0+, e Flax. Siga as instruções de instalação abaixo para a biblioteca de deep learning que você está utilizando: +🤗 Diffusers é testado no Python 3.8+ e PyTorch 1.7.0+. Siga as instruções de instalação abaixo para a biblioteca de deep learning que você está utilizando: - [PyTorch](https://pytorch.org/get-started/locally/) instruções de instalação -- [Flax](https://flax.readthedocs.io/en/latest/) instruções de instalação ## Instalação com pip @@ -43,11 +42,6 @@ Recomenda-se a instalação do 🤗 Transformers porque 🤗 Diffusers depende d pip install diffusers["torch"] transformers ``` - -```bash -pip install diffusers["flax"] transformers -``` - ## Instalação a partir do código fonte @@ -93,11 +87,6 @@ cd diffusers pip install -e ".[torch]" ``` - -```bash -pip install -e ".[flax]" -``` - Esses comandos irão vincular a pasta que você clonou o repositório e os caminhos das suas bibliotecas Python. @@ -131,7 +120,7 @@ Para mais detalhes de como gerenciar e limpar o cache, olhe o guia de [caching]( ## Telemetria Nossa biblioteca coleta informações de telemetria durante as requisições [`~DiffusionPipeline.from_pretrained`]. -O dado coletado inclui a versão do 🤗 Diffusers e PyTorch/Flax, o modelo ou classe de pipeline requisitado, +O dado coletado inclui a versão do 🤗 Diffusers e PyTorch, o modelo ou classe de pipeline requisitado, e o caminho para um checkpoint pré-treinado se ele estiver hospedado no Hugging Face Hub. Esse dado de uso nos ajuda a debugar problemas e priorizar novas funcionalidades. Telemetria é enviada apenas quando é carregado modelos e pipelines do Hub, diff --git a/docs/source/zh/installation.md b/docs/source/zh/installation.md index 9941ed24aea4..d68fa12e81db 100644 --- a/docs/source/zh/installation.md +++ b/docs/source/zh/installation.md @@ -14,10 +14,9 @@ specific language governing permissions and limitations under the License. 在你正在使用的任意深度学习框架中安装 🤗 Diffusers 。 -🤗 Diffusers已在Python 3.8+、PyTorch 1.7.0+和Flax上进行了测试。按照下面的安装说明，针对你正在使用的深度学习框架进行安装： +🤗 Diffusers已在Python 3.8+和PyTorch 1.7.0+上进行了测试。按照下面的安装说明，针对你正在使用的深度学习框架进行安装： - [PyTorch](https://pytorch.org/get-started/locally/) installation instructions. -- [Flax](https://flax.readthedocs.io/en/latest/) installation instructions. ## 使用pip安装 @@ -47,12 +46,6 @@ source .env/bin/activate pip install diffusers["torch"] ``` -**Flax** - -```bash -pip install diffusers["flax"] -``` - ## 从源代码安装 在从源代码安装 `diffusers` 之前，确保你已经安装了 `torch` 和 `accelerate`。 @@ -99,12 +92,6 @@ cd diffusers pip install -e ".[torch]" ``` -**Flax** - -```sh -pip install -e ".[flax]" -``` - 这些命令将连接到你克隆的版本库和你的 Python 库路径。 现在，不只是在通常的库路径，Python 还会在你克隆的文件夹内寻找包。 例如，如果你的 Python 包通常安装在 `~/anaconda3/envs/main/lib/python3.10/Site-packages/`，Python 也会搜索你克隆到的文件夹。`~/diffusers/`。 @@ -124,7 +111,7 @@ git pull ## 注意 Telemetry 日志 -我们的库会在使用`from_pretrained()`请求期间收集 telemetry 信息。这些数据包括Diffusers和PyTorch/Flax的版本，请求的模型或管道类，以及预训练检查点的路径（如果它被托管在Hub上的话）。 +我们的库会在使用`from_pretrained()`请求期间收集 telemetry 信息。这些数据包括Diffusers和PyTorch的版本，请求的模型或管道类，以及预训练检查点的路径（如果它被托管在Hub上的话）。 这些使用数据有助于我们调试问题并确定新功能的开发优先级。 Telemetry 数据仅在从 HuggingFace Hub 中加载模型和管道时发送，而不会在本地使用期间收集。 diff --git a/docs/source/zh/training/controlnet.md b/docs/source/zh/training/controlnet.md index 84bc3263a842..bab0bbcf7b52 100644 --- a/docs/source/zh/training/controlnet.md +++ b/docs/source/zh/training/controlnet.md @@ -14,7 +14,7 @@ specific language governing permissions and limitations under the License. [ControlNet](https://hf.co/papers/2302.05543) 是一种基于预训练模型的适配器架构。它通过额外输入的条件图像（如边缘检测图、深度图、人体姿态图等），实现对生成图像的精细化控制。 -在显存有限的GPU上训练时，建议启用训练命令中的 `gradient_checkpointing`（梯度检查点）、`gradient_accumulation_steps`（梯度累积步数）和 `mixed_precision`（混合精度）参数。还可使用 [xFormers](../optimization/xformers) 的内存高效注意力机制进一步降低显存占用。虽然JAX/Flax训练支持在TPU和GPU上高效运行，但不支持梯度检查点和xFormers。若需通过Flax加速训练，建议使用显存大于30GB的GPU。 +在显存有限的GPU上训练时，建议启用训练命令中的 `gradient_checkpointing`（梯度检查点）、`gradient_accumulation_steps`（梯度累积步数）和 `mixed_precision`（混合精度）参数。还可使用 [xFormers](../optimization/xformers) 的内存高效注意力机制进一步降低显存占用。 本指南将解析 [train_controlnet.py](https://github.com/huggingface/diffusers/blob/main/examples/controlnet/train_controlnet.py) 训练脚本，帮助您理解其逻辑并适配自定义需求。 @@ -34,37 +34,6 @@ pip install . cd examples/controlnet pip install -r requirements.txt ``` - - - -若可访问TPU设备，Flax训练脚本将运行得更快！以下是在 [Google Cloud TPU VM](https://cloud.google.com/tpu/docs/run-calculation-jax) 上的配置流程。创建单个TPU v4-8虚拟机并连接： - -```bash -ZONE=us-central2-b -TPU_TYPE=v4-8 -VM_NAME=hg_flax - -gcloud alpha compute tpus tpu-vm create $VM_NAME \ - --zone $ZONE \ - --accelerator-type $TPU_TYPE \ - --version tpu-vm-v4-base - -gcloud alpha compute tpus tpu-vm ssh $VM_NAME --zone $ZONE -- \ -``` - -安装JAX 0.4.5： - -```bash -pip install "jax[tpu]==0.4.5" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html -``` - -然后安装Flax脚本的依赖： - -```bash -cd examples/controlnet -pip install -r requirements_flax.txt -``` - @@ -114,7 +83,7 @@ accelerate launch train_controlnet.py \ ### Min-SNR加权策略 -[Min-SNR](https://huggingface.co/papers/2303.09556) 加权策略通过重新平衡损失函数加速模型收敛。虽然训练脚本支持预测 `epsilon`（噪声）或 `v_prediction`，但Min-SNR对两种预测类型均兼容。该策略仅适用于PyTorch版本，Flax训练脚本暂不支持。 +[Min-SNR](https://huggingface.co/papers/2303.09556) 加权策略通过重新平衡损失函数加速模型收敛。虽然训练脚本支持预测 `epsilon`（噪声）或 `v_prediction`，但Min-SNR对两种预测类型均兼容。 推荐值设为5.0： @@ -283,41 +252,6 @@ accelerate launch train_controlnet.py \ --push_to_hub ``` - - - -Flax版本支持通过 `--profile_steps==5` 参数进行性能分析： - -```bash -pip install tensorflow tensorboard-plugin-profile -tensorboard --logdir runs/fill-circle-100steps-20230411_165612/ -``` - -在 [http://localhost:6006/#profile](http://localhost:6006/#profile) 查看分析结果。 - -> [!WARNING] -> 若遇到插件版本冲突，建议重新安装TensorFlow和Tensorboard。注意性能分析插件仍处实验阶段，部分视图可能不完整。`trace_viewer` 会截断超过1M的事件记录，在编译步骤分析时可能导致设备轨迹丢失。 - -```bash -python3 train_controlnet_flax.py \ - --pretrained_model_name_or_path=$MODEL_DIR \ - --output_dir=$OUTPUT_DIR \ - --dataset_name=fusing/fill50k \ - --resolution=512 \ - --learning_rate=1e-5 \ - --validation_image "./conditioning_image_1.png" "./conditioning_image_2.png" \ - --validation_prompt "red circle with blue background" "cyan circle with brown floral background" \ - --validation_steps=1000 \ - --train_batch_size=2 \ - --revision="non-ema" \ - --from_pt \ - --report_to="wandb" \ - --tracker_project_name=$HUB_MODEL_ID \ - --num_train_epochs=11 \ - --push_to_hub \ - --hub_model_id=$HUB_MODEL_ID -``` - diff --git a/docs/source/zh/training/dreambooth.md b/docs/source/zh/training/dreambooth.md index cae5e30be011..00f94e46e8e1 100644 --- a/docs/source/zh/training/dreambooth.md +++ b/docs/source/zh/training/dreambooth.md @@ -11,7 +11,7 @@ http://www.apache.org/licenses/LICENSE-2.0 [DreamBooth](https://huggingface.co/papers/2208.12242) 是一种训练技术，通过仅训练少数主题或风格的图像来更新整个扩散模型。它通过在提示中关联一个特殊词与示例图像来工作。 -如果您在 vRAM 有限的 GPU 上训练，应尝试在训练命令中启用 `gradient_checkpointing` 和 `mixed_precision` 参数。您还可以通过使用 [xFormers](../optimization/xformers) 的内存高效注意力来减少内存占用。JAX/Flax 训练也支持在 TPU 和 GPU 上进行高效训练，但不支持梯度检查点或 xFormers。如果您想使用 Flax 更快地训练，应拥有内存 >30GB 的 GPU。 +如果您在 vRAM 有限的 GPU 上训练，应尝试在训练命令中启用 `gradient_checkpointing` 和 `mixed_precision` 参数。您还可以通过使用 [xFormers](../optimization/xformers) 的内存高效注意力来减少内存占用。 本指南将探索 [train_dreambooth.py](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py) 脚本，帮助您更熟悉它，以及如何根据您的用例进行适配。 @@ -33,14 +33,6 @@ cd examples/dreambooth pip install -r requirements.txt ``` - - - -```bash -cd examples/dreambooth -pip install -r requirements_flax.txt -``` - @@ -99,7 +91,7 @@ accelerate launch train_dreambooth.py \ ### Min-SNR 加权 -[Min-SNR](https://huggingface.co/papers/2303.09556) 加权策略可以通过重新平衡损失来帮助训练，以实现更快的收敛。训练脚本支持预测 `epsilon`（噪声）或 `v_prediction`，但 Min-SNR 与两种预测类型都兼容。此加权策略仅由 PyTorch 支持，在 Flax 训练脚本中不可用。 +[Min-SNR](https://huggingface.co/papers/2303.09556) 加权策略可以通过重新平衡损失来帮助训练，以实现更快的收敛。训练脚本支持预测 `epsilon`（噪声）或 `v_prediction`，但 Min-SNR 与两种预测类型都兼容。 添加 `--snr_gamma` 参数并将其设置为推荐值 5.0： @@ -325,26 +317,6 @@ accelerate launch train_dreambooth.py \ --push_to_hub ``` - - - -```bash -export MODEL_NAME="duongna/stable-diffusion-v1-4-flax" -export INSTANCE_DIR="./dog" -export OUTPUT_DIR="path-to-save-model" - -python train_dreambooth_flax.py \ - --pretrained_model_name_or_path=$MODEL_NAME \ - --instance_data_dir=$INSTANCE_DIR \ - --output_dir=$OUTPUT_DIR \ - --instance_prompt="a photo of sks dog" \ - --resolution=512 \ - --train_batch_size=1 \ - --learning_rate=5e-6 \ - --max_train_steps=400 \ - --push_to_hub -``` - @@ -383,37 +355,6 @@ image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guida image.save("dog-bucket.png") ``` - - - -```py -import jax -import numpy as np -from flax.jax_utils import replicate -from flax.training.common_utils import shard -from diffusers import FlaxStableDiffusionPipeline - -pipeline, params = FlaxStableDiffusionPipeline.from_pretrained("path-to-your-trained-model", dtype=jax.numpy.bfloat16) - -prompt = "A photo of sks dog in a bucket" -prng_seed = jax.random.PRNGKey(0) -num_inference_steps = 50 - -num_samples = jax.device_count() -prompt = num_samples * [prompt] -prompt_ids = pipeline.prepare_inputs(prompt) - -# 分片输入和随机数生成器 -params = replicate(params) -prng_seed = jax.random.split(prng_seed, jax.device_count()) -prompt_ids = shard(prompt_ids) - -images = pipeline(prompt_ids, params, prng_seed, num_inference_ -steps, jit=True).images -images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]))) -image.save("dog-bucket.png") -``` - diff --git a/docs/source/zh/training/kandinsky.md b/docs/source/zh/training/kandinsky.md index 8ef3524ee7c4..94fa5020efab 100644 --- a/docs/source/zh/training/kandinsky.md +++ b/docs/source/zh/training/kandinsky.md @@ -77,7 +77,7 @@ accelerate launch train_text_to_image_prior.py \ ### Min-SNR 加权 -[Min-SNR](https://huggingface.co/papers/2303.09556) 加权策略可以通过重新平衡损失来帮助训练，实现更快的收敛。训练脚本支持预测 `epsilon`（噪声）或 `v_prediction`，但 Min-SNR 与两种预测类型都兼容。此加权策略仅由 PyTorch 支持，在 Flax 训练脚本中不可用。 +[Min-SNR](https://huggingface.co/papers/2303.09556) 加权策略可以通过重新平衡损失来帮助训练，实现更快的收敛。训练脚本支持预测 `epsilon`（噪声）或 `v_prediction`，但 Min-SNR 与两种预测类型都兼容。 添加 `--snr_gamma` 参数并将其设置为推荐值 5.0： diff --git a/docs/source/zh/training/lora.md b/docs/source/zh/training/lora.md index ce29365450bd..85bd195be27f 100644 --- a/docs/source/zh/training/lora.md +++ b/docs/source/zh/training/lora.md @@ -40,14 +40,6 @@ cd examples/text_to_image pip install -r requirements.txt ``` - - - -```bash -cd examples/text_to_image -pip install -r requirements_flax.txt -``` - diff --git a/docs/source/zh/training/overview.md b/docs/source/zh/training/overview.md index ebf814aefe44..c0fa829ea06a 100644 --- a/docs/source/zh/training/overview.md +++ b/docs/source/zh/training/overview.md @@ -20,18 +20,18 @@ http://www.apache.org/licenses/LICENSE-2.0 当前提供的训练脚本包括： -| 训练类型 | 支持SDXL | 支持LoRA | 支持Flax | -|---|---|---|---| -| [unconditional image generation](https://github.com/huggingface/diffusers/tree/main/examples/unconditional_image_generation) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/training_example.ipynb) | | | | -| [text-to-image](https://github.com/huggingface/diffusers/tree/main/examples/text_to_image) | 👍 | 👍 | 👍 | -| [textual inversion](https://github.com/huggingface/diffusers/tree/main/examples/textual_inversion) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/sd_textual_inversion_training.ipynb) | | | 👍 | -| [DreamBooth](https://github.com/huggingface/diffusers/tree/main/examples/dreambooth) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/sd_dreambooth_training.ipynb) | 👍 | 👍 | 👍 | -| [ControlNet](https://github.com/huggingface/diffusers/tree/main/examples/controlnet) | 👍 | | 👍 | -| [InstructPix2Pix](https://github.com/huggingface/diffusers/tree/main/examples/instruct_pix2pix) | 👍 | | | -| [Custom Diffusion](https://github.com/huggingface/diffusers/tree/main/examples/custom_diffusion) | | | | -| [T2I-Adapters](https://github.com/huggingface/diffusers/tree/main/examples/t2i_adapter) | 👍 | | | -| [Kandinsky 2.2](https://github.com/huggingface/diffusers/tree/main/examples/kandinsky2_2/text_to_image) | | 👍 | | -| [Wuerstchen](https://github.com/huggingface/diffusers/tree/main/examples/wuerstchen/text_to_image) | | 👍 | | +| 训练类型 | 支持SDXL | 支持LoRA | +|---|---|---| +| [unconditional image generation](https://github.com/huggingface/diffusers/tree/main/examples/unconditional_image_generation) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/training_example.ipynb) | | | +| [text-to-image](https://github.com/huggingface/diffusers/tree/main/examples/text_to_image) | 👍 | 👍 | +| [textual inversion](https://github.com/huggingface/diffusers/tree/main/examples/textual_inversion) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/sd_textual_inversion_training.ipynb) | | | +| [DreamBooth](https://github.com/huggingface/diffusers/tree/main/examples/dreambooth) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/sd_dreambooth_training.ipynb) | 👍 | 👍 | +| [ControlNet](https://github.com/huggingface/diffusers/tree/main/examples/controlnet) | 👍 | | +| [InstructPix2Pix](https://github.com/huggingface/diffusers/tree/main/examples/instruct_pix2pix) | 👍 | | +| [Custom Diffusion](https://github.com/huggingface/diffusers/tree/main/examples/custom_diffusion) | | | +| [T2I-Adapters](https://github.com/huggingface/diffusers/tree/main/examples/t2i_adapter) | 👍 | | +| [Kandinsky 2.2](https://github.com/huggingface/diffusers/tree/main/examples/kandinsky2_2/text_to_image) | | 👍 | +| [Wuerstchen](https://github.com/huggingface/diffusers/tree/main/examples/wuerstchen/text_to_image) | | 👍 | 这些示例处于**积极维护**状态，如果遇到问题请随时提交issue。如果您认为应该添加其他训练示例，欢迎创建[功能请求](https://github.com/huggingface/diffusers/issues/new?assignees=&labels=&template=feature_request.md&title=)与我们讨论，我们将评估其是否符合独立完整、易于调整、新手友好和单一用途的标准。 @@ -45,7 +45,7 @@ cd diffusers pip install . ``` -然后进入具体训练脚本目录（例如[DreamBooth](https://github.com/huggingface/diffusers/tree/main/examples/dreambooth)），安装对应的`requirements.txt`文件。部分脚本针对SDXL、LoRA或Flax有特定要求文件，使用时请确保安装对应文件。 +然后进入具体训练脚本目录（例如[DreamBooth](https://github.com/huggingface/diffusers/tree/main/examples/dreambooth)），安装对应的`requirements.txt`文件。部分脚本针对SDXL或LoRA有特定要求文件，使用时请确保安装对应文件。 ```bash cd examples/dreambooth diff --git a/docs/source/zh/training/text2image.md b/docs/source/zh/training/text2image.md index 4465adbe2ad7..d96f0cec200c 100644 --- a/docs/source/zh/training/text2image.md +++ b/docs/source/zh/training/text2image.md @@ -17,7 +17,7 @@ specific language governing permissions and limitations under the License. Stable Diffusion 等文生图模型能够根据文本提示生成对应图像。 -模型训练对硬件要求较高，但启用 `gradient_checkpointing` 和 `mixed_precision` 后，可在单块24GB显存GPU上完成训练。如需更大批次或更快训练速度，建议使用30GB以上显存的GPU设备。通过启用 [xFormers](../optimization/xformers) 内存高效注意力机制可降低显存占用。JAX/Flax 训练方案也支持TPU/GPU高效训练，但不支持梯度检查点、梯度累积和xFormers。使用Flax训练时建议配备30GB以上显存GPU或TPU v3。 +模型训练对硬件要求较高，但启用 `gradient_checkpointing` 和 `mixed_precision` 后，可在单块24GB显存GPU上完成训练。如需更大批次或更快训练速度，建议使用30GB以上显存的GPU设备。通过启用 [xFormers](../optimization/xformers) 内存高效注意力机制可降低显存占用。 本指南将详解 [train_text_to_image.py](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py) 训练脚本，助您掌握其原理并适配自定义需求。 @@ -38,12 +38,6 @@ cd examples/text_to_image pip install -r requirements.txt ``` - -```bash -cd examples/text_to_image -pip install -r requirements_flax.txt -``` - > [!TIP] @@ -97,7 +91,7 @@ accelerate launch train_text_to_image.py \ ### Min-SNR加权策略 -[Min-SNR](https://huggingface.co/papers/2303.09556) 加权策略通过重新平衡损失函数加速模型收敛。训练脚本支持预测 `epsilon`（噪声）或 `v_prediction`，而Min-SNR兼容两种预测类型。该策略仅限PyTorch版本，Flax训练脚本不支持。 +[Min-SNR](https://huggingface.co/papers/2303.09556) 加权策略通过重新平衡损失函数加速模型收敛。训练脚本支持预测 `epsilon`（噪声）或 `v_prediction`，而Min-SNR兼容两种预测类型。 添加 `--snr_gamma` 参数并设为推荐值5.0： @@ -175,32 +169,6 @@ accelerate launch --mixed_precision="fp16" train_text_to_image.py \ --push_to_hub ``` - - - -Flax训练方案在TPU/GPU上效率更高（由 [@duongna211](https://github.com/duongna21) 实现），TPU性能更优但GPU表现同样出色。 - -设置环境变量 `MODEL_NAME` 和 `dataset_name` 指定模型和数据集（Hub或本地路径）。 - -> [!TIP] -> 使用本地数据集时，设置 `TRAIN_DIR` 和 `OUTPUT_DIR` 环境变量为数据集路径和模型保存路径。 - -```bash -export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5" -export dataset_name="lambdalabs/naruto-blip-captions" - -python train_text_to_image_flax.py \ - --pretrained_model_name_or_path=$MODEL_NAME \ - --dataset_name=$dataset_name \ - --resolution=512 --center_crop --random_flip \ - --train_batch_size=1 \ - --max_train_steps=15000 \ - --learning_rate=1e-05 \ - --max_grad_norm=1 \ - --output_dir="sd-naruto-model" \ - --push_to_hub -``` - @@ -219,36 +187,6 @@ image = pipeline(prompt="yoda").images[0] image.save("yoda-naruto.png") ``` - - - -```py -import jax -import numpy as np -from flax.jax_utils import replicate -from flax.training.common_utils import shard -from diffusers import FlaxStableDiffusionPipeline - -pipeline, params = FlaxStableDiffusionPipeline.from_pretrained("path/to/saved_model", dtype=jax.numpy.bfloat16) - -prompt = "yoda naruto" -prng_seed = jax.random.PRNGKey(0) -num_inference_steps = 50 - -num_samples = jax.device_count() -prompt = num_samples * [prompt] -prompt_ids = pipeline.prepare_inputs(prompt) - -# 分片输入和随机数 -params = replicate(params) -prng_seed = jax.random.split(prng_seed, jax.device_count()) -prompt_ids = shard(prompt_ids) - -images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images -images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]))) -image.save("yoda-naruto.png") -``` - diff --git a/docs/source/zh/training/text_inversion.md b/docs/source/zh/training/text_inversion.md index eda9f911441b..d41b0d5ed039 100644 --- a/docs/source/zh/training/text_inversion.md +++ b/docs/source/zh/training/text_inversion.md @@ -12,7 +12,7 @@ http://www.apache.org/licenses/LICENSE-2.0 [文本反转](https://hf.co/papers/2208.01618)是一种训练技术，仅需少量示例图像即可个性化图像生成模型。该技术通过学习和更新文本嵌入（新嵌入会绑定到提示中必须使用的特殊词汇）来匹配您提供的示例图像。 -如果在显存有限的GPU上训练，建议在训练命令中启用`gradient_checkpointing`和`mixed_precision`参数。您还可以通过[xFormers](../optimization/xformers)使用内存高效注意力机制来减少内存占用。JAX/Flax训练也支持在TPU和GPU上进行高效训练，但不支持梯度检查点或xFormers。在配置与PyTorch相同的情况下，Flax训练脚本的速度至少应快70%！ +如果在显存有限的GPU上训练，建议在训练命令中启用`gradient_checkpointing`和`mixed_precision`参数。您还可以通过[xFormers](../optimization/xformers)使用内存高效注意力机制来减少内存占用。 本指南将探索[textual_inversion.py](https://github.com/huggingface/diffusers/blob/main/examples/textual_inversion/textual_inversion.py)脚本，帮助您更熟悉其工作原理，并了解如何根据自身需求进行调整。 @@ -34,14 +34,6 @@ cd examples/textual_inversion pip install -r requirements.txt ``` - - - -```bash -cd examples/textual_inversion -pip install -r requirements_flax.txt -``` - @@ -203,28 +195,6 @@ accelerate launch textual_inversion.py \ --push_to_hub ``` - - - -```bash -export MODEL_NAME="duongna/stable-diffusion-v1-4-flax" -export DATA_DIR="./cat" - -python textual_inversion_flax.py \ - --pretrained_model_name_or_path=$MODEL_NAME \ - --train_data_dir=$DATA_DIR \ - --learnable_property="object" \ - --placeholder_token="" \ - --initializer_token="toy" \ - --resolution=512 \ - --train_batch_size=1 \ - --max_train_steps=3000 \ - --learning_rate=5.0e-04 \ - --scale_lr \ - --output_dir="textual_inversion_cat" \ - --push_to_hub -``` - @@ -243,39 +213,6 @@ image = pipeline("A train", num_inference_steps=50).images[0] image.save("cat-train.png") ``` - - - -Flax不支持[`~loaders.TextualInversionLoaderMixin.load_textual_inversion`]方法，但textual_inversion_flax.py脚本会在训练后[保存](https://github.com/huggingface/diffusers/blob/c0f058265161178f2a88849e92b37ffdc81f1dcc/examples/textual_inversion/textual_inversion_flax.py#L636C2-L636C2)学习到的嵌入作为模型的一部分。这意味着您可以像使用其他Flax模型一样进行推理： - -```py -import jax -import numpy as np -from flax.jax_utils import replicate -from flax.training.common_utils import shard -from diffusers import FlaxStableDiffusionPipeline - -model_path = "path-to-your-trained-model" -pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(model_path, dtype=jax.numpy.bfloat16) - -prompt = "A train" -prng_seed = jax.random.PRNGKey(0) -num_inference_steps = 50 - -num_samples = jax.device_count() -prompt = num_samples * [prompt] -prompt_ids = pipeline.prepare_inputs(prompt) - -# 分片输入和随机数生成器 -params = replicate(params) -prng_seed = jax.random.split(prng_seed, jax.device_count()) -prompt_ids = shard(prompt_ids) - -images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images -images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]))) -image.save("cat-train.png") -``` - diff --git a/docs/source/zh/using-diffusers/schedulers.md b/docs/source/zh/using-diffusers/schedulers.md index 8032c1a98904..d71e4150c165 100644 --- a/docs/source/zh/using-diffusers/schedulers.md +++ b/docs/source/zh/using-diffusers/schedulers.md @@ -165,53 +165,6 @@ image 多数生成图像质量相近，实际选择需根据具体场景测试多种调度器进行比较。 -### Flax调度器 - -对比Flax调度器时，需额外将调度器状态加载到模型参数中。例如将[`FlaxStableDiffusionPipeline`]的默认调度器切换为超高效的[`FlaxDPMSolverMultistepScheduler`]： - -> [!警告] -> [`FlaxLMSDiscreteScheduler`]和[`FlaxDDPMScheduler`]目前暂不兼容[`FlaxStableDiffusionPipeline`]。 - -```python -import jax -import numpy as np -from flax.jax_utils import replicate -from flax.training.common_utils import shard -from diffusers import FlaxStableDiffusionPipeline, FlaxDPMSolverMultistepScheduler - -scheduler, scheduler_state = FlaxDPMSolverMultistepScheduler.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", - subfolder="scheduler" -) -pipeline, params = FlaxStableDiffusionPipeline.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", - scheduler=scheduler, - variant="bf16", - dtype=jax.numpy.bfloat16, -) -params["scheduler"] = scheduler_state -``` - -利用Flax对TPU的兼容性实现并行图像生成。需为每个设备复制模型参数，并分配输入数据： - -```python -# 每个并行设备生成1张图像（TPUv2-8/TPUv3-8支持8设备并行） -prompt = "一张宇航员在火星上骑马的高清照片，高分辨率，高画质。" -num_samples = jax.device_count() -prompt_ids = pipeline.prepare_inputs([prompt] * num_samples) - -prng_seed = jax.random.PRNGKey(0) -num_inference_steps = 25 - -# 分配输入和随机种子 -params = replicate(params) -prng_seed = jax.random.split(prng_seed, jax.device_count()) -prompt_ids = shard(prompt_ids) - -images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images -images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]))) -``` - ## 模型加载 通过[`ModelMixin.from_pretrained`]方法加载模型，该方法会下载并缓存模型权重和配置的最新版本。若本地缓存已存在最新文件，则直接复用缓存而非重复下载。 diff --git a/examples/controlnet/README.md b/examples/controlnet/README.md index 99767617394c..388368a8778f 100644 --- a/examples/controlnet/README.md +++ b/examples/controlnet/README.md @@ -292,6 +292,9 @@ image.save("./output.png") ## Training with Flax/JAX +> [!WARNING] +> JAX/Flax support was removed from the `diffusers` library in v0.40.0. These scripts require `diffusers<=0.39.x` to run. + For faster training on TPUs and GPUs you can leverage the flax training example. Follow the instructions above to get the model and dataset before running the script. ### Running on Google Cloud TPU diff --git a/examples/dreambooth/README.md b/examples/dreambooth/README.md index 006e583e9f16..73ea8238953f 100644 --- a/examples/dreambooth/README.md +++ b/examples/dreambooth/README.md @@ -439,6 +439,9 @@ Note that the use of [`StableDiffusionLoraLoaderMixin.load_lora_weights`](https: ## Training with Flax/JAX +> [!WARNING] +> JAX/Flax support was removed from the `diffusers` library in v0.40.0. These scripts require `diffusers<=0.39.x` to run. + For faster training on TPUs and GPUs you can leverage the flax training example. Follow the instructions above to get the model and dataset before running the script. ____Note: The flax example don't yet support features like gradient checkpoint, gradient accumulation etc, so to use flax for faster training we will need >30GB cards.___ diff --git a/examples/model_search/pipeline_easy.py b/examples/model_search/pipeline_easy.py index ee5dced817ec..d4c306b76b38 100644 --- a/examples/model_search/pipeline_easy.py +++ b/examples/model_search/pipeline_easy.py @@ -1243,8 +1243,7 @@ def from_huggingface(cls, pretrained_model_link_or_path, **kwargs): class). The overwritten components are passed directly to the pipelines `__init__` method. See example below for more information. variant (`str`, *optional*): - Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. + Load weights from a specified variant filename such as `"fp16"` or `"ema"`. > [!TIP] > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with @@ -1495,8 +1494,7 @@ def from_huggingface(cls, pretrained_model_link_or_path, **kwargs): class). The overwritten components are passed directly to the pipelines `__init__` method. See example below for more information. variant (`str`, *optional*): - Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. + Load weights from a specified variant filename such as `"fp16"` or `"ema"`. > [!TIP] > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with @@ -1748,8 +1746,7 @@ def from_huggingface(cls, pretrained_model_link_or_path, **kwargs): class). The overwritten components are passed directly to the pipelines `__init__` method. See example below for more information. variant (`str`, *optional*): - Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. + Load weights from a specified variant filename such as `"fp16"` or `"ema"`. > [!TIP] > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with diff --git a/examples/research_projects/multi_token_textual_inversion/README.md b/examples/research_projects/multi_token_textual_inversion/README.md index 7d80c0beee37..5e5d4991cb5b 100644 --- a/examples/research_projects/multi_token_textual_inversion/README.md +++ b/examples/research_projects/multi_token_textual_inversion/README.md @@ -114,6 +114,9 @@ image.save("cat-backpack.png") ## Training with Flax/JAX +> [!WARNING] +> JAX/Flax support was removed from the `diffusers` library in v0.40.0. These scripts require `diffusers<=0.39.x` to run. + For faster training on TPUs and GPUs you can leverage the flax training example. Follow the instructions above to get the model and dataset before running the script. Before running the scripts, make sure to install the library's training dependencies: diff --git a/examples/research_projects/sdxl_flax/README.md b/examples/research_projects/sdxl_flax/README.md index e2c02d7bb617..c14911fc718a 100644 --- a/examples/research_projects/sdxl_flax/README.md +++ b/examples/research_projects/sdxl_flax/README.md @@ -1,5 +1,8 @@ # Stable Diffusion XL for JAX + TPUv5e +> [!WARNING] +> JAX/Flax support was removed from the `diffusers` library in v0.40.0. This project requires `diffusers<=0.39.x` to run. + [TPU v5e](https://cloud.google.com/blog/products/compute/how-cloud-tpu-v5e-accelerates-large-scale-ai-inference) is a new generation of TPUs from Google Cloud. It is the most cost-effective, versatile, and scalable Cloud TPU to date. This makes them ideal for serving and scaling large diffusion models. [JAX](https://github.com/google/jax) is a high-performance numerical computation library that is well-suited to develop and deploy diffusion models: diff --git a/examples/text_to_image/README.md b/examples/text_to_image/README.md index ebbf0a96becc..4718d59fcfcb 100644 --- a/examples/text_to_image/README.md +++ b/examples/text_to_image/README.md @@ -277,6 +277,9 @@ pipe = StableDiffusionPipeline.from_pretrained(base_model_id, torch_dtype=torch. ## Training with Flax/JAX +> [!WARNING] +> JAX/Flax support was removed from the `diffusers` library in v0.40.0. These scripts require `diffusers<=0.39.x` to run. + For faster training on TPUs and GPUs you can leverage the flax training example. Follow the instructions above to get the model and dataset before running the script. **___Note: The flax example doesn't yet support features like gradient checkpoint, gradient accumulation etc, so to use flax for faster training we will need >30GB cards or TPU v3.___** diff --git a/examples/textual_inversion/README.md b/examples/textual_inversion/README.md index 06e22dbcd804..528c7962e38f 100644 --- a/examples/textual_inversion/README.md +++ b/examples/textual_inversion/README.md @@ -122,6 +122,9 @@ image.save("cat-backpack.png") ## Training with Flax/JAX +> [!WARNING] +> JAX/Flax support was removed from the `diffusers` library in v0.40.0. These scripts require `diffusers<=0.39.x` to run. + For faster training on TPUs and GPUs you can leverage the flax training example. Follow the instructions above to get the model and dataset before running the script. Before running the scripts, make sure to install the library's training dependencies: diff --git a/setup.py b/setup.py index 955f6feb09aa..99ccb52a598e 100644 --- a/setup.py +++ b/setup.py @@ -84,7 +84,6 @@ you need to go back to main before executing this. """ -import os import re import sys @@ -99,7 +98,6 @@ "accelerate>=0.31.0", "datasets", "filelock", - "flax>=0.4.1", "ftfy", "hf-doc-builder>=0.3.0", "httpx<1.0.0", @@ -108,8 +106,6 @@ "importlib_metadata", "invisible-watermark>=0.2.0", "isort>=5.5.4", - "jax>=0.4.1", - "jaxlib>=0.4.1", "Jinja2", "torchsde", "note_seq", @@ -257,14 +253,7 @@ def run(self): extras["nvidia_modelopt"] = deps_list("nvidia_modelopt[hf]") extras["flashpack"] = deps_list("flashpack") -if os.name == "nt": # windows - extras["flax"] = [] # jax is not supported on windows -else: - extras["flax"] = deps_list("jax", "jaxlib", "flax") - -extras["dev"] = ( - extras["quality"] + extras["test"] + extras["training"] + extras["docs"] + extras["torch"] + extras["flax"] -) +extras["dev"] = extras["quality"] + extras["test"] + extras["training"] + extras["docs"] + extras["torch"] install_requires = [ deps["importlib_metadata"], @@ -283,10 +272,10 @@ def run(self): setup( name="diffusers", version="0.40.0.dev0", # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots) - description="State-of-the-art diffusion in PyTorch and JAX.", + description="State-of-the-art diffusion in PyTorch.", long_description=open("README.md", "r", encoding="utf-8").read(), long_description_content_type="text/markdown", - keywords="deep learning diffusion jax pytorch stable diffusion audioldm", + keywords="deep learning diffusion pytorch stable diffusion audioldm", license="Apache 2.0 License", author="The Hugging Face team (past and future) with the help of all our contributors (https://github.com/huggingface/diffusers/graphs/contributors)", author_email="diffusers@huggingface.co", diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 86b9e96c3e97..4d9f7c5aeeba 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -9,7 +9,6 @@ is_accelerate_available, is_auto_round_available, is_bitsandbytes_available, - is_flax_available, is_gguf_available, is_librosa_available, is_note_seq_available, @@ -23,7 +22,6 @@ is_torchao_available, is_torchsde_available, is_transformers_available, - is_transformers_flax_compatible, is_transformers_version, ) @@ -48,7 +46,6 @@ "schedulers": [], "utils": [ "OptionalDependencyNotAvailable", - "is_flax_available", "is_inflect_available", "is_invisible_watermark_available", "is_librosa_available", @@ -900,60 +897,6 @@ else: _import_structure["pipelines"].extend(["SpectrogramDiffusionPipeline"]) -try: - if not is_flax_available(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_flax_objects # noqa F403 - - _import_structure["utils.dummy_flax_objects"] = [ - name for name in dir(dummy_flax_objects) if not name.startswith("_") - ] - - -else: - _import_structure["models.controlnets.controlnet_flax"] = ["FlaxControlNetModel"] - _import_structure["models.modeling_flax_utils"] = ["FlaxModelMixin"] - _import_structure["models.unets.unet_2d_condition_flax"] = ["FlaxUNet2DConditionModel"] - _import_structure["models.vae_flax"] = ["FlaxAutoencoderKL"] - _import_structure["schedulers"].extend( - [ - "FlaxDDIMScheduler", - "FlaxDDPMScheduler", - "FlaxDPMSolverMultistepScheduler", - "FlaxEulerDiscreteScheduler", - "FlaxKarrasVeScheduler", - "FlaxLMSDiscreteScheduler", - "FlaxPNDMScheduler", - "FlaxSchedulerMixin", - "FlaxScoreSdeVeScheduler", - ] - ) - - -try: - if not (is_flax_available() and is_transformers_available() and is_transformers_flax_compatible()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_flax_and_transformers_objects # noqa F403 - - _import_structure["utils.dummy_flax_and_transformers_objects"] = [ - name for name in dir(dummy_flax_and_transformers_objects) if not name.startswith("_") - ] - - -else: - _import_structure["pipelines"].extend( - [ - "FlaxDiffusionPipeline", - "FlaxStableDiffusionControlNetPipeline", - "FlaxStableDiffusionImg2ImgPipeline", - "FlaxStableDiffusionInpaintPipeline", - "FlaxStableDiffusionPipeline", - "FlaxStableDiffusionXLPipeline", - ] - ) - try: if not (is_note_seq_available()): raise OptionalDependencyNotAvailable() @@ -1714,43 +1657,6 @@ else: from .pipelines import SpectrogramDiffusionPipeline - try: - if not is_flax_available(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_flax_objects import * # noqa F403 - else: - from .models.controlnets.controlnet_flax import FlaxControlNetModel - from .models.modeling_flax_utils import FlaxModelMixin - from .models.unets.unet_2d_condition_flax import FlaxUNet2DConditionModel - from .models.vae_flax import FlaxAutoencoderKL - from .schedulers import ( - FlaxDDIMScheduler, - FlaxDDPMScheduler, - FlaxDPMSolverMultistepScheduler, - FlaxEulerDiscreteScheduler, - FlaxKarrasVeScheduler, - FlaxLMSDiscreteScheduler, - FlaxPNDMScheduler, - FlaxSchedulerMixin, - FlaxScoreSdeVeScheduler, - ) - - try: - if not (is_flax_available() and is_transformers_available() and is_transformers_flax_compatible()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_flax_and_transformers_objects import * # noqa F403 - else: - from .pipelines import ( - FlaxDiffusionPipeline, - FlaxStableDiffusionControlNetPipeline, - FlaxStableDiffusionImg2ImgPipeline, - FlaxStableDiffusionInpaintPipeline, - FlaxStableDiffusionPipeline, - FlaxStableDiffusionXLPipeline, - ) - try: if not (is_note_seq_available()): raise OptionalDependencyNotAvailable() diff --git a/src/diffusers/commands/env.py b/src/diffusers/commands/env.py index 3dc937b38bfc..29ae8bc49373 100644 --- a/src/diffusers/commands/env.py +++ b/src/diffusers/commands/env.py @@ -23,7 +23,6 @@ from ..utils import ( is_accelerate_available, is_bitsandbytes_available, - is_flax_available, is_gguf_available, is_google_colab, is_nvidia_modelopt_available, @@ -75,20 +74,6 @@ def run(self) -> dict: pt_version = torch.__version__ pt_cuda_available = torch.cuda.is_available() - flax_version = "not installed" - jax_version = "not installed" - jaxlib_version = "not installed" - jax_backend = "NA" - if is_flax_available(): - import flax - import jax - import jaxlib - - flax_version = flax.__version__ - jax_version = jax.__version__ - jaxlib_version = jaxlib.__version__ - jax_backend = jax.lib.xla_bridge.get_backend().platform - transformers_version = "not installed" if is_transformers_available(): import transformers @@ -173,9 +158,6 @@ def run(self) -> dict: "Running on Google Colab?": is_google_colab_str, "Python version": platform.python_version(), "PyTorch version (GPU?)": f"{pt_version} ({pt_cuda_available})", - "Flax version (CPU?/GPU?/TPU?)": f"{flax_version} ({jax_backend})", - "Jax version": jax_version, - "JaxLib version": jaxlib_version, "Huggingface_hub version": hub_version, "Transformers version": transformers_version, "Accelerate version": accelerate_version, diff --git a/src/diffusers/configuration_utils.py b/src/diffusers/configuration_utils.py index 7a95ce20aaff..f16871b9f56f 100644 --- a/src/diffusers/configuration_utils.py +++ b/src/diffusers/configuration_utils.py @@ -15,7 +15,6 @@ # limitations under the License. """ConfigMixin base class and utilities.""" -import dataclasses import functools import importlib import inspect @@ -510,10 +509,6 @@ def extract_init_dict(cls, config_dict, **kwargs): # remove general kwargs if present in dict if "kwargs" in expected_keys: expected_keys.remove("kwargs") - # remove flax internal keys - if hasattr(cls, "_flax_internal_args"): - for arg in cls._flax_internal_args: - expected_keys.remove(arg) # 2. Remove attributes that cannot be expected from expected config attributes # remove keys to be ignored @@ -737,54 +732,6 @@ def inner_init(self, *args, **kwargs): return inner_init -def flax_register_to_config(cls): - original_init = cls.__init__ - - @functools.wraps(original_init) - def init(self, *args, **kwargs): - if not isinstance(self, ConfigMixin): - raise RuntimeError( - f"`@register_for_config` was applied to {self.__class__.__name__} init method, but this class does " - "not inherit from `ConfigMixin`." - ) - - # Ignore private kwargs in the init. Retrieve all passed attributes - init_kwargs = dict(kwargs.items()) - - # Retrieve default values - fields = dataclasses.fields(self) - default_kwargs = {} - for field in fields: - # ignore flax specific attributes - if field.name in self._flax_internal_args: - continue - if type(field.default) == dataclasses._MISSING_TYPE: - default_kwargs[field.name] = None - else: - default_kwargs[field.name] = getattr(self, field.name) - - # Make sure init_kwargs override default kwargs - new_kwargs = {**default_kwargs, **init_kwargs} - # dtype should be part of `init_kwargs`, but not `new_kwargs` - if "dtype" in new_kwargs: - new_kwargs.pop("dtype") - - # Get positional arguments aligned with kwargs - for i, arg in enumerate(args): - name = fields[i].name - new_kwargs[name] = arg - - # Take note of the parameters that were not present in the loaded config - if len(set(new_kwargs.keys()) - set(init_kwargs)) > 0: - new_kwargs["_use_default_values"] = list(set(new_kwargs.keys()) - set(init_kwargs)) - - getattr(self, "register_to_config")(**new_kwargs) - original_init(self, *args, **kwargs) - - cls.__init__ = init - return cls - - class LegacyConfigMixin(ConfigMixin): r""" A subclass of `ConfigMixin` to resolve class mapping from legacy classes (like `Transformer2DModel`) to more diff --git a/src/diffusers/dependency_versions_table.py b/src/diffusers/dependency_versions_table.py index fcced60419d0..34f1757eb480 100644 --- a/src/diffusers/dependency_versions_table.py +++ b/src/diffusers/dependency_versions_table.py @@ -6,7 +6,6 @@ "accelerate": "accelerate>=0.31.0", "datasets": "datasets", "filelock": "filelock", - "flax": "flax>=0.4.1", "ftfy": "ftfy", "hf-doc-builder": "hf-doc-builder>=0.3.0", "httpx": "httpx<1.0.0", @@ -15,8 +14,6 @@ "importlib_metadata": "importlib_metadata", "invisible-watermark": "invisible-watermark>=0.2.0", "isort": "isort>=5.5.4", - "jax": "jax>=0.4.1", - "jaxlib": "jaxlib>=0.4.1", "Jinja2": "Jinja2", "torchsde": "torchsde", "note_seq": "note_seq", diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 1746db347d32..167ee7a534de 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -17,7 +17,6 @@ from ..utils import ( DIFFUSERS_SLOW_IMPORT, _LazyModule, - is_flax_available, is_torch_available, ) @@ -156,11 +155,6 @@ _import_structure["unets.unet_stable_cascade"] = ["StableCascadeUNet"] _import_structure["unets.uvit_2d"] = ["UVit2DModel"] -if is_flax_available(): - _import_structure["controlnets.controlnet_flax"] = ["FlaxControlNetModel"] - _import_structure["unets.unet_2d_condition_flax"] = ["FlaxUNet2DConditionModel"] - _import_structure["vae_flax"] = ["FlaxAutoencoderKL"] - if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: if is_torch_available(): @@ -303,11 +297,6 @@ UVit2DModel, ) - if is_flax_available(): - from .controlnets import FlaxControlNetModel - from .unets import FlaxUNet2DConditionModel - from .vae_flax import FlaxAutoencoderKL - else: import sys diff --git a/src/diffusers/models/adapter.py b/src/diffusers/models/adapter.py index d387859ffbd1..09f88f285a29 100644 --- a/src/diffusers/models/adapter.py +++ b/src/diffusers/models/adapter.py @@ -187,8 +187,7 @@ def from_pretrained(cls, pretrained_model_path: str | os.PathLike | None, **kwar model. This is only supported when torch version >= 1.9.0. If you are using an older version of torch, setting this argument to `True` will raise an error. variant (`str`, *optional*): - If specified, load weights from a `variant` file (*e.g.* pytorch_model..bin). `variant` will - be ignored when using `from_flax`. + If specified, load weights from a `variant` file (*e.g.* pytorch_model..bin). use_safetensors (`bool`, *optional*, defaults to `None`): If `None`, the `safetensors` weights will be downloaded if available **and** if`safetensors` library is installed. If `True`, the model will be forcibly loaded from`safetensors` weights. If `False`, diff --git a/src/diffusers/models/attention_flax.py b/src/diffusers/models/attention_flax.py deleted file mode 100644 index 2009703d5b43..000000000000 --- a/src/diffusers/models/attention_flax.py +++ /dev/null @@ -1,524 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import functools -import math - -import flax.linen as nn -import jax -import jax.numpy as jnp - -from ..utils import logging - - -logger = logging.get_logger(__name__) - - -def _query_chunk_attention(query, key, value, precision, key_chunk_size: int = 4096): - """Multi-head dot product attention with a limited number of queries.""" - num_kv, num_heads, k_features = key.shape[-3:] - v_features = value.shape[-1] - key_chunk_size = min(key_chunk_size, num_kv) - query = query / jnp.sqrt(k_features) - - @functools.partial(jax.checkpoint, prevent_cse=False) - def summarize_chunk(query, key, value): - attn_weights = jnp.einsum("...qhd,...khd->...qhk", query, key, precision=precision) - - max_score = jnp.max(attn_weights, axis=-1, keepdims=True) - max_score = jax.lax.stop_gradient(max_score) - exp_weights = jnp.exp(attn_weights - max_score) - - exp_values = jnp.einsum("...vhf,...qhv->...qhf", value, exp_weights, precision=precision) - max_score = jnp.einsum("...qhk->...qh", max_score) - - return (exp_values, exp_weights.sum(axis=-1), max_score) - - def chunk_scanner(chunk_idx): - # julienne key array - key_chunk = jax.lax.dynamic_slice( - operand=key, - start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0], # [...,k,h,d] - slice_sizes=list(key.shape[:-3]) + [key_chunk_size, num_heads, k_features], # [...,k,h,d] - ) - - # julienne value array - value_chunk = jax.lax.dynamic_slice( - operand=value, - start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0], # [...,v,h,d] - slice_sizes=list(value.shape[:-3]) + [key_chunk_size, num_heads, v_features], # [...,v,h,d] - ) - - return summarize_chunk(query, key_chunk, value_chunk) - - chunk_values, chunk_weights, chunk_max = jax.lax.map(f=chunk_scanner, xs=jnp.arange(0, num_kv, key_chunk_size)) - - global_max = jnp.max(chunk_max, axis=0, keepdims=True) - max_diffs = jnp.exp(chunk_max - global_max) - - chunk_values *= jnp.expand_dims(max_diffs, axis=-1) - chunk_weights *= max_diffs - - all_values = chunk_values.sum(axis=0) - all_weights = jnp.expand_dims(chunk_weights, -1).sum(axis=0) - - return all_values / all_weights - - -def jax_memory_efficient_attention( - query, key, value, precision=jax.lax.Precision.HIGHEST, query_chunk_size: int = 1024, key_chunk_size: int = 4096 -): - r""" - Flax Memory-efficient multi-head dot product attention. https://huggingface.co/papers/2112.05682v2 - https://github.com/AminRezaei0x443/memory-efficient-attention - - Args: - query (`jnp.ndarray`): (batch..., query_length, head, query_key_depth_per_head) - key (`jnp.ndarray`): (batch..., key_value_length, head, query_key_depth_per_head) - value (`jnp.ndarray`): (batch..., key_value_length, head, value_depth_per_head) - precision (`jax.lax.Precision`, *optional*, defaults to `jax.lax.Precision.HIGHEST`): - numerical precision for computation - query_chunk_size (`int`, *optional*, defaults to 1024): - chunk size to divide query array value must divide query_length equally without remainder - key_chunk_size (`int`, *optional*, defaults to 4096): - chunk size to divide key and value array value must divide key_value_length equally without remainder - - Returns: - (`jnp.ndarray`) with shape of (batch..., query_length, head, value_depth_per_head) - """ - num_q, num_heads, q_features = query.shape[-3:] - - def chunk_scanner(chunk_idx, _): - # julienne query array - query_chunk = jax.lax.dynamic_slice( - operand=query, - start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0], # [...,q,h,d] - slice_sizes=list(query.shape[:-3]) + [min(query_chunk_size, num_q), num_heads, q_features], # [...,q,h,d] - ) - - return ( - chunk_idx + query_chunk_size, # unused ignore it - _query_chunk_attention( - query=query_chunk, key=key, value=value, precision=precision, key_chunk_size=key_chunk_size - ), - ) - - _, res = jax.lax.scan( - f=chunk_scanner, - init=0, - xs=None, - length=math.ceil(num_q / query_chunk_size), # start counter # stop counter - ) - - return jnp.concatenate(res, axis=-3) # fuse the chunked result back - - -class FlaxAttention(nn.Module): - r""" - A Flax multi-head attention module as described in: https://huggingface.co/papers/1706.03762 - - Parameters: - query_dim (:obj:`int`): - Input hidden states dimension - heads (:obj:`int`, *optional*, defaults to 8): - Number of heads - dim_head (:obj:`int`, *optional*, defaults to 64): - Hidden states dimension inside each head - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - use_memory_efficient_attention (`bool`, *optional*, defaults to `False`): - enable memory efficient attention https://huggingface.co/papers/2112.05682 - split_head_dim (`bool`, *optional*, defaults to `False`): - Whether to split the head dimension into a new axis for the self-attention computation. In most cases, - enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL. - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - - """ - - query_dim: int - heads: int = 8 - dim_head: int = 64 - dropout: float = 0.0 - use_memory_efficient_attention: bool = False - split_head_dim: bool = False - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - inner_dim = self.dim_head * self.heads - self.scale = self.dim_head**-0.5 - - # Weights were exported with old names {to_q, to_k, to_v, to_out} - self.query = nn.Dense(inner_dim, use_bias=False, dtype=self.dtype, name="to_q") - self.key = nn.Dense(inner_dim, use_bias=False, dtype=self.dtype, name="to_k") - self.value = nn.Dense(inner_dim, use_bias=False, dtype=self.dtype, name="to_v") - - self.proj_attn = nn.Dense(self.query_dim, dtype=self.dtype, name="to_out_0") - self.dropout_layer = nn.Dropout(rate=self.dropout) - - def reshape_heads_to_batch_dim(self, tensor): - batch_size, seq_len, dim = tensor.shape - head_size = self.heads - tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size) - tensor = jnp.transpose(tensor, (0, 2, 1, 3)) - tensor = tensor.reshape(batch_size * head_size, seq_len, dim // head_size) - return tensor - - def reshape_batch_dim_to_heads(self, tensor): - batch_size, seq_len, dim = tensor.shape - head_size = self.heads - tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim) - tensor = jnp.transpose(tensor, (0, 2, 1, 3)) - tensor = tensor.reshape(batch_size // head_size, seq_len, dim * head_size) - return tensor - - def __call__(self, hidden_states, context=None, deterministic=True): - context = hidden_states if context is None else context - - query_proj = self.query(hidden_states) - key_proj = self.key(context) - value_proj = self.value(context) - - if self.split_head_dim: - b = hidden_states.shape[0] - query_states = jnp.reshape(query_proj, (b, -1, self.heads, self.dim_head)) - key_states = jnp.reshape(key_proj, (b, -1, self.heads, self.dim_head)) - value_states = jnp.reshape(value_proj, (b, -1, self.heads, self.dim_head)) - else: - query_states = self.reshape_heads_to_batch_dim(query_proj) - key_states = self.reshape_heads_to_batch_dim(key_proj) - value_states = self.reshape_heads_to_batch_dim(value_proj) - - if self.use_memory_efficient_attention: - query_states = query_states.transpose(1, 0, 2) - key_states = key_states.transpose(1, 0, 2) - value_states = value_states.transpose(1, 0, 2) - - # this if statement create a chunk size for each layer of the unet - # the chunk size is equal to the query_length dimension of the deepest layer of the unet - - flatten_latent_dim = query_states.shape[-3] - if flatten_latent_dim % 64 == 0: - query_chunk_size = int(flatten_latent_dim / 64) - elif flatten_latent_dim % 16 == 0: - query_chunk_size = int(flatten_latent_dim / 16) - elif flatten_latent_dim % 4 == 0: - query_chunk_size = int(flatten_latent_dim / 4) - else: - query_chunk_size = int(flatten_latent_dim) - - hidden_states = jax_memory_efficient_attention( - query_states, key_states, value_states, query_chunk_size=query_chunk_size, key_chunk_size=4096 * 4 - ) - hidden_states = hidden_states.transpose(1, 0, 2) - hidden_states = self.reshape_batch_dim_to_heads(hidden_states) - else: - # compute attentions - if self.split_head_dim: - attention_scores = jnp.einsum("b t n h, b f n h -> b n f t", key_states, query_states) - else: - attention_scores = jnp.einsum("b i d, b j d->b i j", query_states, key_states) - - attention_scores = attention_scores * self.scale - attention_probs = nn.softmax(attention_scores, axis=-1 if self.split_head_dim else 2) - - # attend to values - if self.split_head_dim: - hidden_states = jnp.einsum("b n f t, b t n h -> b f n h", attention_probs, value_states) - b = hidden_states.shape[0] - hidden_states = jnp.reshape(hidden_states, (b, -1, self.heads * self.dim_head)) - else: - hidden_states = jnp.einsum("b i j, b j d -> b i d", attention_probs, value_states) - hidden_states = self.reshape_batch_dim_to_heads(hidden_states) - - hidden_states = self.proj_attn(hidden_states) - return self.dropout_layer(hidden_states, deterministic=deterministic) - - -class FlaxBasicTransformerBlock(nn.Module): - r""" - A Flax transformer block layer with `GLU` (Gated Linear Unit) activation function as described in: - https://huggingface.co/papers/1706.03762 - - - Parameters: - dim (:obj:`int`): - Inner hidden states dimension - n_heads (:obj:`int`): - Number of heads - d_head (:obj:`int`): - Hidden states dimension inside each head - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - only_cross_attention (`bool`, defaults to `False`): - Whether to only apply cross attention. - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - use_memory_efficient_attention (`bool`, *optional*, defaults to `False`): - enable memory efficient attention https://huggingface.co/papers/2112.05682 - split_head_dim (`bool`, *optional*, defaults to `False`): - Whether to split the head dimension into a new axis for the self-attention computation. In most cases, - enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL. - """ - - dim: int - n_heads: int - d_head: int - dropout: float = 0.0 - only_cross_attention: bool = False - dtype: jnp.dtype = jnp.float32 - use_memory_efficient_attention: bool = False - split_head_dim: bool = False - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - # self attention (or cross_attention if only_cross_attention is True) - self.attn1 = FlaxAttention( - self.dim, - self.n_heads, - self.d_head, - self.dropout, - self.use_memory_efficient_attention, - self.split_head_dim, - dtype=self.dtype, - ) - # cross attention - self.attn2 = FlaxAttention( - self.dim, - self.n_heads, - self.d_head, - self.dropout, - self.use_memory_efficient_attention, - self.split_head_dim, - dtype=self.dtype, - ) - self.ff = FlaxFeedForward(dim=self.dim, dropout=self.dropout, dtype=self.dtype) - self.norm1 = nn.LayerNorm(epsilon=1e-5, dtype=self.dtype) - self.norm2 = nn.LayerNorm(epsilon=1e-5, dtype=self.dtype) - self.norm3 = nn.LayerNorm(epsilon=1e-5, dtype=self.dtype) - self.dropout_layer = nn.Dropout(rate=self.dropout) - - def __call__(self, hidden_states, context, deterministic=True): - # self attention - residual = hidden_states - if self.only_cross_attention: - hidden_states = self.attn1(self.norm1(hidden_states), context, deterministic=deterministic) - else: - hidden_states = self.attn1(self.norm1(hidden_states), deterministic=deterministic) - hidden_states = hidden_states + residual - - # cross attention - residual = hidden_states - hidden_states = self.attn2(self.norm2(hidden_states), context, deterministic=deterministic) - hidden_states = hidden_states + residual - - # feed forward - residual = hidden_states - hidden_states = self.ff(self.norm3(hidden_states), deterministic=deterministic) - hidden_states = hidden_states + residual - - return self.dropout_layer(hidden_states, deterministic=deterministic) - - -class FlaxTransformer2DModel(nn.Module): - r""" - A Spatial Transformer layer with Gated Linear Unit (GLU) activation function as described in: - https://huggingface.co/papers/1506.02025 - - - Parameters: - in_channels (:obj:`int`): - Input number of channels - n_heads (:obj:`int`): - Number of heads - d_head (:obj:`int`): - Hidden states dimension inside each head - depth (:obj:`int`, *optional*, defaults to 1): - Number of transformers block - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - use_linear_projection (`bool`, defaults to `False`): tbd - only_cross_attention (`bool`, defaults to `False`): tbd - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - use_memory_efficient_attention (`bool`, *optional*, defaults to `False`): - enable memory efficient attention https://huggingface.co/papers/2112.05682 - split_head_dim (`bool`, *optional*, defaults to `False`): - Whether to split the head dimension into a new axis for the self-attention computation. In most cases, - enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL. - """ - - in_channels: int - n_heads: int - d_head: int - depth: int = 1 - dropout: float = 0.0 - use_linear_projection: bool = False - only_cross_attention: bool = False - dtype: jnp.dtype = jnp.float32 - use_memory_efficient_attention: bool = False - split_head_dim: bool = False - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - self.norm = nn.GroupNorm(num_groups=32, epsilon=1e-5) - - inner_dim = self.n_heads * self.d_head - if self.use_linear_projection: - self.proj_in = nn.Dense(inner_dim, dtype=self.dtype) - else: - self.proj_in = nn.Conv( - inner_dim, - kernel_size=(1, 1), - strides=(1, 1), - padding="VALID", - dtype=self.dtype, - ) - - self.transformer_blocks = [ - FlaxBasicTransformerBlock( - inner_dim, - self.n_heads, - self.d_head, - dropout=self.dropout, - only_cross_attention=self.only_cross_attention, - dtype=self.dtype, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - ) - for _ in range(self.depth) - ] - - if self.use_linear_projection: - self.proj_out = nn.Dense(inner_dim, dtype=self.dtype) - else: - self.proj_out = nn.Conv( - inner_dim, - kernel_size=(1, 1), - strides=(1, 1), - padding="VALID", - dtype=self.dtype, - ) - - self.dropout_layer = nn.Dropout(rate=self.dropout) - - def __call__(self, hidden_states, context, deterministic=True): - batch, height, width, channels = hidden_states.shape - residual = hidden_states - hidden_states = self.norm(hidden_states) - if self.use_linear_projection: - hidden_states = hidden_states.reshape(batch, height * width, channels) - hidden_states = self.proj_in(hidden_states) - else: - hidden_states = self.proj_in(hidden_states) - hidden_states = hidden_states.reshape(batch, height * width, channels) - - for transformer_block in self.transformer_blocks: - hidden_states = transformer_block(hidden_states, context, deterministic=deterministic) - - if self.use_linear_projection: - hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.reshape(batch, height, width, channels) - else: - hidden_states = hidden_states.reshape(batch, height, width, channels) - hidden_states = self.proj_out(hidden_states) - - hidden_states = hidden_states + residual - return self.dropout_layer(hidden_states, deterministic=deterministic) - - -class FlaxFeedForward(nn.Module): - r""" - Flax module that encapsulates two Linear layers separated by a non-linearity. It is the counterpart of PyTorch's - [`FeedForward`] class, with the following simplifications: - - The activation function is currently hardcoded to a gated linear unit from: - https://huggingface.co/papers/2002.05202 - - `dim_out` is equal to `dim`. - - The number of hidden dimensions is hardcoded to `dim * 4` in [`FlaxGELU`]. - - Parameters: - dim (:obj:`int`): - Inner hidden states dimension - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - dim: int - dropout: float = 0.0 - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - # The second linear layer needs to be called - # net_2 for now to match the index of the Sequential layer - self.net_0 = FlaxGEGLU(self.dim, self.dropout, self.dtype) - self.net_2 = nn.Dense(self.dim, dtype=self.dtype) - - def __call__(self, hidden_states, deterministic=True): - hidden_states = self.net_0(hidden_states, deterministic=deterministic) - hidden_states = self.net_2(hidden_states) - return hidden_states - - -class FlaxGEGLU(nn.Module): - r""" - Flax implementation of a Linear layer followed by the variant of the gated linear unit activation function from - https://huggingface.co/papers/2002.05202. - - Parameters: - dim (:obj:`int`): - Input hidden states dimension - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - dim: int - dropout: float = 0.0 - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - inner_dim = self.dim * 4 - self.proj = nn.Dense(inner_dim * 2, dtype=self.dtype) - self.dropout_layer = nn.Dropout(rate=self.dropout) - - def __call__(self, hidden_states, deterministic=True): - hidden_states = self.proj(hidden_states) - hidden_linear, hidden_gelu = jnp.split(hidden_states, 2, axis=2) - return self.dropout_layer(hidden_linear * nn.gelu(hidden_gelu), deterministic=deterministic) diff --git a/src/diffusers/models/auto_model.py b/src/diffusers/models/auto_model.py index d3010266ef70..81dda81bc6ae 100644 --- a/src/diffusers/models/auto_model.py +++ b/src/diffusers/models/auto_model.py @@ -192,8 +192,6 @@ def from_pretrained(cls, pretrained_model_or_path: str | os.PathLike | None = No revision (`str`, *optional*, defaults to `"main"`): The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier allowed by Git. - from_flax (`bool`, *optional*, defaults to `False`): - Load the model weights from a Flax checkpoint save file. subfolder (`str`, *optional*, defaults to `""`): The subfolder location of a model file within a larger model repository on the Hub or locally. mirror (`str`, *optional*): @@ -223,8 +221,7 @@ def from_pretrained(cls, pretrained_model_or_path: str | os.PathLike | None = No Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this argument to `True` will raise an error. variant (`str`, *optional*): - Load weights from a specified `variant` filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. + Load weights from a specified `variant` filename such as `"fp16"` or `"ema"`. use_safetensors (`bool`, *optional*, defaults to `None`): If set to `None`, the `safetensors` weights are downloaded if they're available **and** if the `safetensors` library is installed. If set to `True`, the model is forcibly loaded from `safetensors` diff --git a/src/diffusers/models/controlnets/__init__.py b/src/diffusers/models/controlnets/__init__.py index 853a2207f903..3f9c7337e0f2 100644 --- a/src/diffusers/models/controlnets/__init__.py +++ b/src/diffusers/models/controlnets/__init__.py @@ -1,4 +1,4 @@ -from ...utils import is_flax_available, is_torch_available +from ...utils import is_torch_available if is_torch_available(): @@ -23,6 +23,3 @@ from .controlnet_z_image import ZImageControlNetModel from .multicontrolnet import MultiControlNetModel from .multicontrolnet_union import MultiControlNetUnionModel - -if is_flax_available(): - from .controlnet_flax import FlaxControlNetModel diff --git a/src/diffusers/models/controlnets/controlnet_flax.py b/src/diffusers/models/controlnets/controlnet_flax.py deleted file mode 100644 index 61f750d6f02c..000000000000 --- a/src/diffusers/models/controlnets/controlnet_flax.py +++ /dev/null @@ -1,406 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import flax -import flax.linen as nn -import jax -import jax.numpy as jnp -from flax.core.frozen_dict import FrozenDict - -from ...configuration_utils import ConfigMixin, flax_register_to_config -from ...utils import BaseOutput, logging -from ..embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps -from ..modeling_flax_utils import FlaxModelMixin -from ..unets.unet_2d_blocks_flax import ( - FlaxCrossAttnDownBlock2D, - FlaxDownBlock2D, - FlaxUNetMidBlock2DCrossAttn, -) - - -logger = logging.get_logger(__name__) - - -@flax.struct.dataclass -class FlaxControlNetOutput(BaseOutput): - """ - The output of [`FlaxControlNetModel`]. - - Args: - down_block_res_samples (`jnp.ndarray`): - mid_block_res_sample (`jnp.ndarray`): - """ - - down_block_res_samples: jnp.ndarray - mid_block_res_sample: jnp.ndarray - - -class FlaxControlNetConditioningEmbedding(nn.Module): - conditioning_embedding_channels: int - block_out_channels: tuple[int, ...] = (16, 32, 96, 256) - dtype: jnp.dtype = jnp.float32 - - def setup(self) -> None: - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - self.conv_in = nn.Conv( - self.block_out_channels[0], - kernel_size=(3, 3), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - blocks = [] - for i in range(len(self.block_out_channels) - 1): - channel_in = self.block_out_channels[i] - channel_out = self.block_out_channels[i + 1] - conv1 = nn.Conv( - channel_in, - kernel_size=(3, 3), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - blocks.append(conv1) - conv2 = nn.Conv( - channel_out, - kernel_size=(3, 3), - strides=(2, 2), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - blocks.append(conv2) - self.blocks = blocks - - self.conv_out = nn.Conv( - self.conditioning_embedding_channels, - kernel_size=(3, 3), - padding=((1, 1), (1, 1)), - kernel_init=nn.initializers.zeros_init(), - bias_init=nn.initializers.zeros_init(), - dtype=self.dtype, - ) - - def __call__(self, conditioning: jnp.ndarray) -> jnp.ndarray: - embedding = self.conv_in(conditioning) - embedding = nn.silu(embedding) - - for block in self.blocks: - embedding = block(embedding) - embedding = nn.silu(embedding) - - embedding = self.conv_out(embedding) - - return embedding - - -@flax_register_to_config -class FlaxControlNetModel(nn.Module, FlaxModelMixin, ConfigMixin): - r""" - A ControlNet model. - - This model inherits from [`FlaxModelMixin`]. Check the superclass documentation for it’s generic methods - implemented for all models (such as downloading or saving). - - This model is also a Flax Linen [`flax.linen.Module`](https://flax.readthedocs.io/en/latest/flax.linen.html#module) - subclass. Use it as a regular Flax Linen module and refer to the Flax documentation for all matters related to its - general usage and behavior. - - Inherent JAX features such as the following are supported: - - - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit) - - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation) - - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap) - - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap) - - Parameters: - sample_size (`int`, *optional*): - The size of the input sample. - in_channels (`int`, *optional*, defaults to 4): - The number of channels in the input sample. - down_block_types (`tuple[str]`, *optional*, defaults to `("FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxDownBlock2D")`): - The tuple of downsample blocks to use. - block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): - The tuple of output channels for each block. - layers_per_block (`int`, *optional*, defaults to 2): - The number of layers per block. - attention_head_dim (`int` or `tuple[int]`, *optional*, defaults to 8): - The dimension of the attention heads. - num_attention_heads (`int` or `tuple[int]`, *optional*): - The number of attention heads. - cross_attention_dim (`int`, *optional*, defaults to 768): - The dimension of the cross attention features. - dropout (`float`, *optional*, defaults to 0): - Dropout probability for down, up and bottleneck blocks. - flip_sin_to_cos (`bool`, *optional*, defaults to `True`): - Whether to flip the sin to cos in the time embedding. - freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding. - controlnet_conditioning_channel_order (`str`, *optional*, defaults to `rgb`): - The channel order of conditional image. Will convert to `rgb` if it's `bgr`. - conditioning_embedding_out_channels (`tuple`, *optional*, defaults to `(16, 32, 96, 256)`): - The tuple of output channel for each block in the `conditioning_embedding` layer. - """ - - sample_size: int = 32 - in_channels: int = 4 - down_block_types: tuple[str, ...] = ( - "CrossAttnDownBlock2D", - "CrossAttnDownBlock2D", - "CrossAttnDownBlock2D", - "DownBlock2D", - ) - only_cross_attention: bool | tuple[bool, ...] = False - block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280) - layers_per_block: int = 2 - attention_head_dim: int | tuple[int, ...] = 8 - num_attention_heads: int | tuple[int, ...] | None = None - cross_attention_dim: int = 1280 - dropout: float = 0.0 - use_linear_projection: bool = False - dtype: jnp.dtype = jnp.float32 - flip_sin_to_cos: bool = True - freq_shift: int = 0 - controlnet_conditioning_channel_order: str = "rgb" - conditioning_embedding_out_channels: tuple[int, ...] = (16, 32, 96, 256) - - def init_weights(self, rng: jax.Array) -> FrozenDict: - # init input tensors - sample_shape = (1, self.in_channels, self.sample_size, self.sample_size) - sample = jnp.zeros(sample_shape, dtype=jnp.float32) - timesteps = jnp.ones((1,), dtype=jnp.int32) - encoder_hidden_states = jnp.zeros((1, 1, self.cross_attention_dim), dtype=jnp.float32) - controlnet_cond_shape = (1, 3, self.sample_size * 8, self.sample_size * 8) - controlnet_cond = jnp.zeros(controlnet_cond_shape, dtype=jnp.float32) - - params_rng, dropout_rng = jax.random.split(rng) - rngs = {"params": params_rng, "dropout": dropout_rng} - - return self.init(rngs, sample, timesteps, encoder_hidden_states, controlnet_cond)["params"] - - def setup(self) -> None: - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - block_out_channels = self.block_out_channels - time_embed_dim = block_out_channels[0] * 4 - - # If `num_attention_heads` is not defined (which is the case for most models) - # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. - # The reason for this behavior is to correct for incorrectly named variables that were introduced - # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 - # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking - # which is why we correct for the naming here. - num_attention_heads = self.num_attention_heads or self.attention_head_dim - - # input - self.conv_in = nn.Conv( - block_out_channels[0], - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - # time - self.time_proj = FlaxTimesteps( - block_out_channels[0], flip_sin_to_cos=self.flip_sin_to_cos, freq_shift=self.config.freq_shift - ) - self.time_embedding = FlaxTimestepEmbedding(time_embed_dim, dtype=self.dtype) - - self.controlnet_cond_embedding = FlaxControlNetConditioningEmbedding( - conditioning_embedding_channels=block_out_channels[0], - block_out_channels=self.conditioning_embedding_out_channels, - ) - - only_cross_attention = self.only_cross_attention - if isinstance(only_cross_attention, bool): - only_cross_attention = (only_cross_attention,) * len(self.down_block_types) - - if isinstance(num_attention_heads, int): - num_attention_heads = (num_attention_heads,) * len(self.down_block_types) - - # down - down_blocks = [] - controlnet_down_blocks = [] - - output_channel = block_out_channels[0] - - controlnet_block = nn.Conv( - output_channel, - kernel_size=(1, 1), - padding="VALID", - kernel_init=nn.initializers.zeros_init(), - bias_init=nn.initializers.zeros_init(), - dtype=self.dtype, - ) - controlnet_down_blocks.append(controlnet_block) - - for i, down_block_type in enumerate(self.down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - if down_block_type == "CrossAttnDownBlock2D": - down_block = FlaxCrossAttnDownBlock2D( - in_channels=input_channel, - out_channels=output_channel, - dropout=self.dropout, - num_layers=self.layers_per_block, - num_attention_heads=num_attention_heads[i], - add_downsample=not is_final_block, - use_linear_projection=self.use_linear_projection, - only_cross_attention=only_cross_attention[i], - dtype=self.dtype, - ) - else: - down_block = FlaxDownBlock2D( - in_channels=input_channel, - out_channels=output_channel, - dropout=self.dropout, - num_layers=self.layers_per_block, - add_downsample=not is_final_block, - dtype=self.dtype, - ) - - down_blocks.append(down_block) - - for _ in range(self.layers_per_block): - controlnet_block = nn.Conv( - output_channel, - kernel_size=(1, 1), - padding="VALID", - kernel_init=nn.initializers.zeros_init(), - bias_init=nn.initializers.zeros_init(), - dtype=self.dtype, - ) - controlnet_down_blocks.append(controlnet_block) - - if not is_final_block: - controlnet_block = nn.Conv( - output_channel, - kernel_size=(1, 1), - padding="VALID", - kernel_init=nn.initializers.zeros_init(), - bias_init=nn.initializers.zeros_init(), - dtype=self.dtype, - ) - controlnet_down_blocks.append(controlnet_block) - - self.down_blocks = down_blocks - self.controlnet_down_blocks = controlnet_down_blocks - - # mid - mid_block_channel = block_out_channels[-1] - self.mid_block = FlaxUNetMidBlock2DCrossAttn( - in_channels=mid_block_channel, - dropout=self.dropout, - num_attention_heads=num_attention_heads[-1], - use_linear_projection=self.use_linear_projection, - dtype=self.dtype, - ) - - self.controlnet_mid_block = nn.Conv( - mid_block_channel, - kernel_size=(1, 1), - padding="VALID", - kernel_init=nn.initializers.zeros_init(), - bias_init=nn.initializers.zeros_init(), - dtype=self.dtype, - ) - - def __call__( - self, - sample: jnp.ndarray, - timesteps: jnp.ndarray | float | int, - encoder_hidden_states: jnp.ndarray, - controlnet_cond: jnp.ndarray, - conditioning_scale: float = 1.0, - return_dict: bool = True, - train: bool = False, - ) -> FlaxControlNetOutput | tuple[tuple[jnp.ndarray, ...], jnp.ndarray]: - r""" - Args: - sample (`jnp.ndarray`): (batch, channel, height, width) noisy inputs tensor - timestep (`jnp.ndarray` or `float` or `int`): timesteps - encoder_hidden_states (`jnp.ndarray`): (batch_size, sequence_length, hidden_size) encoder hidden states - controlnet_cond (`jnp.ndarray`): (batch, channel, height, width) the conditional input tensor - conditioning_scale (`float`, *optional*, defaults to `1.0`): the scale factor for controlnet outputs - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`models.unets.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] instead of - a plain tuple. - train (`bool`, *optional*, defaults to `False`): - Use deterministic functions and disable dropout when not training. - - Returns: - [`~models.unets.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] or `tuple`: - [`~models.unets.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] if `return_dict` is True, otherwise - a `tuple`. When returning a tuple, the first element is the sample tensor. - """ - channel_order = self.controlnet_conditioning_channel_order - if channel_order == "bgr": - controlnet_cond = jnp.flip(controlnet_cond, axis=1) - - # 1. time - if not isinstance(timesteps, jnp.ndarray): - timesteps = jnp.array([timesteps], dtype=jnp.int32) - elif isinstance(timesteps, jnp.ndarray) and len(timesteps.shape) == 0: - timesteps = timesteps.astype(dtype=jnp.float32) - timesteps = jnp.expand_dims(timesteps, 0) - - t_emb = self.time_proj(timesteps) - t_emb = self.time_embedding(t_emb) - - # 2. pre-process - sample = jnp.transpose(sample, (0, 2, 3, 1)) - sample = self.conv_in(sample) - - controlnet_cond = jnp.transpose(controlnet_cond, (0, 2, 3, 1)) - controlnet_cond = self.controlnet_cond_embedding(controlnet_cond) - sample += controlnet_cond - - # 3. down - down_block_res_samples = (sample,) - for down_block in self.down_blocks: - if isinstance(down_block, FlaxCrossAttnDownBlock2D): - sample, res_samples = down_block(sample, t_emb, encoder_hidden_states, deterministic=not train) - else: - sample, res_samples = down_block(sample, t_emb, deterministic=not train) - down_block_res_samples += res_samples - - # 4. mid - sample = self.mid_block(sample, t_emb, encoder_hidden_states, deterministic=not train) - - # 5. contronet blocks - controlnet_down_block_res_samples = () - for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks): - down_block_res_sample = controlnet_block(down_block_res_sample) - controlnet_down_block_res_samples += (down_block_res_sample,) - - down_block_res_samples = controlnet_down_block_res_samples - - mid_block_res_sample = self.controlnet_mid_block(sample) - - # 6. scaling - down_block_res_samples = [sample * conditioning_scale for sample in down_block_res_samples] - mid_block_res_sample *= conditioning_scale - - if not return_dict: - return (down_block_res_samples, mid_block_res_sample) - - return FlaxControlNetOutput( - down_block_res_samples=down_block_res_samples, mid_block_res_sample=mid_block_res_sample - ) diff --git a/src/diffusers/models/controlnets/multicontrolnet.py b/src/diffusers/models/controlnets/multicontrolnet.py index 995637c965d3..308592ce9bc2 100644 --- a/src/diffusers/models/controlnets/multicontrolnet.py +++ b/src/diffusers/models/controlnets/multicontrolnet.py @@ -184,8 +184,7 @@ def from_pretrained(cls, pretrained_model_path: str | os.PathLike | None, **kwar model. This is only supported when torch version >= 1.9.0. If you are using an older version of torch, setting this argument to `True` will raise an error. variant (`str`, *optional*): - If specified load weights from `variant` filename, *e.g.* pytorch_model..bin. `variant` is - ignored when using `from_flax`. + If specified load weights from `variant` filename, *e.g.* pytorch_model..bin. use_safetensors (`bool`, *optional*, defaults to `None`): If set to `None`, the `safetensors` weights will be downloaded if they're available **and** if the `safetensors` library is installed. If set to `True`, the model will be forcibly loaded from diff --git a/src/diffusers/models/controlnets/multicontrolnet_union.py b/src/diffusers/models/controlnets/multicontrolnet_union.py index 577492074243..2c50c0184b2e 100644 --- a/src/diffusers/models/controlnets/multicontrolnet_union.py +++ b/src/diffusers/models/controlnets/multicontrolnet_union.py @@ -201,8 +201,7 @@ def from_pretrained(cls, pretrained_model_path: str | os.PathLike | None, **kwar model. This is only supported when torch version >= 1.9.0. If you are using an older version of torch, setting this argument to `True` will raise an error. variant (`str`, *optional*): - If specified load weights from `variant` filename, *e.g.* pytorch_model..bin. `variant` is - ignored when using `from_flax`. + If specified load weights from `variant` filename, *e.g.* pytorch_model..bin. use_safetensors (`bool`, *optional*, defaults to `None`): If set to `None`, the `safetensors` weights will be downloaded if they're available **and** if the `safetensors` library is installed. If set to `True`, the model will be forcibly loaded from diff --git a/src/diffusers/models/embeddings_flax.py b/src/diffusers/models/embeddings_flax.py deleted file mode 100644 index 7bf7f63c438b..000000000000 --- a/src/diffusers/models/embeddings_flax.py +++ /dev/null @@ -1,126 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import math - -import flax.linen as nn -import jax.numpy as jnp - -from ..utils import logging - - -logger = logging.get_logger(__name__) - - -def get_sinusoidal_embeddings( - timesteps: jnp.ndarray, - embedding_dim: int, - freq_shift: float = 1, - min_timescale: float = 1, - max_timescale: float = 1.0e4, - flip_sin_to_cos: bool = False, - scale: float = 1.0, -) -> jnp.ndarray: - """Returns the positional encoding (same as Tensor2Tensor). - - Args: - timesteps (`jnp.ndarray` of shape `(N,)`): - A 1-D array of N indices, one per batch element. These may be fractional. - embedding_dim (`int`): - The number of output channels. - freq_shift (`float`, *optional*, defaults to `1`): - Shift applied to the frequency scaling of the embeddings. - min_timescale (`float`, *optional*, defaults to `1`): - The smallest time unit used in the sinusoidal calculation (should probably be 0.0). - max_timescale (`float`, *optional*, defaults to `1.0e4`): - The largest time unit used in the sinusoidal calculation. - flip_sin_to_cos (`bool`, *optional*, defaults to `False`): - Whether to flip the order of sinusoidal components to cosine first. - scale (`float`, *optional*, defaults to `1.0`): - A scaling factor applied to the positional embeddings. - - Returns: - a Tensor of timing signals [N, num_channels] - """ - assert timesteps.ndim == 1, "Timesteps should be a 1d-array" - assert embedding_dim % 2 == 0, f"Embedding dimension {embedding_dim} should be even" - num_timescales = float(embedding_dim // 2) - log_timescale_increment = math.log(max_timescale / min_timescale) / (num_timescales - freq_shift) - inv_timescales = min_timescale * jnp.exp(jnp.arange(num_timescales, dtype=jnp.float32) * -log_timescale_increment) - emb = jnp.expand_dims(timesteps, 1) * jnp.expand_dims(inv_timescales, 0) - - # scale embeddings - scaled_time = scale * emb - - if flip_sin_to_cos: - signal = jnp.concatenate([jnp.cos(scaled_time), jnp.sin(scaled_time)], axis=1) - else: - signal = jnp.concatenate([jnp.sin(scaled_time), jnp.cos(scaled_time)], axis=1) - signal = jnp.reshape(signal, [jnp.shape(timesteps)[0], embedding_dim]) - return signal - - -class FlaxTimestepEmbedding(nn.Module): - r""" - Time step Embedding Module. Learns embeddings for input time steps. - - Args: - time_embed_dim (`int`, *optional*, defaults to `32`): - Time step embedding dimension. - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - The data type for the embedding parameters. - """ - - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - time_embed_dim: int = 32 - dtype: jnp.dtype = jnp.float32 - - @nn.compact - def __call__(self, temb): - temb = nn.Dense(self.time_embed_dim, dtype=self.dtype, name="linear_1")(temb) - temb = nn.silu(temb) - temb = nn.Dense(self.time_embed_dim, dtype=self.dtype, name="linear_2")(temb) - return temb - - -class FlaxTimesteps(nn.Module): - r""" - Wrapper Module for sinusoidal Time step Embeddings as described in https://huggingface.co/papers/2006.11239 - - Args: - dim (`int`, *optional*, defaults to `32`): - Time step embedding dimension. - flip_sin_to_cos (`bool`, *optional*, defaults to `False`): - Whether to flip the sinusoidal function from sine to cosine. - freq_shift (`float`, *optional*, defaults to `1`): - Frequency shift applied to the sinusoidal embeddings. - """ - - dim: int = 32 - flip_sin_to_cos: bool = False - freq_shift: float = 1 - - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - @nn.compact - def __call__(self, timesteps): - return get_sinusoidal_embeddings( - timesteps, embedding_dim=self.dim, flip_sin_to_cos=self.flip_sin_to_cos, freq_shift=self.freq_shift - ) diff --git a/src/diffusers/models/modeling_flax_pytorch_utils.py b/src/diffusers/models/modeling_flax_pytorch_utils.py deleted file mode 100644 index d64c48a9601e..000000000000 --- a/src/diffusers/models/modeling_flax_pytorch_utils.py +++ /dev/null @@ -1,135 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""PyTorch - Flax general utilities.""" - -import re - -import jax.numpy as jnp -from flax.traverse_util import flatten_dict, unflatten_dict -from jax.random import PRNGKey - -from ..utils import logging - - -logger = logging.get_logger(__name__) - - -def rename_key(key): - regex = r"\w+[.]\d+" - pats = re.findall(regex, key) - for pat in pats: - key = key.replace(pat, "_".join(pat.split("."))) - return key - - -##################### -# PyTorch => Flax # -##################### - - -# Adapted from https://github.com/huggingface/transformers/blob/c603c80f46881ae18b2ca50770ef65fa4033eacd/src/transformers/modeling_flax_pytorch_utils.py#L69 -# and https://github.com/patil-suraj/stable-diffusion-jax/blob/main/stable_diffusion_jax/convert_diffusers_to_jax.py -def rename_key_and_reshape_tensor(pt_tuple_key, pt_tensor, random_flax_state_dict): - """Rename PT weight names to corresponding Flax weight names and reshape tensor if necessary""" - # conv norm or layer norm - renamed_pt_tuple_key = pt_tuple_key[:-1] + ("scale",) - - # rename attention layers - if len(pt_tuple_key) > 1: - for rename_from, rename_to in ( - ("to_out_0", "proj_attn"), - ("to_k", "key"), - ("to_v", "value"), - ("to_q", "query"), - ): - if pt_tuple_key[-2] == rename_from: - weight_name = pt_tuple_key[-1] - weight_name = "kernel" if weight_name == "weight" else weight_name - renamed_pt_tuple_key = pt_tuple_key[:-2] + (rename_to, weight_name) - if renamed_pt_tuple_key in random_flax_state_dict: - assert random_flax_state_dict[renamed_pt_tuple_key].shape == pt_tensor.T.shape - return renamed_pt_tuple_key, pt_tensor.T - - if ( - any("norm" in str_ for str_ in pt_tuple_key) - and (pt_tuple_key[-1] == "bias") - and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict) - and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict) - ): - renamed_pt_tuple_key = pt_tuple_key[:-1] + ("scale",) - return renamed_pt_tuple_key, pt_tensor - elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict: - renamed_pt_tuple_key = pt_tuple_key[:-1] + ("scale",) - return renamed_pt_tuple_key, pt_tensor - - # embedding - if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict: - pt_tuple_key = pt_tuple_key[:-1] + ("embedding",) - return renamed_pt_tuple_key, pt_tensor - - # conv layer - renamed_pt_tuple_key = pt_tuple_key[:-1] + ("kernel",) - if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4: - pt_tensor = pt_tensor.transpose(2, 3, 1, 0) - return renamed_pt_tuple_key, pt_tensor - - # linear layer - renamed_pt_tuple_key = pt_tuple_key[:-1] + ("kernel",) - if pt_tuple_key[-1] == "weight": - pt_tensor = pt_tensor.T - return renamed_pt_tuple_key, pt_tensor - - # old PyTorch layer norm weight - renamed_pt_tuple_key = pt_tuple_key[:-1] + ("weight",) - if pt_tuple_key[-1] == "gamma": - return renamed_pt_tuple_key, pt_tensor - - # old PyTorch layer norm bias - renamed_pt_tuple_key = pt_tuple_key[:-1] + ("bias",) - if pt_tuple_key[-1] == "beta": - return renamed_pt_tuple_key, pt_tensor - - return pt_tuple_key, pt_tensor - - -def convert_pytorch_state_dict_to_flax(pt_state_dict, flax_model, init_key=42): - # Step 1: Convert pytorch tensor to numpy - pt_state_dict = {k: v.numpy() for k, v in pt_state_dict.items()} - - # Step 2: Since the model is stateless, get random Flax params - random_flax_params = flax_model.init_weights(PRNGKey(init_key)) - - random_flax_state_dict = flatten_dict(random_flax_params) - flax_state_dict = {} - - # Need to change some parameters name to match Flax names - for pt_key, pt_tensor in pt_state_dict.items(): - renamed_pt_key = rename_key(pt_key) - pt_tuple_key = tuple(renamed_pt_key.split(".")) - - # Correctly rename weight parameters - flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, pt_tensor, random_flax_state_dict) - - if flax_key in random_flax_state_dict: - if flax_tensor.shape != random_flax_state_dict[flax_key].shape: - raise ValueError( - f"PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape " - f"{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}." - ) - - # also add unexpected weight so that warning is thrown - flax_state_dict[flax_key] = jnp.asarray(flax_tensor) - - return unflatten_dict(flax_state_dict) diff --git a/src/diffusers/models/modeling_flax_utils.py b/src/diffusers/models/modeling_flax_utils.py deleted file mode 100644 index 3bc68172a23b..000000000000 --- a/src/diffusers/models/modeling_flax_utils.py +++ /dev/null @@ -1,558 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import os -from pickle import UnpicklingError -from typing import Any - -import jax -import jax.numpy as jnp -import msgpack.exceptions -from flax.core.frozen_dict import FrozenDict, unfreeze -from flax.serialization import from_bytes, to_bytes -from flax.traverse_util import flatten_dict, unflatten_dict -from huggingface_hub import create_repo, hf_hub_download -from huggingface_hub.utils import ( - EntryNotFoundError, - HfHubHTTPError, - RepositoryNotFoundError, - RevisionNotFoundError, - validate_hf_hub_args, -) - -from .. import __version__, is_torch_available -from ..utils import ( - CONFIG_NAME, - FLAX_WEIGHTS_NAME, - HUGGINGFACE_CO_RESOLVE_ENDPOINT, - WEIGHTS_NAME, - PushToHubMixin, - logging, -) -from .modeling_flax_pytorch_utils import convert_pytorch_state_dict_to_flax - - -logger = logging.get_logger(__name__) - - -class FlaxModelMixin(PushToHubMixin): - r""" - Base class for all Flax models. - - [`FlaxModelMixin`] takes care of storing the model configuration and provides methods for loading, downloading and - saving models. - - - **config_name** ([`str`]) -- Filename to save a model to when calling [`~FlaxModelMixin.save_pretrained`]. - """ - - config_name = CONFIG_NAME - _automatically_saved_args = ["_diffusers_version", "_class_name", "_name_or_path"] - _flax_internal_args = ["name", "parent", "dtype"] - - @classmethod - def _from_config(cls, config, **kwargs): - """ - All context managers that the model should be initialized under go here. - """ - return cls(config, **kwargs) - - def _cast_floating_to(self, params: dict | FrozenDict, dtype: jnp.dtype, mask: Any = None) -> Any: - """ - Helper method to cast floating-point values of given parameter `PyTree` to given `dtype`. - """ - - # taken from https://github.com/deepmind/jmp/blob/3a8318abc3292be38582794dbf7b094e6583b192/jmp/_src/policy.py#L27 - def conditional_cast(param): - if isinstance(param, jnp.ndarray) and jnp.issubdtype(param.dtype, jnp.floating): - param = param.astype(dtype) - return param - - if mask is None: - return jax.tree_map(conditional_cast, params) - - flat_params = flatten_dict(params) - flat_mask, _ = jax.tree_flatten(mask) - - for masked, key in zip(flat_mask, flat_params.keys()): - if masked: - param = flat_params[key] - flat_params[key] = conditional_cast(param) - - return unflatten_dict(flat_params) - - def to_bf16(self, params: dict | FrozenDict, mask: Any = None): - r""" - Cast the floating-point `params` to `jax.numpy.bfloat16`. This returns a new `params` tree and does not cast - the `params` in place. - - This method can be used on a TPU to explicitly convert the model parameters to bfloat16 precision to do full - half-precision training or to save weights in bfloat16 for inference in order to save memory and improve speed. - - Arguments: - params (`dict | FrozenDict`): - A `PyTree` of model parameters. - mask (`dict | FrozenDict`): - A `PyTree` with same structure as the `params` tree. The leaves should be booleans. It should be `True` - for params you want to cast, and `False` for those you want to skip. - - Examples: - - ```python - >>> from diffusers import FlaxUNet2DConditionModel - - >>> # load model - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") - >>> # By default, the model parameters will be in fp32 precision, to cast these to bfloat16 precision - >>> params = model.to_bf16(params) - >>> # If you don't want to cast certain parameters (for example layer norm bias and scale) - >>> # then pass the mask as follows - >>> from flax import traverse_util - - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") - >>> flat_params = traverse_util.flatten_dict(params) - >>> mask = { - ... path: (path[-2] != ("LayerNorm", "bias") and path[-2:] != ("LayerNorm", "scale")) - ... for path in flat_params - ... } - >>> mask = traverse_util.unflatten_dict(mask) - >>> params = model.to_bf16(params, mask) - ```""" - return self._cast_floating_to(params, jnp.bfloat16, mask) - - def to_fp32(self, params: dict | FrozenDict, mask: Any = None): - r""" - Cast the floating-point `params` to `jax.numpy.float32`. This method can be used to explicitly convert the - model parameters to fp32 precision. This returns a new `params` tree and does not cast the `params` in place. - - Arguments: - params (`dict | FrozenDict`): - A `PyTree` of model parameters. - mask (`dict | FrozenDict`): - A `PyTree` with same structure as the `params` tree. The leaves should be booleans. It should be `True` - for params you want to cast, and `False` for those you want to skip. - - Examples: - - ```python - >>> from diffusers import FlaxUNet2DConditionModel - - >>> # Download model and configuration from huggingface.co - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") - >>> # By default, the model params will be in fp32, to illustrate the use of this method, - >>> # we'll first cast to fp16 and back to fp32 - >>> params = model.to_f16(params) - >>> # now cast back to fp32 - >>> params = model.to_fp32(params) - ```""" - return self._cast_floating_to(params, jnp.float32, mask) - - def to_fp16(self, params: dict | FrozenDict, mask: Any = None): - r""" - Cast the floating-point `params` to `jax.numpy.float16`. This returns a new `params` tree and does not cast the - `params` in place. - - This method can be used on a GPU to explicitly convert the model parameters to float16 precision to do full - half-precision training or to save weights in float16 for inference in order to save memory and improve speed. - - Arguments: - params (`dict | FrozenDict`): - A `PyTree` of model parameters. - mask (`dict | FrozenDict`): - A `PyTree` with same structure as the `params` tree. The leaves should be booleans. It should be `True` - for params you want to cast, and `False` for those you want to skip. - - Examples: - - ```python - >>> from diffusers import FlaxUNet2DConditionModel - - >>> # load model - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") - >>> # By default, the model params will be in fp32, to cast these to float16 - >>> params = model.to_fp16(params) - >>> # If you want don't want to cast certain parameters (for example layer norm bias and scale) - >>> # then pass the mask as follows - >>> from flax import traverse_util - - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") - >>> flat_params = traverse_util.flatten_dict(params) - >>> mask = { - ... path: (path[-2] != ("LayerNorm", "bias") and path[-2:] != ("LayerNorm", "scale")) - ... for path in flat_params - ... } - >>> mask = traverse_util.unflatten_dict(mask) - >>> params = model.to_fp16(params, mask) - ```""" - return self._cast_floating_to(params, jnp.float16, mask) - - def init_weights(self, rng: jax.Array) -> dict: - raise NotImplementedError(f"init_weights method has to be implemented for {self}") - - @classmethod - @validate_hf_hub_args - def from_pretrained( - cls, - pretrained_model_name_or_path: str | os.PathLike, - dtype: jnp.dtype = jnp.float32, - *model_args, - **kwargs, - ): - r""" - Instantiate a pretrained Flax model from a pretrained model configuration. - - Parameters: - pretrained_model_name_or_path (`str` or `os.PathLike`): - Can be either: - - - A string, the *model id* (for example `stable-diffusion-v1-5/stable-diffusion-v1-5`) of a - pretrained model hosted on the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - using [`~FlaxModelMixin.save_pretrained`]. - dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`): - The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and - `jax.numpy.bfloat16` (on TPUs). - - This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If - specified, all the computation will be performed with the given `dtype`. - - > [!TIP] > This only specifies the dtype of the *computation* and does not influence the dtype of model - > parameters. > > If you wish to change the dtype of the model parameters, see - [`~FlaxModelMixin.to_fp16`] and > [`~FlaxModelMixin.to_bf16`]. - - model_args (sequence of positional arguments, *optional*): - All remaining positional arguments are passed to the underlying model's `__init__` method. - cache_dir (`str | os.PathLike`, *optional*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - local_files_only(`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - from_pt (`bool`, *optional*, defaults to `False`): - Load the model weights from a PyTorch checkpoint save file. - kwargs (remaining dictionary of keyword arguments, *optional*): - Can be used to update the configuration object (after it is loaded) and initiate the model (for - example, `output_attentions=True`). Behaves differently depending on whether a `config` is provided or - automatically loaded: - - - If a configuration is provided with `config`, `kwargs` are directly passed to the underlying - model's `__init__` method (we assume all relevant updates to the configuration have already been - done). - - If a configuration is not provided, `kwargs` are first passed to the configuration class - initialization function [`~ConfigMixin.from_config`]. Each key of the `kwargs` that corresponds - to a configuration attribute is used to override said attribute with the supplied `kwargs` value. - Remaining keys that do not correspond to any configuration attribute are passed to the underlying - model's `__init__` function. - - Examples: - - ```python - >>> from diffusers import FlaxUNet2DConditionModel - - >>> # Download model and configuration from huggingface.co and cache. - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") - >>> # Model was saved using *save_pretrained('./test/saved_model/')* (for example purposes, not runnable). - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("./test/saved_model/") - ``` - - If you get the error message below, you need to finetune the weights for your downstream task: - - ```bash - Some weights of UNet2DConditionModel were not initialized from the model checkpoint at stable-diffusion-v1-5/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: - - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated - You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. - ``` - """ - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - config = kwargs.pop("config", None) - cache_dir = kwargs.pop("cache_dir", None) - force_download = kwargs.pop("force_download", False) - from_pt = kwargs.pop("from_pt", False) - proxies = kwargs.pop("proxies", None) - local_files_only = kwargs.pop("local_files_only", False) - token = kwargs.pop("token", None) - revision = kwargs.pop("revision", None) - subfolder = kwargs.pop("subfolder", None) - - user_agent = { - "diffusers": __version__, - "file_type": "model", - "framework": "flax", - } - - # Load config if we don't provide one - if config is None: - config, unused_kwargs = cls.load_config( - pretrained_model_name_or_path, - cache_dir=cache_dir, - return_unused_kwargs=True, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - **kwargs, - ) - - model, model_kwargs = cls.from_config(config, dtype=dtype, return_unused_kwargs=True, **unused_kwargs) - - # Load model - pretrained_path_with_subfolder = ( - pretrained_model_name_or_path - if subfolder is None - else os.path.join(pretrained_model_name_or_path, subfolder) - ) - if os.path.isdir(pretrained_path_with_subfolder): - if from_pt: - if not os.path.isfile(os.path.join(pretrained_path_with_subfolder, WEIGHTS_NAME)): - raise EnvironmentError( - f"Error no file named {WEIGHTS_NAME} found in directory {pretrained_path_with_subfolder} " - ) - model_file = os.path.join(pretrained_path_with_subfolder, WEIGHTS_NAME) - elif os.path.isfile(os.path.join(pretrained_path_with_subfolder, FLAX_WEIGHTS_NAME)): - # Load from a Flax checkpoint - model_file = os.path.join(pretrained_path_with_subfolder, FLAX_WEIGHTS_NAME) - # Check if pytorch weights exist instead - elif os.path.isfile(os.path.join(pretrained_path_with_subfolder, WEIGHTS_NAME)): - raise EnvironmentError( - f"{WEIGHTS_NAME} file found in directory {pretrained_path_with_subfolder}. Please load the model" - " using `from_pt=True`." - ) - else: - raise EnvironmentError( - f"Error no file named {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME} found in directory " - f"{pretrained_path_with_subfolder}." - ) - else: - try: - model_file = hf_hub_download( - pretrained_model_name_or_path, - filename=FLAX_WEIGHTS_NAME if not from_pt else WEIGHTS_NAME, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - user_agent=user_agent, - subfolder=subfolder, - revision=revision, - ) - - except RepositoryNotFoundError: - raise EnvironmentError( - f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier " - "listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a " - "token having permission to this repo with `token` or log in with `hf auth login`." - ) - except RevisionNotFoundError: - raise EnvironmentError( - f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for " - "this model name. Check the model page at " - f"'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions." - ) - except EntryNotFoundError: - raise EnvironmentError( - f"{pretrained_model_name_or_path} does not appear to have a file named {FLAX_WEIGHTS_NAME}." - ) - except HfHubHTTPError as err: - raise EnvironmentError( - f"There was a specific connection error when trying to load {pretrained_model_name_or_path}:\n" - f"{err}" - ) - except ValueError: - raise EnvironmentError( - f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it" - f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a" - f" directory containing a file named {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME}.\nCheckout your" - " internet connection or see how to run the library in offline mode at" - " 'https://huggingface.co/docs/transformers/installation#offline-mode'." - ) - except EnvironmentError: - raise EnvironmentError( - f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it from " - "'https://huggingface.co/models', make sure you don't have a local directory with the same name. " - f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory " - f"containing a file named {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME}." - ) - - if from_pt: - if is_torch_available(): - from .modeling_utils import load_state_dict - else: - raise EnvironmentError( - "Can't load the model in PyTorch format because PyTorch is not installed. " - "Please, install PyTorch or use native Flax weights." - ) - - # Step 1: Get the pytorch file - pytorch_model_file = load_state_dict(model_file) - - # Step 2: Convert the weights - state = convert_pytorch_state_dict_to_flax(pytorch_model_file, model) - else: - try: - with open(model_file, "rb") as state_f: - state = from_bytes(cls, state_f.read()) - except (UnpicklingError, msgpack.exceptions.ExtraData) as e: - try: - with open(model_file) as f: - if f.read().startswith("version"): - raise OSError( - "You seem to have cloned a repository without having git-lfs installed. Please" - " install git-lfs and run `git lfs install` followed by `git lfs pull` in the" - " folder you cloned." - ) - else: - raise ValueError from e - except (UnicodeDecodeError, ValueError): - raise EnvironmentError(f"Unable to convert {model_file} to Flax deserializable object. ") - # make sure all arrays are stored as jnp.ndarray - # NOTE: This is to prevent a bug this will be fixed in Flax >= v0.3.4: - # https://github.com/google/flax/issues/1261 - state = jax.tree_util.tree_map(lambda x: jax.device_put(x, jax.local_devices(backend="cpu")[0]), state) - - # flatten dicts - state = flatten_dict(state) - - params_shape_tree = jax.eval_shape(model.init_weights, rng=jax.random.PRNGKey(0)) - required_params = set(flatten_dict(unfreeze(params_shape_tree)).keys()) - - shape_state = flatten_dict(unfreeze(params_shape_tree)) - - missing_keys = required_params - set(state.keys()) - unexpected_keys = set(state.keys()) - required_params - - if missing_keys: - logger.warning( - f"The checkpoint {pretrained_model_name_or_path} is missing required keys: {missing_keys}. " - "Make sure to call model.init_weights to initialize the missing weights." - ) - cls._missing_keys = missing_keys - - for key in state.keys(): - if key in shape_state and state[key].shape != shape_state[key].shape: - raise ValueError( - f"Trying to load the pretrained weight for {key} failed: checkpoint has shape " - f"{state[key].shape} which is incompatible with the model shape {shape_state[key].shape}. " - ) - - # remove unexpected keys to not be saved again - for unexpected_key in unexpected_keys: - del state[unexpected_key] - - if len(unexpected_keys) > 0: - logger.warning( - f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when" - f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are" - f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or" - " with another architecture." - ) - else: - logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n") - - if len(missing_keys) > 0: - logger.warning( - f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" - f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably" - " TRAIN this model on a down-stream task to be able to use it for predictions and inference." - ) - else: - logger.info( - f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at" - f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the checkpoint" - f" was trained on, you can already use {model.__class__.__name__} for predictions without further" - " training." - ) - - return model, unflatten_dict(state) - - def save_pretrained( - self, - save_directory: str | os.PathLike, - params: dict | FrozenDict, - is_main_process: bool = True, - push_to_hub: bool = False, - **kwargs, - ): - """ - Save a model and its configuration file to a directory so that it can be reloaded using the - [`~FlaxModelMixin.from_pretrained`] class method. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save a model and its configuration file to. Will be created if it doesn't exist. - params (`dict | FrozenDict`): - A `PyTree` of model parameters. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - push_to_hub (`bool`, *optional*, defaults to `False`): - Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the - repository you want to push to with `repo_id` (will default to the name of `save_directory` in your - namespace). - kwargs (`dict[str, Any]`, *optional*): - Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. - """ - if os.path.isfile(save_directory): - logger.error(f"Provided path ({save_directory}) should be a directory, not a file") - return - - os.makedirs(save_directory, exist_ok=True) - - if push_to_hub: - commit_message = kwargs.pop("commit_message", None) - private = kwargs.pop("private", None) - create_pr = kwargs.pop("create_pr", False) - token = kwargs.pop("token", None) - repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) - repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id - - model_to_save = self - - # Attach architecture to the config - # Save the config - if is_main_process: - model_to_save.save_config(save_directory) - - # save model - output_model_file = os.path.join(save_directory, FLAX_WEIGHTS_NAME) - with open(output_model_file, "wb") as f: - model_bytes = to_bytes(params) - f.write(model_bytes) - - logger.info(f"Model weights saved in {output_model_file}") - - if push_to_hub: - self._upload_folder( - save_directory, - repo_id, - token=token, - commit_message=commit_message, - create_pr=create_pr, - ) diff --git a/src/diffusers/models/modeling_pytorch_flax_utils.py b/src/diffusers/models/modeling_pytorch_flax_utils.py deleted file mode 100644 index ada55073dd55..000000000000 --- a/src/diffusers/models/modeling_pytorch_flax_utils.py +++ /dev/null @@ -1,161 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""PyTorch - Flax general utilities.""" - -from pickle import UnpicklingError - -import jax -import jax.numpy as jnp -import numpy as np -from flax.serialization import from_bytes -from flax.traverse_util import flatten_dict - -from ..utils import logging - - -logger = logging.get_logger(__name__) - - -##################### -# Flax => PyTorch # -##################### - - -# from https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_flax_pytorch_utils.py#L224-L352 -def load_flax_checkpoint_in_pytorch_model(pt_model, model_file): - try: - with open(model_file, "rb") as flax_state_f: - flax_state = from_bytes(None, flax_state_f.read()) - except UnpicklingError as e: - try: - with open(model_file) as f: - if f.read().startswith("version"): - raise OSError( - "You seem to have cloned a repository without having git-lfs installed. Please" - " install git-lfs and run `git lfs install` followed by `git lfs pull` in the" - " folder you cloned." - ) - else: - raise ValueError from e - except (UnicodeDecodeError, ValueError): - raise EnvironmentError(f"Unable to convert {model_file} to Flax deserializable object. ") - - return load_flax_weights_in_pytorch_model(pt_model, flax_state) - - -def load_flax_weights_in_pytorch_model(pt_model, flax_state): - """Load flax checkpoints in a PyTorch model""" - - try: - import torch # noqa: F401 - except ImportError: - logger.error( - "Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see" - " https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation" - " instructions." - ) - raise - - # check if we have bf16 weights - is_type_bf16 = flatten_dict(jax.tree_util.tree_map(lambda x: x.dtype == jnp.bfloat16, flax_state)).values() - if any(is_type_bf16): - # convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16 - - # and bf16 is not fully supported in PT yet. - logger.warning( - "Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` " - "before loading those in PyTorch model." - ) - flax_state = jax.tree_util.tree_map( - lambda params: params.astype(np.float32) if params.dtype == jnp.bfloat16 else params, flax_state - ) - - pt_model.base_model_prefix = "" - - flax_state_dict = flatten_dict(flax_state, sep=".") - pt_model_dict = pt_model.state_dict() - - # keep track of unexpected & missing keys - unexpected_keys = [] - missing_keys = set(pt_model_dict.keys()) - - for flax_key_tuple, flax_tensor in flax_state_dict.items(): - flax_key_tuple_array = flax_key_tuple.split(".") - - if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4: - flax_key_tuple_array = flax_key_tuple_array[:-1] + ["weight"] - flax_tensor = jnp.transpose(flax_tensor, (3, 2, 0, 1)) - elif flax_key_tuple_array[-1] == "kernel": - flax_key_tuple_array = flax_key_tuple_array[:-1] + ["weight"] - flax_tensor = flax_tensor.T - elif flax_key_tuple_array[-1] == "scale": - flax_key_tuple_array = flax_key_tuple_array[:-1] + ["weight"] - - if "time_embedding" not in flax_key_tuple_array: - for i, flax_key_tuple_string in enumerate(flax_key_tuple_array): - flax_key_tuple_array[i] = ( - flax_key_tuple_string.replace("_0", ".0") - .replace("_1", ".1") - .replace("_2", ".2") - .replace("_3", ".3") - .replace("_4", ".4") - .replace("_5", ".5") - .replace("_6", ".6") - .replace("_7", ".7") - .replace("_8", ".8") - .replace("_9", ".9") - ) - - flax_key = ".".join(flax_key_tuple_array) - - if flax_key in pt_model_dict: - if flax_tensor.shape != pt_model_dict[flax_key].shape: - raise ValueError( - f"Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected " - f"to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}." - ) - else: - # add weight to pytorch dict - flax_tensor = np.asarray(flax_tensor) if not isinstance(flax_tensor, np.ndarray) else flax_tensor - pt_model_dict[flax_key] = torch.from_numpy(flax_tensor) - # remove from missing keys - missing_keys.remove(flax_key) - else: - # weight is not expected by PyTorch model - unexpected_keys.append(flax_key) - - pt_model.load_state_dict(pt_model_dict) - - # re-transform missing_keys to list - missing_keys = list(missing_keys) - - if len(unexpected_keys) > 0: - logger.warning( - "Some weights of the Flax model were not used when initializing the PyTorch model" - f" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing" - f" {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture" - " (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This" - f" IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect" - " to be exactly identical (e.g. initializing a BertForSequenceClassification model from a" - " FlaxBertForSequenceClassification model)." - ) - if len(missing_keys) > 0: - logger.warning( - f"Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly" - f" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to" - " use it for predictions and inference." - ) - - return pt_model diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 650ab4c7fcda..bb8a5689f002 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -43,7 +43,6 @@ from ..utils import ( CONFIG_NAME, FLASHPACK_WEIGHTS_NAME, - FLAX_WEIGHTS_NAME, HF_ENABLE_PARALLEL_LOADING, SAFE_WEIGHTS_INDEX_NAME, SAFETENSORS_WEIGHTS_NAME, @@ -911,8 +910,6 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None revision (`str`, *optional*, defaults to `"main"`): The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier allowed by Git. - from_flax (`bool`, *optional*, defaults to `False`): - Load the model weights from a Flax checkpoint save file. subfolder (`str`, *optional*, defaults to `""`): The subfolder location of a model file within a larger model repository on the Hub or locally. mirror (`str`, *optional*): @@ -971,8 +968,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this argument to `True` will raise an error. variant (`str`, *optional*): - Load weights from a specified `variant` filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. + Load weights from a specified `variant` filename such as `"fp16"` or `"ema"`. use_safetensors (`bool`, *optional*, defaults to `None`): If set to `None`, the `safetensors` weights are downloaded if they're available **and** if the `safetensors` library is installed. If set to `True`, the model is forcibly loaded from `safetensors` @@ -1011,7 +1007,6 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None cache_dir = kwargs.pop("cache_dir", None) ignore_mismatched_sizes = kwargs.pop("ignore_mismatched_sizes", False) force_download = kwargs.pop("force_download", False) - from_flax = kwargs.pop("from_flax", False) proxies = kwargs.pop("proxies", None) output_loading_info = kwargs.pop("output_loading_info", False) local_files_only = kwargs.pop("local_files_only", None) @@ -1159,7 +1154,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None hf_quantizer = None 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) torch_dtype = hf_quantizer.update_torch_dtype(torch_dtype) device_map = hf_quantizer.update_device_map(device_map) @@ -1223,14 +1218,57 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None if index_file is not None and (dduf_entries or index_file.is_file()): is_sharded = True - if is_sharded and from_flax: - raise ValueError("Loading of sharded checkpoints is not supported when `from_flax=True`.") - # load model - if from_flax: + # in the case it is sharded, we have already the index + if is_sharded: + resolved_model_file, sharded_metadata = _get_checkpoint_shard_files( + pretrained_model_name_or_path, + index_file, + cache_dir=cache_dir, + proxies=proxies, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + subfolder=subfolder or "", + dduf_entries=dduf_entries, + ) + else: + if use_flashpack: + weights_name = FLASHPACK_WEIGHTS_NAME + elif use_safetensors: + weights_name = _add_variant(SAFETENSORS_WEIGHTS_NAME, variant) + else: + weights_name = None + if weights_name is not None: + try: + resolved_model_file = _get_model_file( + pretrained_model_name_or_path, + weights_name=weights_name, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + local_files_only=local_files_only, + token=token, + revision=revision, + subfolder=subfolder, + user_agent=user_agent, + commit_hash=commit_hash, + dduf_entries=dduf_entries, + ) + + except IOError as e: + logger.error(f"An error occurred while trying to fetch {pretrained_model_name_or_path}: {e}") + if not allow_pickle: + raise + logger.warning( + "Defaulting to unsafe serialization. Pass `allow_pickle=False` to raise an error instead." + ) + + if resolved_model_file is None and not is_sharded: resolved_model_file = _get_model_file( pretrained_model_name_or_path, - weights_name=FLAX_WEIGHTS_NAME, + weights_name=_add_variant(WEIGHTS_NAME, variant), cache_dir=cache_dir, force_download=force_download, proxies=proxies, @@ -1240,75 +1278,8 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None subfolder=subfolder, user_agent=user_agent, commit_hash=commit_hash, + dduf_entries=dduf_entries, ) - model = cls.from_config(config, **unused_kwargs) - - # Convert the weights - from .modeling_pytorch_flax_utils import load_flax_checkpoint_in_pytorch_model - - model = load_flax_checkpoint_in_pytorch_model(model, resolved_model_file) - else: - # in the case it is sharded, we have already the index - if is_sharded: - resolved_model_file, sharded_metadata = _get_checkpoint_shard_files( - pretrained_model_name_or_path, - index_file, - cache_dir=cache_dir, - proxies=proxies, - local_files_only=local_files_only, - token=token, - user_agent=user_agent, - revision=revision, - subfolder=subfolder or "", - dduf_entries=dduf_entries, - ) - else: - if use_flashpack: - weights_name = FLASHPACK_WEIGHTS_NAME - elif use_safetensors: - weights_name = _add_variant(SAFETENSORS_WEIGHTS_NAME, variant) - else: - weights_name = None - if weights_name is not None: - try: - resolved_model_file = _get_model_file( - pretrained_model_name_or_path, - weights_name=weights_name, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - commit_hash=commit_hash, - dduf_entries=dduf_entries, - ) - - except IOError as e: - logger.error(f"An error occurred while trying to fetch {pretrained_model_name_or_path}: {e}") - if not allow_pickle: - raise - logger.warning( - "Defaulting to unsafe serialization. Pass `allow_pickle=False` to raise an error instead." - ) - - if resolved_model_file is None and not is_sharded: - resolved_model_file = _get_model_file( - pretrained_model_name_or_path, - weights_name=_add_variant(WEIGHTS_NAME, variant), - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - commit_hash=commit_hash, - dduf_entries=dduf_entries, - ) if not isinstance(resolved_model_file, list): resolved_model_file = [resolved_model_file] diff --git a/src/diffusers/models/resnet_flax.py b/src/diffusers/models/resnet_flax.py deleted file mode 100644 index 23bf8ee23c02..000000000000 --- a/src/diffusers/models/resnet_flax.py +++ /dev/null @@ -1,144 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import flax.linen as nn -import jax -import jax.numpy as jnp - -from ..utils import logging - - -logger = logging.get_logger(__name__) - - -class FlaxUpsample2D(nn.Module): - out_channels: int - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - self.conv = nn.Conv( - self.out_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - def __call__(self, hidden_states): - batch, height, width, channels = hidden_states.shape - hidden_states = jax.image.resize( - hidden_states, - shape=(batch, height * 2, width * 2, channels), - method="nearest", - ) - hidden_states = self.conv(hidden_states) - return hidden_states - - -class FlaxDownsample2D(nn.Module): - out_channels: int - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - self.conv = nn.Conv( - self.out_channels, - kernel_size=(3, 3), - strides=(2, 2), - padding=((1, 1), (1, 1)), # padding="VALID", - dtype=self.dtype, - ) - - def __call__(self, hidden_states): - # pad = ((0, 0), (0, 1), (0, 1), (0, 0)) # pad height and width dim - # hidden_states = jnp.pad(hidden_states, pad_width=pad) - hidden_states = self.conv(hidden_states) - return hidden_states - - -class FlaxResnetBlock2D(nn.Module): - in_channels: int - out_channels: int = None - dropout_prob: float = 0.0 - use_nin_shortcut: bool = None - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - out_channels = self.in_channels if self.out_channels is None else self.out_channels - - self.norm1 = nn.GroupNorm(num_groups=32, epsilon=1e-5) - self.conv1 = nn.Conv( - out_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - self.time_emb_proj = nn.Dense(out_channels, dtype=self.dtype) - - self.norm2 = nn.GroupNorm(num_groups=32, epsilon=1e-5) - self.dropout = nn.Dropout(self.dropout_prob) - self.conv2 = nn.Conv( - out_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - use_nin_shortcut = self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut - - self.conv_shortcut = None - if use_nin_shortcut: - self.conv_shortcut = nn.Conv( - out_channels, - kernel_size=(1, 1), - strides=(1, 1), - padding="VALID", - dtype=self.dtype, - ) - - def __call__(self, hidden_states, temb, deterministic=True): - residual = hidden_states - hidden_states = self.norm1(hidden_states) - hidden_states = nn.swish(hidden_states) - hidden_states = self.conv1(hidden_states) - - temb = self.time_emb_proj(nn.swish(temb)) - temb = jnp.expand_dims(jnp.expand_dims(temb, 1), 1) - hidden_states = hidden_states + temb - - hidden_states = self.norm2(hidden_states) - hidden_states = nn.swish(hidden_states) - hidden_states = self.dropout(hidden_states, deterministic) - hidden_states = self.conv2(hidden_states) - - if self.conv_shortcut is not None: - residual = self.conv_shortcut(residual) - - return hidden_states + residual diff --git a/src/diffusers/models/unets/__init__.py b/src/diffusers/models/unets/__init__.py index 394df72261c6..d3b69d6d5e8c 100644 --- a/src/diffusers/models/unets/__init__.py +++ b/src/diffusers/models/unets/__init__.py @@ -1,4 +1,4 @@ -from ...utils import is_flax_available, is_torch_available +from ...utils import is_torch_available if is_torch_available(): @@ -13,7 +13,3 @@ from .unet_spatio_temporal_condition import UNetSpatioTemporalConditionModel from .unet_stable_cascade import StableCascadeUNet from .uvit_2d import UVit2DModel - - -if is_flax_available(): - from .unet_2d_condition_flax import FlaxUNet2DConditionModel diff --git a/src/diffusers/models/unets/unet_2d_blocks_flax.py b/src/diffusers/models/unets/unet_2d_blocks_flax.py deleted file mode 100644 index 959d03e38036..000000000000 --- a/src/diffusers/models/unets/unet_2d_blocks_flax.py +++ /dev/null @@ -1,430 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import flax.linen as nn -import jax.numpy as jnp - -from ...utils import logging -from ..attention_flax import FlaxTransformer2DModel -from ..resnet_flax import FlaxDownsample2D, FlaxResnetBlock2D, FlaxUpsample2D - - -logger = logging.get_logger(__name__) - - -class FlaxCrossAttnDownBlock2D(nn.Module): - r""" - Cross Attention 2D Downsizing block - original architecture from Unet transformers: - https://huggingface.co/papers/2103.06104 - - Parameters: - in_channels (:obj:`int`): - Input channels - out_channels (:obj:`int`): - Output channels - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - num_layers (:obj:`int`, *optional*, defaults to 1): - Number of attention blocks layers - num_attention_heads (:obj:`int`, *optional*, defaults to 1): - Number of attention heads of each spatial transformer block - add_downsample (:obj:`bool`, *optional*, defaults to `True`): - Whether to add downsampling layer before each final output - use_memory_efficient_attention (`bool`, *optional*, defaults to `False`): - enable memory efficient attention https://huggingface.co/papers/2112.05682 - split_head_dim (`bool`, *optional*, defaults to `False`): - Whether to split the head dimension into a new axis for the self-attention computation. In most cases, - enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL. - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - out_channels: int - dropout: float = 0.0 - num_layers: int = 1 - num_attention_heads: int = 1 - add_downsample: bool = True - use_linear_projection: bool = False - only_cross_attention: bool = False - use_memory_efficient_attention: bool = False - split_head_dim: bool = False - dtype: jnp.dtype = jnp.float32 - transformer_layers_per_block: int = 1 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - resnets = [] - attentions = [] - - for i in range(self.num_layers): - in_channels = self.in_channels if i == 0 else self.out_channels - - res_block = FlaxResnetBlock2D( - in_channels=in_channels, - out_channels=self.out_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - resnets.append(res_block) - - attn_block = FlaxTransformer2DModel( - in_channels=self.out_channels, - n_heads=self.num_attention_heads, - d_head=self.out_channels // self.num_attention_heads, - depth=self.transformer_layers_per_block, - use_linear_projection=self.use_linear_projection, - only_cross_attention=self.only_cross_attention, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - attentions.append(attn_block) - - self.resnets = resnets - self.attentions = attentions - - if self.add_downsample: - self.downsamplers_0 = FlaxDownsample2D(self.out_channels, dtype=self.dtype) - - def __call__(self, hidden_states, temb, encoder_hidden_states, deterministic=True): - output_states = () - - for resnet, attn in zip(self.resnets, self.attentions): - hidden_states = resnet(hidden_states, temb, deterministic=deterministic) - hidden_states = attn(hidden_states, encoder_hidden_states, deterministic=deterministic) - output_states += (hidden_states,) - - if self.add_downsample: - hidden_states = self.downsamplers_0(hidden_states) - output_states += (hidden_states,) - - return hidden_states, output_states - - -class FlaxDownBlock2D(nn.Module): - r""" - Flax 2D downsizing block - - Parameters: - in_channels (:obj:`int`): - Input channels - out_channels (:obj:`int`): - Output channels - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - num_layers (:obj:`int`, *optional*, defaults to 1): - Number of attention blocks layers - add_downsample (:obj:`bool`, *optional*, defaults to `True`): - Whether to add downsampling layer before each final output - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - out_channels: int - dropout: float = 0.0 - num_layers: int = 1 - add_downsample: bool = True - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - resnets = [] - - for i in range(self.num_layers): - in_channels = self.in_channels if i == 0 else self.out_channels - - res_block = FlaxResnetBlock2D( - in_channels=in_channels, - out_channels=self.out_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - resnets.append(res_block) - self.resnets = resnets - - if self.add_downsample: - self.downsamplers_0 = FlaxDownsample2D(self.out_channels, dtype=self.dtype) - - def __call__(self, hidden_states, temb, deterministic=True): - output_states = () - - for resnet in self.resnets: - hidden_states = resnet(hidden_states, temb, deterministic=deterministic) - output_states += (hidden_states,) - - if self.add_downsample: - hidden_states = self.downsamplers_0(hidden_states) - output_states += (hidden_states,) - - return hidden_states, output_states - - -class FlaxCrossAttnUpBlock2D(nn.Module): - r""" - Cross Attention 2D Upsampling block - original architecture from Unet transformers: - https://huggingface.co/papers/2103.06104 - - Parameters: - in_channels (:obj:`int`): - Input channels - out_channels (:obj:`int`): - Output channels - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - num_layers (:obj:`int`, *optional*, defaults to 1): - Number of attention blocks layers - num_attention_heads (:obj:`int`, *optional*, defaults to 1): - Number of attention heads of each spatial transformer block - add_upsample (:obj:`bool`, *optional*, defaults to `True`): - Whether to add upsampling layer before each final output - use_memory_efficient_attention (`bool`, *optional*, defaults to `False`): - enable memory efficient attention https://huggingface.co/papers/2112.05682 - split_head_dim (`bool`, *optional*, defaults to `False`): - Whether to split the head dimension into a new axis for the self-attention computation. In most cases, - enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL. - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - out_channels: int - prev_output_channel: int - dropout: float = 0.0 - num_layers: int = 1 - num_attention_heads: int = 1 - add_upsample: bool = True - use_linear_projection: bool = False - only_cross_attention: bool = False - use_memory_efficient_attention: bool = False - split_head_dim: bool = False - dtype: jnp.dtype = jnp.float32 - transformer_layers_per_block: int = 1 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - resnets = [] - attentions = [] - - for i in range(self.num_layers): - res_skip_channels = self.in_channels if (i == self.num_layers - 1) else self.out_channels - resnet_in_channels = self.prev_output_channel if i == 0 else self.out_channels - - res_block = FlaxResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=self.out_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - resnets.append(res_block) - - attn_block = FlaxTransformer2DModel( - in_channels=self.out_channels, - n_heads=self.num_attention_heads, - d_head=self.out_channels // self.num_attention_heads, - depth=self.transformer_layers_per_block, - use_linear_projection=self.use_linear_projection, - only_cross_attention=self.only_cross_attention, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - attentions.append(attn_block) - - self.resnets = resnets - self.attentions = attentions - - if self.add_upsample: - self.upsamplers_0 = FlaxUpsample2D(self.out_channels, dtype=self.dtype) - - def __call__(self, hidden_states, res_hidden_states_tuple, temb, encoder_hidden_states, deterministic=True): - for resnet, attn in zip(self.resnets, self.attentions): - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - hidden_states = jnp.concatenate((hidden_states, res_hidden_states), axis=-1) - - hidden_states = resnet(hidden_states, temb, deterministic=deterministic) - hidden_states = attn(hidden_states, encoder_hidden_states, deterministic=deterministic) - - if self.add_upsample: - hidden_states = self.upsamplers_0(hidden_states) - - return hidden_states - - -class FlaxUpBlock2D(nn.Module): - r""" - Flax 2D upsampling block - - Parameters: - in_channels (:obj:`int`): - Input channels - out_channels (:obj:`int`): - Output channels - prev_output_channel (:obj:`int`): - Output channels from the previous block - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - num_layers (:obj:`int`, *optional*, defaults to 1): - Number of attention blocks layers - add_downsample (:obj:`bool`, *optional*, defaults to `True`): - Whether to add downsampling layer before each final output - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - out_channels: int - prev_output_channel: int - dropout: float = 0.0 - num_layers: int = 1 - add_upsample: bool = True - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - resnets = [] - - for i in range(self.num_layers): - res_skip_channels = self.in_channels if (i == self.num_layers - 1) else self.out_channels - resnet_in_channels = self.prev_output_channel if i == 0 else self.out_channels - - res_block = FlaxResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=self.out_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - resnets.append(res_block) - - self.resnets = resnets - - if self.add_upsample: - self.upsamplers_0 = FlaxUpsample2D(self.out_channels, dtype=self.dtype) - - def __call__(self, hidden_states, res_hidden_states_tuple, temb, deterministic=True): - for resnet in self.resnets: - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - hidden_states = jnp.concatenate((hidden_states, res_hidden_states), axis=-1) - - hidden_states = resnet(hidden_states, temb, deterministic=deterministic) - - if self.add_upsample: - hidden_states = self.upsamplers_0(hidden_states) - - return hidden_states - - -class FlaxUNetMidBlock2DCrossAttn(nn.Module): - r""" - Cross Attention 2D Mid-level block - original architecture from Unet transformers: - https://huggingface.co/papers/2103.06104 - - Parameters: - in_channels (:obj:`int`): - Input channels - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - num_layers (:obj:`int`, *optional*, defaults to 1): - Number of attention blocks layers - num_attention_heads (:obj:`int`, *optional*, defaults to 1): - Number of attention heads of each spatial transformer block - use_memory_efficient_attention (`bool`, *optional*, defaults to `False`): - enable memory efficient attention https://huggingface.co/papers/2112.05682 - split_head_dim (`bool`, *optional*, defaults to `False`): - Whether to split the head dimension into a new axis for the self-attention computation. In most cases, - enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL. - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - dropout: float = 0.0 - num_layers: int = 1 - num_attention_heads: int = 1 - use_linear_projection: bool = False - use_memory_efficient_attention: bool = False - split_head_dim: bool = False - dtype: jnp.dtype = jnp.float32 - transformer_layers_per_block: int = 1 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - # there is always at least one resnet - resnets = [ - FlaxResnetBlock2D( - in_channels=self.in_channels, - out_channels=self.in_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - ] - - attentions = [] - - for _ in range(self.num_layers): - attn_block = FlaxTransformer2DModel( - in_channels=self.in_channels, - n_heads=self.num_attention_heads, - d_head=self.in_channels // self.num_attention_heads, - depth=self.transformer_layers_per_block, - use_linear_projection=self.use_linear_projection, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - attentions.append(attn_block) - - res_block = FlaxResnetBlock2D( - in_channels=self.in_channels, - out_channels=self.in_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - resnets.append(res_block) - - self.resnets = resnets - self.attentions = attentions - - def __call__(self, hidden_states, temb, encoder_hidden_states, deterministic=True): - hidden_states = self.resnets[0](hidden_states, temb) - for attn, resnet in zip(self.attentions, self.resnets[1:]): - hidden_states = attn(hidden_states, encoder_hidden_states, deterministic=deterministic) - hidden_states = resnet(hidden_states, temb, deterministic=deterministic) - - return hidden_states diff --git a/src/diffusers/models/unets/unet_2d_condition_flax.py b/src/diffusers/models/unets/unet_2d_condition_flax.py deleted file mode 100644 index 0ff2f86a5bda..000000000000 --- a/src/diffusers/models/unets/unet_2d_condition_flax.py +++ /dev/null @@ -1,461 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import flax -import flax.linen as nn -import jax -import jax.numpy as jnp -from flax.core.frozen_dict import FrozenDict - -from ...configuration_utils import ConfigMixin, flax_register_to_config -from ...utils import BaseOutput, logging -from ..embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps -from ..modeling_flax_utils import FlaxModelMixin -from .unet_2d_blocks_flax import ( - FlaxCrossAttnDownBlock2D, - FlaxCrossAttnUpBlock2D, - FlaxDownBlock2D, - FlaxUNetMidBlock2DCrossAttn, - FlaxUpBlock2D, -) - - -logger = logging.get_logger(__name__) - - -@flax.struct.dataclass -class FlaxUNet2DConditionOutput(BaseOutput): - """ - The output of [`FlaxUNet2DConditionModel`]. - - Args: - sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)`): - The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model. - """ - - sample: jnp.ndarray - - -@flax_register_to_config -class FlaxUNet2DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin): - r""" - A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample - shaped output. - - This model inherits from [`FlaxModelMixin`]. Check the superclass documentation for it's generic methods - implemented for all models (such as downloading or saving). - - This model is also a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module) - subclass. Use it as a regular Flax Linen module and refer to the Flax documentation for all matters related to its - general usage and behavior. - - Inherent JAX features such as the following are supported: - - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit) - - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation) - - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap) - - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap) - - Parameters: - sample_size (`int`, *optional*): - The size of the input sample. - in_channels (`int`, *optional*, defaults to 4): - The number of channels in the input sample. - out_channels (`int`, *optional*, defaults to 4): - The number of channels in the output. - down_block_types (`tuple[str]`, *optional*, defaults to `("FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxDownBlock2D")`): - The tuple of downsample blocks to use. - up_block_types (`tuple[str]`, *optional*, defaults to `("FlaxUpBlock2D", "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D")`): - The tuple of upsample blocks to use. - mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`): - Block type for middle of UNet, it can be one of `UNetMidBlock2DCrossAttn`. If `None`, the mid block layer - is skipped. - block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): - The tuple of output channels for each block. - layers_per_block (`int`, *optional*, defaults to 2): - The number of layers per block. - attention_head_dim (`int` or `tuple[int]`, *optional*, defaults to 8): - The dimension of the attention heads. - num_attention_heads (`int` or `tuple[int]`, *optional*): - The number of attention heads. - cross_attention_dim (`int`, *optional*, defaults to 768): - The dimension of the cross attention features. - dropout (`float`, *optional*, defaults to 0): - Dropout probability for down, up and bottleneck blocks. - flip_sin_to_cos (`bool`, *optional*, defaults to `True`): - Whether to flip the sin to cos in the time embedding. - freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding. - use_memory_efficient_attention (`bool`, *optional*, defaults to `False`): - Enable memory efficient attention as described [here](https://huggingface.co/papers/2112.05682). - split_head_dim (`bool`, *optional*, defaults to `False`): - Whether to split the head dimension into a new axis for the self-attention computation. In most cases, - enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL. - """ - - sample_size: int = 32 - in_channels: int = 4 - out_channels: int = 4 - down_block_types: tuple[str, ...] = ( - "CrossAttnDownBlock2D", - "CrossAttnDownBlock2D", - "CrossAttnDownBlock2D", - "DownBlock2D", - ) - up_block_types: tuple[str, ...] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D") - mid_block_type: str | None = "UNetMidBlock2DCrossAttn" - only_cross_attention: bool | tuple[bool] = False - block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280) - layers_per_block: int = 2 - attention_head_dim: int | tuple[int, ...] = 8 - num_attention_heads: int | tuple[int, ...] | None = None - cross_attention_dim: int = 1280 - dropout: float = 0.0 - use_linear_projection: bool = False - dtype: jnp.dtype = jnp.float32 - flip_sin_to_cos: bool = True - freq_shift: int = 0 - use_memory_efficient_attention: bool = False - split_head_dim: bool = False - transformer_layers_per_block: int | tuple[int, ...] = 1 - addition_embed_type: str | None = None - addition_time_embed_dim: int | None = None - addition_embed_type_num_heads: int = 64 - projection_class_embeddings_input_dim: int | None = None - - def init_weights(self, rng: jax.Array) -> FrozenDict: - # init input tensors - sample_shape = (1, self.in_channels, self.sample_size, self.sample_size) - sample = jnp.zeros(sample_shape, dtype=jnp.float32) - timesteps = jnp.ones((1,), dtype=jnp.int32) - encoder_hidden_states = jnp.zeros((1, 1, self.cross_attention_dim), dtype=jnp.float32) - - params_rng, dropout_rng = jax.random.split(rng) - rngs = {"params": params_rng, "dropout": dropout_rng} - - added_cond_kwargs = None - if self.addition_embed_type == "text_time": - # we retrieve the expected `text_embeds_dim` by first checking if the architecture is a refiner - # or non-refiner architecture and then by "reverse-computing" from `projection_class_embeddings_input_dim` - is_refiner = ( - 5 * self.config.addition_time_embed_dim + self.config.cross_attention_dim - == self.config.projection_class_embeddings_input_dim - ) - num_micro_conditions = 5 if is_refiner else 6 - - text_embeds_dim = self.config.projection_class_embeddings_input_dim - ( - num_micro_conditions * self.config.addition_time_embed_dim - ) - - time_ids_channels = self.projection_class_embeddings_input_dim - text_embeds_dim - time_ids_dims = time_ids_channels // self.addition_time_embed_dim - added_cond_kwargs = { - "text_embeds": jnp.zeros((1, text_embeds_dim), dtype=jnp.float32), - "time_ids": jnp.zeros((1, time_ids_dims), dtype=jnp.float32), - } - return self.init(rngs, sample, timesteps, encoder_hidden_states, added_cond_kwargs)["params"] - - def setup(self) -> None: - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - block_out_channels = self.block_out_channels - time_embed_dim = block_out_channels[0] * 4 - - if self.num_attention_heads is not None: - raise ValueError( - "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19." - ) - - # If `num_attention_heads` is not defined (which is the case for most models) - # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. - # The reason for this behavior is to correct for incorrectly named variables that were introduced - # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 - # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking - # which is why we correct for the naming here. - num_attention_heads = self.num_attention_heads or self.attention_head_dim - - # input - self.conv_in = nn.Conv( - block_out_channels[0], - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - # time - self.time_proj = FlaxTimesteps( - block_out_channels[0], flip_sin_to_cos=self.flip_sin_to_cos, freq_shift=self.config.freq_shift - ) - self.time_embedding = FlaxTimestepEmbedding(time_embed_dim, dtype=self.dtype) - - only_cross_attention = self.only_cross_attention - if isinstance(only_cross_attention, bool): - only_cross_attention = (only_cross_attention,) * len(self.down_block_types) - - if isinstance(num_attention_heads, int): - num_attention_heads = (num_attention_heads,) * len(self.down_block_types) - - # transformer layers per block - transformer_layers_per_block = self.transformer_layers_per_block - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * len(self.down_block_types) - - # addition embed types - if self.addition_embed_type is None: - self.add_embedding = None - elif self.addition_embed_type == "text_time": - if self.addition_time_embed_dim is None: - raise ValueError( - f"addition_embed_type {self.addition_embed_type} requires `addition_time_embed_dim` to not be None" - ) - self.add_time_proj = FlaxTimesteps(self.addition_time_embed_dim, self.flip_sin_to_cos, self.freq_shift) - self.add_embedding = FlaxTimestepEmbedding(time_embed_dim, dtype=self.dtype) - else: - raise ValueError(f"addition_embed_type: {self.addition_embed_type} must be None or `text_time`.") - - # down - down_blocks = [] - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(self.down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - if down_block_type == "CrossAttnDownBlock2D": - down_block = FlaxCrossAttnDownBlock2D( - in_channels=input_channel, - out_channels=output_channel, - dropout=self.dropout, - num_layers=self.layers_per_block, - transformer_layers_per_block=transformer_layers_per_block[i], - num_attention_heads=num_attention_heads[i], - add_downsample=not is_final_block, - use_linear_projection=self.use_linear_projection, - only_cross_attention=only_cross_attention[i], - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - else: - down_block = FlaxDownBlock2D( - in_channels=input_channel, - out_channels=output_channel, - dropout=self.dropout, - num_layers=self.layers_per_block, - add_downsample=not is_final_block, - dtype=self.dtype, - ) - - down_blocks.append(down_block) - self.down_blocks = down_blocks - - # mid - if self.config.mid_block_type == "UNetMidBlock2DCrossAttn": - self.mid_block = FlaxUNetMidBlock2DCrossAttn( - in_channels=block_out_channels[-1], - dropout=self.dropout, - num_attention_heads=num_attention_heads[-1], - transformer_layers_per_block=transformer_layers_per_block[-1], - use_linear_projection=self.use_linear_projection, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - elif self.config.mid_block_type is None: - self.mid_block = None - else: - raise ValueError(f"Unexpected mid_block_type {self.config.mid_block_type}") - - # up - up_blocks = [] - reversed_block_out_channels = list(reversed(block_out_channels)) - reversed_num_attention_heads = list(reversed(num_attention_heads)) - only_cross_attention = list(reversed(only_cross_attention)) - output_channel = reversed_block_out_channels[0] - reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block)) - for i, up_block_type in enumerate(self.up_block_types): - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] - - is_final_block = i == len(block_out_channels) - 1 - - if up_block_type == "CrossAttnUpBlock2D": - up_block = FlaxCrossAttnUpBlock2D( - in_channels=input_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - num_layers=self.layers_per_block + 1, - transformer_layers_per_block=reversed_transformer_layers_per_block[i], - num_attention_heads=reversed_num_attention_heads[i], - add_upsample=not is_final_block, - dropout=self.dropout, - use_linear_projection=self.use_linear_projection, - only_cross_attention=only_cross_attention[i], - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - else: - up_block = FlaxUpBlock2D( - in_channels=input_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - num_layers=self.layers_per_block + 1, - add_upsample=not is_final_block, - dropout=self.dropout, - dtype=self.dtype, - ) - - up_blocks.append(up_block) - prev_output_channel = output_channel - self.up_blocks = up_blocks - - # out - self.conv_norm_out = nn.GroupNorm(num_groups=32, epsilon=1e-5) - self.conv_out = nn.Conv( - self.out_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - def __call__( - self, - sample: jnp.ndarray, - timesteps: jnp.ndarray | float | int, - encoder_hidden_states: jnp.ndarray, - added_cond_kwargs: dict | FrozenDict | None = None, - down_block_additional_residuals: tuple[jnp.ndarray, ...] | None = None, - mid_block_additional_residual: jnp.ndarray | None = None, - return_dict: bool = True, - train: bool = False, - ) -> FlaxUNet2DConditionOutput | tuple[jnp.ndarray]: - r""" - Args: - sample (`jnp.ndarray`): (batch, channel, height, width) noisy inputs tensor - timestep (`jnp.ndarray` or `float` or `int`): timesteps - encoder_hidden_states (`jnp.ndarray`): (batch_size, sequence_length, hidden_size) encoder hidden states - added_cond_kwargs: (`dict`, *optional*): - A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that - are passed along to the UNet blocks. - down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*): - A tuple of tensors that if specified are added to the residuals of down unet blocks. - mid_block_additional_residual: (`torch.Tensor`, *optional*): - A tensor that if specified is added to the residual of the middle unet block. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`models.unets.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] instead of - a plain tuple. - train (`bool`, *optional*, defaults to `False`): - Use deterministic functions and disable dropout when not training. - - Returns: - [`~models.unets.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] or `tuple`: - [`~models.unets.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is the sample tensor. - """ - # 1. time - if not isinstance(timesteps, jnp.ndarray): - timesteps = jnp.array([timesteps], dtype=jnp.int32) - elif isinstance(timesteps, jnp.ndarray) and len(timesteps.shape) == 0: - timesteps = timesteps.astype(dtype=jnp.float32) - timesteps = jnp.expand_dims(timesteps, 0) - - t_emb = self.time_proj(timesteps) - t_emb = self.time_embedding(t_emb) - - # additional embeddings - aug_emb = None - if self.addition_embed_type == "text_time": - if added_cond_kwargs is None: - raise ValueError( - f"Need to provide argument `added_cond_kwargs` for {self.__class__} when using `addition_embed_type={self.addition_embed_type}`" - ) - text_embeds = added_cond_kwargs.get("text_embeds") - if text_embeds is None: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`" - ) - time_ids = added_cond_kwargs.get("time_ids") - if time_ids is None: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`" - ) - # compute time embeds - time_embeds = self.add_time_proj(jnp.ravel(time_ids)) # (1, 6) => (6,) => (6, 256) - time_embeds = jnp.reshape(time_embeds, (text_embeds.shape[0], -1)) - add_embeds = jnp.concatenate([text_embeds, time_embeds], axis=-1) - aug_emb = self.add_embedding(add_embeds) - - t_emb = t_emb + aug_emb if aug_emb is not None else t_emb - - # 2. pre-process - sample = jnp.transpose(sample, (0, 2, 3, 1)) - sample = self.conv_in(sample) - - # 3. down - down_block_res_samples = (sample,) - for down_block in self.down_blocks: - if isinstance(down_block, FlaxCrossAttnDownBlock2D): - sample, res_samples = down_block(sample, t_emb, encoder_hidden_states, deterministic=not train) - else: - sample, res_samples = down_block(sample, t_emb, deterministic=not train) - down_block_res_samples += res_samples - - if down_block_additional_residuals is not None: - new_down_block_res_samples = () - - for down_block_res_sample, down_block_additional_residual in zip( - down_block_res_samples, down_block_additional_residuals - ): - down_block_res_sample += down_block_additional_residual - new_down_block_res_samples += (down_block_res_sample,) - - down_block_res_samples = new_down_block_res_samples - - # 4. mid - if self.mid_block is not None: - sample = self.mid_block(sample, t_emb, encoder_hidden_states, deterministic=not train) - - if mid_block_additional_residual is not None: - sample += mid_block_additional_residual - - # 5. up - for up_block in self.up_blocks: - res_samples = down_block_res_samples[-(self.layers_per_block + 1) :] - down_block_res_samples = down_block_res_samples[: -(self.layers_per_block + 1)] - if isinstance(up_block, FlaxCrossAttnUpBlock2D): - sample = up_block( - sample, - temb=t_emb, - encoder_hidden_states=encoder_hidden_states, - res_hidden_states_tuple=res_samples, - deterministic=not train, - ) - else: - sample = up_block(sample, temb=t_emb, res_hidden_states_tuple=res_samples, deterministic=not train) - - # 6. post-process - sample = self.conv_norm_out(sample) - sample = nn.silu(sample) - sample = self.conv_out(sample) - sample = jnp.transpose(sample, (0, 3, 1, 2)) - - if not return_dict: - return (sample,) - - return FlaxUNet2DConditionOutput(sample=sample) diff --git a/src/diffusers/models/vae_flax.py b/src/diffusers/models/vae_flax.py deleted file mode 100644 index 8b4e16c16228..000000000000 --- a/src/diffusers/models/vae_flax.py +++ /dev/null @@ -1,927 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# JAX implementation of VQGAN from taming-transformers https://github.com/CompVis/taming-transformers - -import math -from functools import partial - -import flax -import flax.linen as nn -import jax -import jax.numpy as jnp -from flax.core.frozen_dict import FrozenDict - -from ..configuration_utils import ConfigMixin, flax_register_to_config -from ..utils import BaseOutput, logging -from .modeling_flax_utils import FlaxModelMixin - - -logger = logging.get_logger(__name__) - - -@flax.struct.dataclass -class FlaxDecoderOutput(BaseOutput): - """ - Output of decoding method. - - Args: - sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)`): - The decoded output sample from the last layer of the model. - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - The `dtype` of the parameters. - """ - - sample: jnp.ndarray - - -@flax.struct.dataclass -class FlaxAutoencoderKLOutput(BaseOutput): - """ - Output of AutoencoderKL encoding method. - - Args: - latent_dist (`FlaxDiagonalGaussianDistribution`): - Encoded outputs of `Encoder` represented as the mean and logvar of `FlaxDiagonalGaussianDistribution`. - `FlaxDiagonalGaussianDistribution` allows for sampling latents from the distribution. - """ - - latent_dist: "FlaxDiagonalGaussianDistribution" - - -class FlaxUpsample2D(nn.Module): - """ - Flax implementation of 2D Upsample layer - - Args: - in_channels (`int`): - Input channels - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - self.conv = nn.Conv( - self.in_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - def __call__(self, hidden_states): - batch, height, width, channels = hidden_states.shape - hidden_states = jax.image.resize( - hidden_states, - shape=(batch, height * 2, width * 2, channels), - method="nearest", - ) - hidden_states = self.conv(hidden_states) - return hidden_states - - -class FlaxDownsample2D(nn.Module): - """ - Flax implementation of 2D Downsample layer - - Args: - in_channels (`int`): - Input channels - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - self.conv = nn.Conv( - self.in_channels, - kernel_size=(3, 3), - strides=(2, 2), - padding="VALID", - dtype=self.dtype, - ) - - def __call__(self, hidden_states): - pad = ((0, 0), (0, 1), (0, 1), (0, 0)) # pad height and width dim - hidden_states = jnp.pad(hidden_states, pad_width=pad) - hidden_states = self.conv(hidden_states) - return hidden_states - - -class FlaxResnetBlock2D(nn.Module): - """ - Flax implementation of 2D Resnet Block. - - Args: - in_channels (`int`): - Input channels - out_channels (`int`): - Output channels - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - groups (:obj:`int`, *optional*, defaults to `32`): - The number of groups to use for group norm. - use_nin_shortcut (:obj:`bool`, *optional*, defaults to `None`): - Whether to use `nin_shortcut`. This activates a new layer inside ResNet block - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - out_channels: int = None - dropout: float = 0.0 - groups: int = 32 - use_nin_shortcut: bool = None - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - out_channels = self.in_channels if self.out_channels is None else self.out_channels - - self.norm1 = nn.GroupNorm(num_groups=self.groups, epsilon=1e-6) - self.conv1 = nn.Conv( - out_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - self.norm2 = nn.GroupNorm(num_groups=self.groups, epsilon=1e-6) - self.dropout_layer = nn.Dropout(self.dropout) - self.conv2 = nn.Conv( - out_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - use_nin_shortcut = self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut - - self.conv_shortcut = None - if use_nin_shortcut: - self.conv_shortcut = nn.Conv( - out_channels, - kernel_size=(1, 1), - strides=(1, 1), - padding="VALID", - dtype=self.dtype, - ) - - def __call__(self, hidden_states, deterministic=True): - residual = hidden_states - hidden_states = self.norm1(hidden_states) - hidden_states = nn.swish(hidden_states) - hidden_states = self.conv1(hidden_states) - - hidden_states = self.norm2(hidden_states) - hidden_states = nn.swish(hidden_states) - hidden_states = self.dropout_layer(hidden_states, deterministic) - hidden_states = self.conv2(hidden_states) - - if self.conv_shortcut is not None: - residual = self.conv_shortcut(residual) - - return hidden_states + residual - - -class FlaxAttentionBlock(nn.Module): - r""" - Flax Convolutional based multi-head attention block for diffusion-based VAE. - - Parameters: - channels (:obj:`int`): - Input channels - num_head_channels (:obj:`int`, *optional*, defaults to `None`): - Number of attention heads - num_groups (:obj:`int`, *optional*, defaults to `32`): - The number of groups to use for group norm - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - - """ - - channels: int - num_head_channels: int = None - num_groups: int = 32 - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - self.num_heads = self.channels // self.num_head_channels if self.num_head_channels is not None else 1 - - dense = partial(nn.Dense, self.channels, dtype=self.dtype) - - self.group_norm = nn.GroupNorm(num_groups=self.num_groups, epsilon=1e-6) - self.query, self.key, self.value = dense(), dense(), dense() - self.proj_attn = dense() - - def transpose_for_scores(self, projection): - new_projection_shape = projection.shape[:-1] + (self.num_heads, -1) - # move heads to 2nd position (B, T, H * D) -> (B, T, H, D) - new_projection = projection.reshape(new_projection_shape) - # (B, T, H, D) -> (B, H, T, D) - new_projection = jnp.transpose(new_projection, (0, 2, 1, 3)) - return new_projection - - def __call__(self, hidden_states): - residual = hidden_states - batch, height, width, channels = hidden_states.shape - - hidden_states = self.group_norm(hidden_states) - - hidden_states = hidden_states.reshape((batch, height * width, channels)) - - query = self.query(hidden_states) - key = self.key(hidden_states) - value = self.value(hidden_states) - - # transpose - query = self.transpose_for_scores(query) - key = self.transpose_for_scores(key) - value = self.transpose_for_scores(value) - - # compute attentions - scale = 1 / math.sqrt(math.sqrt(self.channels / self.num_heads)) - attn_weights = jnp.einsum("...qc,...kc->...qk", query * scale, key * scale) - attn_weights = nn.softmax(attn_weights, axis=-1) - - # attend to values - hidden_states = jnp.einsum("...kc,...qk->...qc", value, attn_weights) - - hidden_states = jnp.transpose(hidden_states, (0, 2, 1, 3)) - new_hidden_states_shape = hidden_states.shape[:-2] + (self.channels,) - hidden_states = hidden_states.reshape(new_hidden_states_shape) - - hidden_states = self.proj_attn(hidden_states) - hidden_states = hidden_states.reshape((batch, height, width, channels)) - hidden_states = hidden_states + residual - return hidden_states - - -class FlaxDownEncoderBlock2D(nn.Module): - r""" - Flax Resnet blocks-based Encoder block for diffusion-based VAE. - - Parameters: - in_channels (:obj:`int`): - Input channels - out_channels (:obj:`int`): - Output channels - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - num_layers (:obj:`int`, *optional*, defaults to 1): - Number of Resnet layer block - resnet_groups (:obj:`int`, *optional*, defaults to `32`): - The number of groups to use for the Resnet block group norm - add_downsample (:obj:`bool`, *optional*, defaults to `True`): - Whether to add downsample layer - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - out_channels: int - dropout: float = 0.0 - num_layers: int = 1 - resnet_groups: int = 32 - add_downsample: bool = True - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - resnets = [] - for i in range(self.num_layers): - in_channels = self.in_channels if i == 0 else self.out_channels - - res_block = FlaxResnetBlock2D( - in_channels=in_channels, - out_channels=self.out_channels, - dropout=self.dropout, - groups=self.resnet_groups, - dtype=self.dtype, - ) - resnets.append(res_block) - self.resnets = resnets - - if self.add_downsample: - self.downsamplers_0 = FlaxDownsample2D(self.out_channels, dtype=self.dtype) - - def __call__(self, hidden_states, deterministic=True): - for resnet in self.resnets: - hidden_states = resnet(hidden_states, deterministic=deterministic) - - if self.add_downsample: - hidden_states = self.downsamplers_0(hidden_states) - - return hidden_states - - -class FlaxUpDecoderBlock2D(nn.Module): - r""" - Flax Resnet blocks-based Decoder block for diffusion-based VAE. - - Parameters: - in_channels (:obj:`int`): - Input channels - out_channels (:obj:`int`): - Output channels - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - num_layers (:obj:`int`, *optional*, defaults to 1): - Number of Resnet layer block - resnet_groups (:obj:`int`, *optional*, defaults to `32`): - The number of groups to use for the Resnet block group norm - add_upsample (:obj:`bool`, *optional*, defaults to `True`): - Whether to add upsample layer - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - out_channels: int - dropout: float = 0.0 - num_layers: int = 1 - resnet_groups: int = 32 - add_upsample: bool = True - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - resnets = [] - for i in range(self.num_layers): - in_channels = self.in_channels if i == 0 else self.out_channels - res_block = FlaxResnetBlock2D( - in_channels=in_channels, - out_channels=self.out_channels, - dropout=self.dropout, - groups=self.resnet_groups, - dtype=self.dtype, - ) - resnets.append(res_block) - - self.resnets = resnets - - if self.add_upsample: - self.upsamplers_0 = FlaxUpsample2D(self.out_channels, dtype=self.dtype) - - def __call__(self, hidden_states, deterministic=True): - for resnet in self.resnets: - hidden_states = resnet(hidden_states, deterministic=deterministic) - - if self.add_upsample: - hidden_states = self.upsamplers_0(hidden_states) - - return hidden_states - - -class FlaxUNetMidBlock2D(nn.Module): - r""" - Flax Unet Mid-Block module. - - Parameters: - in_channels (:obj:`int`): - Input channels - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - num_layers (:obj:`int`, *optional*, defaults to 1): - Number of Resnet layer block - resnet_groups (:obj:`int`, *optional*, defaults to `32`): - The number of groups to use for the Resnet and Attention block group norm - num_attention_heads (:obj:`int`, *optional*, defaults to `1`): - Number of attention heads for each attention block - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - dropout: float = 0.0 - num_layers: int = 1 - resnet_groups: int = 32 - num_attention_heads: int = 1 - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - resnet_groups = self.resnet_groups if self.resnet_groups is not None else min(self.in_channels // 4, 32) - - # there is always at least one resnet - resnets = [ - FlaxResnetBlock2D( - in_channels=self.in_channels, - out_channels=self.in_channels, - dropout=self.dropout, - groups=resnet_groups, - dtype=self.dtype, - ) - ] - - attentions = [] - - for _ in range(self.num_layers): - attn_block = FlaxAttentionBlock( - channels=self.in_channels, - num_head_channels=self.num_attention_heads, - num_groups=resnet_groups, - dtype=self.dtype, - ) - attentions.append(attn_block) - - res_block = FlaxResnetBlock2D( - in_channels=self.in_channels, - out_channels=self.in_channels, - dropout=self.dropout, - groups=resnet_groups, - dtype=self.dtype, - ) - resnets.append(res_block) - - self.resnets = resnets - self.attentions = attentions - - def __call__(self, hidden_states, deterministic=True): - hidden_states = self.resnets[0](hidden_states, deterministic=deterministic) - for attn, resnet in zip(self.attentions, self.resnets[1:]): - hidden_states = attn(hidden_states) - hidden_states = resnet(hidden_states, deterministic=deterministic) - - return hidden_states - - -class FlaxEncoder(nn.Module): - r""" - Flax Implementation of VAE Encoder. - - This model is a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module) - subclass. Use it as a regular Flax linen Module and refer to the Flax documentation for all matter related to - general usage and behavior. - - Finally, this model supports inherent JAX features such as: - - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit) - - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation) - - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap) - - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap) - - Parameters: - in_channels (:obj:`int`, *optional*, defaults to 3): - Input channels - out_channels (:obj:`int`, *optional*, defaults to 3): - Output channels - down_block_types (:obj:`tuple[str]`, *optional*, defaults to `(DownEncoderBlock2D)`): - DownEncoder block type - block_out_channels (:obj:`tuple[str]`, *optional*, defaults to `(64,)`): - tuple[ containing the number of output channels for each block - layers_per_block (:obj:`int`, *optional*, defaults to `2`): - Number of Resnet layer for each block - norm_num_groups (:obj:`int`, *optional*, defaults to `32`): - norm num group - act_fn (:obj:`str`, *optional*, defaults to `silu`): - Activation function - double_z (:obj:`bool`, *optional*, defaults to `False`): - Whether to double the last output channels - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int = 3 - out_channels: int = 3 - down_block_types: tuple[str, ...] = ("DownEncoderBlock2D",) - block_out_channels: tuple[int, ...] = (64,) - layers_per_block: int = 2 - norm_num_groups: int = 32 - act_fn: str = "silu" - double_z: bool = False - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - block_out_channels = self.block_out_channels - # in - self.conv_in = nn.Conv( - block_out_channels[0], - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - # downsampling - down_blocks = [] - output_channel = block_out_channels[0] - for i, _ in enumerate(self.down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - down_block = FlaxDownEncoderBlock2D( - in_channels=input_channel, - out_channels=output_channel, - num_layers=self.layers_per_block, - resnet_groups=self.norm_num_groups, - add_downsample=not is_final_block, - dtype=self.dtype, - ) - down_blocks.append(down_block) - self.down_blocks = down_blocks - - # middle - self.mid_block = FlaxUNetMidBlock2D( - in_channels=block_out_channels[-1], - resnet_groups=self.norm_num_groups, - num_attention_heads=None, - dtype=self.dtype, - ) - - # end - conv_out_channels = 2 * self.out_channels if self.double_z else self.out_channels - self.conv_norm_out = nn.GroupNorm(num_groups=self.norm_num_groups, epsilon=1e-6) - self.conv_out = nn.Conv( - conv_out_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - def __call__(self, sample, deterministic: bool = True): - # in - sample = self.conv_in(sample) - - # downsampling - for block in self.down_blocks: - sample = block(sample, deterministic=deterministic) - - # middle - sample = self.mid_block(sample, deterministic=deterministic) - - # end - sample = self.conv_norm_out(sample) - sample = nn.swish(sample) - sample = self.conv_out(sample) - - return sample - - -class FlaxDecoder(nn.Module): - r""" - Flax Implementation of VAE Decoder. - - This model is a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module) - subclass. Use it as a regular Flax linen Module and refer to the Flax documentation for all matter related to - general usage and behavior. - - Finally, this model supports inherent JAX features such as: - - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit) - - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation) - - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap) - - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap) - - Parameters: - in_channels (:obj:`int`, *optional*, defaults to 3): - Input channels - out_channels (:obj:`int`, *optional*, defaults to 3): - Output channels - up_block_types (:obj:`tuple[str]`, *optional*, defaults to `(UpDecoderBlock2D)`): - UpDecoder block type - block_out_channels (:obj:`tuple[str]`, *optional*, defaults to `(64,)`): - tuple[ containing the number of output channels for each block - layers_per_block (:obj:`int`, *optional*, defaults to `2`): - Number of Resnet layer for each block - norm_num_groups (:obj:`int`, *optional*, defaults to `32`): - norm num group - act_fn (:obj:`str`, *optional*, defaults to `silu`): - Activation function - double_z (:obj:`bool`, *optional*, defaults to `False`): - Whether to double the last output channels - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - parameters `dtype` - """ - - in_channels: int = 3 - out_channels: int = 3 - up_block_types: tuple[str, ...] = ("UpDecoderBlock2D",) - block_out_channels: tuple[int, ...] = (64,) - layers_per_block: int = 2 - norm_num_groups: int = 32 - act_fn: str = "silu" - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - block_out_channels = self.block_out_channels - - # z to block_in - self.conv_in = nn.Conv( - block_out_channels[-1], - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - # middle - self.mid_block = FlaxUNetMidBlock2D( - in_channels=block_out_channels[-1], - resnet_groups=self.norm_num_groups, - num_attention_heads=None, - dtype=self.dtype, - ) - - # upsampling - reversed_block_out_channels = list(reversed(block_out_channels)) - output_channel = reversed_block_out_channels[0] - up_blocks = [] - for i, _ in enumerate(self.up_block_types): - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - - is_final_block = i == len(block_out_channels) - 1 - - up_block = FlaxUpDecoderBlock2D( - in_channels=prev_output_channel, - out_channels=output_channel, - num_layers=self.layers_per_block + 1, - resnet_groups=self.norm_num_groups, - add_upsample=not is_final_block, - dtype=self.dtype, - ) - up_blocks.append(up_block) - prev_output_channel = output_channel - - self.up_blocks = up_blocks - - # end - self.conv_norm_out = nn.GroupNorm(num_groups=self.norm_num_groups, epsilon=1e-6) - self.conv_out = nn.Conv( - self.out_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - def __call__(self, sample, deterministic: bool = True): - # z to block_in - sample = self.conv_in(sample) - - # middle - sample = self.mid_block(sample, deterministic=deterministic) - - # upsampling - for block in self.up_blocks: - sample = block(sample, deterministic=deterministic) - - sample = self.conv_norm_out(sample) - sample = nn.swish(sample) - sample = self.conv_out(sample) - - return sample - - -class FlaxDiagonalGaussianDistribution(object): - def __init__(self, parameters, deterministic=False): - # Last axis to account for channels-last - self.mean, self.logvar = jnp.split(parameters, 2, axis=-1) - self.logvar = jnp.clip(self.logvar, -30.0, 20.0) - self.deterministic = deterministic - self.std = jnp.exp(0.5 * self.logvar) - self.var = jnp.exp(self.logvar) - if self.deterministic: - self.var = self.std = jnp.zeros_like(self.mean) - - def sample(self, key): - return self.mean + self.std * jax.random.normal(key, self.mean.shape) - - def kl(self, other=None): - if self.deterministic: - return jnp.array([0.0]) - - if other is None: - return 0.5 * jnp.sum(self.mean**2 + self.var - 1.0 - self.logvar, axis=[1, 2, 3]) - - return 0.5 * jnp.sum( - jnp.square(self.mean - other.mean) / other.var + self.var / other.var - 1.0 - self.logvar + other.logvar, - axis=[1, 2, 3], - ) - - def nll(self, sample, axis=[1, 2, 3]): - if self.deterministic: - return jnp.array([0.0]) - - logtwopi = jnp.log(2.0 * jnp.pi) - return 0.5 * jnp.sum(logtwopi + self.logvar + jnp.square(sample - self.mean) / self.var, axis=axis) - - def mode(self): - return self.mean - - -@flax_register_to_config -class FlaxAutoencoderKL(nn.Module, FlaxModelMixin, ConfigMixin): - r""" - Flax implementation of a VAE model with KL loss for decoding latent representations. - - This model inherits from [`FlaxModelMixin`]. Check the superclass documentation for it's generic methods - implemented for all models (such as downloading or saving). - - This model is a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module) - subclass. Use it as a regular Flax Linen module and refer to the Flax documentation for all matter related to its - general usage and behavior. - - Inherent JAX features such as the following are supported: - - - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit) - - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation) - - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap) - - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap) - - Parameters: - in_channels (`int`, *optional*, defaults to 3): - Number of channels in the input image. - out_channels (`int`, *optional*, defaults to 3): - Number of channels in the output. - down_block_types (`tuple[str]`, *optional*, defaults to `(DownEncoderBlock2D)`): - tuple[ of downsample block types. - up_block_types (`tuple[str]`, *optional*, defaults to `(UpDecoderBlock2D)`): - tuple[ of upsample block types. - block_out_channels (`tuple[str]`, *optional*, defaults to `(64,)`): - tuple[ of block output channels. - layers_per_block (`int`, *optional*, defaults to `2`): - Number of ResNet layer for each block. - act_fn (`str`, *optional*, defaults to `silu`): - The activation function to use. - latent_channels (`int`, *optional*, defaults to `4`): - Number of channels in the latent space. - norm_num_groups (`int`, *optional*, defaults to `32`): - The number of groups for normalization. - sample_size (`int`, *optional*, defaults to 32): - Sample input size. - scaling_factor (`float`, *optional*, defaults to 0.18215): - The component-wise standard deviation of the trained latent space computed using the first batch of the - training set. This is used to scale the latent space to have unit variance when training the diffusion - model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the - diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1 - / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image - Synthesis with Latent Diffusion Models](https://huggingface.co/papers/2112.10752) paper. - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - The `dtype` of the parameters. - """ - - in_channels: int = 3 - out_channels: int = 3 - down_block_types: tuple[str, ...] = ("DownEncoderBlock2D",) - up_block_types: tuple[str, ...] = ("UpDecoderBlock2D",) - block_out_channels: tuple[int, ...] = (64,) - layers_per_block: int = 1 - act_fn: str = "silu" - latent_channels: int = 4 - norm_num_groups: int = 32 - sample_size: int = 32 - scaling_factor: float = 0.18215 - dtype: jnp.dtype = jnp.float32 - - def setup(self): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - self.encoder = FlaxEncoder( - in_channels=self.config.in_channels, - out_channels=self.config.latent_channels, - down_block_types=self.config.down_block_types, - block_out_channels=self.config.block_out_channels, - layers_per_block=self.config.layers_per_block, - act_fn=self.config.act_fn, - norm_num_groups=self.config.norm_num_groups, - double_z=True, - dtype=self.dtype, - ) - self.decoder = FlaxDecoder( - in_channels=self.config.latent_channels, - out_channels=self.config.out_channels, - up_block_types=self.config.up_block_types, - block_out_channels=self.config.block_out_channels, - layers_per_block=self.config.layers_per_block, - norm_num_groups=self.config.norm_num_groups, - act_fn=self.config.act_fn, - dtype=self.dtype, - ) - self.quant_conv = nn.Conv( - 2 * self.config.latent_channels, - kernel_size=(1, 1), - strides=(1, 1), - padding="VALID", - dtype=self.dtype, - ) - self.post_quant_conv = nn.Conv( - self.config.latent_channels, - kernel_size=(1, 1), - strides=(1, 1), - padding="VALID", - dtype=self.dtype, - ) - - def init_weights(self, rng: jax.Array) -> FrozenDict: - # init input tensors - sample_shape = (1, self.in_channels, self.sample_size, self.sample_size) - sample = jnp.zeros(sample_shape, dtype=jnp.float32) - - params_rng, dropout_rng, gaussian_rng = jax.random.split(rng, 3) - rngs = {"params": params_rng, "dropout": dropout_rng, "gaussian": gaussian_rng} - - return self.init(rngs, sample)["params"] - - def encode(self, sample, deterministic: bool = True, return_dict: bool = True): - sample = jnp.transpose(sample, (0, 2, 3, 1)) - - hidden_states = self.encoder(sample, deterministic=deterministic) - moments = self.quant_conv(hidden_states) - posterior = FlaxDiagonalGaussianDistribution(moments) - - if not return_dict: - return (posterior,) - - return FlaxAutoencoderKLOutput(latent_dist=posterior) - - def decode(self, latents, deterministic: bool = True, return_dict: bool = True): - if latents.shape[-1] != self.config.latent_channels: - latents = jnp.transpose(latents, (0, 2, 3, 1)) - - hidden_states = self.post_quant_conv(latents) - hidden_states = self.decoder(hidden_states, deterministic=deterministic) - - hidden_states = jnp.transpose(hidden_states, (0, 3, 1, 2)) - - if not return_dict: - return (hidden_states,) - - return FlaxDecoderOutput(sample=hidden_states) - - def __call__(self, sample, sample_posterior=False, deterministic: bool = True, return_dict: bool = True): - posterior = self.encode(sample, deterministic=deterministic, return_dict=return_dict) - if sample_posterior: - rng = self.make_rng("gaussian") - hidden_states = posterior.latent_dist.sample(rng) - else: - hidden_states = posterior.latent_dist.mode() - - sample = self.decode(hidden_states, return_dict=return_dict).sample - - if not return_dict: - return (sample,) - - return FlaxDecoderOutput(sample=sample) diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 7b0c43727975..6c0c8667aab5 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -13,7 +13,6 @@ is_torch_available, is_torch_npu_available, is_transformers_available, - is_transformers_flax_compatible, is_transformers_version, ) @@ -526,37 +525,6 @@ else: _import_structure["consisid"] = ["ConsisIDPipeline"] -try: - if not is_transformers_flax_compatible(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ..utils import dummy_flax_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_flax_objects)) -else: - _import_structure["pipeline_flax_utils"] = ["FlaxDiffusionPipeline"] -try: - if not is_transformers_flax_compatible(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ..utils import dummy_flax_and_transformers_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_flax_and_transformers_objects)) -else: - _import_structure["controlnet"].extend(["FlaxStableDiffusionControlNetPipeline"]) - _import_structure["stable_diffusion"].extend( - [ - "FlaxStableDiffusionImg2ImgPipeline", - "FlaxStableDiffusionInpaintPipeline", - "FlaxStableDiffusionPipeline", - ] - ) - _import_structure["stable_diffusion_xl"].extend( - [ - "FlaxStableDiffusionXLPipeline", - ] - ) - if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: if not is_torch_available(): @@ -984,30 +952,6 @@ else: from .consisid import ConsisIDPipeline - try: - if not is_transformers_flax_compatible(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_flax_objects import * # noqa F403 - else: - from .pipeline_flax_utils import FlaxDiffusionPipeline - - try: - if not is_transformers_flax_compatible(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_flax_and_transformers_objects import * - else: - from .controlnet import FlaxStableDiffusionControlNetPipeline - from .stable_diffusion import ( - FlaxStableDiffusionImg2ImgPipeline, - FlaxStableDiffusionInpaintPipeline, - FlaxStableDiffusionPipeline, - ) - from .stable_diffusion_xl import ( - FlaxStableDiffusionXLPipeline, - ) - try: if not (is_transformers_available() and is_torch_available() and is_note_seq_available()): raise OptionalDependencyNotAvailable() diff --git a/src/diffusers/pipelines/auto_pipeline.py b/src/diffusers/pipelines/auto_pipeline.py index 15c23039c1cb..e8f4a8fa8ae4 100644 --- a/src/diffusers/pipelines/auto_pipeline.py +++ b/src/diffusers/pipelines/auto_pipeline.py @@ -480,8 +480,7 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs): class). The overwritten components are passed directly to the pipelines `__init__` method. See example below for more information. variant (`str`, *optional*): - Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. + Load weights from a specified variant filename such as `"fp16"` or `"ema"`. > [!TIP] > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with `hf > auth login`. @@ -771,8 +770,7 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs): class). The overwritten components are passed directly to the pipelines `__init__` method. See example below for more information. variant (`str`, *optional*): - Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. + Load weights from a specified variant filename such as `"fp16"` or `"ema"`. > [!TIP] > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with `hf > auth login`. @@ -1076,8 +1074,7 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs): class). The overwritten components are passed directly to the pipelines `__init__` method. See example below for more information. variant (`str`, *optional*): - Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. + Load weights from a specified variant filename such as `"fp16"` or `"ema"`. > [!TIP] > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with `hf > auth login`. @@ -1366,8 +1363,7 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs): class). The overwritten components are passed directly to the pipelines `__init__` method. See example below for more information. variant (`str`, *optional*): - Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. + Load weights from a specified variant filename such as `"fp16"` or `"ema"`. > [!TIP] > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with `hf > auth login`. diff --git a/src/diffusers/pipelines/controlnet/__init__.py b/src/diffusers/pipelines/controlnet/__init__.py index cd94327bb0b7..3fb1b8571b43 100644 --- a/src/diffusers/pipelines/controlnet/__init__.py +++ b/src/diffusers/pipelines/controlnet/__init__.py @@ -7,7 +7,6 @@ get_objects_from_module, is_torch_available, is_transformers_available, - is_transformers_flax_compatible, ) @@ -33,15 +32,6 @@ _import_structure["pipeline_controlnet_union_inpaint_sd_xl"] = ["StableDiffusionXLControlNetUnionInpaintPipeline"] _import_structure["pipeline_controlnet_union_sd_xl"] = ["StableDiffusionXLControlNetUnionPipeline"] _import_structure["pipeline_controlnet_union_sd_xl_img2img"] = ["StableDiffusionXLControlNetUnionImg2ImgPipeline"] -try: - if not is_transformers_flax_compatible(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils import dummy_flax_and_transformers_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_flax_and_transformers_objects)) -else: - _import_structure["pipeline_flax_controlnet"] = ["FlaxStableDiffusionControlNetPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: @@ -64,14 +54,6 @@ from .pipeline_controlnet_union_sd_xl import StableDiffusionXLControlNetUnionPipeline from .pipeline_controlnet_union_sd_xl_img2img import StableDiffusionXLControlNetUnionImg2ImgPipeline - try: - if not is_transformers_flax_compatible(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_flax_and_transformers_objects import * # noqa F403 - else: - from .pipeline_flax_controlnet import FlaxStableDiffusionControlNetPipeline - else: import sys diff --git a/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py b/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py deleted file mode 100644 index a76f86035b81..000000000000 --- a/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py +++ /dev/null @@ -1,528 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import warnings -from functools import partial - -import jax -import jax.numpy as jnp -import numpy as np -from flax.core.frozen_dict import FrozenDict -from flax.jax_utils import unreplicate -from flax.training.common_utils import shard -from PIL import Image -from transformers import CLIPImageProcessor, CLIPTokenizer, FlaxCLIPTextModel - -from ...models import FlaxAutoencoderKL, FlaxControlNetModel, FlaxUNet2DConditionModel -from ...schedulers import ( - FlaxDDIMScheduler, - FlaxDPMSolverMultistepScheduler, - FlaxLMSDiscreteScheduler, - FlaxPNDMScheduler, -) -from ...utils import PIL_INTERPOLATION, logging, replace_example_docstring -from ..pipeline_flax_utils import FlaxDiffusionPipeline -from ..stable_diffusion import FlaxStableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker_flax import FlaxStableDiffusionSafetyChecker - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -# Set to True to use python for loop instead of jax.fori_loop for easier debugging -DEBUG = False - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import jax - >>> import numpy as np - >>> import jax.numpy as jnp - >>> from flax.jax_utils import replicate - >>> from flax.training.common_utils import shard - >>> from diffusers.utils import load_image, make_image_grid - >>> from PIL import Image - >>> from diffusers import FlaxStableDiffusionControlNetPipeline, FlaxControlNetModel - - - >>> def create_key(seed=0): - ... return jax.random.PRNGKey(seed) - - - >>> rng = create_key(0) - - >>> # get canny image - >>> canny_image = load_image( - ... "https://huggingface.co/datasets/YiYiXu/test-doc-assets/resolve/main/blog_post_cell_10_output_0.jpeg" - ... ) - - >>> prompts = "best quality, extremely detailed" - >>> negative_prompts = "monochrome, lowres, bad anatomy, worst quality, low quality" - - >>> # load control net and stable diffusion v1-5 - >>> controlnet, controlnet_params = FlaxControlNetModel.from_pretrained( - ... "lllyasviel/sd-controlnet-canny", from_pt=True, dtype=jnp.float32 - ... ) - >>> pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", - ... controlnet=controlnet, - ... revision="flax", - ... dtype=jnp.float32, - ... ) - >>> params["controlnet"] = controlnet_params - - >>> num_samples = jax.device_count() - >>> rng = jax.random.split(rng, jax.device_count()) - - >>> prompt_ids = pipe.prepare_text_inputs([prompts] * num_samples) - >>> negative_prompt_ids = pipe.prepare_text_inputs([negative_prompts] * num_samples) - >>> processed_image = pipe.prepare_image_inputs([canny_image] * num_samples) - - >>> p_params = replicate(params) - >>> prompt_ids = shard(prompt_ids) - >>> negative_prompt_ids = shard(negative_prompt_ids) - >>> processed_image = shard(processed_image) - - >>> output = pipe( - ... prompt_ids=prompt_ids, - ... image=processed_image, - ... params=p_params, - ... prng_seed=rng, - ... num_inference_steps=50, - ... neg_prompt_ids=negative_prompt_ids, - ... jit=True, - ... ).images - - >>> output_images = pipe.numpy_to_pil(np.asarray(output.reshape((num_samples,) + output.shape[-3:]))) - >>> output_images = make_image_grid(output_images, num_samples // 4, 4) - >>> output_images.save("generated_image.png") - ``` -""" - - -class FlaxStableDiffusionControlNetPipeline(FlaxDiffusionPipeline): - r""" - Flax-based pipeline for text-to-image generation using Stable Diffusion with ControlNet Guidance. - - This model inherits from [`FlaxDiffusionPipeline`]. Check the superclass documentation for the generic methods - implemented for all pipelines (downloading, saving, running on a particular device, etc.). - - Args: - vae ([`FlaxAutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.FlaxCLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`FlaxUNet2DConditionModel`]): - A `FlaxUNet2DConditionModel` to denoise the encoded image latents. - controlnet ([`FlaxControlNetModel`]: - Provides additional conditioning to the `unet` during the denoising process. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of - [`FlaxDDIMScheduler`], [`FlaxLMSDiscreteScheduler`], [`FlaxPNDMScheduler`], or - [`FlaxDPMSolverMultistepScheduler`]. - safety_checker ([`FlaxStableDiffusionSafetyChecker`]): - Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - def __init__( - self, - vae: FlaxAutoencoderKL, - text_encoder: FlaxCLIPTextModel, - tokenizer: CLIPTokenizer, - unet: FlaxUNet2DConditionModel, - controlnet: FlaxControlNetModel, - scheduler: FlaxDDIMScheduler | FlaxPNDMScheduler | FlaxLMSDiscreteScheduler | FlaxDPMSolverMultistepScheduler, - safety_checker: FlaxStableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - dtype: jnp.dtype = jnp.float32, - ): - super().__init__() - self.dtype = dtype - - if safety_checker is None: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - - def prepare_text_inputs(self, prompt: str | list[str]): - if not isinstance(prompt, (str, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - text_input = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="np", - ) - - return text_input.input_ids - - def prepare_image_inputs(self, image: Image.Image | list[Image.Image]): - if not isinstance(image, (Image.Image, list)): - raise ValueError(f"image has to be of type `PIL.Image.Image` or list but is {type(image)}") - - if isinstance(image, Image.Image): - image = [image] - - processed_images = jnp.concatenate([preprocess(img, jnp.float32) for img in image]) - - return processed_images - - def _get_has_nsfw_concepts(self, features, params): - has_nsfw_concepts = self.safety_checker(features, params) - return has_nsfw_concepts - - def _run_safety_checker(self, images, safety_model_params, jit=False): - # safety_model_params should already be replicated when jit is True - pil_images = [Image.fromarray(image) for image in images] - features = self.feature_extractor(pil_images, return_tensors="np").pixel_values - - if jit: - features = shard(features) - has_nsfw_concepts = _p_get_has_nsfw_concepts(self, features, safety_model_params) - has_nsfw_concepts = unshard(has_nsfw_concepts) - safety_model_params = unreplicate(safety_model_params) - else: - has_nsfw_concepts = self._get_has_nsfw_concepts(features, safety_model_params) - - images_was_copied = False - for idx, has_nsfw_concept in enumerate(has_nsfw_concepts): - if has_nsfw_concept: - if not images_was_copied: - images_was_copied = True - images = images.copy() - - images[idx] = np.zeros(images[idx].shape, dtype=np.uint8) # black image - - if any(has_nsfw_concepts): - warnings.warn( - "Potential NSFW content was detected in one or more images. A black image will be returned" - " instead. Try again with a different prompt and/or seed." - ) - - return images, has_nsfw_concepts - - def _generate( - self, - prompt_ids: jnp.ndarray, - image: jnp.ndarray, - params: dict | FrozenDict, - prng_seed: jax.Array, - num_inference_steps: int, - guidance_scale: float, - latents: jnp.ndarray | None = None, - neg_prompt_ids: jnp.ndarray | None = None, - controlnet_conditioning_scale: float = 1.0, - ): - height, width = image.shape[-2:] - if height % 64 != 0 or width % 64 != 0: - raise ValueError(f"`height` and `width` have to be divisible by 64 but are {height} and {width}.") - - # get prompt text embeddings - prompt_embeds = self.text_encoder(prompt_ids, params=params["text_encoder"])[0] - - # TODO: currently it is assumed `do_classifier_free_guidance = guidance_scale > 1.0` - # implement this conditional `do_classifier_free_guidance = guidance_scale > 1.0` - batch_size = prompt_ids.shape[0] - - max_length = prompt_ids.shape[-1] - - if neg_prompt_ids is None: - uncond_input = self.tokenizer( - [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="np" - ).input_ids - else: - uncond_input = neg_prompt_ids - negative_prompt_embeds = self.text_encoder(uncond_input, params=params["text_encoder"])[0] - context = jnp.concatenate([negative_prompt_embeds, prompt_embeds]) - - image = jnp.concatenate([image] * 2) - - latents_shape = ( - batch_size, - self.unet.config.in_channels, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - ) - if latents is None: - latents = jax.random.normal(prng_seed, shape=latents_shape, dtype=jnp.float32) - else: - if latents.shape != latents_shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") - - def loop_body(step, args): - latents, scheduler_state = args - # For classifier free guidance, we need to do two forward passes. - # Here we concatenate the unconditional and text embeddings into a single batch - # to avoid doing two forward passes - latents_input = jnp.concatenate([latents] * 2) - - t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step] - timestep = jnp.broadcast_to(t, latents_input.shape[0]) - - latents_input = self.scheduler.scale_model_input(scheduler_state, latents_input, t) - - down_block_res_samples, mid_block_res_sample = self.controlnet.apply( - {"params": params["controlnet"]}, - jnp.array(latents_input), - jnp.array(timestep, dtype=jnp.int32), - encoder_hidden_states=context, - controlnet_cond=image, - conditioning_scale=controlnet_conditioning_scale, - return_dict=False, - ) - - # predict the noise residual - noise_pred = self.unet.apply( - {"params": params["unet"]}, - jnp.array(latents_input), - jnp.array(timestep, dtype=jnp.int32), - encoder_hidden_states=context, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - ).sample - - # perform guidance - noise_pred_uncond, noise_prediction_text = jnp.split(noise_pred, 2, axis=0) - noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents, scheduler_state = self.scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() - return latents, scheduler_state - - scheduler_state = self.scheduler.set_timesteps( - params["scheduler"], num_inference_steps=num_inference_steps, shape=latents_shape - ) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * params["scheduler"].init_noise_sigma - - if DEBUG: - # run with python for loop - for i in range(num_inference_steps): - latents, scheduler_state = loop_body(i, (latents, scheduler_state)) - else: - latents, _ = jax.lax.fori_loop(0, num_inference_steps, loop_body, (latents, scheduler_state)) - - # scale and decode the image latents with vae - latents = 1 / self.vae.config.scaling_factor * latents - image = self.vae.apply({"params": params["vae"]}, latents, method=self.vae.decode).sample - - image = (image / 2 + 0.5).clip(0, 1).transpose(0, 2, 3, 1) - return image - - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt_ids: jnp.ndarray, - image: jnp.ndarray, - params: dict | FrozenDict, - prng_seed: jax.Array, - num_inference_steps: int = 50, - guidance_scale: float | jnp.ndarray = 7.5, - latents: jnp.ndarray = None, - neg_prompt_ids: jnp.ndarray = None, - controlnet_conditioning_scale: float | jnp.ndarray = 1.0, - return_dict: bool = True, - jit: bool = False, - ): - r""" - The call function to the pipeline for generation. - - Args: - prompt_ids (`jnp.ndarray`): - The prompt or prompts to guide the image generation. - image (`jnp.ndarray`): - Array representing the ControlNet input condition to provide guidance to the `unet` for generation. - params (`Dict` or `FrozenDict`): - Dictionary containing the model parameters/weights. - prng_seed (`jax.Array`): - Array containing random number generator key. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - guidance_scale (`float`, *optional*, defaults to 7.5): - A higher guidance scale value encourages the model to generate images closely linked to the text - `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - latents (`jnp.ndarray`, *optional*): - Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - array is generated by sampling using the supplied random `generator`. - controlnet_conditioning_scale (`float` or `jnp.ndarray`, *optional*, defaults to 1.0): - The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added - to the residual in the original `unet`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] instead of - a plain tuple. - jit (`bool`, defaults to `False`): - Whether to run `pmap` versions of the generation and safety scoring functions. - - > [!WARNING] > This argument exists because `__call__` is not yet end-to-end pmap-able. It will be - removed in a > future release. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] is - returned, otherwise a `tuple` is returned where the first element is a list with the generated images - and the second element is a list of `bool`s indicating whether the corresponding generated image - contains "not-safe-for-work" (nsfw) content. - """ - - height, width = image.shape[-2:] - - if isinstance(guidance_scale, float): - # Convert to a tensor so each device gets a copy. Follow the prompt_ids for - # shape information, as they may be sharded (when `jit` is `True`), or not. - guidance_scale = jnp.array([guidance_scale] * prompt_ids.shape[0]) - if len(prompt_ids.shape) > 2: - # Assume sharded - guidance_scale = guidance_scale[:, None] - - if isinstance(controlnet_conditioning_scale, float): - # Convert to a tensor so each device gets a copy. Follow the prompt_ids for - # shape information, as they may be sharded (when `jit` is `True`), or not. - controlnet_conditioning_scale = jnp.array([controlnet_conditioning_scale] * prompt_ids.shape[0]) - if len(prompt_ids.shape) > 2: - # Assume sharded - controlnet_conditioning_scale = controlnet_conditioning_scale[:, None] - - if jit: - images = _p_generate( - self, - prompt_ids, - image, - params, - prng_seed, - num_inference_steps, - guidance_scale, - latents, - neg_prompt_ids, - controlnet_conditioning_scale, - ) - else: - images = self._generate( - prompt_ids, - image, - params, - prng_seed, - num_inference_steps, - guidance_scale, - latents, - neg_prompt_ids, - controlnet_conditioning_scale, - ) - - if self.safety_checker is not None: - safety_params = params["safety_checker"] - images_uint8_casted = (images * 255).round().astype("uint8") - num_devices, batch_size = images.shape[:2] - - images_uint8_casted = np.asarray(images_uint8_casted).reshape(num_devices * batch_size, height, width, 3) - images_uint8_casted, has_nsfw_concept = self._run_safety_checker(images_uint8_casted, safety_params, jit) - images = np.array(images) - - # block images - if any(has_nsfw_concept): - for i, is_nsfw in enumerate(has_nsfw_concept): - if is_nsfw: - images[i] = np.asarray(images_uint8_casted[i]) - - images = images.reshape(num_devices, batch_size, height, width, 3) - else: - images = np.asarray(images) - has_nsfw_concept = False - - if not return_dict: - return (images, has_nsfw_concept) - - return FlaxStableDiffusionPipelineOutput(images=images, nsfw_content_detected=has_nsfw_concept) - - -# Static argnums are pipe, num_inference_steps. A change would trigger recompilation. -# Non-static args are (sharded) input tensors mapped over their first dimension (hence, `0`). -@partial( - jax.pmap, - in_axes=(None, 0, 0, 0, 0, None, 0, 0, 0, 0), - static_broadcasted_argnums=(0, 5), -) -def _p_generate( - pipe, - prompt_ids, - image, - params, - prng_seed, - num_inference_steps, - guidance_scale, - latents, - neg_prompt_ids, - controlnet_conditioning_scale, -): - return pipe._generate( - prompt_ids, - image, - params, - prng_seed, - num_inference_steps, - guidance_scale, - latents, - neg_prompt_ids, - controlnet_conditioning_scale, - ) - - -@partial(jax.pmap, static_broadcasted_argnums=(0,)) -def _p_get_has_nsfw_concepts(pipe, features, params): - return pipe._get_has_nsfw_concepts(features, params) - - -def unshard(x: jnp.ndarray): - # einops.rearrange(x, 'd b ... -> (d b) ...') - num_devices, batch_size = x.shape[:2] - rest = x.shape[2:] - return x.reshape(num_devices * batch_size, *rest) - - -def preprocess(image, dtype): - image = image.convert("RGB") - w, h = image.size - w, h = (x - x % 64 for x in (w, h)) # resize to integer multiple of 64 - image = image.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]) - image = jnp.array(image).astype(dtype) / 255.0 - image = image[None].transpose(0, 3, 1, 2) - return image diff --git a/src/diffusers/pipelines/controlnet_sd3/__init__.py b/src/diffusers/pipelines/controlnet_sd3/__init__.py index e647706aa2f9..9488506af75d 100644 --- a/src/diffusers/pipelines/controlnet_sd3/__init__.py +++ b/src/diffusers/pipelines/controlnet_sd3/__init__.py @@ -7,7 +7,6 @@ get_objects_from_module, is_torch_available, is_transformers_available, - is_transformers_flax_compatible, ) @@ -38,12 +37,6 @@ from .pipeline_stable_diffusion_3_controlnet import StableDiffusion3ControlNetPipeline from .pipeline_stable_diffusion_3_controlnet_inpainting import StableDiffusion3ControlNetInpaintingPipeline - try: - if not is_transformers_flax_compatible(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_flax_and_transformers_objects import * # noqa F403 - else: import sys diff --git a/src/diffusers/pipelines/deprecated/controlnet_xs/__init__.py b/src/diffusers/pipelines/deprecated/controlnet_xs/__init__.py index cbd8c7468f29..8d469eed6b12 100644 --- a/src/diffusers/pipelines/deprecated/controlnet_xs/__init__.py +++ b/src/diffusers/pipelines/deprecated/controlnet_xs/__init__.py @@ -7,7 +7,6 @@ get_objects_from_module, is_torch_available, is_transformers_available, - is_transformers_flax_compatible, ) @@ -24,15 +23,6 @@ else: _import_structure["pipeline_controlnet_xs"] = ["StableDiffusionControlNetXSPipeline"] _import_structure["pipeline_controlnet_xs_sd_xl"] = ["StableDiffusionXLControlNetXSPipeline"] -try: - if not is_transformers_flax_compatible(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ....utils import dummy_flax_and_transformers_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_flax_and_transformers_objects)) -else: - pass # _import_structure["pipeline_flax_controlnet"] = ["FlaxStableDiffusionControlNetPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: @@ -46,14 +36,6 @@ from .pipeline_controlnet_xs import StableDiffusionControlNetXSPipeline from .pipeline_controlnet_xs_sd_xl import StableDiffusionXLControlNetXSPipeline - try: - if not is_transformers_flax_compatible(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ....utils.dummy_flax_and_transformers_objects import * # noqa F403 - else: - pass # from .pipeline_flax_controlnet import FlaxStableDiffusionControlNetPipeline - else: import sys diff --git a/src/diffusers/pipelines/pag/__init__.py b/src/diffusers/pipelines/pag/__init__.py index 176efe3adef6..9d7f0f34f529 100644 --- a/src/diffusers/pipelines/pag/__init__.py +++ b/src/diffusers/pipelines/pag/__init__.py @@ -5,7 +5,6 @@ OptionalDependencyNotAvailable, _LazyModule, get_objects_from_module, - is_flax_available, is_torch_available, is_transformers_available, ) diff --git a/src/diffusers/pipelines/pipeline_flax_utils.py b/src/diffusers/pipelines/pipeline_flax_utils.py deleted file mode 100644 index 8c29db8aa45e..000000000000 --- a/src/diffusers/pipelines/pipeline_flax_utils.py +++ /dev/null @@ -1,611 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import importlib -import inspect -import os -from typing import Any - -import flax -import numpy as np -import PIL.Image -from flax.core.frozen_dict import FrozenDict -from huggingface_hub import create_repo, snapshot_download -from huggingface_hub.utils import validate_hf_hub_args -from PIL import Image -from tqdm.auto import tqdm - -from ..configuration_utils import ConfigMixin -from ..models.modeling_flax_utils import FLAX_WEIGHTS_NAME, FlaxModelMixin -from ..schedulers.scheduling_utils_flax import SCHEDULER_CONFIG_NAME, FlaxSchedulerMixin -from ..utils import ( - CONFIG_NAME, - BaseOutput, - PushToHubMixin, - http_user_agent, - is_transformers_flax_compatible, - logging, -) - - -if is_transformers_flax_compatible(): - from transformers import FlaxPreTrainedModel - -INDEX_FILE = "diffusion_flax_model.bin" - - -logger = logging.get_logger(__name__) - - -LOADABLE_CLASSES = { - "diffusers": { - "FlaxModelMixin": ["save_pretrained", "from_pretrained"], - "FlaxSchedulerMixin": ["save_pretrained", "from_pretrained"], - "FlaxDiffusionPipeline": ["save_pretrained", "from_pretrained"], - }, - "transformers": { - "PreTrainedTokenizer": ["save_pretrained", "from_pretrained"], - "PreTrainedTokenizerFast": ["save_pretrained", "from_pretrained"], - "FlaxPreTrainedModel": ["save_pretrained", "from_pretrained"], - "FeatureExtractionMixin": ["save_pretrained", "from_pretrained"], - "ProcessorMixin": ["save_pretrained", "from_pretrained"], - "ImageProcessingMixin": ["save_pretrained", "from_pretrained"], - }, -} - -ALL_IMPORTABLE_CLASSES = {} -for library in LOADABLE_CLASSES: - ALL_IMPORTABLE_CLASSES.update(LOADABLE_CLASSES[library]) - - -def import_flax_or_no_model(module, class_name): - try: - # 1. First make sure that if a Flax object is present, import this one - class_obj = getattr(module, "Flax" + class_name) - except AttributeError: - # 2. If this doesn't work, it's not a model and we don't append "Flax" - class_obj = getattr(module, class_name) - except AttributeError: - raise ValueError(f"Neither Flax{class_name} nor {class_name} exist in {module}") - - return class_obj - - -@flax.struct.dataclass -class FlaxImagePipelineOutput(BaseOutput): - """ - Output class for image pipelines. - - Args: - images (`list[PIL.Image.Image]` or `np.ndarray`) - list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, - num_channels)`. - """ - - images: list[PIL.Image.Image] | np.ndarray - - -class FlaxDiffusionPipeline(ConfigMixin, PushToHubMixin): - r""" - Base class for Flax-based pipelines. - - [`FlaxDiffusionPipeline`] stores all components (models, schedulers, and processors) for diffusion pipelines and - provides methods for loading, downloading and saving models. It also includes methods to: - - - enable/disable the progress bar for the denoising iteration - - Class attributes: - - - **config_name** ([`str`]) -- The configuration filename that stores the class and module names of all the - diffusion pipeline's components. - """ - - config_name = "model_index.json" - - def register_modules(self, **kwargs): - # import it here to avoid circular import - from diffusers import pipelines - - for name, module in kwargs.items(): - if module is None: - register_dict = {name: (None, None)} - else: - # retrieve library - library = module.__module__.split(".")[0] - - # check if the module is a pipeline module - pipeline_dir = module.__module__.split(".")[-2] - path = module.__module__.split(".") - is_pipeline_module = pipeline_dir in path and hasattr(pipelines, pipeline_dir) - - # if library is not in LOADABLE_CLASSES, then it is a custom module. - # Or if it's a pipeline module, then the module is inside the pipeline - # folder so we set the library to module name. - if library not in LOADABLE_CLASSES or is_pipeline_module: - library = pipeline_dir - - # retrieve class_name - class_name = module.__class__.__name__ - - register_dict = {name: (library, class_name)} - - # save model index config - self.register_to_config(**register_dict) - - # set models - setattr(self, name, module) - - def save_pretrained( - self, - save_directory: str | os.PathLike, - params: dict | FrozenDict, - push_to_hub: bool = False, - **kwargs, - ): - # TODO: handle inference_state - """ - Save all saveable variables of the pipeline to a directory. A pipeline variable can be saved and loaded if its - class implements both a save and loading method. The pipeline is easily reloaded using the - [`~FlaxDiffusionPipeline.from_pretrained`] class method. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to which to save. Will be created if it doesn't exist. - push_to_hub (`bool`, *optional*, defaults to `False`): - Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the - repository you want to push to with `repo_id` (will default to the name of `save_directory` in your - namespace). - kwargs (`dict[str, Any]`, *optional*): - Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. - """ - self.save_config(save_directory) - - model_index_dict = dict(self.config) - model_index_dict.pop("_class_name") - model_index_dict.pop("_diffusers_version") - model_index_dict.pop("_module", None) - - if push_to_hub: - commit_message = kwargs.pop("commit_message", None) - private = kwargs.pop("private", None) - create_pr = kwargs.pop("create_pr", False) - token = kwargs.pop("token", None) - repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) - repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id - - for pipeline_component_name in model_index_dict.keys(): - sub_model = getattr(self, pipeline_component_name) - if sub_model is None: - # edge case for saving a pipeline with safety_checker=None - continue - - model_cls = sub_model.__class__ - - save_method_name = None - # search for the model's base class in LOADABLE_CLASSES - for library_name, library_classes in LOADABLE_CLASSES.items(): - library = importlib.import_module(library_name) - for base_class, save_load_methods in library_classes.items(): - class_candidate = getattr(library, base_class, None) - if class_candidate is not None and issubclass(model_cls, class_candidate): - # if we found a suitable base class in LOADABLE_CLASSES then grab its save method - save_method_name = save_load_methods[0] - break - if save_method_name is not None: - break - - save_method = getattr(sub_model, save_method_name) - expects_params = "params" in set(inspect.signature(save_method).parameters.keys()) - - if expects_params: - save_method( - os.path.join(save_directory, pipeline_component_name), params=params[pipeline_component_name] - ) - else: - save_method(os.path.join(save_directory, pipeline_component_name)) - - if push_to_hub: - self._upload_folder( - save_directory, - repo_id, - token=token, - commit_message=commit_message, - create_pr=create_pr, - ) - - @classmethod - @validate_hf_hub_args - def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None, **kwargs): - r""" - Instantiate a Flax-based diffusion pipeline from pretrained pipeline weights. - - The pipeline is set in evaluation mode (`model.eval()) by default and dropout modules are deactivated. - - If you get the error message below, you need to finetune the weights for your downstream task: - - ``` - Some weights of FlaxUNet2DConditionModel were not initialized from the model checkpoint at stable-diffusion-v1-5/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: - ``` - - Parameters: - pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): - Can be either: - - - A string, the *repo id* (for example `stable-diffusion-v1-5/stable-diffusion-v1-5`) of a - pretrained pipeline hosted on the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - using [`~FlaxDiffusionPipeline.save_pretrained`]. - dtype (`jnp.dtype`, *optional*): - Override the default `jnp.dtype` and load the model under this dtype. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - output_loading_info(`bool`, *optional*, defaults to `False`): - Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - mirror (`str`, *optional*): - Mirror source to resolve accessibility issues if you're downloading a model in China. We do not - guarantee the timeliness or safety of the source, and you should refer to the mirror site for more - information. - kwargs (remaining dictionary of keyword arguments, *optional*): - Can be used to overwrite load and saveable variables (the pipeline components) of the specific pipeline - class. The overwritten components are passed directly to the pipelines `__init__` method. - - > [!TIP] > To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in - with `hf > auth login`. - - Examples: - - ```py - >>> from diffusers import FlaxDiffusionPipeline - - >>> # Download pipeline from huggingface.co and cache. - >>> # Requires to be logged in to Hugging Face hub, - >>> # see more in [the documentation](https://huggingface.co/docs/hub/security-tokens) - >>> pipeline, params = FlaxDiffusionPipeline.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", - ... variant="bf16", - ... dtype=jnp.bfloat16, - ... ) - - >>> # Download pipeline, but use a different scheduler - >>> from diffusers import FlaxDPMSolverMultistepScheduler - - >>> model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5" - >>> dpmpp, dpmpp_state = FlaxDPMSolverMultistepScheduler.from_pretrained( - ... model_id, - ... subfolder="scheduler", - ... ) - - >>> dpm_pipe, dpm_params = FlaxStableDiffusionPipeline.from_pretrained( - ... model_id, variant="bf16", dtype=jnp.bfloat16, scheduler=dpmpp - ... ) - >>> dpm_params["scheduler"] = dpmpp_state - ``` - """ - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - cache_dir = kwargs.pop("cache_dir", None) - proxies = kwargs.pop("proxies", None) - local_files_only = kwargs.pop("local_files_only", False) - token = kwargs.pop("token", None) - revision = kwargs.pop("revision", None) - from_pt = kwargs.pop("from_pt", False) - use_memory_efficient_attention = kwargs.pop("use_memory_efficient_attention", False) - split_head_dim = kwargs.pop("split_head_dim", False) - dtype = kwargs.pop("dtype", None) - - # 1. Download the checkpoints and configs - # use snapshot download here to get it working from from_pretrained - if not os.path.isdir(pretrained_model_name_or_path): - config_dict = cls.load_config( - pretrained_model_name_or_path, - cache_dir=cache_dir, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - ) - # make sure we only download sub-folders and `diffusers` filenames - folder_names = [k for k in config_dict.keys() if not k.startswith("_")] - allow_patterns = [os.path.join(k, "*") for k in folder_names] - allow_patterns += [FLAX_WEIGHTS_NAME, SCHEDULER_CONFIG_NAME, CONFIG_NAME, cls.config_name] - - ignore_patterns = ["*.bin", "*.safetensors"] if not from_pt else [] - ignore_patterns += ["*.onnx", "*.onnx_data", "*.xml", "*.pb"] - - if cls != FlaxDiffusionPipeline: - requested_pipeline_class = cls.__name__ - else: - requested_pipeline_class = config_dict.get("_class_name", cls.__name__) - requested_pipeline_class = ( - requested_pipeline_class - if requested_pipeline_class.startswith("Flax") - else "Flax" + requested_pipeline_class - ) - - user_agent = {"pipeline_class": requested_pipeline_class} - user_agent = http_user_agent(user_agent) - - # download all allow_patterns - cached_folder = snapshot_download( - pretrained_model_name_or_path, - cache_dir=cache_dir, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - allow_patterns=allow_patterns, - ignore_patterns=ignore_patterns, - user_agent=user_agent, - ) - else: - cached_folder = pretrained_model_name_or_path - - config_dict = cls.load_config(cached_folder) - - # 2. Load the pipeline class, if using custom module then load it from the hub - # if we load from explicit class, let's use it - if cls != FlaxDiffusionPipeline: - pipeline_class = cls - else: - diffusers_module = importlib.import_module(cls.__module__.split(".")[0]) - class_name = ( - config_dict["_class_name"] - if config_dict["_class_name"].startswith("Flax") - else "Flax" + config_dict["_class_name"] - ) - pipeline_class = getattr(diffusers_module, class_name) - - # some modules can be passed directly to the init - # in this case they are already instantiated in `kwargs` - # extract them here - expected_modules, optional_kwargs = cls._get_signature_keys(pipeline_class) - passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs} - passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs} - - init_dict, unused_kwargs, _ = pipeline_class.extract_init_dict(config_dict, **kwargs) - - # define init kwargs - init_kwargs = {k: init_dict.pop(k) for k in optional_kwargs if k in init_dict} - init_kwargs = {**init_kwargs, **passed_pipe_kwargs} - - # remove `null` components - def load_module(name, value): - if value[0] is None: - return False - if name in passed_class_obj and passed_class_obj[name] is None: - return False - return True - - init_dict = {k: v for k, v in init_dict.items() if load_module(k, v)} - - # Throw nice warnings / errors for fast accelerate loading - if len(unused_kwargs) > 0: - logger.warning( - f"Keyword arguments {unused_kwargs} are not expected by {pipeline_class.__name__} and will be ignored." - ) - - # inference_params - params = {} - - # import it here to avoid circular import - from diffusers import pipelines - - # 3. Load each module in the pipeline - for name, (library_name, class_name) in init_dict.items(): - if class_name is None: - # edge case for when the pipeline was saved with safety_checker=None - init_kwargs[name] = None - continue - - is_pipeline_module = hasattr(pipelines, library_name) - loaded_sub_model = None - sub_model_should_be_defined = True - - # if the model is in a pipeline module, then we load it from the pipeline - if name in passed_class_obj: - # 1. check that passed_class_obj has correct parent class - if not is_pipeline_module: - library = importlib.import_module(library_name) - class_obj = getattr(library, class_name) - importable_classes = LOADABLE_CLASSES[library_name] - class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()} - - expected_class_obj = None - for class_name, class_candidate in class_candidates.items(): - if class_candidate is not None and issubclass(class_obj, class_candidate): - expected_class_obj = class_candidate - - if not issubclass(passed_class_obj[name].__class__, expected_class_obj): - raise ValueError( - f"{passed_class_obj[name]} is of type: {type(passed_class_obj[name])}, but should be" - f" {expected_class_obj}" - ) - elif passed_class_obj[name] is None: - logger.warning( - f"You have passed `None` for {name} to disable its functionality in {pipeline_class}. Note" - f" that this might lead to problems when using {pipeline_class} and is not recommended." - ) - sub_model_should_be_defined = False - else: - logger.warning( - f"You have passed a non-standard module {passed_class_obj[name]}. We cannot verify whether it" - " has the correct type" - ) - - # set passed class object - loaded_sub_model = passed_class_obj[name] - elif is_pipeline_module: - pipeline_module = getattr(pipelines, library_name) - class_obj = import_flax_or_no_model(pipeline_module, class_name) - - importable_classes = ALL_IMPORTABLE_CLASSES - class_candidates = dict.fromkeys(importable_classes.keys(), class_obj) - else: - # else we just import it from the library. - library = importlib.import_module(library_name) - class_obj = import_flax_or_no_model(library, class_name) - - importable_classes = LOADABLE_CLASSES[library_name] - class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()} - - if loaded_sub_model is None and sub_model_should_be_defined: - load_method_name = None - for class_name, class_candidate in class_candidates.items(): - if class_candidate is not None and issubclass(class_obj, class_candidate): - load_method_name = importable_classes[class_name][1] - - load_method = getattr(class_obj, load_method_name) - - # check if the module is in a subdirectory - if os.path.isdir(os.path.join(cached_folder, name)): - loadable_folder = os.path.join(cached_folder, name) - else: - loaded_sub_model = cached_folder - - if issubclass(class_obj, FlaxModelMixin): - loaded_sub_model, loaded_params = load_method( - loadable_folder, - from_pt=from_pt, - use_memory_efficient_attention=use_memory_efficient_attention, - split_head_dim=split_head_dim, - dtype=dtype, - ) - params[name] = loaded_params - elif is_transformers_flax_compatible() and issubclass(class_obj, FlaxPreTrainedModel): - if from_pt: - # TODO(Suraj): Fix this in Transformers. We should be able to use `_do_init=False` here - loaded_sub_model = load_method(loadable_folder, from_pt=from_pt) - loaded_params = loaded_sub_model.params - del loaded_sub_model._params - else: - loaded_sub_model, loaded_params = load_method(loadable_folder, _do_init=False) - params[name] = loaded_params - elif issubclass(class_obj, FlaxSchedulerMixin): - loaded_sub_model, scheduler_state = load_method(loadable_folder) - params[name] = scheduler_state - else: - loaded_sub_model = load_method(loadable_folder) - - init_kwargs[name] = loaded_sub_model # UNet(...), # DiffusionSchedule(...) - - # 4. Potentially add passed objects if expected - missing_modules = set(expected_modules) - set(init_kwargs.keys()) - passed_modules = list(passed_class_obj.keys()) - - if len(missing_modules) > 0 and missing_modules <= set(passed_modules): - for module in missing_modules: - init_kwargs[module] = passed_class_obj.get(module, None) - elif len(missing_modules) > 0: - passed_modules = set(list(init_kwargs.keys()) + list(passed_class_obj.keys())) - optional_kwargs - raise ValueError( - f"Pipeline {pipeline_class} expected {expected_modules}, but only {passed_modules} were passed." - ) - - model = pipeline_class(**init_kwargs, dtype=dtype) - return model, params - - @classmethod - def _get_signature_keys(cls, obj): - parameters = inspect.signature(obj.__init__).parameters - required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty} - optional_parameters = set({k for k, v in parameters.items() if v.default != inspect._empty}) - expected_modules = set(required_parameters.keys()) - {"self"} - - return expected_modules, optional_parameters - - @property - def components(self) -> dict[str, Any]: - r""" - - The `self.components` property can be useful to run different pipelines with the same weights and - configurations to not have to re-allocate memory. - - Examples: - - ```py - >>> from diffusers import ( - ... FlaxStableDiffusionPipeline, - ... FlaxStableDiffusionImg2ImgPipeline, - ... ) - - >>> text2img = FlaxStableDiffusionPipeline.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", variant="bf16", dtype=jnp.bfloat16 - ... ) - >>> img2img = FlaxStableDiffusionImg2ImgPipeline(**text2img.components) - ``` - - Returns: - A dictionary containing all the modules needed to initialize the pipeline. - """ - expected_modules, optional_parameters = self._get_signature_keys(self) - components = { - k: getattr(self, k) for k in self.config.keys() if not k.startswith("_") and k not in optional_parameters - } - - if set(components.keys()) != expected_modules: - raise ValueError( - f"{self} has been incorrectly initialized or {self.__class__} is incorrectly implemented. Expected" - f" {expected_modules} to be defined, but {components} are defined." - ) - - return components - - @staticmethod - def numpy_to_pil(images): - """ - Convert a NumPy image or a batch of images to a PIL image. - """ - if images.ndim == 3: - images = images[None, ...] - images = (images * 255).round().astype("uint8") - if images.shape[-1] == 1: - # special case for grayscale (single channel) images - pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images] - else: - pil_images = [Image.fromarray(image) for image in images] - - return pil_images - - # TODO: make it compatible with jax.lax - def progress_bar(self, iterable): - if not hasattr(self, "_progress_bar_config"): - self._progress_bar_config = {} - elif not isinstance(self._progress_bar_config, dict): - raise ValueError( - f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}." - ) - - return tqdm(iterable, **self._progress_bar_config) - - def set_progress_bar_config(self, **kwargs): - self._progress_bar_config = kwargs diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index 90cbffc5b69d..e0c708bfc433 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/src/diffusers/pipelines/pipeline_loading_utils.py @@ -29,7 +29,6 @@ from .. import __version__ from ..utils import ( FLASHPACK_WEIGHTS_NAME, - FLAX_WEIGHTS_NAME, ONNX_EXTERNAL_WEIGHTS_NAME, ONNX_WEIGHTS_NAME, SAFETENSORS_WEIGHTS_NAME, @@ -53,8 +52,6 @@ from transformers.utils import SAFE_WEIGHTS_NAME as TRANSFORMERS_SAFE_WEIGHTS_NAME from transformers.utils import WEIGHTS_NAME as TRANSFORMERS_WEIGHTS_NAME - if is_transformers_version("<=", "4.56.2"): - from transformers.utils import FLAX_WEIGHTS_NAME as TRANSFORMERS_FLAX_WEIGHTS_NAME if is_accelerate_available(): import accelerate @@ -124,15 +121,12 @@ def is_safetensors_compatible(filenames, passed_components=None, folder_names=No weight_names = [ WEIGHTS_NAME, SAFETENSORS_WEIGHTS_NAME, - FLAX_WEIGHTS_NAME, ONNX_WEIGHTS_NAME, ONNX_EXTERNAL_WEIGHTS_NAME, ] if is_transformers_available(): weight_names += [TRANSFORMERS_WEIGHTS_NAME, TRANSFORMERS_SAFE_WEIGHTS_NAME] - if is_transformers_version("<=", "4.56.2"): - weight_names += [TRANSFORMERS_FLAX_WEIGHTS_NAME] # model_pytorch, diffusion_model_pytorch, ... weight_prefixes = [w.split(".")[0] for w in weight_names] @@ -207,7 +201,6 @@ def filter_model_files(filenames): weight_names = [ WEIGHTS_NAME, SAFETENSORS_WEIGHTS_NAME, - FLAX_WEIGHTS_NAME, ONNX_WEIGHTS_NAME, ONNX_EXTERNAL_WEIGHTS_NAME, FLASHPACK_WEIGHTS_NAME, @@ -215,8 +208,6 @@ def filter_model_files(filenames): if is_transformers_available(): weight_names += [TRANSFORMERS_WEIGHTS_NAME, TRANSFORMERS_SAFE_WEIGHTS_NAME] - if is_transformers_version("<=", "4.56.2"): - weight_names += [TRANSFORMERS_FLAX_WEIGHTS_NAME] allowed_extensions = [wn.split(".")[-1] for wn in weight_names] @@ -231,15 +222,12 @@ def variant_compatible_siblings(filenames, variant=None, ignore_patterns=None) - weight_names = [ WEIGHTS_NAME, SAFETENSORS_WEIGHTS_NAME, - FLAX_WEIGHTS_NAME, ONNX_WEIGHTS_NAME, ONNX_EXTERNAL_WEIGHTS_NAME, ] if is_transformers_available(): weight_names += [TRANSFORMERS_WEIGHTS_NAME, TRANSFORMERS_SAFE_WEIGHTS_NAME] - if is_transformers_version("<=", "4.56.2"): - weight_names += [TRANSFORMERS_FLAX_WEIGHTS_NAME] # model_pytorch, diffusion_model_pytorch, ... weight_prefixes = [w.split(".")[0] for w in weight_names] @@ -536,8 +524,6 @@ def _get_pipeline_class( "The class name could not be found in the configuration file. Please make sure to pass the correct `class_name`." ) - class_name = class_name[4:] if class_name.startswith("Flax") else class_name - pipeline_cls = getattr(diffusers_module, class_name) if load_connected_pipeline: @@ -682,9 +668,6 @@ def _get_final_device_map(device_map, pipeline_class, passed_class_obj, init_dic # Load each module in the pipeline on a meta device so that we can derive the device map. init_empty_modules = {} for name, (library_name, class_name) in init_dict.items(): - if class_name.startswith("Flax"): - raise ValueError("Flax pipelines are not supported with `device_map`.") - # Define all importable classes is_pipeline_module = hasattr(pipelines, library_name) importable_classes = ALL_IMPORTABLE_CLASSES @@ -785,7 +768,6 @@ def load_sub_model( offload_state_dict: bool, model_variants: dict[str, str], name: str, - from_flax: bool, variant: str, low_cpu_mem_usage: bool, cached_folder: str | os.PathLike, @@ -878,9 +860,6 @@ def load_sub_model( if is_diffusers_model: loading_kwargs["use_flashpack"] = use_flashpack - if from_flax: - loading_kwargs["from_flax"] = True - # the following can be deleted once the minimum required `transformers` version # is higher than 4.27 if ( @@ -894,11 +873,7 @@ def load_sub_model( elif is_transformers_model and loading_kwargs["variant"] is None: loading_kwargs.pop("variant") - # if `from_flax` and model is transformer model, can currently not load with `low_cpu_mem_usage` - if not (from_flax and is_transformers_model): - loading_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage - else: - loading_kwargs["low_cpu_mem_usage"] = False + loading_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage if is_diffusers_model: loading_kwargs["disable_mmap"] = disable_mmap @@ -1129,7 +1104,6 @@ def _get_ignore_patterns( model_folder_names: list[str], model_filenames: list[str], use_safetensors: bool, - from_flax: bool, allow_pickle: bool, use_onnx: bool, is_onnx: bool, @@ -1147,10 +1121,7 @@ def _get_ignore_patterns( f"Could not find the necessary `safetensors` weights in {model_filenames} (variant={variant})" ) - if from_flax: - ignore_patterns = ["*.bin", "*.safetensors", "*.onnx", "*.pb"] - - elif use_safetensors and is_safetensors_compatible( + if use_safetensors and is_safetensors_compatible( model_filenames, passed_components=passed_components, folder_names=model_folder_names, variant=variant ): ignore_patterns = ["*.bin", "*.msgpack"] diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 7b653caa0657..64893ce80c42 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -727,8 +727,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike, **kwa class). The overwritten components are passed directly to the pipelines `__init__` method. See example below for more information. variant (`str`, *optional*): - Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. + Load weights from a specified variant filename such as `"fp16"` or `"ema"`. dduf_file(`str`, *optional*): Load weights from the specified dduf file. disable_mmap ('bool', *optional*, defaults to 'False'): @@ -767,7 +766,6 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike, **kwa local_files_only = kwargs.pop("local_files_only", None) token = kwargs.pop("token", None) revision = kwargs.pop("revision", None) - from_flax = kwargs.pop("from_flax", False) torch_dtype = kwargs.pop("torch_dtype", None) custom_pipeline = kwargs.pop("custom_pipeline", None) custom_revision = kwargs.pop("custom_revision", None) @@ -864,7 +862,6 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike, **kwa local_files_only=local_files_only, token=token, revision=revision, - from_flax=from_flax, use_safetensors=use_safetensors, use_onnx=use_onnx, custom_pipeline=custom_pipeline, @@ -975,14 +972,6 @@ def load_module(name, value): init_dict = {k: v for k, v in init_dict.items() if load_module(k, v)} - # Special case: safety_checker must be loaded separately when using `from_flax` - if from_flax and "safety_checker" in init_dict and "safety_checker" not in passed_class_obj: - raise NotImplementedError( - "The safety checker cannot be automatically loaded when loading weights `from_flax`." - " Please, pass `safety_checker=None` to `from_pretrained`, and load the safety checker" - " separately if you need it." - ) - # 5. Throw nice warnings / errors for fast accelerate loading if len(unused_kwargs) > 0: logger.warning( @@ -1030,15 +1019,12 @@ def load_module(name, value): elif isinstance(final_device_map, str): current_device_map = final_device_map - # 7.2 - now that JAX/Flax is an official framework of the library, we might load from Flax names - class_name = class_name[4:] if class_name.startswith("Flax") else class_name - - # 7.3 Define all importable classes + # 7.2 Define all importable classes is_pipeline_module = hasattr(pipelines, library_name) importable_classes = ALL_IMPORTABLE_CLASSES loaded_sub_model = None - # 7.4 Use passed sub model or load class_name from library_name + # 7.3 Use passed sub model or load class_name from library_name if name in passed_class_obj: # if the model is in a pipeline module, then we load it from the pipeline # check that passed_class_obj has correct parent class @@ -1070,7 +1056,6 @@ def load_module(name, value): offload_state_dict=offload_state_dict, model_variants=model_variants, name=name, - from_flax=from_flax, variant=variant, low_cpu_mem_usage=low_cpu_mem_usage, cached_folder=cached_folder, @@ -1569,8 +1554,7 @@ def download(cls, pretrained_model_name, **kwargs) -> str | os.PathLike: guarantee the timeliness or safety of the source, and you should refer to the mirror site for more information. variant (`str`, *optional*): - Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. + Load weights from a specified variant filename such as `"fp16"` or `"ema"`. dduf_file(`str`, *optional*): Load weights from the specified DDUF file. use_safetensors (`bool`, *optional*, defaults to `None`): @@ -1604,7 +1588,6 @@ def download(cls, pretrained_model_name, **kwargs) -> str | os.PathLike: local_files_only = kwargs.pop("local_files_only", None) token = kwargs.pop("token", None) revision = kwargs.pop("revision", None) - from_flax = kwargs.pop("from_flax", False) custom_pipeline = kwargs.pop("custom_pipeline", None) custom_revision = kwargs.pop("custom_revision", None) variant = kwargs.pop("variant", None) @@ -1715,7 +1698,6 @@ def download(cls, pretrained_model_name, **kwargs) -> str | os.PathLike: model_folder_names, filenames, use_safetensors, - from_flax, allow_pickle, use_onnx, pipeline_class._is_onnx, @@ -1794,7 +1776,6 @@ def download(cls, pretrained_model_name, **kwargs) -> str | os.PathLike: ) cls_name = cls.load_config(os.path.join(cached_folder, "model_index.json")).get("_class_name", None) - cls_name = cls_name[4:] if isinstance(cls_name, str) and cls_name.startswith("Flax") else cls_name diffusers_module = importlib.import_module(__name__.split(".")[0]) pipeline_class = getattr(diffusers_module, cls_name, None) if isinstance(cls_name, str) else None diff --git a/src/diffusers/pipelines/stable_diffusion/__init__.py b/src/diffusers/pipelines/stable_diffusion/__init__.py index 8acdd219423a..263457a76489 100644 --- a/src/diffusers/pipelines/stable_diffusion/__init__.py +++ b/src/diffusers/pipelines/stable_diffusion/__init__.py @@ -8,7 +8,6 @@ is_onnx_available, is_torch_available, is_transformers_available, - is_transformers_flax_compatible, is_transformers_version, ) @@ -17,8 +16,6 @@ _additional_imports = {} _import_structure = {"pipeline_output": ["StableDiffusionPipelineOutput"]} -if is_transformers_flax_compatible(): - _import_structure["pipeline_output"].extend(["FlaxStableDiffusionPipelineOutput"]) try: if not (is_transformers_available() and is_torch_available()): raise OptionalDependencyNotAvailable() @@ -82,15 +79,6 @@ _import_structure["pipeline_onnx_stable_diffusion_inpaint_legacy"] = ["OnnxStableDiffusionInpaintPipelineLegacy"] _import_structure["pipeline_onnx_stable_diffusion_upscale"] = ["OnnxStableDiffusionUpscalePipeline"] -if is_transformers_flax_compatible(): - from ...schedulers.scheduling_pndm_flax import PNDMSchedulerState - - _additional_imports.update({"PNDMSchedulerState": PNDMSchedulerState}) - _import_structure["pipeline_flax_stable_diffusion"] = ["FlaxStableDiffusionPipeline"] - _import_structure["pipeline_flax_stable_diffusion_img2img"] = ["FlaxStableDiffusionImg2ImgPipeline"] - _import_structure["pipeline_flax_stable_diffusion_inpaint"] = ["FlaxStableDiffusionInpaintPipeline"] - _import_structure["safety_checker_flax"] = ["FlaxStableDiffusionSafetyChecker"] - if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: if not (is_transformers_available() and is_torch_available()): @@ -161,22 +149,6 @@ OnnxStableDiffusionUpscalePipeline, ) - try: - if not is_transformers_flax_compatible(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_flax_objects import * - else: - from .pipeline_flax_stable_diffusion import FlaxStableDiffusionPipeline - from .pipeline_flax_stable_diffusion_img2img import ( - FlaxStableDiffusionImg2ImgPipeline, - ) - from .pipeline_flax_stable_diffusion_inpaint import ( - FlaxStableDiffusionInpaintPipeline, - ) - from .pipeline_output import FlaxStableDiffusionPipelineOutput - from .safety_checker_flax import FlaxStableDiffusionSafetyChecker - else: import sys diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py deleted file mode 100644 index d4be6b384c1c..000000000000 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py +++ /dev/null @@ -1,470 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import warnings -from functools import partial - -import jax -import jax.numpy as jnp -import numpy as np -from flax.core.frozen_dict import FrozenDict -from flax.jax_utils import unreplicate -from flax.training.common_utils import shard -from packaging import version -from PIL import Image -from transformers import CLIPImageProcessor, CLIPTokenizer, FlaxCLIPTextModel - -from ...models import FlaxAutoencoderKL, FlaxUNet2DConditionModel -from ...schedulers import ( - FlaxDDIMScheduler, - FlaxDPMSolverMultistepScheduler, - FlaxLMSDiscreteScheduler, - FlaxPNDMScheduler, -) -from ...utils import deprecate, logging, replace_example_docstring -from ..pipeline_flax_utils import FlaxDiffusionPipeline -from .pipeline_output import FlaxStableDiffusionPipelineOutput -from .safety_checker_flax import FlaxStableDiffusionSafetyChecker - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -# Set to True to use python for loop instead of jax.fori_loop for easier debugging -DEBUG = False - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import jax - >>> import numpy as np - >>> from flax.jax_utils import replicate - >>> from flax.training.common_utils import shard - - >>> from diffusers import FlaxStableDiffusionPipeline - - >>> pipeline, params = FlaxStableDiffusionPipeline.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", variant="bf16", dtype=jax.numpy.bfloat16 - ... ) - - >>> prompt = "a photo of an astronaut riding a horse on mars" - - >>> prng_seed = jax.random.PRNGKey(0) - >>> num_inference_steps = 50 - - >>> num_samples = jax.device_count() - >>> prompt = num_samples * [prompt] - >>> prompt_ids = pipeline.prepare_inputs(prompt) - # shard inputs and rng - - >>> params = replicate(params) - >>> prng_seed = jax.random.split(prng_seed, jax.device_count()) - >>> prompt_ids = shard(prompt_ids) - - >>> images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images - >>> images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]))) - ``` -""" - - -class FlaxStableDiffusionPipeline(FlaxDiffusionPipeline): - r""" - Flax-based pipeline for text-to-image generation using Stable Diffusion. - - This model inherits from [`FlaxDiffusionPipeline`]. Check the superclass documentation for the generic methods - implemented for all pipelines (downloading, saving, running on a particular device, etc.). - - Args: - vae ([`FlaxAutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.FlaxCLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`FlaxUNet2DConditionModel`]): - A `FlaxUNet2DConditionModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of - [`FlaxDDIMScheduler`], [`FlaxLMSDiscreteScheduler`], [`FlaxPNDMScheduler`], or - [`FlaxDPMSolverMultistepScheduler`]. - safety_checker ([`FlaxStableDiffusionSafetyChecker`]): - Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - def __init__( - self, - vae: FlaxAutoencoderKL, - text_encoder: FlaxCLIPTextModel, - tokenizer: CLIPTokenizer, - unet: FlaxUNet2DConditionModel, - scheduler: FlaxDDIMScheduler | FlaxPNDMScheduler | FlaxLMSDiscreteScheduler | FlaxDPMSolverMultistepScheduler, - safety_checker: FlaxStableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - dtype: jnp.dtype = jnp.float32, - ): - super().__init__() - self.dtype = dtype - - if safety_checker is None: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" - " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_config) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - - def prepare_inputs(self, prompt: str | list[str]): - if not isinstance(prompt, (str, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - text_input = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="np", - ) - return text_input.input_ids - - def _get_has_nsfw_concepts(self, features, params): - has_nsfw_concepts = self.safety_checker(features, params) - return has_nsfw_concepts - - def _run_safety_checker(self, images, safety_model_params, jit=False): - # safety_model_params should already be replicated when jit is True - pil_images = [Image.fromarray(image) for image in images] - features = self.feature_extractor(pil_images, return_tensors="np").pixel_values - - if jit: - features = shard(features) - has_nsfw_concepts = _p_get_has_nsfw_concepts(self, features, safety_model_params) - has_nsfw_concepts = unshard(has_nsfw_concepts) - safety_model_params = unreplicate(safety_model_params) - else: - has_nsfw_concepts = self._get_has_nsfw_concepts(features, safety_model_params) - - images_was_copied = False - for idx, has_nsfw_concept in enumerate(has_nsfw_concepts): - if has_nsfw_concept: - if not images_was_copied: - images_was_copied = True - images = images.copy() - - images[idx] = np.zeros(images[idx].shape, dtype=np.uint8) # black image - - if any(has_nsfw_concepts): - warnings.warn( - "Potential NSFW content was detected in one or more images. A black image will be returned" - " instead. Try again with a different prompt and/or seed." - ) - - return images, has_nsfw_concepts - - def _generate( - self, - prompt_ids: jnp.array, - params: dict | FrozenDict, - prng_seed: jax.Array, - num_inference_steps: int, - height: int, - width: int, - guidance_scale: float, - latents: jnp.ndarray | None = None, - neg_prompt_ids: jnp.ndarray | None = None, - ): - if height % 8 != 0 or width % 8 != 0: - raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") - - # get prompt text embeddings - prompt_embeds = self.text_encoder(prompt_ids, params=params["text_encoder"])[0] - - # TODO: currently it is assumed `do_classifier_free_guidance = guidance_scale > 1.0` - # implement this conditional `do_classifier_free_guidance = guidance_scale > 1.0` - batch_size = prompt_ids.shape[0] - - max_length = prompt_ids.shape[-1] - - if neg_prompt_ids is None: - uncond_input = self.tokenizer( - [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="np" - ).input_ids - else: - uncond_input = neg_prompt_ids - negative_prompt_embeds = self.text_encoder(uncond_input, params=params["text_encoder"])[0] - context = jnp.concatenate([negative_prompt_embeds, prompt_embeds]) - - # Ensure model output will be `float32` before going into the scheduler - guidance_scale = jnp.array([guidance_scale], dtype=jnp.float32) - - latents_shape = ( - batch_size, - self.unet.config.in_channels, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - ) - if latents is None: - latents = jax.random.normal(prng_seed, shape=latents_shape, dtype=jnp.float32) - else: - if latents.shape != latents_shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") - - def loop_body(step, args): - latents, scheduler_state = args - # For classifier free guidance, we need to do two forward passes. - # Here we concatenate the unconditional and text embeddings into a single batch - # to avoid doing two forward passes - latents_input = jnp.concatenate([latents] * 2) - - t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step] - timestep = jnp.broadcast_to(t, latents_input.shape[0]) - - latents_input = self.scheduler.scale_model_input(scheduler_state, latents_input, t) - - # predict the noise residual - noise_pred = self.unet.apply( - {"params": params["unet"]}, - jnp.array(latents_input), - jnp.array(timestep, dtype=jnp.int32), - encoder_hidden_states=context, - ).sample - # perform guidance - noise_pred_uncond, noise_prediction_text = jnp.split(noise_pred, 2, axis=0) - noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents, scheduler_state = self.scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() - return latents, scheduler_state - - scheduler_state = self.scheduler.set_timesteps( - params["scheduler"], num_inference_steps=num_inference_steps, shape=latents.shape - ) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * params["scheduler"].init_noise_sigma - - if DEBUG: - # run with python for loop - for i in range(num_inference_steps): - latents, scheduler_state = loop_body(i, (latents, scheduler_state)) - else: - latents, _ = jax.lax.fori_loop(0, num_inference_steps, loop_body, (latents, scheduler_state)) - - # scale and decode the image latents with vae - latents = 1 / self.vae.config.scaling_factor * latents - image = self.vae.apply({"params": params["vae"]}, latents, method=self.vae.decode).sample - - image = (image / 2 + 0.5).clip(0, 1).transpose(0, 2, 3, 1) - return image - - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt_ids: jnp.array, - params: dict | FrozenDict, - prng_seed: jax.Array, - num_inference_steps: int = 50, - height: int | None = None, - width: int | None = None, - guidance_scale: float | jnp.ndarray = 7.5, - latents: jnp.ndarray = None, - neg_prompt_ids: jnp.ndarray = None, - return_dict: bool = True, - jit: bool = False, - ): - r""" - The call function to the pipeline for generation. - - Args: - prompt (`str` or `list[str]`, *optional*): - The prompt or prompts to guide image generation. - height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The width in pixels of the generated image. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - guidance_scale (`float`, *optional*, defaults to 7.5): - A higher guidance scale value encourages the model to generate images closely linked to the text - `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - latents (`jnp.ndarray`, *optional*): - Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - array is generated by sampling using the supplied random `generator`. - jit (`bool`, defaults to `False`): - Whether to run `pmap` versions of the generation and safety scoring functions. - - > [!WARNING] > This argument exists because `__call__` is not yet end-to-end pmap-able. It will be - removed in a > future release. - - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] instead of - a plain tuple. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] is - returned, otherwise a `tuple` is returned where the first element is a list with the generated images - and the second element is a list of `bool`s indicating whether the corresponding generated image - contains "not-safe-for-work" (nsfw) content. - """ - # 0. Default height and width to unet - height = height or self.unet.config.sample_size * self.vae_scale_factor - width = width or self.unet.config.sample_size * self.vae_scale_factor - - if isinstance(guidance_scale, float): - # Convert to a tensor so each device gets a copy. Follow the prompt_ids for - # shape information, as they may be sharded (when `jit` is `True`), or not. - guidance_scale = jnp.array([guidance_scale] * prompt_ids.shape[0]) - if len(prompt_ids.shape) > 2: - # Assume sharded - guidance_scale = guidance_scale[:, None] - - if jit: - images = _p_generate( - self, - prompt_ids, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - ) - else: - images = self._generate( - prompt_ids, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - ) - - if self.safety_checker is not None: - safety_params = params["safety_checker"] - images_uint8_casted = (images * 255).round().astype("uint8") - num_devices, batch_size = images.shape[:2] - - images_uint8_casted = np.asarray(images_uint8_casted).reshape(num_devices * batch_size, height, width, 3) - images_uint8_casted, has_nsfw_concept = self._run_safety_checker(images_uint8_casted, safety_params, jit) - images = np.asarray(images).copy() - - # block images - if any(has_nsfw_concept): - for i, is_nsfw in enumerate(has_nsfw_concept): - if is_nsfw: - images[i, 0] = np.asarray(images_uint8_casted[i]) - - images = images.reshape(num_devices, batch_size, height, width, 3) - else: - images = np.asarray(images) - has_nsfw_concept = False - - if not return_dict: - return (images, has_nsfw_concept) - - return FlaxStableDiffusionPipelineOutput(images=images, nsfw_content_detected=has_nsfw_concept) - - -# Static argnums are pipe, num_inference_steps, height, width. A change would trigger recompilation. -# Non-static args are (sharded) input tensors mapped over their first dimension (hence, `0`). -@partial( - jax.pmap, - in_axes=(None, 0, 0, 0, None, None, None, 0, 0, 0), - static_broadcasted_argnums=(0, 4, 5, 6), -) -def _p_generate( - pipe, - prompt_ids, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, -): - return pipe._generate( - prompt_ids, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - ) - - -@partial(jax.pmap, static_broadcasted_argnums=(0,)) -def _p_get_has_nsfw_concepts(pipe, features, params): - return pipe._get_has_nsfw_concepts(features, params) - - -def unshard(x: jnp.ndarray): - # einops.rearrange(x, 'd b ... -> (d b) ...') - num_devices, batch_size = x.shape[:2] - rest = x.shape[2:] - return x.reshape(num_devices * batch_size, *rest) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py deleted file mode 100644 index 291d6187f294..000000000000 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py +++ /dev/null @@ -1,525 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import warnings -from functools import partial - -import jax -import jax.numpy as jnp -import numpy as np -from flax.core.frozen_dict import FrozenDict -from flax.jax_utils import unreplicate -from flax.training.common_utils import shard -from PIL import Image -from transformers import CLIPImageProcessor, CLIPTokenizer, FlaxCLIPTextModel - -from ...models import FlaxAutoencoderKL, FlaxUNet2DConditionModel -from ...schedulers import ( - FlaxDDIMScheduler, - FlaxDPMSolverMultistepScheduler, - FlaxLMSDiscreteScheduler, - FlaxPNDMScheduler, -) -from ...utils import PIL_INTERPOLATION, logging, replace_example_docstring -from ..pipeline_flax_utils import FlaxDiffusionPipeline -from .pipeline_output import FlaxStableDiffusionPipelineOutput -from .safety_checker_flax import FlaxStableDiffusionSafetyChecker - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -# Set to True to use python for loop instead of jax.fori_loop for easier debugging -DEBUG = False - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import jax - >>> import numpy as np - >>> import jax.numpy as jnp - >>> from flax.jax_utils import replicate - >>> from flax.training.common_utils import shard - >>> import requests - >>> from io import BytesIO - >>> from PIL import Image - >>> from diffusers import FlaxStableDiffusionImg2ImgPipeline - - - >>> def create_key(seed=0): - ... return jax.random.PRNGKey(seed) - - - >>> rng = create_key(0) - - >>> url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" - >>> response = requests.get(url) - >>> init_img = Image.open(BytesIO(response.content)).convert("RGB") - >>> init_img = init_img.resize((768, 512)) - - >>> prompts = "A fantasy landscape, trending on artstation" - - >>> pipeline, params = FlaxStableDiffusionImg2ImgPipeline.from_pretrained( - ... "CompVis/stable-diffusion-v1-4", - ... revision="flax", - ... dtype=jnp.bfloat16, - ... ) - - >>> num_samples = jax.device_count() - >>> rng = jax.random.split(rng, jax.device_count()) - >>> prompt_ids, processed_image = pipeline.prepare_inputs( - ... prompt=[prompts] * num_samples, image=[init_img] * num_samples - ... ) - >>> p_params = replicate(params) - >>> prompt_ids = shard(prompt_ids) - >>> processed_image = shard(processed_image) - - >>> output = pipeline( - ... prompt_ids=prompt_ids, - ... image=processed_image, - ... params=p_params, - ... prng_seed=rng, - ... strength=0.75, - ... num_inference_steps=50, - ... jit=True, - ... height=512, - ... width=768, - ... ).images - - >>> output_images = pipeline.numpy_to_pil(np.asarray(output.reshape((num_samples,) + output.shape[-3:]))) - ``` -""" - - -class FlaxStableDiffusionImg2ImgPipeline(FlaxDiffusionPipeline): - r""" - Flax-based pipeline for text-guided image-to-image generation using Stable Diffusion. - - This model inherits from [`FlaxDiffusionPipeline`]. Check the superclass documentation for the generic methods - implemented for all pipelines (downloading, saving, running on a particular device, etc.). - - Args: - vae ([`FlaxAutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.FlaxCLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`FlaxUNet2DConditionModel`]): - A `FlaxUNet2DConditionModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of - [`FlaxDDIMScheduler`], [`FlaxLMSDiscreteScheduler`], [`FlaxPNDMScheduler`], or - [`FlaxDPMSolverMultistepScheduler`]. - safety_checker ([`FlaxStableDiffusionSafetyChecker`]): - Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - def __init__( - self, - vae: FlaxAutoencoderKL, - text_encoder: FlaxCLIPTextModel, - tokenizer: CLIPTokenizer, - unet: FlaxUNet2DConditionModel, - scheduler: FlaxDDIMScheduler | FlaxPNDMScheduler | FlaxLMSDiscreteScheduler | FlaxDPMSolverMultistepScheduler, - safety_checker: FlaxStableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - dtype: jnp.dtype = jnp.float32, - ): - super().__init__() - self.dtype = dtype - - if safety_checker is None: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - - def prepare_inputs(self, prompt: str | list[str], image: Image.Image | list[Image.Image]): - if not isinstance(prompt, (str, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - if not isinstance(image, (Image.Image, list)): - raise ValueError(f"image has to be of type `PIL.Image.Image` or list but is {type(image)}") - - if isinstance(image, Image.Image): - image = [image] - - processed_images = jnp.concatenate([preprocess(img, jnp.float32) for img in image]) - - text_input = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="np", - ) - return text_input.input_ids, processed_images - - def _get_has_nsfw_concepts(self, features, params): - has_nsfw_concepts = self.safety_checker(features, params) - return has_nsfw_concepts - - def _run_safety_checker(self, images, safety_model_params, jit=False): - # safety_model_params should already be replicated when jit is True - pil_images = [Image.fromarray(image) for image in images] - features = self.feature_extractor(pil_images, return_tensors="np").pixel_values - - if jit: - features = shard(features) - has_nsfw_concepts = _p_get_has_nsfw_concepts(self, features, safety_model_params) - has_nsfw_concepts = unshard(has_nsfw_concepts) - safety_model_params = unreplicate(safety_model_params) - else: - has_nsfw_concepts = self._get_has_nsfw_concepts(features, safety_model_params) - - images_was_copied = False - for idx, has_nsfw_concept in enumerate(has_nsfw_concepts): - if has_nsfw_concept: - if not images_was_copied: - images_was_copied = True - images = images.copy() - - images[idx] = np.zeros(images[idx].shape, dtype=np.uint8) # black image - - if any(has_nsfw_concepts): - warnings.warn( - "Potential NSFW content was detected in one or more images. A black image will be returned" - " instead. Try again with a different prompt and/or seed." - ) - - return images, has_nsfw_concepts - - def get_timestep_start(self, num_inference_steps, strength): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - - return t_start - - def _generate( - self, - prompt_ids: jnp.ndarray, - image: jnp.ndarray, - params: dict | FrozenDict, - prng_seed: jax.Array, - start_timestep: int, - num_inference_steps: int, - height: int, - width: int, - guidance_scale: float, - noise: jnp.ndarray | None = None, - neg_prompt_ids: jnp.ndarray | None = None, - ): - if height % 8 != 0 or width % 8 != 0: - raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") - - # get prompt text embeddings - prompt_embeds = self.text_encoder(prompt_ids, params=params["text_encoder"])[0] - - # TODO: currently it is assumed `do_classifier_free_guidance = guidance_scale > 1.0` - # implement this conditional `do_classifier_free_guidance = guidance_scale > 1.0` - batch_size = prompt_ids.shape[0] - - max_length = prompt_ids.shape[-1] - - if neg_prompt_ids is None: - uncond_input = self.tokenizer( - [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="np" - ).input_ids - else: - uncond_input = neg_prompt_ids - negative_prompt_embeds = self.text_encoder(uncond_input, params=params["text_encoder"])[0] - context = jnp.concatenate([negative_prompt_embeds, prompt_embeds]) - - latents_shape = ( - batch_size, - self.unet.config.in_channels, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - ) - if noise is None: - noise = jax.random.normal(prng_seed, shape=latents_shape, dtype=jnp.float32) - else: - if noise.shape != latents_shape: - raise ValueError(f"Unexpected latents shape, got {noise.shape}, expected {latents_shape}") - - # Create init_latents - init_latent_dist = self.vae.apply({"params": params["vae"]}, image, method=self.vae.encode).latent_dist - init_latents = init_latent_dist.sample(key=prng_seed).transpose((0, 3, 1, 2)) - init_latents = self.vae.config.scaling_factor * init_latents - - def loop_body(step, args): - latents, scheduler_state = args - # For classifier free guidance, we need to do two forward passes. - # Here we concatenate the unconditional and text embeddings into a single batch - # to avoid doing two forward passes - latents_input = jnp.concatenate([latents] * 2) - - t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step] - timestep = jnp.broadcast_to(t, latents_input.shape[0]) - - latents_input = self.scheduler.scale_model_input(scheduler_state, latents_input, t) - - # predict the noise residual - noise_pred = self.unet.apply( - {"params": params["unet"]}, - jnp.array(latents_input), - jnp.array(timestep, dtype=jnp.int32), - encoder_hidden_states=context, - ).sample - # perform guidance - noise_pred_uncond, noise_prediction_text = jnp.split(noise_pred, 2, axis=0) - noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents, scheduler_state = self.scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() - return latents, scheduler_state - - scheduler_state = self.scheduler.set_timesteps( - params["scheduler"], num_inference_steps=num_inference_steps, shape=latents_shape - ) - - latent_timestep = scheduler_state.timesteps[start_timestep : start_timestep + 1].repeat(batch_size) - - latents = self.scheduler.add_noise(params["scheduler"], init_latents, noise, latent_timestep) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * params["scheduler"].init_noise_sigma - - if DEBUG: - # run with python for loop - for i in range(start_timestep, num_inference_steps): - latents, scheduler_state = loop_body(i, (latents, scheduler_state)) - else: - latents, _ = jax.lax.fori_loop(start_timestep, num_inference_steps, loop_body, (latents, scheduler_state)) - - # scale and decode the image latents with vae - latents = 1 / self.vae.config.scaling_factor * latents - image = self.vae.apply({"params": params["vae"]}, latents, method=self.vae.decode).sample - - image = (image / 2 + 0.5).clip(0, 1).transpose(0, 2, 3, 1) - return image - - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt_ids: jnp.ndarray, - image: jnp.ndarray, - params: dict | FrozenDict, - prng_seed: jax.Array, - strength: float = 0.8, - num_inference_steps: int = 50, - height: int | None = None, - width: int | None = None, - guidance_scale: float | jnp.ndarray = 7.5, - noise: jnp.ndarray = None, - neg_prompt_ids: jnp.ndarray = None, - return_dict: bool = True, - jit: bool = False, - ): - r""" - The call function to the pipeline for generation. - - Args: - prompt_ids (`jnp.ndarray`): - The prompt or prompts to guide image generation. - image (`jnp.ndarray`): - Array representing an image batch to be used as the starting point. - params (`Dict` or `FrozenDict`): - Dictionary containing the model parameters/weights. - prng_seed (`jax.Array` or `jax.Array`): - Array containing random number generator key. - strength (`float`, *optional*, defaults to 0.8): - Indicates extent to transform the reference `image`. Must be between 0 and 1. `image` is used as a - starting point and more noise is added the higher the `strength`. The number of denoising steps depends - on the amount of noise initially added. When `strength` is 1, added noise is maximum and the denoising - process runs for the full number of iterations specified in `num_inference_steps`. A value of 1 - essentially ignores `image`. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. This parameter is modulated by `strength`. - height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The width in pixels of the generated image. - guidance_scale (`float`, *optional*, defaults to 7.5): - A higher guidance scale value encourages the model to generate images closely linked to the text - `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - noise (`jnp.ndarray`, *optional*): - Pre-generated noisy latents sampled from a Gaussian distribution to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. The array is generated by - sampling using the supplied random `generator`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] instead of - a plain tuple. - jit (`bool`, defaults to `False`): - Whether to run `pmap` versions of the generation and safety scoring functions. - - > [!WARNING] > This argument exists because `__call__` is not yet end-to-end pmap-able. It will be - removed in a > future release. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] is - returned, otherwise a `tuple` is returned where the first element is a list with the generated images - and the second element is a list of `bool`s indicating whether the corresponding generated image - contains "not-safe-for-work" (nsfw) content. - """ - # 0. Default height and width to unet - height = height or self.unet.config.sample_size * self.vae_scale_factor - width = width or self.unet.config.sample_size * self.vae_scale_factor - - if isinstance(guidance_scale, float): - # Convert to a tensor so each device gets a copy. Follow the prompt_ids for - # shape information, as they may be sharded (when `jit` is `True`), or not. - guidance_scale = jnp.array([guidance_scale] * prompt_ids.shape[0]) - if len(prompt_ids.shape) > 2: - # Assume sharded - guidance_scale = guidance_scale[:, None] - - start_timestep = self.get_timestep_start(num_inference_steps, strength) - - if jit: - images = _p_generate( - self, - prompt_ids, - image, - params, - prng_seed, - start_timestep, - num_inference_steps, - height, - width, - guidance_scale, - noise, - neg_prompt_ids, - ) - else: - images = self._generate( - prompt_ids, - image, - params, - prng_seed, - start_timestep, - num_inference_steps, - height, - width, - guidance_scale, - noise, - neg_prompt_ids, - ) - - if self.safety_checker is not None: - safety_params = params["safety_checker"] - images_uint8_casted = (images * 255).round().astype("uint8") - num_devices, batch_size = images.shape[:2] - - images_uint8_casted = np.asarray(images_uint8_casted).reshape(num_devices * batch_size, height, width, 3) - images_uint8_casted, has_nsfw_concept = self._run_safety_checker(images_uint8_casted, safety_params, jit) - images = np.asarray(images) - - # block images - if any(has_nsfw_concept): - for i, is_nsfw in enumerate(has_nsfw_concept): - if is_nsfw: - images[i] = np.asarray(images_uint8_casted[i]) - - images = images.reshape(num_devices, batch_size, height, width, 3) - else: - images = np.asarray(images) - has_nsfw_concept = False - - if not return_dict: - return (images, has_nsfw_concept) - - return FlaxStableDiffusionPipelineOutput(images=images, nsfw_content_detected=has_nsfw_concept) - - -# Static argnums are pipe, start_timestep, num_inference_steps, height, width. A change would trigger recompilation. -# Non-static args are (sharded) input tensors mapped over their first dimension (hence, `0`). -@partial( - jax.pmap, - in_axes=(None, 0, 0, 0, 0, None, None, None, None, 0, 0, 0), - static_broadcasted_argnums=(0, 5, 6, 7, 8), -) -def _p_generate( - pipe, - prompt_ids, - image, - params, - prng_seed, - start_timestep, - num_inference_steps, - height, - width, - guidance_scale, - noise, - neg_prompt_ids, -): - return pipe._generate( - prompt_ids, - image, - params, - prng_seed, - start_timestep, - num_inference_steps, - height, - width, - guidance_scale, - noise, - neg_prompt_ids, - ) - - -@partial(jax.pmap, static_broadcasted_argnums=(0,)) -def _p_get_has_nsfw_concepts(pipe, features, params): - return pipe._get_has_nsfw_concepts(features, params) - - -def unshard(x: jnp.ndarray): - # einops.rearrange(x, 'd b ... -> (d b) ...') - num_devices, batch_size = x.shape[:2] - rest = x.shape[2:] - return x.reshape(num_devices * batch_size, *rest) - - -def preprocess(image, dtype): - w, h = image.size - w, h = (x - x % 32 for x in (w, h)) # resize to integer multiple of 32 - image = image.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]) - image = jnp.array(image).astype(dtype) / 255.0 - image = image[None].transpose(0, 3, 1, 2) - return 2.0 * image - 1.0 diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py deleted file mode 100644 index b1ffffb5a62c..000000000000 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py +++ /dev/null @@ -1,582 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import warnings -from functools import partial - -import jax -import jax.numpy as jnp -import numpy as np -from flax.core.frozen_dict import FrozenDict -from flax.jax_utils import unreplicate -from flax.training.common_utils import shard -from packaging import version -from PIL import Image -from transformers import CLIPImageProcessor, CLIPTokenizer, FlaxCLIPTextModel - -from ...models import FlaxAutoencoderKL, FlaxUNet2DConditionModel -from ...schedulers import ( - FlaxDDIMScheduler, - FlaxDPMSolverMultistepScheduler, - FlaxLMSDiscreteScheduler, - FlaxPNDMScheduler, -) -from ...utils import PIL_INTERPOLATION, deprecate, logging, replace_example_docstring -from ..pipeline_flax_utils import FlaxDiffusionPipeline -from .pipeline_output import FlaxStableDiffusionPipelineOutput -from .safety_checker_flax import FlaxStableDiffusionSafetyChecker - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -# Set to True to use python for loop instead of jax.fori_loop for easier debugging -DEBUG = False - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import jax - >>> import numpy as np - >>> from flax.jax_utils import replicate - >>> from flax.training.common_utils import shard - >>> import PIL - >>> import requests - >>> from io import BytesIO - >>> from diffusers import FlaxStableDiffusionInpaintPipeline - - - >>> def download_image(url): - ... response = requests.get(url) - ... return PIL.Image.open(BytesIO(response.content)).convert("RGB") - - - >>> img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png" - >>> mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png" - - >>> init_image = download_image(img_url).resize((512, 512)) - >>> mask_image = download_image(mask_url).resize((512, 512)) - - >>> pipeline, params = FlaxStableDiffusionInpaintPipeline.from_pretrained( - ... "xvjiarui/stable-diffusion-2-inpainting" - ... ) - - >>> prompt = "Face of a yellow cat, high resolution, sitting on a park bench" - >>> prng_seed = jax.random.PRNGKey(0) - >>> num_inference_steps = 50 - - >>> num_samples = jax.device_count() - >>> prompt = num_samples * [prompt] - >>> init_image = num_samples * [init_image] - >>> mask_image = num_samples * [mask_image] - >>> prompt_ids, processed_masked_images, processed_masks = pipeline.prepare_inputs( - ... prompt, init_image, mask_image - ... ) - # shard inputs and rng - - >>> params = replicate(params) - >>> prng_seed = jax.random.split(prng_seed, jax.device_count()) - >>> prompt_ids = shard(prompt_ids) - >>> processed_masked_images = shard(processed_masked_images) - >>> processed_masks = shard(processed_masks) - - >>> images = pipeline( - ... prompt_ids, processed_masks, processed_masked_images, params, prng_seed, num_inference_steps, jit=True - ... ).images - >>> images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]))) - ``` -""" - - -class FlaxStableDiffusionInpaintPipeline(FlaxDiffusionPipeline): - r""" - Flax-based pipeline for text-guided image inpainting using Stable Diffusion. - - > [!WARNING] > 🧪 This is an experimental feature! - - This model inherits from [`FlaxDiffusionPipeline`]. Check the superclass documentation for the generic methods - implemented for all pipelines (downloading, saving, running on a particular device, etc.). - - Args: - vae ([`FlaxAutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.FlaxCLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`FlaxUNet2DConditionModel`]): - A `FlaxUNet2DConditionModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of - [`FlaxDDIMScheduler`], [`FlaxLMSDiscreteScheduler`], [`FlaxPNDMScheduler`], or - [`FlaxDPMSolverMultistepScheduler`]. - safety_checker ([`FlaxStableDiffusionSafetyChecker`]): - Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - def __init__( - self, - vae: FlaxAutoencoderKL, - text_encoder: FlaxCLIPTextModel, - tokenizer: CLIPTokenizer, - unet: FlaxUNet2DConditionModel, - scheduler: FlaxDDIMScheduler | FlaxPNDMScheduler | FlaxLMSDiscreteScheduler | FlaxDPMSolverMultistepScheduler, - safety_checker: FlaxStableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - dtype: jnp.dtype = jnp.float32, - ): - super().__init__() - self.dtype = dtype - - if safety_checker is None: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" - " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_config) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - - def prepare_inputs( - self, - prompt: str | list[str], - image: Image.Image | list[Image.Image], - mask: Image.Image | list[Image.Image], - ): - if not isinstance(prompt, (str, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - if not isinstance(image, (Image.Image, list)): - raise ValueError(f"image has to be of type `PIL.Image.Image` or list but is {type(image)}") - - if isinstance(image, Image.Image): - image = [image] - - if not isinstance(mask, (Image.Image, list)): - raise ValueError(f"image has to be of type `PIL.Image.Image` or list but is {type(image)}") - - if isinstance(mask, Image.Image): - mask = [mask] - - processed_images = jnp.concatenate([preprocess_image(img, jnp.float32) for img in image]) - processed_masks = jnp.concatenate([preprocess_mask(m, jnp.float32) for m in mask]) - # processed_masks[processed_masks < 0.5] = 0 - processed_masks = processed_masks.at[processed_masks < 0.5].set(0) - # processed_masks[processed_masks >= 0.5] = 1 - processed_masks = processed_masks.at[processed_masks >= 0.5].set(1) - - processed_masked_images = processed_images * (processed_masks < 0.5) - - text_input = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="np", - ) - return text_input.input_ids, processed_masked_images, processed_masks - - def _get_has_nsfw_concepts(self, features, params): - has_nsfw_concepts = self.safety_checker(features, params) - return has_nsfw_concepts - - def _run_safety_checker(self, images, safety_model_params, jit=False): - # safety_model_params should already be replicated when jit is True - pil_images = [Image.fromarray(image) for image in images] - features = self.feature_extractor(pil_images, return_tensors="np").pixel_values - - if jit: - features = shard(features) - has_nsfw_concepts = _p_get_has_nsfw_concepts(self, features, safety_model_params) - has_nsfw_concepts = unshard(has_nsfw_concepts) - safety_model_params = unreplicate(safety_model_params) - else: - has_nsfw_concepts = self._get_has_nsfw_concepts(features, safety_model_params) - - images_was_copied = False - for idx, has_nsfw_concept in enumerate(has_nsfw_concepts): - if has_nsfw_concept: - if not images_was_copied: - images_was_copied = True - images = images.copy() - - images[idx] = np.zeros(images[idx].shape, dtype=np.uint8) # black image - - if any(has_nsfw_concepts): - warnings.warn( - "Potential NSFW content was detected in one or more images. A black image will be returned" - " instead. Try again with a different prompt and/or seed." - ) - - return images, has_nsfw_concepts - - def _generate( - self, - prompt_ids: jnp.ndarray, - mask: jnp.ndarray, - masked_image: jnp.ndarray, - params: dict | FrozenDict, - prng_seed: jax.Array, - num_inference_steps: int, - height: int, - width: int, - guidance_scale: float, - latents: jnp.ndarray | None = None, - neg_prompt_ids: jnp.ndarray | None = None, - ): - if height % 8 != 0 or width % 8 != 0: - raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") - - # get prompt text embeddings - prompt_embeds = self.text_encoder(prompt_ids, params=params["text_encoder"])[0] - - # TODO: currently it is assumed `do_classifier_free_guidance = guidance_scale > 1.0` - # implement this conditional `do_classifier_free_guidance = guidance_scale > 1.0` - batch_size = prompt_ids.shape[0] - - max_length = prompt_ids.shape[-1] - - if neg_prompt_ids is None: - uncond_input = self.tokenizer( - [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="np" - ).input_ids - else: - uncond_input = neg_prompt_ids - negative_prompt_embeds = self.text_encoder(uncond_input, params=params["text_encoder"])[0] - context = jnp.concatenate([negative_prompt_embeds, prompt_embeds]) - - latents_shape = ( - batch_size, - self.vae.config.latent_channels, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - ) - if latents is None: - latents = jax.random.normal(prng_seed, shape=latents_shape, dtype=self.dtype) - else: - if latents.shape != latents_shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") - - prng_seed, mask_prng_seed = jax.random.split(prng_seed) - - masked_image_latent_dist = self.vae.apply( - {"params": params["vae"]}, masked_image, method=self.vae.encode - ).latent_dist - masked_image_latents = masked_image_latent_dist.sample(key=mask_prng_seed).transpose((0, 3, 1, 2)) - masked_image_latents = self.vae.config.scaling_factor * masked_image_latents - del mask_prng_seed - - mask = jax.image.resize(mask, (*mask.shape[:-2], *masked_image_latents.shape[-2:]), method="nearest") - - # 8. Check that sizes of mask, masked image and latents match - num_channels_latents = self.vae.config.latent_channels - num_channels_mask = mask.shape[1] - num_channels_masked_image = masked_image_latents.shape[1] - if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: - raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" - f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" - f" = {num_channels_latents + num_channels_masked_image + num_channels_mask}. Please verify the config of" - " `pipeline.unet` or your `mask_image` or `image` input." - ) - - def loop_body(step, args): - latents, mask, masked_image_latents, scheduler_state = args - # For classifier free guidance, we need to do two forward passes. - # Here we concatenate the unconditional and text embeddings into a single batch - # to avoid doing two forward passes - latents_input = jnp.concatenate([latents] * 2) - mask_input = jnp.concatenate([mask] * 2) - masked_image_latents_input = jnp.concatenate([masked_image_latents] * 2) - - t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step] - timestep = jnp.broadcast_to(t, latents_input.shape[0]) - - latents_input = self.scheduler.scale_model_input(scheduler_state, latents_input, t) - # concat latents, mask, masked_image_latents in the channel dimension - latents_input = jnp.concatenate([latents_input, mask_input, masked_image_latents_input], axis=1) - - # predict the noise residual - noise_pred = self.unet.apply( - {"params": params["unet"]}, - jnp.array(latents_input), - jnp.array(timestep, dtype=jnp.int32), - encoder_hidden_states=context, - ).sample - # perform guidance - noise_pred_uncond, noise_prediction_text = jnp.split(noise_pred, 2, axis=0) - noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents, scheduler_state = self.scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() - return latents, mask, masked_image_latents, scheduler_state - - scheduler_state = self.scheduler.set_timesteps( - params["scheduler"], num_inference_steps=num_inference_steps, shape=latents.shape - ) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * params["scheduler"].init_noise_sigma - - if DEBUG: - # run with python for loop - for i in range(num_inference_steps): - latents, mask, masked_image_latents, scheduler_state = loop_body( - i, (latents, mask, masked_image_latents, scheduler_state) - ) - else: - latents, _, _, _ = jax.lax.fori_loop( - 0, num_inference_steps, loop_body, (latents, mask, masked_image_latents, scheduler_state) - ) - - # scale and decode the image latents with vae - latents = 1 / self.vae.config.scaling_factor * latents - image = self.vae.apply({"params": params["vae"]}, latents, method=self.vae.decode).sample - - image = (image / 2 + 0.5).clip(0, 1).transpose(0, 2, 3, 1) - return image - - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt_ids: jnp.ndarray, - mask: jnp.ndarray, - masked_image: jnp.ndarray, - params: dict | FrozenDict, - prng_seed: jax.Array, - num_inference_steps: int = 50, - height: int | None = None, - width: int | None = None, - guidance_scale: float | jnp.ndarray = 7.5, - latents: jnp.ndarray = None, - neg_prompt_ids: jnp.ndarray = None, - return_dict: bool = True, - jit: bool = False, - ): - r""" - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` or `list[str]`): - The prompt or prompts to guide image generation. - height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The width in pixels of the generated image. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. This parameter is modulated by `strength`. - guidance_scale (`float`, *optional*, defaults to 7.5): - A higher guidance scale value encourages the model to generate images closely linked to the text - `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - latents (`jnp.ndarray`, *optional*): - Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - array is generated by sampling using the supplied random `generator`. - jit (`bool`, defaults to `False`): - Whether to run `pmap` versions of the generation and safety scoring functions. - - > [!WARNING] > This argument exists because `__call__` is not yet end-to-end pmap-able. It will be - removed in a > future release. - - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] instead of - a plain tuple. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] is - returned, otherwise a `tuple` is returned where the first element is a list with the generated images - and the second element is a list of `bool`s indicating whether the corresponding generated image - contains "not-safe-for-work" (nsfw) content. - """ - # 0. Default height and width to unet - height = height or self.unet.config.sample_size * self.vae_scale_factor - width = width or self.unet.config.sample_size * self.vae_scale_factor - - masked_image = jax.image.resize(masked_image, (*masked_image.shape[:-2], height, width), method="bicubic") - mask = jax.image.resize(mask, (*mask.shape[:-2], height, width), method="nearest") - - if isinstance(guidance_scale, float): - # Convert to a tensor so each device gets a copy. Follow the prompt_ids for - # shape information, as they may be sharded (when `jit` is `True`), or not. - guidance_scale = jnp.array([guidance_scale] * prompt_ids.shape[0]) - if len(prompt_ids.shape) > 2: - # Assume sharded - guidance_scale = guidance_scale[:, None] - - if jit: - images = _p_generate( - self, - prompt_ids, - mask, - masked_image, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - ) - else: - images = self._generate( - prompt_ids, - mask, - masked_image, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - ) - - if self.safety_checker is not None: - safety_params = params["safety_checker"] - images_uint8_casted = (images * 255).round().astype("uint8") - num_devices, batch_size = images.shape[:2] - - images_uint8_casted = np.asarray(images_uint8_casted).reshape(num_devices * batch_size, height, width, 3) - images_uint8_casted, has_nsfw_concept = self._run_safety_checker(images_uint8_casted, safety_params, jit) - images = np.asarray(images) - - # block images - if any(has_nsfw_concept): - for i, is_nsfw in enumerate(has_nsfw_concept): - if is_nsfw: - images[i] = np.asarray(images_uint8_casted[i]) - - images = images.reshape(num_devices, batch_size, height, width, 3) - else: - images = np.asarray(images) - has_nsfw_concept = False - - if not return_dict: - return (images, has_nsfw_concept) - - return FlaxStableDiffusionPipelineOutput(images=images, nsfw_content_detected=has_nsfw_concept) - - -# Static argnums are pipe, num_inference_steps, height, width. A change would trigger recompilation. -# Non-static args are (sharded) input tensors mapped over their first dimension (hence, `0`). -@partial( - jax.pmap, - in_axes=(None, 0, 0, 0, 0, 0, None, None, None, 0, 0, 0), - static_broadcasted_argnums=(0, 6, 7, 8), -) -def _p_generate( - pipe, - prompt_ids, - mask, - masked_image, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, -): - return pipe._generate( - prompt_ids, - mask, - masked_image, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - ) - - -@partial(jax.pmap, static_broadcasted_argnums=(0,)) -def _p_get_has_nsfw_concepts(pipe, features, params): - return pipe._get_has_nsfw_concepts(features, params) - - -def unshard(x: jnp.ndarray): - # einops.rearrange(x, 'd b ... -> (d b) ...') - num_devices, batch_size = x.shape[:2] - rest = x.shape[2:] - return x.reshape(num_devices * batch_size, *rest) - - -def preprocess_image(image, dtype): - w, h = image.size - w, h = (x - x % 32 for x in (w, h)) # resize to integer multiple of 32 - image = image.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]) - image = jnp.array(image).astype(dtype) / 255.0 - image = image[None].transpose(0, 3, 1, 2) - return 2.0 * image - 1.0 - - -def preprocess_mask(mask, dtype): - w, h = mask.size - w, h = (x - x % 32 for x in (w, h)) # resize to integer multiple of 32 - mask = mask.resize((w, h)) - mask = jnp.array(mask.convert("L")).astype(dtype) / 255.0 - mask = jnp.expand_dims(mask, axis=(0, 1)) - - return mask diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_output.py b/src/diffusers/pipelines/stable_diffusion/pipeline_output.py index 804e5655d369..22b9433a26d0 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_output.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_output.py @@ -3,7 +3,7 @@ import numpy as np import PIL.Image -from ...utils import BaseOutput, is_flax_available +from ...utils import BaseOutput @dataclass @@ -22,23 +22,3 @@ class StableDiffusionPipelineOutput(BaseOutput): images: list[PIL.Image.Image] | np.ndarray nsfw_content_detected: list[bool] | None - - -if is_flax_available(): - import flax - - @flax.struct.dataclass - class FlaxStableDiffusionPipelineOutput(BaseOutput): - """ - Output class for Flax-based Stable Diffusion pipelines. - - Args: - images (`np.ndarray`): - Denoised images of array shape of `(batch_size, height, width, num_channels)`. - nsfw_content_detected (`list[bool]`): - list indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content - or `None` if safety checking could not be performed. - """ - - images: np.ndarray - nsfw_content_detected: list[bool] diff --git a/src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py b/src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py deleted file mode 100644 index 612eb48313be..000000000000 --- a/src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py +++ /dev/null @@ -1,110 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import jax -import jax.numpy as jnp -from flax import linen as nn -from flax.core.frozen_dict import FrozenDict -from transformers import CLIPConfig, FlaxPreTrainedModel -from transformers.models.clip.modeling_flax_clip import FlaxCLIPVisionModule - - -def jax_cosine_distance(emb_1, emb_2, eps=1e-12): - norm_emb_1 = jnp.divide(emb_1.T, jnp.clip(jnp.linalg.norm(emb_1, axis=1), a_min=eps)).T - norm_emb_2 = jnp.divide(emb_2.T, jnp.clip(jnp.linalg.norm(emb_2, axis=1), a_min=eps)).T - return jnp.matmul(norm_emb_1, norm_emb_2.T) - - -class FlaxStableDiffusionSafetyCheckerModule(nn.Module): - config: CLIPConfig - dtype: jnp.dtype = jnp.float32 - - def setup(self): - self.vision_model = FlaxCLIPVisionModule(self.config.vision_config) - self.visual_projection = nn.Dense(self.config.projection_dim, use_bias=False, dtype=self.dtype) - - self.concept_embeds = self.param("concept_embeds", jax.nn.initializers.ones, (17, self.config.projection_dim)) - self.special_care_embeds = self.param( - "special_care_embeds", jax.nn.initializers.ones, (3, self.config.projection_dim) - ) - - self.concept_embeds_weights = self.param("concept_embeds_weights", jax.nn.initializers.ones, (17,)) - self.special_care_embeds_weights = self.param("special_care_embeds_weights", jax.nn.initializers.ones, (3,)) - - def __call__(self, clip_input): - pooled_output = self.vision_model(clip_input)[1] - image_embeds = self.visual_projection(pooled_output) - - special_cos_dist = jax_cosine_distance(image_embeds, self.special_care_embeds) - cos_dist = jax_cosine_distance(image_embeds, self.concept_embeds) - - # increase this value to create a stronger `nfsw` filter - # at the cost of increasing the possibility of filtering benign image inputs - adjustment = 0.0 - - special_scores = special_cos_dist - self.special_care_embeds_weights[None, :] + adjustment - special_scores = jnp.round(special_scores, 3) - is_special_care = jnp.any(special_scores > 0, axis=1, keepdims=True) - # Use a lower threshold if an image has any special care concept - special_adjustment = is_special_care * 0.01 - - concept_scores = cos_dist - self.concept_embeds_weights[None, :] + special_adjustment - concept_scores = jnp.round(concept_scores, 3) - has_nsfw_concepts = jnp.any(concept_scores > 0, axis=1) - - return has_nsfw_concepts - - -class FlaxStableDiffusionSafetyChecker(FlaxPreTrainedModel): - config_class = CLIPConfig - main_input_name = "clip_input" - module_class = FlaxStableDiffusionSafetyCheckerModule - - def __init__( - self, - config: CLIPConfig, - input_shape: tuple | None = None, - seed: int = 0, - dtype: jnp.dtype = jnp.float32, - _do_init: bool = True, - **kwargs, - ): - if input_shape is None: - input_shape = (1, 224, 224, 3) - module = self.module_class(config=config, dtype=dtype, **kwargs) - super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init) - - def init_weights(self, rng: jax.Array, input_shape: tuple, params: FrozenDict = None) -> FrozenDict: - # init input tensor - clip_input = jax.random.normal(rng, input_shape) - - params_rng, dropout_rng = jax.random.split(rng) - rngs = {"params": params_rng, "dropout": dropout_rng} - - random_params = self.module.init(rngs, clip_input)["params"] - - return random_params - - def __call__( - self, - clip_input, - params: dict = None, - ): - clip_input = jnp.transpose(clip_input, (0, 2, 3, 1)) - - return self.module.apply( - {"params": params or self.params}, - jnp.array(clip_input, dtype=jnp.float32), - rngs={}, - ) diff --git a/src/diffusers/pipelines/stable_diffusion_3/__init__.py b/src/diffusers/pipelines/stable_diffusion_3/__init__.py index b0604589a208..2907dc5f26ee 100644 --- a/src/diffusers/pipelines/stable_diffusion_3/__init__.py +++ b/src/diffusers/pipelines/stable_diffusion_3/__init__.py @@ -5,7 +5,6 @@ OptionalDependencyNotAvailable, _LazyModule, get_objects_from_module, - is_flax_available, is_torch_available, is_transformers_available, ) diff --git a/src/diffusers/pipelines/stable_diffusion_xl/__init__.py b/src/diffusers/pipelines/stable_diffusion_xl/__init__.py index 183a91d85aaf..ab0b7183019b 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/__init__.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/__init__.py @@ -7,7 +7,6 @@ get_objects_from_module, is_torch_available, is_transformers_available, - is_transformers_flax_compatible, ) @@ -15,8 +14,6 @@ _additional_imports = {} _import_structure = {"pipeline_output": ["StableDiffusionXLPipelineOutput"]} -if is_transformers_flax_compatible(): - _import_structure["pipeline_output"].extend(["FlaxStableDiffusionXLPipelineOutput"]) try: if not (is_transformers_available() and is_torch_available()): raise OptionalDependencyNotAvailable() @@ -30,12 +27,6 @@ _import_structure["pipeline_stable_diffusion_xl_inpaint"] = ["StableDiffusionXLInpaintPipeline"] _import_structure["pipeline_stable_diffusion_xl_instruct_pix2pix"] = ["StableDiffusionXLInstructPix2PixPipeline"] -if is_transformers_flax_compatible(): - from ...schedulers.scheduling_pndm_flax import PNDMSchedulerState - - _additional_imports.update({"PNDMSchedulerState": PNDMSchedulerState}) - _import_structure["pipeline_flax_stable_diffusion_xl"] = ["FlaxStableDiffusionXLPipeline"] - if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: @@ -49,17 +40,6 @@ from .pipeline_stable_diffusion_xl_inpaint import StableDiffusionXLInpaintPipeline from .pipeline_stable_diffusion_xl_instruct_pix2pix import StableDiffusionXLInstructPix2PixPipeline - try: - if not is_transformers_flax_compatible(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_flax_objects import * - else: - from .pipeline_flax_stable_diffusion_xl import ( - FlaxStableDiffusionXLPipeline, - ) - from .pipeline_output import FlaxStableDiffusionXLPipelineOutput - else: import sys diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py deleted file mode 100644 index b40fa2607710..000000000000 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py +++ /dev/null @@ -1,307 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from __future__ import annotations - -from functools import partial - -import jax -import jax.numpy as jnp -from flax.core.frozen_dict import FrozenDict -from transformers import CLIPTokenizer, FlaxCLIPTextModel - -from diffusers.utils import logging - -from ...models import FlaxAutoencoderKL, FlaxUNet2DConditionModel -from ...schedulers import ( - FlaxDDIMScheduler, - FlaxDPMSolverMultistepScheduler, - FlaxLMSDiscreteScheduler, - FlaxPNDMScheduler, -) -from ..pipeline_flax_utils import FlaxDiffusionPipeline -from .pipeline_output import FlaxStableDiffusionXLPipelineOutput - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -# Set to True to use python for loop instead of jax.fori_loop for easier debugging -DEBUG = False - - -class FlaxStableDiffusionXLPipeline(FlaxDiffusionPipeline): - def __init__( - self, - text_encoder: FlaxCLIPTextModel, - text_encoder_2: FlaxCLIPTextModel, - vae: FlaxAutoencoderKL, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: FlaxUNet2DConditionModel, - scheduler: FlaxDDIMScheduler | FlaxPNDMScheduler | FlaxLMSDiscreteScheduler | FlaxDPMSolverMultistepScheduler, - dtype: jnp.dtype = jnp.float32, - ): - super().__init__() - self.dtype = dtype - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - scheduler=scheduler, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - - def prepare_inputs(self, prompt: str | list[str]): - if not isinstance(prompt, (str, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - # Assume we have the two encoders - inputs = [] - for tokenizer in [self.tokenizer, self.tokenizer_2]: - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="np", - ) - inputs.append(text_inputs.input_ids) - inputs = jnp.stack(inputs, axis=1) - return inputs - - def __call__( - self, - prompt_ids: jax.Array, - params: dict | FrozenDict, - prng_seed: jax.Array, - num_inference_steps: int = 50, - guidance_scale: float | jax.Array = 7.5, - height: int | None = None, - width: int | None = None, - latents: jnp.array = None, - neg_prompt_ids: jnp.array = None, - return_dict: bool = True, - output_type: str = None, - jit: bool = False, - ): - # 0. Default height and width to unet - height = height or self.unet.config.sample_size * self.vae_scale_factor - width = width or self.unet.config.sample_size * self.vae_scale_factor - - if isinstance(guidance_scale, float) and jit: - # Convert to a tensor so each device gets a copy. - guidance_scale = jnp.array([guidance_scale] * prompt_ids.shape[0]) - guidance_scale = guidance_scale[:, None] - - return_latents = output_type == "latent" - - if jit: - images = _p_generate( - self, - prompt_ids, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - return_latents, - ) - else: - images = self._generate( - prompt_ids, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - return_latents, - ) - - if not return_dict: - return (images,) - - return FlaxStableDiffusionXLPipelineOutput(images=images) - - def get_embeddings(self, prompt_ids: jnp.array, params): - # We assume we have the two encoders - - # bs, encoder_input, seq_length - te_1_inputs = prompt_ids[:, 0, :] - te_2_inputs = prompt_ids[:, 1, :] - - prompt_embeds = self.text_encoder(te_1_inputs, params=params["text_encoder"], output_hidden_states=True) - prompt_embeds = prompt_embeds["hidden_states"][-2] - prompt_embeds_2_out = self.text_encoder_2( - te_2_inputs, params=params["text_encoder_2"], output_hidden_states=True - ) - prompt_embeds_2 = prompt_embeds_2_out["hidden_states"][-2] - text_embeds = prompt_embeds_2_out["text_embeds"] - prompt_embeds = jnp.concatenate([prompt_embeds, prompt_embeds_2], axis=-1) - return prompt_embeds, text_embeds - - def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, bs, dtype): - add_time_ids = list(original_size + crops_coords_top_left + target_size) - add_time_ids = jnp.array([add_time_ids] * bs, dtype=dtype) - return add_time_ids - - def _generate( - self, - prompt_ids: jnp.array, - params: dict | FrozenDict, - prng_seed: jax.Array, - num_inference_steps: int, - height: int, - width: int, - guidance_scale: float, - latents: jnp.array | None = None, - neg_prompt_ids: jnp.array | None = None, - return_latents=False, - ): - if height % 8 != 0 or width % 8 != 0: - raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") - - # Encode input prompt - prompt_embeds, pooled_embeds = self.get_embeddings(prompt_ids, params) - - # Get unconditional embeddings - batch_size = prompt_embeds.shape[0] - if neg_prompt_ids is None: - neg_prompt_embeds = jnp.zeros_like(prompt_embeds) - negative_pooled_embeds = jnp.zeros_like(pooled_embeds) - else: - neg_prompt_embeds, negative_pooled_embeds = self.get_embeddings(neg_prompt_ids, params) - - add_time_ids = self._get_add_time_ids( - (height, width), (0, 0), (height, width), prompt_embeds.shape[0], dtype=prompt_embeds.dtype - ) - - prompt_embeds = jnp.concatenate([neg_prompt_embeds, prompt_embeds], axis=0) # (2, 77, 2048) - add_text_embeds = jnp.concatenate([negative_pooled_embeds, pooled_embeds], axis=0) - add_time_ids = jnp.concatenate([add_time_ids, add_time_ids], axis=0) - - # Ensure model output will be `float32` before going into the scheduler - guidance_scale = jnp.array([guidance_scale], dtype=jnp.float32) - - # Create random latents - latents_shape = ( - batch_size, - self.unet.config.in_channels, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - ) - if latents is None: - latents = jax.random.normal(prng_seed, shape=latents_shape, dtype=jnp.float32) - else: - if latents.shape != latents_shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") - - # Prepare scheduler state - scheduler_state = self.scheduler.set_timesteps( - params["scheduler"], num_inference_steps=num_inference_steps, shape=latents.shape - ) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * scheduler_state.init_noise_sigma - - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - - # Denoising loop - def loop_body(step, args): - latents, scheduler_state = args - # For classifier free guidance, we need to do two forward passes. - # Here we concatenate the unconditional and text embeddings into a single batch - # to avoid doing two forward passes - latents_input = jnp.concatenate([latents] * 2) - - t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step] - timestep = jnp.broadcast_to(t, latents_input.shape[0]) - - latents_input = self.scheduler.scale_model_input(scheduler_state, latents_input, t) - - # predict the noise residual - noise_pred = self.unet.apply( - {"params": params["unet"]}, - jnp.array(latents_input), - jnp.array(timestep, dtype=jnp.int32), - encoder_hidden_states=prompt_embeds, - added_cond_kwargs=added_cond_kwargs, - ).sample - # perform guidance - noise_pred_uncond, noise_prediction_text = jnp.split(noise_pred, 2, axis=0) - noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents, scheduler_state = self.scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() - return latents, scheduler_state - - if DEBUG: - # run with python for loop - for i in range(num_inference_steps): - latents, scheduler_state = loop_body(i, (latents, scheduler_state)) - else: - latents, _ = jax.lax.fori_loop(0, num_inference_steps, loop_body, (latents, scheduler_state)) - - if return_latents: - return latents - - # Decode latents - latents = 1 / self.vae.config.scaling_factor * latents - image = self.vae.apply({"params": params["vae"]}, latents, method=self.vae.decode).sample - - image = (image / 2 + 0.5).clip(0, 1).transpose(0, 2, 3, 1) - return image - - -# Static argnums are pipe, num_inference_steps, height, width, return_latents. A change would trigger recompilation. -# Non-static args are (sharded) input tensors mapped over their first dimension (hence, `0`). -@partial( - jax.pmap, - in_axes=(None, 0, 0, 0, None, None, None, 0, 0, 0, None), - static_broadcasted_argnums=(0, 4, 5, 6, 10), -) -def _p_generate( - pipe, - prompt_ids, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - return_latents, -): - return pipe._generate( - prompt_ids, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - return_latents, - ) diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_output.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_output.py index 6617e7cd058c..805a2f7d379d 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_output.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_output.py @@ -3,7 +3,7 @@ import numpy as np import PIL.Image -from ...utils import BaseOutput, is_flax_available +from ...utils import BaseOutput @dataclass @@ -18,19 +18,3 @@ class StableDiffusionXLPipelineOutput(BaseOutput): """ images: list[PIL.Image.Image] | np.ndarray - - -if is_flax_available(): - import flax - - @flax.struct.dataclass - class FlaxStableDiffusionXLPipelineOutput(BaseOutput): - """ - Output class for Flax Stable Diffusion XL pipelines. - - Args: - images (`np.ndarray`) - Array of shape `(batch_size, height, width, num_channels)` with images from the diffusion pipeline. - """ - - images: np.ndarray diff --git a/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py b/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py index c6f6ff886a8d..7c5dcbc73d5f 100644 --- a/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py +++ b/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py @@ -72,12 +72,6 @@ def validate_environment(self, *args, **kwargs): "Using `bitsandbytes` 4-bit quantization requires the latest version of bitsandbytes: `pip install -U bitsandbytes`" ) - if kwargs.get("from_flax", False): - raise ValueError( - "Converting into 4-bit weights from flax weights is currently not supported, please make" - " sure the weights are in PyTorch format." - ) - device_map = kwargs.get("device_map", None) if ( device_map is not None @@ -363,12 +357,6 @@ def validate_environment(self, *args, **kwargs): "Using `bitsandbytes` 8-bit quantization requires the latest version of bitsandbytes: `pip install -U bitsandbytes`" ) - if kwargs.get("from_flax", False): - raise ValueError( - "Converting into 8-bit weights from flax weights is currently not supported, please make" - " sure the weights are in PyTorch format." - ) - device_map = kwargs.get("device_map", None) if ( device_map is not None diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index 871f227a8415..9b54c669293c 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -19,7 +19,6 @@ OptionalDependencyNotAvailable, _LazyModule, get_objects_from_module, - is_flax_available, is_scipy_available, is_torch_available, is_torchsde_available, @@ -84,31 +83,6 @@ _import_structure["scheduling_utils"] = ["AysSchedules", "KarrasDiffusionSchedulers", "SchedulerMixin"] _import_structure["scheduling_vq_diffusion"] = ["VQDiffusionScheduler"] -try: - if not is_flax_available(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ..utils import dummy_flax_objects # noqa F403 - - _dummy_modules.update(get_objects_from_module(dummy_flax_objects)) - -else: - _import_structure["scheduling_ddim_flax"] = ["FlaxDDIMScheduler"] - _import_structure["scheduling_ddpm_flax"] = ["FlaxDDPMScheduler"] - _import_structure["scheduling_dpmsolver_multistep_flax"] = ["FlaxDPMSolverMultistepScheduler"] - _import_structure["scheduling_euler_discrete_flax"] = ["FlaxEulerDiscreteScheduler"] - _import_structure["scheduling_karras_ve_flax"] = ["FlaxKarrasVeScheduler"] - _import_structure["scheduling_lms_discrete_flax"] = ["FlaxLMSDiscreteScheduler"] - _import_structure["scheduling_pndm_flax"] = ["FlaxPNDMScheduler"] - _import_structure["scheduling_sde_ve_flax"] = ["FlaxScoreSdeVeScheduler"] - _import_structure["scheduling_utils_flax"] = [ - "FlaxKarrasDiffusionSchedulers", - "FlaxSchedulerMixin", - "FlaxSchedulerOutput", - "broadcast_to_shape_from_left", - ] - - try: if not (is_torch_available() and is_scipy_available()): raise OptionalDependencyNotAvailable() @@ -135,7 +109,6 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: from ..utils import ( OptionalDependencyNotAvailable, - is_flax_available, is_scipy_available, is_torch_available, is_torchsde_available, @@ -192,26 +165,6 @@ from .scheduling_unipc_multistep import UniPCMultistepScheduler from .scheduling_utils import AysSchedules, KarrasDiffusionSchedulers, SchedulerMixin from .scheduling_vq_diffusion import VQDiffusionScheduler - try: - if not is_flax_available(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_flax_objects import * # noqa F403 - else: - from .scheduling_ddim_flax import FlaxDDIMScheduler - from .scheduling_ddpm_flax import FlaxDDPMScheduler - from .scheduling_dpmsolver_multistep_flax import FlaxDPMSolverMultistepScheduler - from .scheduling_euler_discrete_flax import FlaxEulerDiscreteScheduler - from .scheduling_karras_ve_flax import FlaxKarrasVeScheduler - from .scheduling_lms_discrete_flax import FlaxLMSDiscreteScheduler - from .scheduling_pndm_flax import FlaxPNDMScheduler - from .scheduling_sde_ve_flax import FlaxScoreSdeVeScheduler - from .scheduling_utils_flax import ( - FlaxKarrasDiffusionSchedulers, - FlaxSchedulerMixin, - FlaxSchedulerOutput, - broadcast_to_shape_from_left, - ) try: if not (is_torch_available() and is_scipy_available()): diff --git a/src/diffusers/schedulers/scheduling_ddim_flax.py b/src/diffusers/schedulers/scheduling_ddim_flax.py deleted file mode 100644 index 45c173ffbff5..000000000000 --- a/src/diffusers/schedulers/scheduling_ddim_flax.py +++ /dev/null @@ -1,326 +0,0 @@ -# Copyright 2025 Stanford University Team and The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion -# and https://github.com/hojonathanho/diffusion - -from dataclasses import dataclass - -import flax -import jax.numpy as jnp - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import logging -from .scheduling_utils_flax import ( - CommonSchedulerState, - FlaxKarrasDiffusionSchedulers, - FlaxSchedulerMixin, - FlaxSchedulerOutput, - add_noise_common, - get_velocity_common, -) - - -logger = logging.get_logger(__name__) - - -@flax.struct.dataclass -class DDIMSchedulerState: - common: CommonSchedulerState - final_alpha_cumprod: jnp.ndarray - - # setable values - init_noise_sigma: jnp.ndarray - timesteps: jnp.ndarray - num_inference_steps: int = None - - @classmethod - def create( - cls, - common: CommonSchedulerState, - final_alpha_cumprod: jnp.ndarray, - init_noise_sigma: jnp.ndarray, - timesteps: jnp.ndarray, - ): - return cls( - common=common, - final_alpha_cumprod=final_alpha_cumprod, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - ) - - -@dataclass -class FlaxDDIMSchedulerOutput(FlaxSchedulerOutput): - state: DDIMSchedulerState - - -class FlaxDDIMScheduler(FlaxSchedulerMixin, ConfigMixin): - """ - Denoising diffusion implicit models is a scheduler that extends the denoising procedure introduced in denoising - diffusion probabilistic models (DDPMs) with non-Markovian guidance. - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - For more details, see the original paper: https://huggingface.co/papers/2010.02502 - - Args: - num_train_timesteps (`int`): number of diffusion steps used to train the model. - beta_start (`float`): the starting `beta` value of inference. - beta_end (`float`): the final `beta` value. - beta_schedule (`str`): - the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from - `linear`, `scaled_linear`, or `squaredcos_cap_v2`. - trained_betas (`jnp.ndarray`, optional): - option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. - clip_sample (`bool`, default `True`): - option to clip predicted sample between for numerical stability. The clip range is determined by - `clip_sample_range`. - clip_sample_range (`float`, default `1.0`): - the maximum magnitude for sample clipping. Valid only when `clip_sample=True`. - set_alpha_to_one (`bool`, default `True`): - each diffusion step uses the value of alphas product at that step and at the previous one. For the final - step there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`, - otherwise it uses the value of alpha at step 0. - steps_offset (`int`, default `0`): - An offset added to the inference steps, as required by some model families. - prediction_type (`str`, default `epsilon`): - indicates whether the model predicts the noise (epsilon), or the samples. One of `epsilon`, `sample`. - `v-prediction` is not supported for this scheduler. - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - the `dtype` used for params and computation. - """ - - _compatibles = [e.name for e in FlaxKarrasDiffusionSchedulers] - - dtype: jnp.dtype - - @property - def has_state(self): - return True - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.0001, - beta_end: float = 0.02, - beta_schedule: str = "linear", - trained_betas: jnp.ndarray | None = None, - clip_sample: bool = True, - clip_sample_range: float = 1.0, - set_alpha_to_one: bool = True, - steps_offset: int = 0, - prediction_type: str = "epsilon", - dtype: jnp.dtype = jnp.float32, - ): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - self.dtype = dtype - - def create_state(self, common: CommonSchedulerState | None = None) -> DDIMSchedulerState: - if common is None: - common = CommonSchedulerState.create(self) - - # At every step in ddim, we are looking into the previous alphas_cumprod - # For the final step, there is no previous alphas_cumprod because we are already at 0 - # `set_alpha_to_one` decides whether we set this parameter simply to one or - # whether we use the final alpha of the "non-previous" one. - final_alpha_cumprod = ( - jnp.array(1.0, dtype=self.dtype) if self.config.set_alpha_to_one else common.alphas_cumprod[0] - ) - - # standard deviation of the initial noise distribution - init_noise_sigma = jnp.array(1.0, dtype=self.dtype) - - timesteps = jnp.arange(0, self.config.num_train_timesteps).round()[::-1] - - return DDIMSchedulerState.create( - common=common, - final_alpha_cumprod=final_alpha_cumprod, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - ) - - def scale_model_input( - self, - state: DDIMSchedulerState, - sample: jnp.ndarray, - timestep: int | None = None, - ) -> jnp.ndarray: - """ - Args: - state (`PNDMSchedulerState`): the `FlaxPNDMScheduler` state data class instance. - sample (`jnp.ndarray`): input sample - timestep (`int`, optional): current timestep - - Returns: - `jnp.ndarray`: scaled input sample - """ - return sample - - def set_timesteps( - self, state: DDIMSchedulerState, num_inference_steps: int, shape: tuple = () - ) -> DDIMSchedulerState: - """ - Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference. - - Args: - state (`DDIMSchedulerState`): - the `FlaxDDIMScheduler` state data class instance. - num_inference_steps (`int`): - the number of diffusion steps used when generating samples with a pre-trained model. - """ - step_ratio = self.config.num_train_timesteps // num_inference_steps - # creates integer timesteps by multiplying by ratio - # rounding to avoid issues when num_inference_step is power of 3 - timesteps = (jnp.arange(0, num_inference_steps) * step_ratio).round()[::-1] + self.config.steps_offset - - return state.replace( - num_inference_steps=num_inference_steps, - timesteps=timesteps, - ) - - def _get_variance(self, state: DDIMSchedulerState, timestep, prev_timestep): - alpha_prod_t = state.common.alphas_cumprod[timestep] - alpha_prod_t_prev = jnp.where( - prev_timestep >= 0, - state.common.alphas_cumprod[prev_timestep], - state.final_alpha_cumprod, - ) - beta_prod_t = 1 - alpha_prod_t - beta_prod_t_prev = 1 - alpha_prod_t_prev - - variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev) - - return variance - - def step( - self, - state: DDIMSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - eta: float = 0.0, - return_dict: bool = True, - ) -> FlaxDDIMSchedulerOutput | tuple: - """ - Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion - process from the learned model outputs (most often the predicted noise). - - Args: - state (`DDIMSchedulerState`): the `FlaxDDIMScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - return_dict (`bool`): option for returning tuple rather than FlaxDDIMSchedulerOutput class - - Returns: - [`FlaxDDIMSchedulerOutput`] or `tuple`: [`FlaxDDIMSchedulerOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - if state.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - # See formulas (12) and (16) of DDIM paper https://huggingface.co/papers/2010.02502 - # Ideally, read DDIM paper in-detail understanding - - # Notation ( -> - # - pred_noise_t -> e_theta(x_t, t) - # - pred_original_sample -> f_theta(x_t, t) or x_0 - # - std_dev_t -> sigma_t - # - eta -> η - # - pred_sample_direction -> "direction pointing to x_t" - # - pred_prev_sample -> "x_t-1" - - # 1. get previous step value (=t-1) - prev_timestep = timestep - self.config.num_train_timesteps // state.num_inference_steps - - alphas_cumprod = state.common.alphas_cumprod - final_alpha_cumprod = state.final_alpha_cumprod - - # 2. compute alphas, betas - alpha_prod_t = alphas_cumprod[timestep] - alpha_prod_t_prev = jnp.where(prev_timestep >= 0, alphas_cumprod[prev_timestep], final_alpha_cumprod) - - beta_prod_t = 1 - alpha_prod_t - - # 3. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (12) from https://huggingface.co/papers/2010.02502 - if self.config.prediction_type == "epsilon": - pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) - pred_epsilon = model_output - elif self.config.prediction_type == "sample": - pred_original_sample = model_output - pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) - elif self.config.prediction_type == "v_prediction": - pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output - pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" - " `v_prediction`" - ) - - # 4. Clip or threshold "predicted x_0" - if self.config.clip_sample: - pred_original_sample = pred_original_sample.clip( - -self.config.clip_sample_range, self.config.clip_sample_range - ) - - # 4. compute variance: "sigma_t(η)" -> see formula (16) - # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1) - variance = self._get_variance(state, timestep, prev_timestep) - std_dev_t = eta * variance ** (0.5) - - # 5. compute "direction pointing to x_t" of formula (12) from https://huggingface.co/papers/2010.02502 - pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * pred_epsilon - - # 6. compute x_t without "random noise" of formula (12) from https://huggingface.co/papers/2010.02502 - prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction - - if not return_dict: - return (prev_sample, state) - - return FlaxDDIMSchedulerOutput(prev_sample=prev_sample, state=state) - - def add_noise( - self, - state: DDIMSchedulerState, - original_samples: jnp.ndarray, - noise: jnp.ndarray, - timesteps: jnp.ndarray, - ) -> jnp.ndarray: - return add_noise_common(state.common, original_samples, noise, timesteps) - - def get_velocity( - self, - state: DDIMSchedulerState, - sample: jnp.ndarray, - noise: jnp.ndarray, - timesteps: jnp.ndarray, - ) -> jnp.ndarray: - return get_velocity_common(state.common, sample, noise, timesteps) - - def __len__(self): - return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_ddpm_flax.py b/src/diffusers/schedulers/scheduling_ddpm_flax.py deleted file mode 100644 index b286999bcf3c..000000000000 --- a/src/diffusers/schedulers/scheduling_ddpm_flax.py +++ /dev/null @@ -1,318 +0,0 @@ -# Copyright 2025 UC Berkeley Team and The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim - -from dataclasses import dataclass - -import flax -import jax -import jax.numpy as jnp - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import logging -from .scheduling_utils_flax import ( - CommonSchedulerState, - FlaxKarrasDiffusionSchedulers, - FlaxSchedulerMixin, - FlaxSchedulerOutput, - add_noise_common, - get_velocity_common, -) - - -logger = logging.get_logger(__name__) - - -@flax.struct.dataclass -class DDPMSchedulerState: - common: CommonSchedulerState - - # setable values - init_noise_sigma: jnp.ndarray - timesteps: jnp.ndarray - num_inference_steps: int = None - - @classmethod - def create( - cls, - common: CommonSchedulerState, - init_noise_sigma: jnp.ndarray, - timesteps: jnp.ndarray, - ): - return cls(common=common, init_noise_sigma=init_noise_sigma, timesteps=timesteps) - - -@dataclass -class FlaxDDPMSchedulerOutput(FlaxSchedulerOutput): - state: DDPMSchedulerState - - -class FlaxDDPMScheduler(FlaxSchedulerMixin, ConfigMixin): - """ - Denoising diffusion probabilistic models (DDPMs) explores the connections between denoising score matching and - Langevin dynamics sampling. - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - For more details, see the original paper: https://huggingface.co/papers/2006.11239 - - Args: - num_train_timesteps (`int`): number of diffusion steps used to train the model. - beta_start (`float`): the starting `beta` value of inference. - beta_end (`float`): the final `beta` value. - beta_schedule (`str`): - the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from - `linear`, `scaled_linear`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, optional): - option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. - variance_type (`str`): - options to clip the variance used when adding noise to the denoised sample. Choose from `fixed_small`, - `fixed_small_log`, `fixed_large`, `fixed_large_log`, `learned` or `learned_range`. - clip_sample (`bool`, default `True`): - option to clip predicted sample between -1 and 1 for numerical stability. - prediction_type (`str`, default `epsilon`): - indicates whether the model predicts the noise (epsilon), or the samples. One of `epsilon`, `sample`. - `v-prediction` is not supported for this scheduler. - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - the `dtype` used for params and computation. - """ - - _compatibles = [e.name for e in FlaxKarrasDiffusionSchedulers] - - dtype: jnp.dtype - - @property - def has_state(self): - return True - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.0001, - beta_end: float = 0.02, - beta_schedule: str = "linear", - trained_betas: jnp.ndarray | None = None, - variance_type: str = "fixed_small", - clip_sample: bool = True, - prediction_type: str = "epsilon", - dtype: jnp.dtype = jnp.float32, - ): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - self.dtype = dtype - - def create_state(self, common: CommonSchedulerState | None = None) -> DDPMSchedulerState: - if common is None: - common = CommonSchedulerState.create(self) - - # standard deviation of the initial noise distribution - init_noise_sigma = jnp.array(1.0, dtype=self.dtype) - - timesteps = jnp.arange(0, self.config.num_train_timesteps).round()[::-1] - - return DDPMSchedulerState.create( - common=common, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - ) - - def scale_model_input( - self, - state: DDPMSchedulerState, - sample: jnp.ndarray, - timestep: int | None = None, - ) -> jnp.ndarray: - """ - Args: - state (`PNDMSchedulerState`): the `FlaxPNDMScheduler` state data class instance. - sample (`jnp.ndarray`): input sample - timestep (`int`, optional): current timestep - - Returns: - `jnp.ndarray`: scaled input sample - """ - return sample - - def set_timesteps( - self, state: DDPMSchedulerState, num_inference_steps: int, shape: tuple = () - ) -> DDPMSchedulerState: - """ - Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference. - - Args: - state (`DDIMSchedulerState`): - the `FlaxDDPMScheduler` state data class instance. - num_inference_steps (`int`): - the number of diffusion steps used when generating samples with a pre-trained model. - """ - - step_ratio = self.config.num_train_timesteps // num_inference_steps - # creates integer timesteps by multiplying by ratio - # rounding to avoid issues when num_inference_step is power of 3 - timesteps = (jnp.arange(0, num_inference_steps) * step_ratio).round()[::-1] - - return state.replace( - num_inference_steps=num_inference_steps, - timesteps=timesteps, - ) - - def _get_variance(self, state: DDPMSchedulerState, t, predicted_variance=None, variance_type=None): - alpha_prod_t = state.common.alphas_cumprod[t] - alpha_prod_t_prev = jnp.where(t > 0, state.common.alphas_cumprod[t - 1], jnp.array(1.0, dtype=self.dtype)) - - # For t > 0, compute predicted variance βt (see formula (6) and (7) from https://huggingface.co/papers/2006.11239) - # and sample from it to get previous sample - # x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample - variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * state.common.betas[t] - - if variance_type is None: - variance_type = self.config.variance_type - - # hacks - were probably added for training stability - if variance_type == "fixed_small": - variance = jnp.clip(variance, a_min=1e-20) - # for rl-diffuser https://huggingface.co/papers/2205.09991 - elif variance_type == "fixed_small_log": - variance = jnp.log(jnp.clip(variance, a_min=1e-20)) - elif variance_type == "fixed_large": - variance = state.common.betas[t] - elif variance_type == "fixed_large_log": - # Glide max_log - variance = jnp.log(state.common.betas[t]) - elif variance_type == "learned": - return predicted_variance - elif variance_type == "learned_range": - min_log = variance - max_log = state.common.betas[t] - frac = (predicted_variance + 1) / 2 - variance = frac * max_log + (1 - frac) * min_log - - return variance - - def step( - self, - state: DDPMSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - key: jax.Array | None = None, - return_dict: bool = True, - ) -> FlaxDDPMSchedulerOutput | tuple: - """ - Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion - process from the learned model outputs (most often the predicted noise). - - Args: - state (`DDPMSchedulerState`): the `FlaxDDPMScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - key (`jax.Array`): a PRNG key. - return_dict (`bool`): option for returning tuple rather than FlaxDDPMSchedulerOutput class - - Returns: - [`FlaxDDPMSchedulerOutput`] or `tuple`: [`FlaxDDPMSchedulerOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - t = timestep - - if key is None: - key = jax.random.key(0) - - if ( - len(model_output.shape) > 1 - and model_output.shape[1] == sample.shape[1] * 2 - and self.config.variance_type in ["learned", "learned_range"] - ): - model_output, predicted_variance = jnp.split(model_output, sample.shape[1], axis=1) - else: - predicted_variance = None - - # 1. compute alphas, betas - alpha_prod_t = state.common.alphas_cumprod[t] - alpha_prod_t_prev = jnp.where(t > 0, state.common.alphas_cumprod[t - 1], jnp.array(1.0, dtype=self.dtype)) - beta_prod_t = 1 - alpha_prod_t - beta_prod_t_prev = 1 - alpha_prod_t_prev - - # 2. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (15) from https://huggingface.co/papers/2006.11239 - if self.config.prediction_type == "epsilon": - pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) - elif self.config.prediction_type == "sample": - pred_original_sample = model_output - elif self.config.prediction_type == "v_prediction": - pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` " - " for the FlaxDDPMScheduler." - ) - - # 3. Clip "predicted x_0" - if self.config.clip_sample: - pred_original_sample = jnp.clip(pred_original_sample, -1, 1) - - # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t - # See formula (7) from https://huggingface.co/papers/2006.11239 - pred_original_sample_coeff = (alpha_prod_t_prev ** (0.5) * state.common.betas[t]) / beta_prod_t - current_sample_coeff = state.common.alphas[t] ** (0.5) * beta_prod_t_prev / beta_prod_t - - # 5. Compute predicted previous sample µ_t - # See formula (7) from https://huggingface.co/papers/2006.11239 - pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample - - # 6. Add noise - def random_variance(): - split_key = jax.random.split(key, num=1)[0] - noise = jax.random.normal(split_key, shape=model_output.shape, dtype=self.dtype) - return (self._get_variance(state, t, predicted_variance=predicted_variance) ** 0.5) * noise - - variance = jnp.where(t > 0, random_variance(), jnp.zeros(model_output.shape, dtype=self.dtype)) - - pred_prev_sample = pred_prev_sample + variance - - if not return_dict: - return (pred_prev_sample, state) - - return FlaxDDPMSchedulerOutput(prev_sample=pred_prev_sample, state=state) - - def add_noise( - self, - state: DDPMSchedulerState, - original_samples: jnp.ndarray, - noise: jnp.ndarray, - timesteps: jnp.ndarray, - ) -> jnp.ndarray: - return add_noise_common(state.common, original_samples, noise, timesteps) - - def get_velocity( - self, - state: DDPMSchedulerState, - sample: jnp.ndarray, - noise: jnp.ndarray, - timesteps: jnp.ndarray, - ) -> jnp.ndarray: - return get_velocity_common(state.common, sample, noise, timesteps) - - def __len__(self): - return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py deleted file mode 100644 index e8b5f6673037..000000000000 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py +++ /dev/null @@ -1,670 +0,0 @@ -# Copyright 2025 TSAIL Team and The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver - -from dataclasses import dataclass - -import flax -import jax -import jax.numpy as jnp - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import logging -from .scheduling_utils_flax import ( - CommonSchedulerState, - FlaxKarrasDiffusionSchedulers, - FlaxSchedulerMixin, - FlaxSchedulerOutput, - add_noise_common, -) - - -logger = logging.get_logger(__name__) - - -@flax.struct.dataclass -class DPMSolverMultistepSchedulerState: - common: CommonSchedulerState - alpha_t: jnp.ndarray - sigma_t: jnp.ndarray - lambda_t: jnp.ndarray - - # setable values - init_noise_sigma: jnp.ndarray - timesteps: jnp.ndarray - num_inference_steps: int = None - - # running values - model_outputs: jnp.ndarray | None = None - lower_order_nums: jnp.int32 | None = None - prev_timestep: jnp.int32 | None = None - cur_sample: jnp.ndarray | None = None - - @classmethod - def create( - cls, - common: CommonSchedulerState, - alpha_t: jnp.ndarray, - sigma_t: jnp.ndarray, - lambda_t: jnp.ndarray, - init_noise_sigma: jnp.ndarray, - timesteps: jnp.ndarray, - ): - return cls( - common=common, - alpha_t=alpha_t, - sigma_t=sigma_t, - lambda_t=lambda_t, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - ) - - -@dataclass -class FlaxDPMSolverMultistepSchedulerOutput(FlaxSchedulerOutput): - state: DPMSolverMultistepSchedulerState - - -class FlaxDPMSolverMultistepScheduler(FlaxSchedulerMixin, ConfigMixin): - """ - DPM-Solver (and the improved version DPM-Solver++) is a fast dedicated high-order solver for diffusion ODEs with - the convergence order guarantee. Empirically, sampling by DPM-Solver with only 20 steps can generate high-quality - samples, and it can generate quite good samples even in only 10 steps. - - For more details, see the original paper: https://huggingface.co/papers/2206.00927 and - https://huggingface.co/papers/2211.01095 - - Currently, we support the multistep DPM-Solver for both noise prediction models and data prediction models. We - recommend to use `solver_order=2` for guided sampling, and `solver_order=3` for unconditional sampling. - - We also support the "dynamic thresholding" method in Imagen (https://huggingface.co/papers/2205.11487). For - pixel-space diffusion models, you can set both `algorithm_type="dpmsolver++"` and `thresholding=True` to use the - dynamic thresholding. Note that the thresholding method is unsuitable for latent-space diffusion models (such as - stable-diffusion). - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - For more details, see the original paper: https://huggingface.co/papers/2206.00927 and - https://huggingface.co/papers/2211.01095 - - Args: - num_train_timesteps (`int`): number of diffusion steps used to train the model. - beta_start (`float`): the starting `beta` value of inference. - beta_end (`float`): the final `beta` value. - beta_schedule (`str`): - the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from - `linear`, `scaled_linear`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, optional): - option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. - solver_order (`int`, default `2`): - the order of DPM-Solver; can be `1` or `2` or `3`. We recommend to use `solver_order=2` for guided - sampling, and `solver_order=3` for unconditional sampling. - prediction_type (`str`, default `epsilon`): - indicates whether the model predicts the noise (epsilon), or the data / `x0`. One of `epsilon`, `sample`, - or `v-prediction`. - thresholding (`bool`, default `False`): - whether to use the "dynamic thresholding" method (introduced by Imagen, - https://huggingface.co/papers/2205.11487). For pixel-space diffusion models, you can set both - `algorithm_type=dpmsolver++` and `thresholding=True` to use the dynamic thresholding. Note that the - thresholding method is unsuitable for latent-space diffusion models (such as stable-diffusion). - dynamic_thresholding_ratio (`float`, default `0.995`): - the ratio for the dynamic thresholding method. Default is `0.995`, the same as Imagen - (https://huggingface.co/papers/2205.11487). - sample_max_value (`float`, default `1.0`): - the threshold value for dynamic thresholding. Valid only when `thresholding=True` and - `algorithm_type="dpmsolver++`. - algorithm_type (`str`, default `dpmsolver++`): - the algorithm type for the solver. Either `dpmsolver` or `dpmsolver++`. The `dpmsolver` type implements the - algorithms in https://huggingface.co/papers/2206.00927, and the `dpmsolver++` type implements the - algorithms in https://huggingface.co/papers/2211.01095. We recommend to use `dpmsolver++` with - `solver_order=2` for guided sampling (e.g. stable-diffusion). - solver_type (`str`, default `midpoint`): - the solver type for the second-order solver. Either `midpoint` or `heun`. The solver type slightly affects - the sample quality, especially for small number of steps. We empirically find that `midpoint` solvers are - slightly better, so we recommend to use the `midpoint` type. - lower_order_final (`bool`, default `True`): - whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. We empirically - find this trick can stabilize the sampling of DPM-Solver for steps < 15, especially for steps <= 10. - timestep_spacing (`str`, defaults to `"linspace"`): - The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and - Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - the `dtype` used for params and computation. - """ - - _compatibles = [e.name for e in FlaxKarrasDiffusionSchedulers] - - dtype: jnp.dtype - - @property - def has_state(self): - return True - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.0001, - beta_end: float = 0.02, - beta_schedule: str = "linear", - trained_betas: jnp.ndarray | None = None, - solver_order: int = 2, - prediction_type: str = "epsilon", - thresholding: bool = False, - dynamic_thresholding_ratio: float = 0.995, - sample_max_value: float = 1.0, - algorithm_type: str = "dpmsolver++", - solver_type: str = "midpoint", - lower_order_final: bool = True, - timestep_spacing: str = "linspace", - dtype: jnp.dtype = jnp.float32, - ): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - self.dtype = dtype - - def create_state(self, common: CommonSchedulerState | None = None) -> DPMSolverMultistepSchedulerState: - if common is None: - common = CommonSchedulerState.create(self) - - # Currently we only support VP-type noise schedule - alpha_t = jnp.sqrt(common.alphas_cumprod) - sigma_t = jnp.sqrt(1 - common.alphas_cumprod) - lambda_t = jnp.log(alpha_t) - jnp.log(sigma_t) - - # settings for DPM-Solver - if self.config.algorithm_type not in ["dpmsolver", "dpmsolver++"]: - raise NotImplementedError(f"{self.config.algorithm_type} is not implemented for {self.__class__}") - if self.config.solver_type not in ["midpoint", "heun"]: - raise NotImplementedError(f"{self.config.solver_type} is not implemented for {self.__class__}") - - # standard deviation of the initial noise distribution - init_noise_sigma = jnp.array(1.0, dtype=self.dtype) - - timesteps = jnp.arange(0, self.config.num_train_timesteps).round()[::-1] - - return DPMSolverMultistepSchedulerState.create( - common=common, - alpha_t=alpha_t, - sigma_t=sigma_t, - lambda_t=lambda_t, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - ) - - def set_timesteps( - self, - state: DPMSolverMultistepSchedulerState, - num_inference_steps: int, - shape: tuple, - ) -> DPMSolverMultistepSchedulerState: - """ - Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference. - - Args: - state (`DPMSolverMultistepSchedulerState`): - the `FlaxDPMSolverMultistepScheduler` state data class instance. - num_inference_steps (`int`): - the number of diffusion steps used when generating samples with a pre-trained model. - shape (`tuple`): - the shape of the samples to be generated. - """ - last_timestep = self.config.num_train_timesteps - if self.config.timestep_spacing == "linspace": - timesteps = ( - jnp.linspace(0, last_timestep - 1, num_inference_steps + 1).round()[::-1][:-1].astype(jnp.int32) - ) - elif self.config.timestep_spacing == "leading": - step_ratio = last_timestep // (num_inference_steps + 1) - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = ( - (jnp.arange(0, num_inference_steps + 1) * step_ratio).round()[::-1][:-1].copy().astype(jnp.int32) - ) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = jnp.arange(last_timestep, 0, -step_ratio).round().copy().astype(jnp.int32) - timesteps -= 1 - else: - raise ValueError( - f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'." - ) - - # initial running values - - model_outputs = jnp.zeros((self.config.solver_order,) + shape, dtype=self.dtype) - lower_order_nums = jnp.int32(0) - prev_timestep = jnp.int32(-1) - cur_sample = jnp.zeros(shape, dtype=self.dtype) - - return state.replace( - num_inference_steps=num_inference_steps, - timesteps=timesteps, - model_outputs=model_outputs, - lower_order_nums=lower_order_nums, - prev_timestep=prev_timestep, - cur_sample=cur_sample, - ) - - def convert_model_output( - self, - state: DPMSolverMultistepSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - ) -> jnp.ndarray: - """ - Convert the model output to the corresponding type that the algorithm (DPM-Solver / DPM-Solver++) needs. - - DPM-Solver is designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to - discretize an integral of the data prediction model. So we need to first convert the model output to the - corresponding type to match the algorithm. - - Note that the algorithm type and the model type is decoupled. That is to say, we can use either DPM-Solver or - DPM-Solver++ for both noise prediction model and data prediction model. - - Args: - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - - Returns: - `jnp.ndarray`: the converted model output. - """ - # DPM-Solver++ needs to solve an integral of the data prediction model. - if self.config.algorithm_type == "dpmsolver++": - if self.config.prediction_type == "epsilon": - alpha_t, sigma_t = state.alpha_t[timestep], state.sigma_t[timestep] - x0_pred = (sample - sigma_t * model_output) / alpha_t - elif self.config.prediction_type == "sample": - x0_pred = model_output - elif self.config.prediction_type == "v_prediction": - alpha_t, sigma_t = state.alpha_t[timestep], state.sigma_t[timestep] - x0_pred = alpha_t * sample - sigma_t * model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, " - " or `v_prediction` for the FlaxDPMSolverMultistepScheduler." - ) - - if self.config.thresholding: - # Dynamic thresholding in https://huggingface.co/papers/2205.11487 - dynamic_max_val = jnp.percentile( - jnp.abs(x0_pred), - self.config.dynamic_thresholding_ratio, - axis=tuple(range(1, x0_pred.ndim)), - ) - dynamic_max_val = jnp.maximum( - dynamic_max_val, - self.config.sample_max_value * jnp.ones_like(dynamic_max_val), - ) - x0_pred = jnp.clip(x0_pred, -dynamic_max_val, dynamic_max_val) / dynamic_max_val - return x0_pred - # DPM-Solver needs to solve an integral of the noise prediction model. - elif self.config.algorithm_type == "dpmsolver": - if self.config.prediction_type == "epsilon": - return model_output - elif self.config.prediction_type == "sample": - alpha_t, sigma_t = state.alpha_t[timestep], state.sigma_t[timestep] - epsilon = (sample - alpha_t * model_output) / sigma_t - return epsilon - elif self.config.prediction_type == "v_prediction": - alpha_t, sigma_t = state.alpha_t[timestep], state.sigma_t[timestep] - epsilon = alpha_t * model_output + sigma_t * sample - return epsilon - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, " - " or `v_prediction` for the FlaxDPMSolverMultistepScheduler." - ) - - def dpm_solver_first_order_update( - self, - state: DPMSolverMultistepSchedulerState, - model_output: jnp.ndarray, - timestep: int, - prev_timestep: int, - sample: jnp.ndarray, - ) -> jnp.ndarray: - """ - One step for the first-order DPM-Solver (equivalent to DDIM). - - See https://huggingface.co/papers/2206.00927 for the detailed derivation. - - Args: - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - prev_timestep (`int`): previous discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - - Returns: - `jnp.ndarray`: the sample tensor at the previous timestep. - """ - t, s0 = prev_timestep, timestep - m0 = model_output - lambda_t, lambda_s = state.lambda_t[t], state.lambda_t[s0] - alpha_t, alpha_s = state.alpha_t[t], state.alpha_t[s0] - sigma_t, sigma_s = state.sigma_t[t], state.sigma_t[s0] - h = lambda_t - lambda_s - if self.config.algorithm_type == "dpmsolver++": - x_t = (sigma_t / sigma_s) * sample - (alpha_t * (jnp.exp(-h) - 1.0)) * m0 - elif self.config.algorithm_type == "dpmsolver": - x_t = (alpha_t / alpha_s) * sample - (sigma_t * (jnp.exp(h) - 1.0)) * m0 - return x_t - - def multistep_dpm_solver_second_order_update( - self, - state: DPMSolverMultistepSchedulerState, - model_output_list: jnp.ndarray, - timestep_list: list[int], - prev_timestep: int, - sample: jnp.ndarray, - ) -> jnp.ndarray: - """ - One step for the second-order multistep DPM-Solver. - - Args: - model_output_list (`list[jnp.ndarray]`): - direct outputs from learned diffusion model at current and latter timesteps. - timestep (`int`): current and latter discrete timestep in the diffusion chain. - prev_timestep (`int`): previous discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - - Returns: - `jnp.ndarray`: the sample tensor at the previous timestep. - """ - t, s0, s1 = prev_timestep, timestep_list[-1], timestep_list[-2] - m0, m1 = model_output_list[-1], model_output_list[-2] - lambda_t, lambda_s0, lambda_s1 = ( - state.lambda_t[t], - state.lambda_t[s0], - state.lambda_t[s1], - ) - alpha_t, alpha_s0 = state.alpha_t[t], state.alpha_t[s0] - sigma_t, sigma_s0 = state.sigma_t[t], state.sigma_t[s0] - h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1 - r0 = h_0 / h - D0, D1 = m0, (1.0 / r0) * (m0 - m1) - if self.config.algorithm_type == "dpmsolver++": - # See https://huggingface.co/papers/2211.01095 for detailed derivations - if self.config.solver_type == "midpoint": - x_t = ( - (sigma_t / sigma_s0) * sample - - (alpha_t * (jnp.exp(-h) - 1.0)) * D0 - - 0.5 * (alpha_t * (jnp.exp(-h) - 1.0)) * D1 - ) - elif self.config.solver_type == "heun": - x_t = ( - (sigma_t / sigma_s0) * sample - - (alpha_t * (jnp.exp(-h) - 1.0)) * D0 - + (alpha_t * ((jnp.exp(-h) - 1.0) / h + 1.0)) * D1 - ) - elif self.config.algorithm_type == "dpmsolver": - # See https://huggingface.co/papers/2206.00927 for detailed derivations - if self.config.solver_type == "midpoint": - x_t = ( - (alpha_t / alpha_s0) * sample - - (sigma_t * (jnp.exp(h) - 1.0)) * D0 - - 0.5 * (sigma_t * (jnp.exp(h) - 1.0)) * D1 - ) - elif self.config.solver_type == "heun": - x_t = ( - (alpha_t / alpha_s0) * sample - - (sigma_t * (jnp.exp(h) - 1.0)) * D0 - - (sigma_t * ((jnp.exp(h) - 1.0) / h - 1.0)) * D1 - ) - return x_t - - def multistep_dpm_solver_third_order_update( - self, - state: DPMSolverMultistepSchedulerState, - model_output_list: jnp.ndarray, - timestep_list: list[int], - prev_timestep: int, - sample: jnp.ndarray, - ) -> jnp.ndarray: - """ - One step for the third-order multistep DPM-Solver. - - Args: - model_output_list (`list[jnp.ndarray]`): - direct outputs from learned diffusion model at current and latter timesteps. - timestep (`int`): current and latter discrete timestep in the diffusion chain. - prev_timestep (`int`): previous discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - - Returns: - `jnp.ndarray`: the sample tensor at the previous timestep. - """ - t, s0, s1, s2 = ( - prev_timestep, - timestep_list[-1], - timestep_list[-2], - timestep_list[-3], - ) - m0, m1, m2 = model_output_list[-1], model_output_list[-2], model_output_list[-3] - lambda_t, lambda_s0, lambda_s1, lambda_s2 = ( - state.lambda_t[t], - state.lambda_t[s0], - state.lambda_t[s1], - state.lambda_t[s2], - ) - alpha_t, alpha_s0 = state.alpha_t[t], state.alpha_t[s0] - sigma_t, sigma_s0 = state.sigma_t[t], state.sigma_t[s0] - h, h_0, h_1 = lambda_t - lambda_s0, lambda_s0 - lambda_s1, lambda_s1 - lambda_s2 - r0, r1 = h_0 / h, h_1 / h - D0 = m0 - D1_0, D1_1 = (1.0 / r0) * (m0 - m1), (1.0 / r1) * (m1 - m2) - D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1) - D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1) - if self.config.algorithm_type == "dpmsolver++": - # See https://huggingface.co/papers/2206.00927 for detailed derivations - x_t = ( - (sigma_t / sigma_s0) * sample - - (alpha_t * (jnp.exp(-h) - 1.0)) * D0 - + (alpha_t * ((jnp.exp(-h) - 1.0) / h + 1.0)) * D1 - - (alpha_t * ((jnp.exp(-h) - 1.0 + h) / h**2 - 0.5)) * D2 - ) - elif self.config.algorithm_type == "dpmsolver": - # See https://huggingface.co/papers/2206.00927 for detailed derivations - x_t = ( - (alpha_t / alpha_s0) * sample - - (sigma_t * (jnp.exp(h) - 1.0)) * D0 - - (sigma_t * ((jnp.exp(h) - 1.0) / h - 1.0)) * D1 - - (sigma_t * ((jnp.exp(h) - 1.0 - h) / h**2 - 0.5)) * D2 - ) - return x_t - - def step( - self, - state: DPMSolverMultistepSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - return_dict: bool = True, - ) -> FlaxDPMSolverMultistepSchedulerOutput | tuple: - """ - Predict the sample at the previous timestep by DPM-Solver. Core function to propagate the diffusion process - from the learned model outputs (most often the predicted noise). - - Args: - state (`DPMSolverMultistepSchedulerState`): - the `FlaxDPMSolverMultistepScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - return_dict (`bool`): option for returning tuple rather than FlaxDPMSolverMultistepSchedulerOutput class - - Returns: - [`FlaxDPMSolverMultistepSchedulerOutput`] or `tuple`: [`FlaxDPMSolverMultistepSchedulerOutput`] if - `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - if state.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - (step_index,) = jnp.where(state.timesteps == timestep, size=1) - step_index = step_index[0] - - prev_timestep = jax.lax.select(step_index == len(state.timesteps) - 1, 0, state.timesteps[step_index + 1]) - - model_output = self.convert_model_output(state, model_output, timestep, sample) - - model_outputs_new = jnp.roll(state.model_outputs, -1, axis=0) - model_outputs_new = model_outputs_new.at[-1].set(model_output) - state = state.replace( - model_outputs=model_outputs_new, - prev_timestep=prev_timestep, - cur_sample=sample, - ) - - def step_1(state: DPMSolverMultistepSchedulerState) -> jnp.ndarray: - return self.dpm_solver_first_order_update( - state, - state.model_outputs[-1], - state.timesteps[step_index], - state.prev_timestep, - state.cur_sample, - ) - - def step_23(state: DPMSolverMultistepSchedulerState) -> jnp.ndarray: - def step_2(state: DPMSolverMultistepSchedulerState) -> jnp.ndarray: - timestep_list = jnp.array([state.timesteps[step_index - 1], state.timesteps[step_index]]) - return self.multistep_dpm_solver_second_order_update( - state, - state.model_outputs, - timestep_list, - state.prev_timestep, - state.cur_sample, - ) - - def step_3(state: DPMSolverMultistepSchedulerState) -> jnp.ndarray: - timestep_list = jnp.array( - [ - state.timesteps[step_index - 2], - state.timesteps[step_index - 1], - state.timesteps[step_index], - ] - ) - return self.multistep_dpm_solver_third_order_update( - state, - state.model_outputs, - timestep_list, - state.prev_timestep, - state.cur_sample, - ) - - step_2_output = step_2(state) - step_3_output = step_3(state) - - if self.config.solver_order == 2: - return step_2_output - elif self.config.lower_order_final and len(state.timesteps) < 15: - return jax.lax.select( - state.lower_order_nums < 2, - step_2_output, - jax.lax.select( - step_index == len(state.timesteps) - 2, - step_2_output, - step_3_output, - ), - ) - else: - return jax.lax.select( - state.lower_order_nums < 2, - step_2_output, - step_3_output, - ) - - step_1_output = step_1(state) - step_23_output = step_23(state) - - if self.config.solver_order == 1: - prev_sample = step_1_output - - elif self.config.lower_order_final and len(state.timesteps) < 15: - prev_sample = jax.lax.select( - state.lower_order_nums < 1, - step_1_output, - jax.lax.select( - step_index == len(state.timesteps) - 1, - step_1_output, - step_23_output, - ), - ) - - else: - prev_sample = jax.lax.select( - state.lower_order_nums < 1, - step_1_output, - step_23_output, - ) - - state = state.replace( - lower_order_nums=jnp.minimum(state.lower_order_nums + 1, self.config.solver_order), - ) - - if not return_dict: - return (prev_sample, state) - - return FlaxDPMSolverMultistepSchedulerOutput(prev_sample=prev_sample, state=state) - - def scale_model_input( - self, - state: DPMSolverMultistepSchedulerState, - sample: jnp.ndarray, - timestep: int | None = None, - ) -> jnp.ndarray: - """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. - - Args: - state (`DPMSolverMultistepSchedulerState`): - the `FlaxDPMSolverMultistepScheduler` state data class instance. - sample (`jnp.ndarray`): input sample - timestep (`int`, optional): current timestep - - Returns: - `jnp.ndarray`: scaled input sample - """ - return sample - - def add_noise( - self, - state: DPMSolverMultistepSchedulerState, - original_samples: jnp.ndarray, - noise: jnp.ndarray, - timesteps: jnp.ndarray, - ) -> jnp.ndarray: - return add_noise_common(state.common, original_samples, noise, timesteps) - - def __len__(self): - return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_euler_discrete_flax.py b/src/diffusers/schedulers/scheduling_euler_discrete_flax.py deleted file mode 100644 index 91150c5cddd8..000000000000 --- a/src/diffusers/schedulers/scheduling_euler_discrete_flax.py +++ /dev/null @@ -1,289 +0,0 @@ -# Copyright 2025 Katherine Crowson and The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from dataclasses import dataclass - -import flax -import jax.numpy as jnp - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import logging -from .scheduling_utils_flax import ( - CommonSchedulerState, - FlaxKarrasDiffusionSchedulers, - FlaxSchedulerMixin, - FlaxSchedulerOutput, - broadcast_to_shape_from_left, -) - - -logger = logging.get_logger(__name__) - - -@flax.struct.dataclass -class EulerDiscreteSchedulerState: - common: CommonSchedulerState - - # setable values - init_noise_sigma: jnp.ndarray - timesteps: jnp.ndarray - sigmas: jnp.ndarray - num_inference_steps: int = None - - @classmethod - def create( - cls, - common: CommonSchedulerState, - init_noise_sigma: jnp.ndarray, - timesteps: jnp.ndarray, - sigmas: jnp.ndarray, - ): - return cls( - common=common, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - sigmas=sigmas, - ) - - -@dataclass -class FlaxEulerDiscreteSchedulerOutput(FlaxSchedulerOutput): - state: EulerDiscreteSchedulerState - - -class FlaxEulerDiscreteScheduler(FlaxSchedulerMixin, ConfigMixin): - """ - Euler scheduler (Algorithm 2) from Karras et al. (2022) https://huggingface.co/papers/2206.00364. . Based on the - original k-diffusion implementation by Katherine Crowson: - https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L51 - - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - Args: - num_train_timesteps (`int`): number of diffusion steps used to train the model. - beta_start (`float`): the starting `beta` value of inference. - beta_end (`float`): the final `beta` value. - beta_schedule (`str`): - the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from - `linear` or `scaled_linear`. - trained_betas (`jnp.ndarray`, optional): - option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. - prediction_type (`str`, default `epsilon`, optional): - prediction type of the scheduler function, one of `epsilon` (predicting the noise of the diffusion - process), `sample` (directly predicting the noisy sample`) or `v_prediction` (see section 2.4 - https://huggingface.co/papers/2210.02303) - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - the `dtype` used for params and computation. - """ - - _compatibles = [e.name for e in FlaxKarrasDiffusionSchedulers] - - dtype: jnp.dtype - - @property - def has_state(self): - return True - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.0001, - beta_end: float = 0.02, - beta_schedule: str = "linear", - trained_betas: jnp.ndarray | None = None, - prediction_type: str = "epsilon", - timestep_spacing: str = "linspace", - dtype: jnp.dtype = jnp.float32, - ): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - self.dtype = dtype - - def create_state(self, common: CommonSchedulerState | None = None) -> EulerDiscreteSchedulerState: - if common is None: - common = CommonSchedulerState.create(self) - - timesteps = jnp.arange(0, self.config.num_train_timesteps).round()[::-1] - sigmas = ((1 - common.alphas_cumprod) / common.alphas_cumprod) ** 0.5 - sigmas = jnp.interp(timesteps, jnp.arange(0, len(sigmas)), sigmas) - sigmas = jnp.concatenate([sigmas, jnp.array([0.0], dtype=self.dtype)]) - - # standard deviation of the initial noise distribution - if self.config.timestep_spacing in ["linspace", "trailing"]: - init_noise_sigma = sigmas.max() - else: - init_noise_sigma = (sigmas.max() ** 2 + 1) ** 0.5 - - return EulerDiscreteSchedulerState.create( - common=common, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - sigmas=sigmas, - ) - - def scale_model_input(self, state: EulerDiscreteSchedulerState, sample: jnp.ndarray, timestep: int) -> jnp.ndarray: - """ - Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the Euler algorithm. - - Args: - state (`EulerDiscreteSchedulerState`): - the `FlaxEulerDiscreteScheduler` state data class instance. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - timestep (`int`): - current discrete timestep in the diffusion chain. - - Returns: - `jnp.ndarray`: scaled input sample - """ - (step_index,) = jnp.where(state.timesteps == timestep, size=1) - step_index = step_index[0] - - sigma = state.sigmas[step_index] - sample = sample / ((sigma**2 + 1) ** 0.5) - return sample - - def set_timesteps( - self, - state: EulerDiscreteSchedulerState, - num_inference_steps: int, - shape: tuple = (), - ) -> EulerDiscreteSchedulerState: - """ - Sets the timesteps used for the diffusion chain. Supporting function to be run before inference. - - Args: - state (`EulerDiscreteSchedulerState`): - the `FlaxEulerDiscreteScheduler` state data class instance. - num_inference_steps (`int`): - the number of diffusion steps used when generating samples with a pre-trained model. - """ - - if self.config.timestep_spacing == "linspace": - timesteps = jnp.linspace( - self.config.num_train_timesteps - 1, - 0, - num_inference_steps, - dtype=self.dtype, - ) - elif self.config.timestep_spacing == "leading": - step_ratio = self.config.num_train_timesteps // num_inference_steps - timesteps = (jnp.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(float) - timesteps += 1 - else: - raise ValueError( - f"timestep_spacing must be one of ['linspace', 'leading'], got {self.config.timestep_spacing}" - ) - - sigmas = ((1 - state.common.alphas_cumprod) / state.common.alphas_cumprod) ** 0.5 - sigmas = jnp.interp(timesteps, jnp.arange(0, len(sigmas)), sigmas) - sigmas = jnp.concatenate([sigmas, jnp.array([0.0], dtype=self.dtype)]) - - # standard deviation of the initial noise distribution - if self.config.timestep_spacing in ["linspace", "trailing"]: - init_noise_sigma = sigmas.max() - else: - init_noise_sigma = (sigmas.max() ** 2 + 1) ** 0.5 - - return state.replace( - timesteps=timesteps, - sigmas=sigmas, - num_inference_steps=num_inference_steps, - init_noise_sigma=init_noise_sigma, - ) - - def step( - self, - state: EulerDiscreteSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - return_dict: bool = True, - ) -> FlaxEulerDiscreteSchedulerOutput | tuple: - """ - Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion - process from the learned model outputs (most often the predicted noise). - - Args: - state (`EulerDiscreteSchedulerState`): - the `FlaxEulerDiscreteScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - order: coefficient for multi-step inference. - return_dict (`bool`): option for returning tuple rather than FlaxEulerDiscreteScheduler class - - Returns: - [`FlaxEulerDiscreteScheduler`] or `tuple`: [`FlaxEulerDiscreteScheduler`] if `return_dict` is True, - otherwise a `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - if state.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - (step_index,) = jnp.where(state.timesteps == timestep, size=1) - step_index = step_index[0] - - sigma = state.sigmas[step_index] - - # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise - if self.config.prediction_type == "epsilon": - pred_original_sample = sample - sigma * model_output - elif self.config.prediction_type == "v_prediction": - # * c_out + input * c_skip - pred_original_sample = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1)) - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`" - ) - - # 2. Convert to an ODE derivative - derivative = (sample - pred_original_sample) / sigma - - # dt = sigma_down - sigma - dt = state.sigmas[step_index + 1] - sigma - - prev_sample = sample + derivative * dt - - if not return_dict: - return (prev_sample, state) - - return FlaxEulerDiscreteSchedulerOutput(prev_sample=prev_sample, state=state) - - def add_noise( - self, - state: EulerDiscreteSchedulerState, - original_samples: jnp.ndarray, - noise: jnp.ndarray, - timesteps: jnp.ndarray, - ) -> jnp.ndarray: - sigma = state.sigmas[timesteps].flatten() - sigma = broadcast_to_shape_from_left(sigma, noise.shape) - - noisy_samples = original_samples + noise * sigma - - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_karras_ve_flax.py b/src/diffusers/schedulers/scheduling_karras_ve_flax.py deleted file mode 100644 index af7c17f17cf0..000000000000 --- a/src/diffusers/schedulers/scheduling_karras_ve_flax.py +++ /dev/null @@ -1,243 +0,0 @@ -# Copyright 2025 NVIDIA and The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - - -from dataclasses import dataclass - -import flax -import jax -import jax.numpy as jnp -from jax import random - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput, logging -from .scheduling_utils_flax import FlaxSchedulerMixin - - -logger = logging.get_logger(__name__) - - -@flax.struct.dataclass -class KarrasVeSchedulerState: - # setable values - num_inference_steps: int = None - timesteps: jnp.ndarray | None = None - schedule: jnp.ndarray | None = None # sigma(t_i) - - @classmethod - def create(cls): - return cls() - - -@dataclass -class FlaxKarrasVeOutput(BaseOutput): - """ - Output class for the scheduler's step function output. - - Args: - prev_sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - derivative (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)` for images): - Derivative of predicted original image sample (x_0). - state (`KarrasVeSchedulerState`): the `FlaxKarrasVeScheduler` state data class. - """ - - prev_sample: jnp.ndarray - derivative: jnp.ndarray - state: KarrasVeSchedulerState - - -class FlaxKarrasVeScheduler(FlaxSchedulerMixin, ConfigMixin): - """ - Stochastic sampling from Karras et al. [1] tailored to the Variance-Expanding (VE) models [2]. Use Algorithm 2 and - the VE column of Table 1 from [1] for reference. - - [1] Karras, Tero, et al. "Elucidating the Design Space of Diffusion-Based Generative Models." - https://huggingface.co/papers/2206.00364 [2] Song, Yang, et al. "Score-based generative modeling through stochastic - differential equations." https://huggingface.co/papers/2011.13456 - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - For more details on the parameters, see the original paper's Appendix E.: "Elucidating the Design Space of - Diffusion-Based Generative Models." https://huggingface.co/papers/2206.00364. The grid search values used to find - the optimal {s_noise, s_churn, s_min, s_max} for a specific model are described in Table 5 of the paper. - - Args: - sigma_min (`float`): minimum noise magnitude - sigma_max (`float`): maximum noise magnitude - s_noise (`float`): the amount of additional noise to counteract loss of detail during sampling. - A reasonable range is [1.000, 1.011]. - s_churn (`float`): the parameter controlling the overall amount of stochasticity. - A reasonable range is [0, 100]. - s_min (`float`): the start value of the sigma range where we add noise (enable stochasticity). - A reasonable range is [0, 10]. - s_max (`float`): the end value of the sigma range where we add noise. - A reasonable range is [0.2, 80]. - """ - - @property - def has_state(self): - return True - - @register_to_config - def __init__( - self, - sigma_min: float = 0.02, - sigma_max: float = 100, - s_noise: float = 1.007, - s_churn: float = 80, - s_min: float = 0.05, - s_max: float = 50, - ): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - def create_state(self): - return KarrasVeSchedulerState.create() - - def set_timesteps( - self, state: KarrasVeSchedulerState, num_inference_steps: int, shape: tuple = () - ) -> KarrasVeSchedulerState: - """ - Sets the continuous timesteps used for the diffusion chain. Supporting function to be run before inference. - - Args: - state (`KarrasVeSchedulerState`): - the `FlaxKarrasVeScheduler` state data class. - num_inference_steps (`int`): - the number of diffusion steps used when generating samples with a pre-trained model. - - """ - timesteps = jnp.arange(0, num_inference_steps)[::-1].copy() - schedule = [ - ( - self.config.sigma_max**2 - * (self.config.sigma_min**2 / self.config.sigma_max**2) ** (i / (num_inference_steps - 1)) - ) - for i in timesteps - ] - - return state.replace( - num_inference_steps=num_inference_steps, - schedule=jnp.array(schedule, dtype=jnp.float32), - timesteps=timesteps, - ) - - def add_noise_to_input( - self, - state: KarrasVeSchedulerState, - sample: jnp.ndarray, - sigma: float, - key: jax.Array, - ) -> tuple[jnp.ndarray, float]: - """ - Explicit Langevin-like "churn" step of adding noise to the sample according to a factor gamma_i ≥ 0 to reach a - higher noise level sigma_hat = sigma_i + gamma_i*sigma_i. - - TODO Args: - """ - if self.config.s_min <= sigma <= self.config.s_max: - gamma = min(self.config.s_churn / state.num_inference_steps, 2**0.5 - 1) - else: - gamma = 0 - - # sample eps ~ N(0, S_noise^2 * I) - key = random.split(key, num=1) - eps = self.config.s_noise * random.normal(key=key, shape=sample.shape) - sigma_hat = sigma + gamma * sigma - sample_hat = sample + ((sigma_hat**2 - sigma**2) ** 0.5 * eps) - - return sample_hat, sigma_hat - - def step( - self, - state: KarrasVeSchedulerState, - model_output: jnp.ndarray, - sigma_hat: float, - sigma_prev: float, - sample_hat: jnp.ndarray, - return_dict: bool = True, - ) -> FlaxKarrasVeOutput | tuple: - """ - Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion - process from the learned model outputs (most often the predicted noise). - - Args: - state (`KarrasVeSchedulerState`): the `FlaxKarrasVeScheduler` state data class. - model_output (`torch.Tensor` or `np.ndarray`): direct output from learned diffusion model. - sigma_hat (`float`): TODO - sigma_prev (`float`): TODO - sample_hat (`torch.Tensor` or `np.ndarray`): TODO - return_dict (`bool`): option for returning tuple rather than FlaxKarrasVeOutput class - - Returns: - [`~schedulers.scheduling_karras_ve_flax.FlaxKarrasVeOutput`] or `tuple`: Updated sample in the diffusion - chain and derivative. [`~schedulers.scheduling_karras_ve_flax.FlaxKarrasVeOutput`] if `return_dict` is - True, otherwise a `tuple`. When returning a tuple, the first element is the sample tensor. - """ - - pred_original_sample = sample_hat + sigma_hat * model_output - derivative = (sample_hat - pred_original_sample) / sigma_hat - sample_prev = sample_hat + (sigma_prev - sigma_hat) * derivative - - if not return_dict: - return (sample_prev, derivative, state) - - return FlaxKarrasVeOutput(prev_sample=sample_prev, derivative=derivative, state=state) - - def step_correct( - self, - state: KarrasVeSchedulerState, - model_output: jnp.ndarray, - sigma_hat: float, - sigma_prev: float, - sample_hat: jnp.ndarray, - sample_prev: jnp.ndarray, - derivative: jnp.ndarray, - return_dict: bool = True, - ) -> FlaxKarrasVeOutput | tuple: - """ - Correct the predicted sample based on the output model_output of the network. TODO complete description - - Args: - state (`KarrasVeSchedulerState`): the `FlaxKarrasVeScheduler` state data class. - model_output (`torch.Tensor` or `np.ndarray`): direct output from learned diffusion model. - sigma_hat (`float`): TODO - sigma_prev (`float`): TODO - sample_hat (`torch.Tensor` or `np.ndarray`): TODO - sample_prev (`torch.Tensor` or `np.ndarray`): TODO - derivative (`torch.Tensor` or `np.ndarray`): TODO - return_dict (`bool`): option for returning tuple rather than FlaxKarrasVeOutput class - - Returns: - prev_sample (TODO): updated sample in the diffusion chain. derivative (TODO): TODO - - """ - pred_original_sample = sample_prev + sigma_prev * model_output - derivative_corr = (sample_prev - pred_original_sample) / sigma_prev - sample_prev = sample_hat + (sigma_prev - sigma_hat) * (0.5 * derivative + 0.5 * derivative_corr) - - if not return_dict: - return (sample_prev, derivative, state) - - return FlaxKarrasVeOutput(prev_sample=sample_prev, derivative=derivative, state=state) - - def add_noise(self, state: KarrasVeSchedulerState, original_samples, noise, timesteps): - raise NotImplementedError() diff --git a/src/diffusers/schedulers/scheduling_lms_discrete_flax.py b/src/diffusers/schedulers/scheduling_lms_discrete_flax.py deleted file mode 100644 index c37d8752f7fb..000000000000 --- a/src/diffusers/schedulers/scheduling_lms_discrete_flax.py +++ /dev/null @@ -1,307 +0,0 @@ -# Copyright 2025 Katherine Crowson and The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from dataclasses import dataclass - -import flax -import jax.numpy as jnp -from scipy import integrate - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import logging -from .scheduling_utils_flax import ( - CommonSchedulerState, - FlaxKarrasDiffusionSchedulers, - FlaxSchedulerMixin, - FlaxSchedulerOutput, - broadcast_to_shape_from_left, -) - - -logger = logging.get_logger(__name__) - - -@flax.struct.dataclass -class LMSDiscreteSchedulerState: - common: CommonSchedulerState - - # setable values - init_noise_sigma: jnp.ndarray - timesteps: jnp.ndarray - sigmas: jnp.ndarray - num_inference_steps: int = None - - # running values - derivatives: jnp.ndarray | None = None - - @classmethod - def create( - cls, - common: CommonSchedulerState, - init_noise_sigma: jnp.ndarray, - timesteps: jnp.ndarray, - sigmas: jnp.ndarray, - ): - return cls( - common=common, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - sigmas=sigmas, - ) - - -@dataclass -class FlaxLMSSchedulerOutput(FlaxSchedulerOutput): - state: LMSDiscreteSchedulerState - - -class FlaxLMSDiscreteScheduler(FlaxSchedulerMixin, ConfigMixin): - """ - Linear Multistep Scheduler for discrete beta schedules. Based on the original k-diffusion implementation by - Katherine Crowson: - https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L181 - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - Args: - num_train_timesteps (`int`): number of diffusion steps used to train the model. - beta_start (`float`): the starting `beta` value of inference. - beta_end (`float`): the final `beta` value. - beta_schedule (`str`): - the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from - `linear` or `scaled_linear`. - trained_betas (`jnp.ndarray`, optional): - option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. - prediction_type (`str`, default `epsilon`, optional): - prediction type of the scheduler function, one of `epsilon` (predicting the noise of the diffusion - process), `sample` (directly predicting the noisy sample`) or `v_prediction` (see section 2.4 - https://huggingface.co/papers/2210.02303) - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - the `dtype` used for params and computation. - """ - - _compatibles = [e.name for e in FlaxKarrasDiffusionSchedulers] - - dtype: jnp.dtype - - @property - def has_state(self): - return True - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.0001, - beta_end: float = 0.02, - beta_schedule: str = "linear", - trained_betas: jnp.ndarray | None = None, - prediction_type: str = "epsilon", - dtype: jnp.dtype = jnp.float32, - ): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - self.dtype = dtype - - def create_state(self, common: CommonSchedulerState | None = None) -> LMSDiscreteSchedulerState: - if common is None: - common = CommonSchedulerState.create(self) - - timesteps = jnp.arange(0, self.config.num_train_timesteps).round()[::-1] - sigmas = ((1 - common.alphas_cumprod) / common.alphas_cumprod) ** 0.5 - - # standard deviation of the initial noise distribution - init_noise_sigma = sigmas.max() - - return LMSDiscreteSchedulerState.create( - common=common, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - sigmas=sigmas, - ) - - def scale_model_input(self, state: LMSDiscreteSchedulerState, sample: jnp.ndarray, timestep: int) -> jnp.ndarray: - """ - Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the K-LMS algorithm. - - Args: - state (`LMSDiscreteSchedulerState`): - the `FlaxLMSDiscreteScheduler` state data class instance. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - timestep (`int`): - current discrete timestep in the diffusion chain. - - Returns: - `jnp.ndarray`: scaled input sample - """ - (step_index,) = jnp.where(state.timesteps == timestep, size=1) - step_index = step_index[0] - - sigma = state.sigmas[step_index] - sample = sample / ((sigma**2 + 1) ** 0.5) - return sample - - def get_lms_coefficient(self, state: LMSDiscreteSchedulerState, order, t, current_order): - """ - Compute a linear multistep coefficient. - - Args: - order (TODO): - t (TODO): - current_order (TODO): - """ - - def lms_derivative(tau): - prod = 1.0 - for k in range(order): - if current_order == k: - continue - prod *= (tau - state.sigmas[t - k]) / (state.sigmas[t - current_order] - state.sigmas[t - k]) - return prod - - integrated_coeff = integrate.quad(lms_derivative, state.sigmas[t], state.sigmas[t + 1], epsrel=1e-4)[0] - - return integrated_coeff - - def set_timesteps( - self, - state: LMSDiscreteSchedulerState, - num_inference_steps: int, - shape: tuple = (), - ) -> LMSDiscreteSchedulerState: - """ - Sets the timesteps used for the diffusion chain. Supporting function to be run before inference. - - Args: - state (`LMSDiscreteSchedulerState`): - the `FlaxLMSDiscreteScheduler` state data class instance. - num_inference_steps (`int`): - the number of diffusion steps used when generating samples with a pre-trained model. - """ - - timesteps = jnp.linspace( - self.config.num_train_timesteps - 1, - 0, - num_inference_steps, - dtype=self.dtype, - ) - - low_idx = jnp.floor(timesteps).astype(jnp.int32) - high_idx = jnp.ceil(timesteps).astype(jnp.int32) - - frac = jnp.mod(timesteps, 1.0) - - sigmas = ((1 - state.common.alphas_cumprod) / state.common.alphas_cumprod) ** 0.5 - sigmas = (1 - frac) * sigmas[low_idx] + frac * sigmas[high_idx] - sigmas = jnp.concatenate([sigmas, jnp.array([0.0], dtype=self.dtype)]) - - timesteps = timesteps.astype(jnp.int32) - - # initial running values - derivatives = jnp.zeros((0,) + shape, dtype=self.dtype) - - return state.replace( - timesteps=timesteps, - sigmas=sigmas, - num_inference_steps=num_inference_steps, - derivatives=derivatives, - ) - - def step( - self, - state: LMSDiscreteSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - order: int = 4, - return_dict: bool = True, - ) -> FlaxLMSSchedulerOutput | tuple: - """ - Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion - process from the learned model outputs (most often the predicted noise). - - Args: - state (`LMSDiscreteSchedulerState`): the `FlaxLMSDiscreteScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - order: coefficient for multi-step inference. - return_dict (`bool`): option for returning tuple rather than FlaxLMSSchedulerOutput class - - Returns: - [`FlaxLMSSchedulerOutput`] or `tuple`: [`FlaxLMSSchedulerOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - if state.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - sigma = state.sigmas[timestep] - - # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise - if self.config.prediction_type == "epsilon": - pred_original_sample = sample - sigma * model_output - elif self.config.prediction_type == "v_prediction": - # * c_out + input * c_skip - pred_original_sample = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1)) - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`" - ) - - # 2. Convert to an ODE derivative - derivative = (sample - pred_original_sample) / sigma - state = state.replace(derivatives=jnp.append(state.derivatives, derivative)) - if len(state.derivatives) > order: - state = state.replace(derivatives=jnp.delete(state.derivatives, 0)) - - # 3. Compute linear multistep coefficients - order = min(timestep + 1, order) - lms_coeffs = [self.get_lms_coefficient(state, order, timestep, curr_order) for curr_order in range(order)] - - # 4. Compute previous sample based on the derivatives path - prev_sample = sample + sum( - coeff * derivative for coeff, derivative in zip(lms_coeffs, reversed(state.derivatives)) - ) - - if not return_dict: - return (prev_sample, state) - - return FlaxLMSSchedulerOutput(prev_sample=prev_sample, state=state) - - def add_noise( - self, - state: LMSDiscreteSchedulerState, - original_samples: jnp.ndarray, - noise: jnp.ndarray, - timesteps: jnp.ndarray, - ) -> jnp.ndarray: - sigma = state.sigmas[timesteps].flatten() - sigma = broadcast_to_shape_from_left(sigma, noise.shape) - - noisy_samples = original_samples + noise * sigma - - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_pndm_flax.py b/src/diffusers/schedulers/scheduling_pndm_flax.py deleted file mode 100644 index e18e484c8c4c..000000000000 --- a/src/diffusers/schedulers/scheduling_pndm_flax.py +++ /dev/null @@ -1,528 +0,0 @@ -# Copyright 2025 Zhejiang University Team and The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim - -from dataclasses import dataclass - -import flax -import jax -import jax.numpy as jnp - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import logging -from .scheduling_utils_flax import ( - CommonSchedulerState, - FlaxKarrasDiffusionSchedulers, - FlaxSchedulerMixin, - FlaxSchedulerOutput, - add_noise_common, -) - - -logger = logging.get_logger(__name__) - - -@flax.struct.dataclass -class PNDMSchedulerState: - common: CommonSchedulerState - final_alpha_cumprod: jnp.ndarray - - # setable values - init_noise_sigma: jnp.ndarray - timesteps: jnp.ndarray - num_inference_steps: int = None - prk_timesteps: jnp.ndarray | None = None - plms_timesteps: jnp.ndarray | None = None - - # running values - cur_model_output: jnp.ndarray | None = None - counter: jnp.int32 | None = None - cur_sample: jnp.ndarray | None = None - ets: jnp.ndarray | None = None - - @classmethod - def create( - cls, - common: CommonSchedulerState, - final_alpha_cumprod: jnp.ndarray, - init_noise_sigma: jnp.ndarray, - timesteps: jnp.ndarray, - ): - return cls( - common=common, - final_alpha_cumprod=final_alpha_cumprod, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - ) - - -@dataclass -class FlaxPNDMSchedulerOutput(FlaxSchedulerOutput): - state: PNDMSchedulerState - - -class FlaxPNDMScheduler(FlaxSchedulerMixin, ConfigMixin): - """ - Pseudo numerical methods for diffusion models (PNDM) proposes using more advanced ODE integration techniques, - namely Runge-Kutta method and a linear multi-step method. - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - For more details, see the original paper: https://huggingface.co/papers/2202.09778 - - Args: - num_train_timesteps (`int`): number of diffusion steps used to train the model. - beta_start (`float`): the starting `beta` value of inference. - beta_end (`float`): the final `beta` value. - beta_schedule (`str`): - the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from - `linear`, `scaled_linear`, or `squaredcos_cap_v2`. - trained_betas (`jnp.ndarray`, optional): - option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. - skip_prk_steps (`bool`): - allows the scheduler to skip the Runge-Kutta steps that are defined in the original paper as being required - before plms steps; defaults to `False`. - set_alpha_to_one (`bool`, default `False`): - each diffusion step uses the value of alphas product at that step and at the previous one. For the final - step there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`, - otherwise it uses the value of alpha at step 0. - steps_offset (`int`, default `0`): - An offset added to the inference steps, as required by some model families. - prediction_type (`str`, default `epsilon`, optional): - prediction type of the scheduler function, one of `epsilon` (predicting the noise of the diffusion - process), `sample` (directly predicting the noisy sample`) or `v_prediction` (see section 2.4 - https://huggingface.co/papers/2210.02303) - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - the `dtype` used for params and computation. - """ - - _compatibles = [e.name for e in FlaxKarrasDiffusionSchedulers] - - dtype: jnp.dtype - pndm_order: int - - @property - def has_state(self): - return True - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.0001, - beta_end: float = 0.02, - beta_schedule: str = "linear", - trained_betas: jnp.ndarray | None = None, - skip_prk_steps: bool = False, - set_alpha_to_one: bool = False, - steps_offset: int = 0, - prediction_type: str = "epsilon", - dtype: jnp.dtype = jnp.float32, - ): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - self.dtype = dtype - - # For now we only support F-PNDM, i.e. the runge-kutta method - # For more information on the algorithm please take a look at the paper: https://huggingface.co/papers/2202.09778 - # mainly at formula (9), (12), (13) and the Algorithm 2. - self.pndm_order = 4 - - def create_state(self, common: CommonSchedulerState | None = None) -> PNDMSchedulerState: - if common is None: - common = CommonSchedulerState.create(self) - - # At every step in ddim, we are looking into the previous alphas_cumprod - # For the final step, there is no previous alphas_cumprod because we are already at 0 - # `set_alpha_to_one` decides whether we set this parameter simply to one or - # whether we use the final alpha of the "non-previous" one. - final_alpha_cumprod = ( - jnp.array(1.0, dtype=self.dtype) if self.config.set_alpha_to_one else common.alphas_cumprod[0] - ) - - # standard deviation of the initial noise distribution - init_noise_sigma = jnp.array(1.0, dtype=self.dtype) - - timesteps = jnp.arange(0, self.config.num_train_timesteps).round()[::-1] - - return PNDMSchedulerState.create( - common=common, - final_alpha_cumprod=final_alpha_cumprod, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - ) - - def set_timesteps(self, state: PNDMSchedulerState, num_inference_steps: int, shape: tuple) -> PNDMSchedulerState: - """ - Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference. - - Args: - state (`PNDMSchedulerState`): - the `FlaxPNDMScheduler` state data class instance. - num_inference_steps (`int`): - the number of diffusion steps used when generating samples with a pre-trained model. - shape (`tuple`): - the shape of the samples to be generated. - """ - - step_ratio = self.config.num_train_timesteps // num_inference_steps - # creates integer timesteps by multiplying by ratio - # rounding to avoid issues when num_inference_step is power of 3 - _timesteps = (jnp.arange(0, num_inference_steps) * step_ratio).round() + self.config.steps_offset - - if self.config.skip_prk_steps: - # for some models like stable diffusion the prk steps can/should be skipped to - # produce better results. When using PNDM with `self.config.skip_prk_steps` the implementation - # is based on crowsonkb's PLMS sampler implementation: https://github.com/CompVis/latent-diffusion/pull/51 - - prk_timesteps = jnp.array([], dtype=jnp.int32) - plms_timesteps = jnp.concatenate([_timesteps[:-1], _timesteps[-2:-1], _timesteps[-1:]])[::-1] - - else: - prk_timesteps = _timesteps[-self.pndm_order :].repeat(2) + jnp.tile( - jnp.array( - [0, self.config.num_train_timesteps // num_inference_steps // 2], - dtype=jnp.int32, - ), - self.pndm_order, - ) - - prk_timesteps = (prk_timesteps[:-1].repeat(2)[1:-1])[::-1] - plms_timesteps = _timesteps[:-3][::-1] - - timesteps = jnp.concatenate([prk_timesteps, plms_timesteps]) - - # initial running values - - cur_model_output = jnp.zeros(shape, dtype=self.dtype) - counter = jnp.int32(0) - cur_sample = jnp.zeros(shape, dtype=self.dtype) - ets = jnp.zeros((4,) + shape, dtype=self.dtype) - - return state.replace( - timesteps=timesteps, - num_inference_steps=num_inference_steps, - prk_timesteps=prk_timesteps, - plms_timesteps=plms_timesteps, - cur_model_output=cur_model_output, - counter=counter, - cur_sample=cur_sample, - ets=ets, - ) - - def scale_model_input( - self, - state: PNDMSchedulerState, - sample: jnp.ndarray, - timestep: int | None = None, - ) -> jnp.ndarray: - """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. - - Args: - state (`PNDMSchedulerState`): the `FlaxPNDMScheduler` state data class instance. - sample (`jnp.ndarray`): input sample - timestep (`int`, optional): current timestep - - Returns: - `jnp.ndarray`: scaled input sample - """ - return sample - - def step( - self, - state: PNDMSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - return_dict: bool = True, - ) -> FlaxPNDMSchedulerOutput | tuple: - """ - Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion - process from the learned model outputs (most often the predicted noise). - - This function calls `step_prk()` or `step_plms()` depending on the internal variable `counter`. - - Args: - state (`PNDMSchedulerState`): the `FlaxPNDMScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - return_dict (`bool`): option for returning tuple rather than FlaxPNDMSchedulerOutput class - - Returns: - [`FlaxPNDMSchedulerOutput`] or `tuple`: [`FlaxPNDMSchedulerOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - - if state.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - if self.config.skip_prk_steps: - prev_sample, state = self.step_plms(state, model_output, timestep, sample) - else: - prk_prev_sample, prk_state = self.step_prk(state, model_output, timestep, sample) - plms_prev_sample, plms_state = self.step_plms(state, model_output, timestep, sample) - - cond = state.counter < len(state.prk_timesteps) - - prev_sample = jax.lax.select(cond, prk_prev_sample, plms_prev_sample) - - state = state.replace( - cur_model_output=jax.lax.select(cond, prk_state.cur_model_output, plms_state.cur_model_output), - ets=jax.lax.select(cond, prk_state.ets, plms_state.ets), - cur_sample=jax.lax.select(cond, prk_state.cur_sample, plms_state.cur_sample), - counter=jax.lax.select(cond, prk_state.counter, plms_state.counter), - ) - - if not return_dict: - return (prev_sample, state) - - return FlaxPNDMSchedulerOutput(prev_sample=prev_sample, state=state) - - def step_prk( - self, - state: PNDMSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - ) -> FlaxPNDMSchedulerOutput | tuple: - """ - Step function propagating the sample with the Runge-Kutta method. RK takes 4 forward passes to approximate the - solution to the differential equation. - - Args: - state (`PNDMSchedulerState`): the `FlaxPNDMScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - return_dict (`bool`): option for returning tuple rather than FlaxPNDMSchedulerOutput class - - Returns: - [`FlaxPNDMSchedulerOutput`] or `tuple`: [`FlaxPNDMSchedulerOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - - if state.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - diff_to_prev = jnp.where( - state.counter % 2, - 0, - self.config.num_train_timesteps // state.num_inference_steps // 2, - ) - prev_timestep = timestep - diff_to_prev - timestep = state.prk_timesteps[state.counter // 4 * 4] - - model_output = jax.lax.select( - (state.counter % 4) != 3, - model_output, # remainder 0, 1, 2 - state.cur_model_output + 1 / 6 * model_output, # remainder 3 - ) - - state = state.replace( - cur_model_output=jax.lax.select_n( - state.counter % 4, - state.cur_model_output + 1 / 6 * model_output, # remainder 0 - state.cur_model_output + 1 / 3 * model_output, # remainder 1 - state.cur_model_output + 1 / 3 * model_output, # remainder 2 - jnp.zeros_like(state.cur_model_output), # remainder 3 - ), - ets=jax.lax.select( - (state.counter % 4) == 0, - state.ets.at[0:3].set(state.ets[1:4]).at[3].set(model_output), # remainder 0 - state.ets, # remainder 1, 2, 3 - ), - cur_sample=jax.lax.select( - (state.counter % 4) == 0, - sample, # remainder 0 - state.cur_sample, # remainder 1, 2, 3 - ), - ) - - cur_sample = state.cur_sample - prev_sample = self._get_prev_sample(state, cur_sample, timestep, prev_timestep, model_output) - state = state.replace(counter=state.counter + 1) - - return (prev_sample, state) - - def step_plms( - self, - state: PNDMSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - ) -> FlaxPNDMSchedulerOutput | tuple: - """ - Step function propagating the sample with the linear multi-step method. This has one forward pass with multiple - times to approximate the solution. - - Args: - state (`PNDMSchedulerState`): the `FlaxPNDMScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - return_dict (`bool`): option for returning tuple rather than FlaxPNDMSchedulerOutput class - - Returns: - [`FlaxPNDMSchedulerOutput`] or `tuple`: [`FlaxPNDMSchedulerOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - - if state.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - # NOTE: There is no way to check in the jitted runtime if the prk mode was ran before - - prev_timestep = timestep - self.config.num_train_timesteps // state.num_inference_steps - prev_timestep = jnp.where(prev_timestep > 0, prev_timestep, 0) - - # Reference: - # if state.counter != 1: - # state.ets.append(model_output) - # else: - # prev_timestep = timestep - # timestep = timestep + self.config.num_train_timesteps // state.num_inference_steps - - prev_timestep = jnp.where(state.counter == 1, timestep, prev_timestep) - timestep = jnp.where( - state.counter == 1, - timestep + self.config.num_train_timesteps // state.num_inference_steps, - timestep, - ) - - # Reference: - # if len(state.ets) == 1 and state.counter == 0: - # model_output = model_output - # state.cur_sample = sample - # elif len(state.ets) == 1 and state.counter == 1: - # model_output = (model_output + state.ets[-1]) / 2 - # sample = state.cur_sample - # state.cur_sample = None - # elif len(state.ets) == 2: - # model_output = (3 * state.ets[-1] - state.ets[-2]) / 2 - # elif len(state.ets) == 3: - # model_output = (23 * state.ets[-1] - 16 * state.ets[-2] + 5 * state.ets[-3]) / 12 - # else: - # model_output = (1 / 24) * (55 * state.ets[-1] - 59 * state.ets[-2] + 37 * state.ets[-3] - 9 * state.ets[-4]) - - state = state.replace( - ets=jax.lax.select( - state.counter != 1, - state.ets.at[0:3].set(state.ets[1:4]).at[3].set(model_output), # counter != 1 - state.ets, # counter 1 - ), - cur_sample=jax.lax.select( - state.counter != 1, - sample, # counter != 1 - state.cur_sample, # counter 1 - ), - ) - - state = state.replace( - cur_model_output=jax.lax.select_n( - jnp.clip(state.counter, 0, 4), - model_output, # counter 0 - (model_output + state.ets[-1]) / 2, # counter 1 - (3 * state.ets[-1] - state.ets[-2]) / 2, # counter 2 - (23 * state.ets[-1] - 16 * state.ets[-2] + 5 * state.ets[-3]) / 12, # counter 3 - (1 / 24) - * (55 * state.ets[-1] - 59 * state.ets[-2] + 37 * state.ets[-3] - 9 * state.ets[-4]), # counter >= 4 - ), - ) - - sample = state.cur_sample - model_output = state.cur_model_output - prev_sample = self._get_prev_sample(state, sample, timestep, prev_timestep, model_output) - state = state.replace(counter=state.counter + 1) - - return (prev_sample, state) - - def _get_prev_sample(self, state: PNDMSchedulerState, sample, timestep, prev_timestep, model_output): - # See formula (9) of PNDM paper https://huggingface.co/papers/2202.09778 - # this function computes x_(t−δ) using the formula of (9) - # Note that x_t needs to be added to both sides of the equation - - # Notation ( -> - # alpha_prod_t -> α_t - # alpha_prod_t_prev -> α_(t−δ) - # beta_prod_t -> (1 - α_t) - # beta_prod_t_prev -> (1 - α_(t−δ)) - # sample -> x_t - # model_output -> e_θ(x_t, t) - # prev_sample -> x_(t−δ) - alpha_prod_t = state.common.alphas_cumprod[timestep] - alpha_prod_t_prev = jnp.where( - prev_timestep >= 0, - state.common.alphas_cumprod[prev_timestep], - state.final_alpha_cumprod, - ) - beta_prod_t = 1 - alpha_prod_t - beta_prod_t_prev = 1 - alpha_prod_t_prev - - if self.config.prediction_type == "v_prediction": - model_output = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample - elif self.config.prediction_type != "epsilon": - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon` or `v_prediction`" - ) - - # corresponds to (α_(t−δ) - α_t) divided by - # denominator of x_t in formula (9) and plus 1 - # Note: (α_(t−δ) - α_t) / (sqrt(α_t) * (sqrt(α_(t−δ)) + sqr(α_t))) = - # sqrt(α_(t−δ)) / sqrt(α_t)) - sample_coeff = (alpha_prod_t_prev / alpha_prod_t) ** (0.5) - - # corresponds to denominator of e_θ(x_t, t) in formula (9) - model_output_denom_coeff = alpha_prod_t * beta_prod_t_prev ** (0.5) + ( - alpha_prod_t * beta_prod_t * alpha_prod_t_prev - ) ** (0.5) - - # full formula (9) - prev_sample = ( - sample_coeff * sample - (alpha_prod_t_prev - alpha_prod_t) * model_output / model_output_denom_coeff - ) - - return prev_sample - - def add_noise( - self, - state: PNDMSchedulerState, - original_samples: jnp.ndarray, - noise: jnp.ndarray, - timesteps: jnp.ndarray, - ) -> jnp.ndarray: - return add_noise_common(state.common, original_samples, noise, timesteps) - - def __len__(self): - return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_sde_ve_flax.py b/src/diffusers/schedulers/scheduling_sde_ve_flax.py deleted file mode 100644 index 59c3a1cbaeeb..000000000000 --- a/src/diffusers/schedulers/scheduling_sde_ve_flax.py +++ /dev/null @@ -1,294 +0,0 @@ -# Copyright 2025 Google Brain and The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch - -from dataclasses import dataclass - -import flax -import jax -import jax.numpy as jnp -from jax import random - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import logging -from .scheduling_utils_flax import ( - FlaxSchedulerMixin, - FlaxSchedulerOutput, - broadcast_to_shape_from_left, -) - - -logger = logging.get_logger(__name__) - - -@flax.struct.dataclass -class ScoreSdeVeSchedulerState: - # setable values - timesteps: jnp.ndarray | None = None - discrete_sigmas: jnp.ndarray | None = None - sigmas: jnp.ndarray | None = None - - @classmethod - def create(cls): - return cls() - - -@dataclass -class FlaxSdeVeOutput(FlaxSchedulerOutput): - """ - Output class for the ScoreSdeVeScheduler's step function output. - - Args: - state (`ScoreSdeVeSchedulerState`): - prev_sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - prev_sample_mean (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)` for images): - Mean averaged `prev_sample`. Same as `prev_sample`, only mean-averaged over previous timesteps. - """ - - state: ScoreSdeVeSchedulerState - prev_sample: jnp.ndarray - prev_sample_mean: jnp.ndarray | None = None - - -class FlaxScoreSdeVeScheduler(FlaxSchedulerMixin, ConfigMixin): - """ - The variance exploding stochastic differential equation (SDE) scheduler. - - For more information, see the original paper: https://huggingface.co/papers/2011.13456 - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - Args: - num_train_timesteps (`int`): number of diffusion steps used to train the model. - snr (`float`): - coefficient weighting the step from the model_output sample (from the network) to the random noise. - sigma_min (`float`): - initial noise scale for sigma sequence in sampling procedure. The minimum sigma should mirror the - distribution of the data. - sigma_max (`float`): maximum value used for the range of continuous timesteps passed into the model. - sampling_eps (`float`): the end value of sampling, where timesteps decrease progressively from 1 to - epsilon. - correct_steps (`int`): number of correction steps performed on a produced sample. - """ - - @property - def has_state(self): - return True - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 2000, - snr: float = 0.15, - sigma_min: float = 0.01, - sigma_max: float = 1348.0, - sampling_eps: float = 1e-5, - correct_steps: int = 1, - ): - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - - def create_state(self): - state = ScoreSdeVeSchedulerState.create() - return self.set_sigmas( - state, - self.config.num_train_timesteps, - self.config.sigma_min, - self.config.sigma_max, - self.config.sampling_eps, - ) - - def set_timesteps( - self, - state: ScoreSdeVeSchedulerState, - num_inference_steps: int, - shape: tuple = (), - sampling_eps: float = None, - ) -> ScoreSdeVeSchedulerState: - """ - Sets the continuous timesteps used for the diffusion chain. Supporting function to be run before inference. - - Args: - state (`ScoreSdeVeSchedulerState`): the `FlaxScoreSdeVeScheduler` state data class instance. - num_inference_steps (`int`): - the number of diffusion steps used when generating samples with a pre-trained model. - sampling_eps (`float`, optional): - final timestep value (overrides value given at Scheduler instantiation). - - """ - sampling_eps = sampling_eps if sampling_eps is not None else self.config.sampling_eps - - timesteps = jnp.linspace(1, sampling_eps, num_inference_steps) - return state.replace(timesteps=timesteps) - - def set_sigmas( - self, - state: ScoreSdeVeSchedulerState, - num_inference_steps: int, - sigma_min: float = None, - sigma_max: float = None, - sampling_eps: float = None, - ) -> ScoreSdeVeSchedulerState: - """ - Sets the noise scales used for the diffusion chain. Supporting function to be run before inference. - - The sigmas control the weight of the `drift` and `diffusion` components of sample update. - - Args: - state (`ScoreSdeVeSchedulerState`): the `FlaxScoreSdeVeScheduler` state data class instance. - num_inference_steps (`int`): - the number of diffusion steps used when generating samples with a pre-trained model. - sigma_min (`float`, optional): - initial noise scale value (overrides value given at Scheduler instantiation). - sigma_max (`float`, optional): - final noise scale value (overrides value given at Scheduler instantiation). - sampling_eps (`float`, optional): - final timestep value (overrides value given at Scheduler instantiation). - """ - sigma_min = sigma_min if sigma_min is not None else self.config.sigma_min - sigma_max = sigma_max if sigma_max is not None else self.config.sigma_max - sampling_eps = sampling_eps if sampling_eps is not None else self.config.sampling_eps - if state.timesteps is None: - state = self.set_timesteps(state, num_inference_steps, sampling_eps) - - discrete_sigmas = jnp.exp(jnp.linspace(jnp.log(sigma_min), jnp.log(sigma_max), num_inference_steps)) - sigmas = jnp.array([sigma_min * (sigma_max / sigma_min) ** t for t in state.timesteps]) - - return state.replace(discrete_sigmas=discrete_sigmas, sigmas=sigmas) - - def get_adjacent_sigma(self, state, timesteps, t): - return jnp.where(timesteps == 0, jnp.zeros_like(t), state.discrete_sigmas[timesteps - 1]) - - def step_pred( - self, - state: ScoreSdeVeSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - key: jax.Array, - return_dict: bool = True, - ) -> FlaxSdeVeOutput | tuple: - """ - Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion - process from the learned model outputs (most often the predicted noise). - - Args: - state (`ScoreSdeVeSchedulerState`): the `FlaxScoreSdeVeScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - generator: random number generator. - return_dict (`bool`): option for returning tuple rather than FlaxSdeVeOutput class - - Returns: - [`FlaxSdeVeOutput`] or `tuple`: [`FlaxSdeVeOutput`] if `return_dict` is True, otherwise a `tuple`. When - returning a tuple, the first element is the sample tensor. - - """ - if state.timesteps is None: - raise ValueError( - "`state.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler" - ) - - timestep = timestep * jnp.ones( - sample.shape[0], - ) - timesteps = (timestep * (len(state.timesteps) - 1)).long() - - sigma = state.discrete_sigmas[timesteps] - adjacent_sigma = self.get_adjacent_sigma(state, timesteps, timestep) - drift = jnp.zeros_like(sample) - diffusion = (sigma**2 - adjacent_sigma**2) ** 0.5 - - # equation 6 in the paper: the model_output modeled by the network is grad_x log pt(x) - # also equation 47 shows the analog from SDE models to ancestral sampling methods - diffusion = diffusion.flatten() - diffusion = broadcast_to_shape_from_left(diffusion, sample.shape) - drift = drift - diffusion**2 * model_output - - # equation 6: sample noise for the diffusion term of - key = random.split(key, num=1) - noise = random.normal(key=key, shape=sample.shape) - prev_sample_mean = sample - drift # subtract because `dt` is a small negative timestep - # TODO is the variable diffusion the correct scaling term for the noise? - prev_sample = prev_sample_mean + diffusion * noise # add impact of diffusion field g - - if not return_dict: - return (prev_sample, prev_sample_mean, state) - - return FlaxSdeVeOutput(prev_sample=prev_sample, prev_sample_mean=prev_sample_mean, state=state) - - def step_correct( - self, - state: ScoreSdeVeSchedulerState, - model_output: jnp.ndarray, - sample: jnp.ndarray, - key: jax.Array, - return_dict: bool = True, - ) -> FlaxSdeVeOutput | tuple: - """ - Correct the predicted sample based on the output model_output of the network. This is often run repeatedly - after making the prediction for the previous timestep. - - Args: - state (`ScoreSdeVeSchedulerState`): the `FlaxScoreSdeVeScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - generator: random number generator. - return_dict (`bool`): option for returning tuple rather than FlaxSdeVeOutput class - - Returns: - [`FlaxSdeVeOutput`] or `tuple`: [`FlaxSdeVeOutput`] if `return_dict` is True, otherwise a `tuple`. When - returning a tuple, the first element is the sample tensor. - - """ - if state.timesteps is None: - raise ValueError( - "`state.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler" - ) - - # For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z" - # sample noise for correction - key = random.split(key, num=1) - noise = random.normal(key=key, shape=sample.shape) - - # compute step size from the model_output, the noise, and the snr - grad_norm = jnp.linalg.norm(model_output) - noise_norm = jnp.linalg.norm(noise) - step_size = (self.config.snr * noise_norm / grad_norm) ** 2 * 2 - step_size = step_size * jnp.ones(sample.shape[0]) - - # compute corrected sample: model_output term and noise term - step_size = step_size.flatten() - step_size = broadcast_to_shape_from_left(step_size, sample.shape) - prev_sample_mean = sample + step_size * model_output - prev_sample = prev_sample_mean + ((step_size * 2) ** 0.5) * noise - - if not return_dict: - return (prev_sample, state) - - return FlaxSdeVeOutput(prev_sample=prev_sample, state=state) - - def __len__(self): - return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_utils_flax.py b/src/diffusers/schedulers/scheduling_utils_flax.py deleted file mode 100644 index eea27dec0f81..000000000000 --- a/src/diffusers/schedulers/scheduling_utils_flax.py +++ /dev/null @@ -1,289 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import importlib -import math -import os -from dataclasses import dataclass -from enum import Enum - -import flax -import jax.numpy as jnp -from huggingface_hub.utils import validate_hf_hub_args - -from ..utils import BaseOutput, PushToHubMixin, logging - - -logger = logging.get_logger(__name__) - -SCHEDULER_CONFIG_NAME = "scheduler_config.json" - - -# NOTE: We make this type an enum because it simplifies usage in docs and prevents -# circular imports when used for `_compatibles` within the schedulers module. -# When it's used as a type in pipelines, it really is a Union because the actual -# scheduler instance is passed in. -class FlaxKarrasDiffusionSchedulers(Enum): - FlaxDDIMScheduler = 1 - FlaxDDPMScheduler = 2 - FlaxPNDMScheduler = 3 - FlaxLMSDiscreteScheduler = 4 - FlaxDPMSolverMultistepScheduler = 5 - FlaxEulerDiscreteScheduler = 6 - - -@dataclass -class FlaxSchedulerOutput(BaseOutput): - """ - Base class for the scheduler's step function output. - - Args: - prev_sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - """ - - prev_sample: jnp.ndarray - - -class FlaxSchedulerMixin(PushToHubMixin): - """ - Mixin containing common functions for the schedulers. - - Class attributes: - - **_compatibles** (`list[str]`) -- A list of classes that are compatible with the parent class, so that - `from_config` can be used from a class different than the one used to save the config (should be overridden - by parent class). - """ - - config_name = SCHEDULER_CONFIG_NAME - ignore_for_config = ["dtype"] - _compatibles = [] - has_compatibles = True - - @classmethod - @validate_hf_hub_args - def from_pretrained( - cls, - pretrained_model_name_or_path: str | os.PathLike | None = None, - subfolder: str | None = None, - return_unused_kwargs=False, - **kwargs, - ): - r""" - Instantiate a Scheduler class from a pre-defined JSON-file. - - Parameters: - pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): - Can be either: - - - A string, the *model id* of a model repo on huggingface.co. Valid model ids should have an - organization name, like `google/ddpm-celebahq-256`. - - A path to a *directory* containing model weights saved using [`~SchedulerMixin.save_pretrained`], - e.g., `./my_model_directory/`. - subfolder (`str`, *optional*): - In case the relevant files are located inside a subfolder of the model repo (either remote in - huggingface.co or downloaded locally), you can specify the folder name here. - return_unused_kwargs (`bool`, *optional*, defaults to `False`): - Whether kwargs that are not consumed by the Python class should be returned or not. - - cache_dir (`str | os.PathLike`, *optional*): - Path to a directory in which a downloaded pretrained model configuration should be cached if the - standard cache should not be used. - force_download (`bool`, *optional*, defaults to `False`): - Whether or not to force the (re-)download of the model weights and configuration files, overriding the - cached versions if they exist. - - proxies (`dict[str, str]`, *optional*): - A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128', - 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. - output_loading_info(`bool`, *optional*, defaults to `False`): - Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. - local_files_only(`bool`, *optional*, defaults to `False`): - Whether or not to only look at local files (i.e., do not try to download the model). - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated - when running `transformers-cli login` (stored in `~/.huggingface`). - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a - git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any - identifier allowed by git. - - > [!TIP] > It is required to be logged in (`hf auth login`) when you want to use private or [gated > - models](https://huggingface.co/docs/hub/models-gated#gated-models). - - > [!TIP] > Activate the special - ["offline-mode"](https://huggingface.co/transformers/installation.html#offline-mode) to > use this method in a - firewalled environment. - - """ - logger.warning( - "Flax classes are deprecated and will be removed in Diffusers v0.40.0. We " - "recommend migrating to PyTorch classes or pinning your version of Diffusers." - ) - config, kwargs = cls.load_config( - pretrained_model_name_or_path=pretrained_model_name_or_path, - subfolder=subfolder, - return_unused_kwargs=True, - **kwargs, - ) - scheduler, unused_kwargs = cls.from_config(config, return_unused_kwargs=True, **kwargs) - - if hasattr(scheduler, "create_state") and getattr(scheduler, "has_state", False): - state = scheduler.create_state() - - if return_unused_kwargs: - return scheduler, state, unused_kwargs - - return scheduler, state - - def save_pretrained(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs): - """ - Save a scheduler configuration object to the directory `save_directory`, so that it can be re-loaded using the - [`~FlaxSchedulerMixin.from_pretrained`] class method. - - Args: - save_directory (`str` or `os.PathLike`): - Directory where the configuration JSON file will be saved (will be created if it does not exist). - push_to_hub (`bool`, *optional*, defaults to `False`): - Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the - repository you want to push to with `repo_id` (will default to the name of `save_directory` in your - namespace). - kwargs (`dict[str, Any]`, *optional*): - Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. - """ - self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs) - - @property - def compatibles(self): - """ - Returns all schedulers that are compatible with this scheduler - - Returns: - `list[SchedulerMixin]`: list of compatible schedulers - """ - return self._get_compatibles() - - @classmethod - def _get_compatibles(cls): - compatible_classes_str = list(set([cls.__name__] + cls._compatibles)) - diffusers_library = importlib.import_module(__name__.split(".")[0]) - compatible_classes = [ - getattr(diffusers_library, c) for c in compatible_classes_str if hasattr(diffusers_library, c) - ] - return compatible_classes - - -def broadcast_to_shape_from_left(x: jnp.ndarray, shape: tuple[int]) -> jnp.ndarray: - assert len(shape) >= x.ndim - return jnp.broadcast_to(x.reshape(x.shape + (1,) * (len(shape) - x.ndim)), shape) - - -def betas_for_alpha_bar(num_diffusion_timesteps: int, max_beta=0.999, dtype=jnp.float32) -> jnp.ndarray: - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - - Returns: - betas (`jnp.ndarray`): the betas used by the scheduler to step the model outputs - """ - - def alpha_bar(time_step): - return math.cos((time_step + 0.008) / 1.008 * math.pi / 2) ** 2 - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) - return jnp.array(betas, dtype=dtype) - - -@flax.struct.dataclass -class CommonSchedulerState: - alphas: jnp.ndarray - betas: jnp.ndarray - alphas_cumprod: jnp.ndarray - - @classmethod - def create(cls, scheduler): - config = scheduler.config - - if config.trained_betas is not None: - betas = jnp.asarray(config.trained_betas, dtype=scheduler.dtype) - elif config.beta_schedule == "linear": - betas = jnp.linspace(config.beta_start, config.beta_end, config.num_train_timesteps, dtype=scheduler.dtype) - elif config.beta_schedule == "scaled_linear": - # this schedule is very specific to the latent diffusion model. - betas = ( - jnp.linspace( - config.beta_start**0.5, config.beta_end**0.5, config.num_train_timesteps, dtype=scheduler.dtype - ) - ** 2 - ) - elif config.beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - betas = betas_for_alpha_bar(config.num_train_timesteps, dtype=scheduler.dtype) - else: - raise NotImplementedError( - f"beta_schedule {config.beta_schedule} is not implemented for scheduler {scheduler.__class__.__name__}" - ) - - alphas = 1.0 - betas - - alphas_cumprod = jnp.cumprod(alphas, axis=0) - - return cls( - alphas=alphas, - betas=betas, - alphas_cumprod=alphas_cumprod, - ) - - -def get_sqrt_alpha_prod( - state: CommonSchedulerState, original_samples: jnp.ndarray, noise: jnp.ndarray, timesteps: jnp.ndarray -): - alphas_cumprod = state.alphas_cumprod - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - sqrt_alpha_prod = broadcast_to_shape_from_left(sqrt_alpha_prod, original_samples.shape) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - sqrt_one_minus_alpha_prod = broadcast_to_shape_from_left(sqrt_one_minus_alpha_prod, original_samples.shape) - - return sqrt_alpha_prod, sqrt_one_minus_alpha_prod - - -def add_noise_common( - state: CommonSchedulerState, original_samples: jnp.ndarray, noise: jnp.ndarray, timesteps: jnp.ndarray -): - sqrt_alpha_prod, sqrt_one_minus_alpha_prod = get_sqrt_alpha_prod(state, original_samples, noise, timesteps) - noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise - return noisy_samples - - -def get_velocity_common(state: CommonSchedulerState, sample: jnp.ndarray, noise: jnp.ndarray, timesteps: jnp.ndarray): - sqrt_alpha_prod, sqrt_one_minus_alpha_prod = get_sqrt_alpha_prod(state, sample, noise, timesteps) - velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample - return velocity diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index 5cd6885e0364..7257518a22a8 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -26,7 +26,6 @@ DIFFUSERS_LOAD_ID_FIELDS, FLASHPACK_FILE_EXTENSION, FLASHPACK_WEIGHTS_NAME, - FLAX_WEIGHTS_NAME, GGUF_FILE_EXTENSION, HF_ENABLE_PARALLEL_LOADING, HF_MODULES_CACHE, @@ -58,7 +57,6 @@ DIFFUSERS_SLOW_IMPORT, ENV_VARS_TRUE_AND_AUTO_VALUES, ENV_VARS_TRUE_VALUES, - USE_JAX, USE_TF, USE_TORCH, DummyObject, @@ -80,7 +78,6 @@ is_flash_attn_available, is_flash_attn_version, is_flashpack_available, - is_flax_available, is_ftfy_available, is_gguf_available, is_gguf_version, @@ -125,7 +122,6 @@ is_torchsde_available, is_torchvision_available, is_transformers_available, - is_transformers_flax_compatible, is_transformers_version, is_unidecode_available, is_wandb_available, diff --git a/src/diffusers/utils/constants.py b/src/diffusers/utils/constants.py index cbfe2da0d32a..edca44d7a2dd 100644 --- a/src/diffusers/utils/constants.py +++ b/src/diffusers/utils/constants.py @@ -29,7 +29,6 @@ CONFIG_NAME = "config.json" WEIGHTS_NAME = "diffusion_pytorch_model.bin" WEIGHTS_INDEX_NAME = "diffusion_pytorch_model.bin.index.json" -FLAX_WEIGHTS_NAME = "diffusion_flax_model.msgpack" ONNX_WEIGHTS_NAME = "model.onnx" SAFETENSORS_WEIGHTS_NAME = "diffusion_pytorch_model.safetensors" SAFE_WEIGHTS_INDEX_NAME = "diffusion_pytorch_model.safetensors.index.json" diff --git a/src/diffusers/utils/dummy_flax_and_transformers_objects.py b/src/diffusers/utils/dummy_flax_and_transformers_objects.py deleted file mode 100644 index 49d34c251e2d..000000000000 --- a/src/diffusers/utils/dummy_flax_and_transformers_objects.py +++ /dev/null @@ -1,92 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class FlaxDiffusionPipeline(metaclass=DummyObject): - _backends = ["flax", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - -class FlaxStableDiffusionControlNetPipeline(metaclass=DummyObject): - _backends = ["flax", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - -class FlaxStableDiffusionImg2ImgPipeline(metaclass=DummyObject): - _backends = ["flax", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - -class FlaxStableDiffusionInpaintPipeline(metaclass=DummyObject): - _backends = ["flax", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - -class FlaxStableDiffusionPipeline(metaclass=DummyObject): - _backends = ["flax", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - -class FlaxStableDiffusionXLPipeline(metaclass=DummyObject): - _backends = ["flax", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) diff --git a/src/diffusers/utils/dummy_flax_objects.py b/src/diffusers/utils/dummy_flax_objects.py deleted file mode 100644 index 181dfea9459f..000000000000 --- a/src/diffusers/utils/dummy_flax_objects.py +++ /dev/null @@ -1,197 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class FlaxControlNetModel(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxModelMixin(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxUNet2DConditionModel(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxAutoencoderKL(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxDDIMScheduler(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxDDPMScheduler(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxDPMSolverMultistepScheduler(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxEulerDiscreteScheduler(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxKarrasVeScheduler(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxLMSDiscreteScheduler(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxPNDMScheduler(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxSchedulerMixin(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxScoreSdeVeScheduler(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) diff --git a/src/diffusers/utils/dummy_transformers_flax_objects.py b/src/diffusers/utils/dummy_transformers_flax_objects.py deleted file mode 100644 index 4e0e9ef94bc4..000000000000 --- a/src/diffusers/utils/dummy_transformers_flax_objects.py +++ /dev/null @@ -1,92 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class FlaxDiffusionPipeline(metaclass=DummyObject): - _backends = ["transformers_flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["transformers_flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["transformers_flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["transformers_flax"]) - - -class FlaxStableDiffusionControlNetPipeline(metaclass=DummyObject): - _backends = ["transformers_flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["transformers_flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["transformers_flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["transformers_flax"]) - - -class FlaxStableDiffusionImg2ImgPipeline(metaclass=DummyObject): - _backends = ["transformers_flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["transformers_flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["transformers_flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["transformers_flax"]) - - -class FlaxStableDiffusionInpaintPipeline(metaclass=DummyObject): - _backends = ["transformers_flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["transformers_flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["transformers_flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["transformers_flax"]) - - -class FlaxStableDiffusionPipeline(metaclass=DummyObject): - _backends = ["transformers_flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["transformers_flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["transformers_flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["transformers_flax"]) - - -class FlaxStableDiffusionXLPipeline(metaclass=DummyObject): - _backends = ["transformers_flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["transformers_flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["transformers_flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["transformers_flax"]) diff --git a/src/diffusers/utils/hub_utils.py b/src/diffusers/utils/hub_utils.py index b5eb9ab2e17f..64292f336d7a 100644 --- a/src/diffusers/utils/hub_utils.py +++ b/src/diffusers/utils/hub_utils.py @@ -54,11 +54,8 @@ ) from .import_utils import ( ENV_VARS_TRUE_VALUES, - _flax_version, - _jax_version, _onnxruntime_version, _torch_version, - is_flax_available, is_onnx_available, is_torch_available, ) @@ -80,9 +77,6 @@ def http_user_agent(user_agent: dict | str | None = None) -> str: return ua + "; telemetry/off" if is_torch_available(): ua += f"; torch/{_torch_version}" - if is_flax_available(): - ua += f"; jax/{_jax_version}" - ua += f"; flax/{_flax_version}" if is_onnx_available(): ua += f"; onnxruntime/{_onnxruntime_version}" # CI will set this value to True diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index 18c189aa26f2..30e00f74bc49 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -49,7 +49,6 @@ USE_TF = os.environ.get("USE_TF", "AUTO").upper() USE_TORCH = os.environ.get("USE_TORCH", "AUTO").upper() -USE_JAX = os.environ.get("USE_FLAX", "AUTO").upper() USE_SAFETENSORS = os.environ.get("USE_SAFETENSORS", "AUTO").upper() DIFFUSERS_SLOW_IMPORT = os.environ.get("DIFFUSERS_SLOW_IMPORT", "FALSE").upper() DIFFUSERS_SLOW_IMPORT = DIFFUSERS_SLOW_IMPORT in ENV_VARS_TRUE_VALUES @@ -103,20 +102,6 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> tuple[b _torch_available = False _torch_version = "N/A" -_jax_version = "N/A" -_flax_version = "N/A" -if USE_JAX in ENV_VARS_TRUE_AND_AUTO_VALUES: - _flax_available = importlib.util.find_spec("jax") is not None and importlib.util.find_spec("flax") is not None - if _flax_available: - try: - _jax_version = importlib_metadata.version("jax") - _flax_version = importlib_metadata.version("flax") - logger.info(f"JAX version {_jax_version}, Flax version {_flax_version} available.") - except importlib_metadata.PackageNotFoundError: - _flax_available = False -else: - _flax_available = False - if USE_SAFETENSORS in ENV_VARS_TRUE_AND_AUTO_VALUES: _safetensors_available, _safetensors_version = _is_package_available("safetensors") @@ -257,30 +242,10 @@ def is_torch_neuronx_available(): return _torch_neuronx_available -def is_flax_available(): - return _flax_available - - def is_transformers_available(): return _transformers_available -def is_transformers_flax_compatible(): - # Flax classes (e.g. FlaxCLIPTextModel, FlaxPreTrainedModel) were removed from - # transformers main on the path to its v5 release. Gate Flax pipeline registration - # on transformers still shipping them so `import diffusers` doesn't crash. - # Name avoids the `is_*_available()` pattern so utils/check_dummies.py keeps - # generating the `flax_and_transformers` backend group when this is combined with - # the legacy is_flax_available()/is_transformers_available() pair. - if not (_transformers_available and _flax_available): - return False - try: - import transformers - except ImportError: - return False - return hasattr(transformers, "FlaxPreTrainedModel") - - def is_inflect_available(): return _inflect_available @@ -457,12 +422,6 @@ def is_av_available(): return _av_available -# docstyle-ignore -FLAX_IMPORT_ERROR = """ -{0} requires the FLAX library but it was not found in your environment. Checkout the instructions on the -installation page: https://github.com/google/flax and follow the ones that match your environment. -""" - # docstyle-ignore INFLECT_IMPORT_ERROR = """ {0} requires the inflect library but it was not found in your environment. You can install it with pip: `pip install @@ -623,7 +582,6 @@ def is_av_available(): BACKENDS_MAPPING = OrderedDict( [ ("bs4", (is_bs4_available, BS4_IMPORT_ERROR)), - ("flax", (is_flax_available, FLAX_IMPORT_ERROR)), ("inflect", (is_inflect_available, INFLECT_IMPORT_ERROR)), ("onnx", (is_onnx_available, ONNX_IMPORT_ERROR)), ("opencv", (is_opencv_available, OPENCV_IMPORT_ERROR)), diff --git a/src/diffusers/utils/testing_utils.py b/src/diffusers/utils/testing_utils.py index eefe52c477a6..1de14ae8539e 100644 --- a/src/diffusers/utils/testing_utils.py +++ b/src/diffusers/utils/testing_utils.py @@ -35,7 +35,6 @@ is_accelerate_available, is_bitsandbytes_available, is_compel_available, - is_flax_available, is_gguf_available, is_kernels_available, is_note_seq_available, @@ -470,13 +469,6 @@ def skip_mps(test_case): return unittest.skipUnless(torch_device != "mps", "test requires non 'mps' device")(test_case) -def require_flax(test_case): - """ - Decorator marking a test that requires JAX & Flax. These tests are skipped when one / both are not installed - """ - return unittest.skipUnless(is_flax_available(), "test requires JAX & Flax")(test_case) - - def require_compel(test_case): """ Decorator marking a test that requires compel: https://github.com/damian0815/compel. These tests are skipped when diff --git a/tests/others/test_check_dummies.py b/tests/others/test_check_dummies.py index cf0bed3c34a8..b9ddc2764465 100644 --- a/tests/others/test_check_dummies.py +++ b/tests/others/test_check_dummies.py @@ -54,14 +54,11 @@ def test_read_init(self): # We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects self.assertIn("torch", objects) self.assertIn("torch_and_transformers", objects) - self.assertIn("flax_and_transformers", objects) self.assertIn("torch_and_transformers_and_onnx", objects) # Likewise, we can't assert on the exact content of a key self.assertIn("UNet2DModel", objects["torch"]) - self.assertIn("FlaxUNet2DConditionModel", objects["flax"]) self.assertIn("StableDiffusionPipeline", objects["torch_and_transformers"]) - self.assertIn("FlaxStableDiffusionPipeline", objects["flax_and_transformers"]) self.assertIn("LMSDiscreteScheduler", objects["torch_and_scipy"]) self.assertIn("OnnxStableDiffusionPipeline", objects["torch_and_transformers_and_onnx"]) diff --git a/tests/pipelines/test_pipeline_utils.py b/tests/pipelines/test_pipeline_utils.py index e0c73b96b8f9..b5630780e44f 100644 --- a/tests/pipelines/test_pipeline_utils.py +++ b/tests/pipelines/test_pipeline_utils.py @@ -629,17 +629,6 @@ def test_download_onnx_models(self): ) assert model_filenames == set(filenames) - def test_download_flax_models(self): - ignore_patterns = ["*.safetensors", "*.bin"] - filenames = [ - "vae/diffusion_flax_model.msgpack", - "unet/diffusion_flax_model.msgpack", - ] - model_filenames, variant_filenames = variant_compatible_siblings( - filenames, variant=None, ignore_patterns=ignore_patterns - ) - assert model_filenames == set(filenames) - class ProgressBarTests(unittest.TestCase): def get_dummy_components_image_generation(self): diff --git a/tests/pipelines/test_pipelines.py b/tests/pipelines/test_pipelines.py index a193fb85182d..7e1af7c1dc44 100644 --- a/tests/pipelines/test_pipelines.py +++ b/tests/pipelines/test_pipelines.py @@ -61,7 +61,6 @@ UniPCMultistepScheduler, logging, ) -from diffusers.pipelines.pipeline_utils import _get_pipeline_class from diffusers.schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME from diffusers.utils import CONFIG_NAME, WEIGHTS_NAME, is_transformers_version from diffusers.utils.torch_utils import is_compiled_module @@ -77,7 +76,6 @@ load_numpy, nightly, require_compel, - require_flax, require_hf_hub_version_greater, require_onnxruntime, require_peft_backend, @@ -1018,14 +1016,6 @@ def test_download_dduf_with_connected_pipeline_raises_error(self): "DDUF/tiny-flux-dev-pipe-dduf", dduf_file="fluxpipeline.dduf", load_connected_pipeline=True ) - def test_get_pipeline_class_from_flax(self): - flax_config = {"_class_name": "FlaxStableDiffusionPipeline"} - config = {"_class_name": "StableDiffusionPipeline"} - - # when loading a PyTorch Pipeline from a FlaxPipeline `model_index.json`, e.g.: https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-lms-pipe/blob/7a9063578b325779f0f1967874a6771caa973cad/model_index.json#L2 - # we need to make sure that we don't load the Flax Pipeline class, but instead the PyTorch pipeline class - assert _get_pipeline_class(DiffusionPipeline, flax_config) == _get_pipeline_class(DiffusionPipeline, config) - class CustomPipelineTests(unittest.TestCase): def test_load_custom_pipeline(self): @@ -2229,47 +2219,6 @@ def test_output_format(self): assert isinstance(images, list) assert isinstance(images[0], PIL.Image.Image) - @require_flax - def test_from_flax_from_pt(self): - pipe_pt = StableDiffusionPipeline.from_pretrained( - "hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None - ) - pipe_pt.to(torch_device) - - from diffusers import FlaxStableDiffusionPipeline - - with tempfile.TemporaryDirectory() as tmpdirname: - pipe_pt.save_pretrained(tmpdirname) - - pipe_flax, params = FlaxStableDiffusionPipeline.from_pretrained( - tmpdirname, safety_checker=None, from_pt=True - ) - - with tempfile.TemporaryDirectory() as tmpdirname: - pipe_flax.save_pretrained(tmpdirname, params=params) - pipe_pt_2 = StableDiffusionPipeline.from_pretrained(tmpdirname, safety_checker=None, from_flax=True) - pipe_pt_2.to(torch_device) - - prompt = "Hello" - - generator = torch.manual_seed(0) - image_0 = pipe_pt( - [prompt], - generator=generator, - num_inference_steps=2, - output_type="np", - ).images[0] - - generator = torch.manual_seed(0) - image_1 = pipe_pt_2( - [prompt], - generator=generator, - num_inference_steps=2, - output_type="np", - ).images[0] - - assert np.abs(image_0 - image_1).sum() < 1e-5, "Models don't give the same forward pass" - @require_compel def test_weighted_prompts_compel(self): from compel import Compel diff --git a/tests/quantization/bnb/README.md b/tests/quantization/bnb/README.md index f1585581597d..c4bdd53b765e 100644 --- a/tests/quantization/bnb/README.md +++ b/tests/quantization/bnb/README.md @@ -28,9 +28,6 @@ They were conducted on the `audace` machine, using a single RTX 4090. Below is ` - Running on Google Colab?: No - Python version: 3.10.12 - PyTorch version (GPU?): 2.5.0.dev20240818+cu124 (True) -- Flax version (CPU?/GPU?/TPU?): not installed (NA) -- Jax version: not installed -- JaxLib version: not installed - Huggingface_hub version: 0.24.5 - Transformers version: 4.44.2 - Accelerate version: 0.34.0.dev0 diff --git a/tests/quantization/torchao/README.md b/tests/quantization/torchao/README.md index 373593091ac0..1b06be1b83e0 100644 --- a/tests/quantization/torchao/README.md +++ b/tests/quantization/torchao/README.md @@ -40,9 +40,6 @@ HF_XET_HIGH_PERFORMANCE=1 RUN_SLOW=1 pytest -s tests/quantization/torchao/test_t - Running on Google Colab?: No - Python version: 3.10.14 - PyTorch version (GPU?): 2.6.0.dev20241112+cu121 (False) -- Flax version (CPU?/GPU?/TPU?): not installed (NA) -- Jax version: not installed -- JaxLib version: not installed - Huggingface_hub version: 0.26.2 - Transformers version: 4.46.3 - Accelerate version: 1.1.1 diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 86887d7af6e9..0fb03e92028f 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -36,7 +36,6 @@ is_bitsandbytes_available, is_compel_available, is_flashpack_available, - is_flax_available, is_gguf_available, is_kernels_available, is_note_seq_available, @@ -673,13 +672,6 @@ def skip_mps(test_case): return pytest.mark.skipif(torch_device == "mps", reason="test requires non 'mps' device")(test_case) -def require_flax(test_case): - """ - Decorator marking a test that requires JAX & Flax. These tests are skipped when one / both are not installed - """ - return pytest.mark.skipif(not is_flax_available(), reason="test requires JAX & Flax")(test_case) - - def require_compel(test_case): """ Decorator marking a test that requires compel: https://github.com/damian0815/compel. These tests are skipped when diff --git a/utils/check_inits.py b/utils/check_inits.py index 8208fa634186..7679a444861e 100644 --- a/utils/check_inits.py +++ b/utils/check_inits.py @@ -267,7 +267,6 @@ def get_transformers_submodules(): IGNORE_SUBMODULES = [ "convert_pytorch_checkpoint_to_tf2", - "modeling_flax_pytorch_utils", ] diff --git a/utils/check_repo.py b/utils/check_repo.py index 14bdbe60adf0..7eb79f6a394a 100644 --- a/utils/check_repo.py +++ b/utils/check_repo.py @@ -23,7 +23,7 @@ from pathlib import Path from diffusers.models.auto import get_values -from diffusers.utils import ENV_VARS_TRUE_VALUES, is_flax_available, is_torch_available +from diffusers.utils import ENV_VARS_TRUE_VALUES, is_torch_available # All paths are set with the intent you should run this script from the root of the repo with the command @@ -270,7 +270,7 @@ def get_model_modules(): def get_models(module, include_pretrained=False): """Get the objects in module that are models.""" models = [] - model_classes = (diffusers.ModelMixin, diffusers.TFModelMixin, diffusers.FlaxModelMixin) + model_classes = (diffusers.ModelMixin, diffusers.TFModelMixin) for attr_name in dir(module): if not include_pretrained and ("Pretrained" in attr_name or "PreTrained" in attr_name): continue @@ -421,10 +421,6 @@ def get_all_auto_configured_models(): for attr_name in dir(diffusers.models.auto.modeling_auto): if attr_name.startswith("MODEL_") and attr_name.endswith("MAPPING_NAMES"): result = result | set(get_values(getattr(diffusers.models.auto.modeling_auto, attr_name))) - if is_flax_available(): - for attr_name in dir(diffusers.models.auto.modeling_flax_auto): - if attr_name.startswith("FLAX_MODEL_") and attr_name.endswith("MAPPING_NAMES"): - result = result | set(get_values(getattr(diffusers.models.auto.modeling_flax_auto, attr_name))) return list(result) @@ -458,8 +454,6 @@ def check_all_models_are_auto_configured(): missing_backends = [] if not is_torch_available(): missing_backends.append("PyTorch") - if not is_flax_available(): - missing_backends.append("Flax") if len(missing_backends) > 0: missing = ", ".join(missing_backends) if os.getenv("TRANSFORMERS_IS_CI", "").upper() in ENV_VARS_TRUE_VALUES: diff --git a/utils/tests_fetcher.py b/utils/tests_fetcher.py index 59ef5499cc2a..d487efb40518 100644 --- a/utils/tests_fetcher.py +++ b/utils/tests_fetcher.py @@ -703,7 +703,7 @@ def init_test_examples_dependencies() -> Tuple[Dict[str, List[str]], List[str]]: """ test_example_deps = {} all_examples = [] - for framework in ["flax", "pytorch", "tensorflow"]: + for framework in ["pytorch", "tensorflow"]: test_files = list((PATH_TO_EXAMPLES / framework).glob("test_*.py")) all_examples.extend(test_files) # Remove the files at the root of examples/framework since they are not proper examples (they are either utils