Suraj Patil e7120bae95 [UNet2DConditionModel] add gradient checkpointing (#461) ...

* add grad ckpt to downsample blocks

* make it work

* don't pass gradient_checkpointing to upsample block

* add tests for UNet2DConditionModel

* add test_gradient_checkpointing

* add gradient_checkpointing for up and down blocks

* add functions to enable and disable grad ckpt

* remove the forward argument

* better naming

* make supports_gradient_checkpointing private

2022-09-22 15:36:47 +02:00

..

__init__.py

fix issues with loading, add test for pipeline

2022-06-07 15:40:36 +02:00

conftest.py

[Tests] Upload custom test artifacts (#572)

2022-09-19 17:08:29 +02:00

test_config.py

[Tests] Make sure tests are on GPU (#269)

2022-08-29 15:58:11 +02:00

test_layers_utils.py

[Tests] Fix spatial transformer tests on GPU (#531)

2022-09-16 16:04:37 +02:00

test_modeling_common.py

[UNet2DConditionModel] add gradient checkpointing (#461)

2022-09-22 15:36:47 +02:00

test_models_unet.py

[UNet2DConditionModel] add gradient checkpointing (#461)

2022-09-22 15:36:47 +02:00

test_models_vae.py

Fix vae tests for cpu and gpu (#480)

2022-09-13 19:14:20 +02:00

test_models_vq.py

update expected results of slow tests (#268)

2022-09-12 15:49:39 +02:00

test_outputs.py

Fix BaseOutput initialization from dict (#570)

2022-09-20 18:32:16 +02:00

test_pipelines.py

[ONNX] Collate the external weights, speed up loading from the hub (#610)

2022-09-21 22:26:30 +02:00

test_scheduler.py

Unify offset configuration in DDIM and PNDM schedulers (#479)

2022-09-17 14:07:43 +02:00

test_training.py

Karras VE, DDIM and DDPM flax schedulers (#508)

2022-09-15 15:55:48 +02:00