From 4cb27b02762985bd445531e34324a6cb448b19a2 Mon Sep 17 00:00:00 2001 From: Prathik Rao Date: Thu, 22 Jun 2023 22:26:34 +0000 Subject: [PATCH 1/6] add default to unet output to prevent it from being a required arg --- src/diffusers/models/unet_2d_condition.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/models/unet_2d_condition.py b/src/diffusers/models/unet_2d_condition.py index 7bca5c336c57..0f9ff8d36be6 100644 --- a/src/diffusers/models/unet_2d_condition.py +++ b/src/diffusers/models/unet_2d_condition.py @@ -55,7 +55,7 @@ class UNet2DConditionOutput(BaseOutput): Hidden states conditioned on `encoder_hidden_states` input. Output of last layer of model. """ - sample: torch.FloatTensor + sample: torch.FloatTensor = None class UNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin): From 508b1992d4ccec7ffedea6b0f71d9151177179fb Mon Sep 17 00:00:00 2001 From: Prathik Rao Date: Wed, 28 Jun 2023 21:30:52 +0000 Subject: [PATCH 2/6] add unit test --- tests/models/test_models_unet_2d_condition.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/models/test_models_unet_2d_condition.py b/tests/models/test_models_unet_2d_condition.py index 24da508227d2..b863e0d8a8a7 100644 --- a/tests/models/test_models_unet_2d_condition.py +++ b/tests/models/test_models_unet_2d_condition.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy +from dataclasses import dataclass import gc import os import tempfile @@ -25,6 +27,7 @@ from diffusers import UNet2DConditionModel from diffusers.models.attention_processor import CustomDiffusionAttnProcessor, LoRAAttnProcessor from diffusers.utils import ( + BaseOutput, floats_tensor, load_hf_numpy, logging, @@ -1088,3 +1091,13 @@ def test_stabilityai_sd_v2_fp16(self, seed, timestep, expected_slice): expected_output_slice = torch.tensor(expected_slice) assert torch_all_close(output_slice, expected_output_slice, atol=5e-3) + + def test_pickle(): + @dataclass + class NetParams(BaseOutput): + sample: torch.FloatTensor + + m = NetParams(sample=torch.randn(1, 10)) + n = copy.copy(m) + + assert m == n \ No newline at end of file From 61595ef9ecf000da7845d8a31cd94ac48f5cf922 Mon Sep 17 00:00:00 2001 From: root Date: Wed, 28 Jun 2023 21:33:10 +0000 Subject: [PATCH 3/6] make style --- tests/models/test_models_unet_2d_condition.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/models/test_models_unet_2d_condition.py b/tests/models/test_models_unet_2d_condition.py index b863e0d8a8a7..9b5f8b8ea9ff 100644 --- a/tests/models/test_models_unet_2d_condition.py +++ b/tests/models/test_models_unet_2d_condition.py @@ -14,11 +14,11 @@ # limitations under the License. import copy -from dataclasses import dataclass import gc import os import tempfile import unittest +from dataclasses import dataclass import torch from parameterized import parameterized @@ -1100,4 +1100,4 @@ class NetParams(BaseOutput): m = NetParams(sample=torch.randn(1, 10)) n = copy.copy(m) - assert m == n \ No newline at end of file + assert m == n From 50467888fa246866b8bc2758ac5da01f079b037a Mon Sep 17 00:00:00 2001 From: root Date: Wed, 28 Jun 2023 21:38:52 +0000 Subject: [PATCH 4/6] adjust unit test --- tests/models/test_models_unet_2d_condition.py | 21 +++++++++++-------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/tests/models/test_models_unet_2d_condition.py b/tests/models/test_models_unet_2d_condition.py index 9b5f8b8ea9ff..603165b92a47 100644 --- a/tests/models/test_models_unet_2d_condition.py +++ b/tests/models/test_models_unet_2d_condition.py @@ -18,7 +18,6 @@ import os import tempfile import unittest -from dataclasses import dataclass import torch from parameterized import parameterized @@ -27,7 +26,6 @@ from diffusers import UNet2DConditionModel from diffusers.models.attention_processor import CustomDiffusionAttnProcessor, LoRAAttnProcessor from diffusers.utils import ( - BaseOutput, floats_tensor, load_hf_numpy, logging, @@ -1092,12 +1090,17 @@ def test_stabilityai_sd_v2_fp16(self, seed, timestep, expected_slice): assert torch_all_close(output_slice, expected_output_slice, atol=5e-3) - def test_pickle(): - @dataclass - class NetParams(BaseOutput): - sample: torch.FloatTensor + @require_torch_gpu + def test_pickle(self, seed, timestep): + model = self.get_unet_model(model_id="CompVis/stable-diffusion-v1-4") + latents = self.get_latents(seed) + encoder_hidden_states = self.get_encoder_hidden_states(seed) + + timestep = torch.tensor([timestep], dtype=torch.long, device=torch_device) + + with torch.no_grad(): + sample = model(latents, timestep=timestep, encoder_hidden_states=encoder_hidden_states).sample - m = NetParams(sample=torch.randn(1, 10)) - n = copy.copy(m) + sample_copy = copy.copy(sample) - assert m == n + assert sample == sample_copy From d7a565674578bf28cdd7dadd61c32b6188f8195f Mon Sep 17 00:00:00 2001 From: Prathik Rao Date: Wed, 5 Jul 2023 18:07:50 +0000 Subject: [PATCH 5/6] mark as fast test --- tests/models/test_models_unet_2d_condition.py | 31 ++++++++++--------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/tests/models/test_models_unet_2d_condition.py b/tests/models/test_models_unet_2d_condition.py index 603165b92a47..932cb90c5818 100644 --- a/tests/models/test_models_unet_2d_condition.py +++ b/tests/models/test_models_unet_2d_condition.py @@ -783,6 +783,22 @@ def test_custom_diffusion_xformers_on_off(self): assert (sample - on_sample).abs().max() < 1e-4 assert (sample - off_sample).abs().max() < 1e-4 + def test_pickle(self): + # enable deterministic behavior for gradient checkpointing + init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + + init_dict["attention_head_dim"] = (8, 16) + + model = self.model_class(**init_dict) + model.to(torch_device) + + with torch.no_grad(): + sample = model(**inputs_dict).sample + + sample_copy = copy.copy(sample) + + assert sample == sample_copy + @slow class UNet2DConditionModelIntegrationTests(unittest.TestCase): @@ -1089,18 +1105,3 @@ def test_stabilityai_sd_v2_fp16(self, seed, timestep, expected_slice): expected_output_slice = torch.tensor(expected_slice) assert torch_all_close(output_slice, expected_output_slice, atol=5e-3) - - @require_torch_gpu - def test_pickle(self, seed, timestep): - model = self.get_unet_model(model_id="CompVis/stable-diffusion-v1-4") - latents = self.get_latents(seed) - encoder_hidden_states = self.get_encoder_hidden_states(seed) - - timestep = torch.tensor([timestep], dtype=torch.long, device=torch_device) - - with torch.no_grad(): - sample = model(latents, timestep=timestep, encoder_hidden_states=encoder_hidden_states).sample - - sample_copy = copy.copy(sample) - - assert sample == sample_copy From 684b080379000a7f7cd64204cf3aeb4d7ff07341 Mon Sep 17 00:00:00 2001 From: Prathik Rao Date: Wed, 5 Jul 2023 18:13:19 +0000 Subject: [PATCH 6/6] adjust assert statement in test --- tests/models/test_models_unet_2d_condition.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/test_models_unet_2d_condition.py b/tests/models/test_models_unet_2d_condition.py index 932cb90c5818..4eeb1b926bec 100644 --- a/tests/models/test_models_unet_2d_condition.py +++ b/tests/models/test_models_unet_2d_condition.py @@ -797,7 +797,7 @@ def test_pickle(self): sample_copy = copy.copy(sample) - assert sample == sample_copy + assert (sample - sample_copy).abs().max() < 1e-4 @slow