|
| 1 | +# coding=utf-8 |
| 2 | +# Copyright 2024 HuggingFace Inc. |
| 3 | +# |
| 4 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +# you may not use this file except in compliance with the License. |
| 6 | +# You may obtain a copy of the License at |
| 7 | +# |
| 8 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +# |
| 10 | +# Unless required by applicable law or agreed to in writing, software |
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +# See the License for the specific language governing permissions and |
| 14 | +# limitations under the License. |
| 15 | + |
| 16 | +import gc |
| 17 | +import unittest |
| 18 | + |
| 19 | +import torch |
| 20 | +from parameterized import parameterized |
| 21 | + |
| 22 | +from diffusers import AsymmetricAutoencoderKL |
| 23 | +from diffusers.utils.import_utils import is_xformers_available |
| 24 | +from diffusers.utils.testing_utils import ( |
| 25 | + backend_empty_cache, |
| 26 | + enable_full_determinism, |
| 27 | + floats_tensor, |
| 28 | + load_hf_numpy, |
| 29 | + require_torch_accelerator, |
| 30 | + require_torch_gpu, |
| 31 | + skip_mps, |
| 32 | + slow, |
| 33 | + torch_all_close, |
| 34 | + torch_device, |
| 35 | +) |
| 36 | + |
| 37 | +from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin |
| 38 | + |
| 39 | + |
| 40 | +enable_full_determinism() |
| 41 | + |
| 42 | + |
| 43 | +class AutoencoderKLTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): |
| 44 | + model_class = AsymmetricAutoencoderKL |
| 45 | + main_input_name = "sample" |
| 46 | + base_precision = 1e-2 |
| 47 | + |
| 48 | + def get_asym_autoencoder_kl_config(self, block_out_channels=None, norm_num_groups=None): |
| 49 | + block_out_channels = block_out_channels or [2, 4] |
| 50 | + norm_num_groups = norm_num_groups or 2 |
| 51 | + init_dict = { |
| 52 | + "in_channels": 3, |
| 53 | + "out_channels": 3, |
| 54 | + "down_block_types": ["DownEncoderBlock2D"] * len(block_out_channels), |
| 55 | + "down_block_out_channels": block_out_channels, |
| 56 | + "layers_per_down_block": 1, |
| 57 | + "up_block_types": ["UpDecoderBlock2D"] * len(block_out_channels), |
| 58 | + "up_block_out_channels": block_out_channels, |
| 59 | + "layers_per_up_block": 1, |
| 60 | + "act_fn": "silu", |
| 61 | + "latent_channels": 4, |
| 62 | + "norm_num_groups": norm_num_groups, |
| 63 | + "sample_size": 32, |
| 64 | + "scaling_factor": 0.18215, |
| 65 | + } |
| 66 | + return init_dict |
| 67 | + |
| 68 | + @property |
| 69 | + def dummy_input(self): |
| 70 | + batch_size = 4 |
| 71 | + num_channels = 3 |
| 72 | + sizes = (32, 32) |
| 73 | + |
| 74 | + image = floats_tensor((batch_size, num_channels) + sizes).to(torch_device) |
| 75 | + mask = torch.ones((batch_size, 1) + sizes).to(torch_device) |
| 76 | + |
| 77 | + return {"sample": image, "mask": mask} |
| 78 | + |
| 79 | + @property |
| 80 | + def input_shape(self): |
| 81 | + return (3, 32, 32) |
| 82 | + |
| 83 | + @property |
| 84 | + def output_shape(self): |
| 85 | + return (3, 32, 32) |
| 86 | + |
| 87 | + def prepare_init_args_and_inputs_for_common(self): |
| 88 | + init_dict = self.get_asym_autoencoder_kl_config() |
| 89 | + inputs_dict = self.dummy_input |
| 90 | + return init_dict, inputs_dict |
| 91 | + |
| 92 | + @unittest.skip("Unsupported test.") |
| 93 | + def test_forward_with_norm_groups(self): |
| 94 | + pass |
| 95 | + |
| 96 | + |
| 97 | +@slow |
| 98 | +class AsymmetricAutoencoderKLIntegrationTests(unittest.TestCase): |
| 99 | + def get_file_format(self, seed, shape): |
| 100 | + return f"gaussian_noise_s={seed}_shape={'_'.join([str(s) for s in shape])}.npy" |
| 101 | + |
| 102 | + def tearDown(self): |
| 103 | + # clean up the VRAM after each test |
| 104 | + super().tearDown() |
| 105 | + gc.collect() |
| 106 | + backend_empty_cache(torch_device) |
| 107 | + |
| 108 | + def get_sd_image(self, seed=0, shape=(4, 3, 512, 512), fp16=False): |
| 109 | + dtype = torch.float16 if fp16 else torch.float32 |
| 110 | + image = torch.from_numpy(load_hf_numpy(self.get_file_format(seed, shape))).to(torch_device).to(dtype) |
| 111 | + return image |
| 112 | + |
| 113 | + def get_sd_vae_model(self, model_id="cross-attention/asymmetric-autoencoder-kl-x-1-5", fp16=False): |
| 114 | + revision = "main" |
| 115 | + torch_dtype = torch.float32 |
| 116 | + |
| 117 | + model = AsymmetricAutoencoderKL.from_pretrained( |
| 118 | + model_id, |
| 119 | + torch_dtype=torch_dtype, |
| 120 | + revision=revision, |
| 121 | + ) |
| 122 | + model.to(torch_device).eval() |
| 123 | + |
| 124 | + return model |
| 125 | + |
| 126 | + def get_generator(self, seed=0): |
| 127 | + generator_device = "cpu" if not torch_device.startswith("cuda") else "cuda" |
| 128 | + if torch_device != "mps": |
| 129 | + return torch.Generator(device=generator_device).manual_seed(seed) |
| 130 | + return torch.manual_seed(seed) |
| 131 | + |
| 132 | + @parameterized.expand( |
| 133 | + [ |
| 134 | + # fmt: off |
| 135 | + [ |
| 136 | + 33, |
| 137 | + [-0.0336, 0.3011, 0.1764, 0.0087, -0.3401, 0.3645, -0.1247, 0.1205], |
| 138 | + [-0.1603, 0.9878, -0.0495, -0.0790, -0.2709, 0.8375, -0.2060, -0.0824], |
| 139 | + ], |
| 140 | + [ |
| 141 | + 47, |
| 142 | + [0.4400, 0.0543, 0.2873, 0.2946, 0.0553, 0.0839, -0.1585, 0.2529], |
| 143 | + [-0.2376, 0.1168, 0.1332, -0.4840, -0.2508, -0.0791, -0.0493, -0.4089], |
| 144 | + ], |
| 145 | + # fmt: on |
| 146 | + ] |
| 147 | + ) |
| 148 | + def test_stable_diffusion(self, seed, expected_slice, expected_slice_mps): |
| 149 | + model = self.get_sd_vae_model() |
| 150 | + image = self.get_sd_image(seed) |
| 151 | + generator = self.get_generator(seed) |
| 152 | + |
| 153 | + with torch.no_grad(): |
| 154 | + sample = model(image, generator=generator, sample_posterior=True).sample |
| 155 | + |
| 156 | + assert sample.shape == image.shape |
| 157 | + |
| 158 | + output_slice = sample[-1, -2:, -2:, :2].flatten().float().cpu() |
| 159 | + expected_output_slice = torch.tensor(expected_slice_mps if torch_device == "mps" else expected_slice) |
| 160 | + |
| 161 | + assert torch_all_close(output_slice, expected_output_slice, atol=5e-3) |
| 162 | + |
| 163 | + @parameterized.expand( |
| 164 | + [ |
| 165 | + # fmt: off |
| 166 | + [ |
| 167 | + 33, |
| 168 | + [-0.0340, 0.2870, 0.1698, -0.0105, -0.3448, 0.3529, -0.1321, 0.1097], |
| 169 | + [-0.0344, 0.2912, 0.1687, -0.0137, -0.3462, 0.3552, -0.1337, 0.1078], |
| 170 | + ], |
| 171 | + [ |
| 172 | + 47, |
| 173 | + [0.4397, 0.0550, 0.2873, 0.2946, 0.0567, 0.0855, -0.1580, 0.2531], |
| 174 | + [0.4397, 0.0550, 0.2873, 0.2946, 0.0567, 0.0855, -0.1580, 0.2531], |
| 175 | + ], |
| 176 | + # fmt: on |
| 177 | + ] |
| 178 | + ) |
| 179 | + def test_stable_diffusion_mode(self, seed, expected_slice, expected_slice_mps): |
| 180 | + model = self.get_sd_vae_model() |
| 181 | + image = self.get_sd_image(seed) |
| 182 | + |
| 183 | + with torch.no_grad(): |
| 184 | + sample = model(image).sample |
| 185 | + |
| 186 | + assert sample.shape == image.shape |
| 187 | + |
| 188 | + output_slice = sample[-1, -2:, -2:, :2].flatten().float().cpu() |
| 189 | + expected_output_slice = torch.tensor(expected_slice_mps if torch_device == "mps" else expected_slice) |
| 190 | + |
| 191 | + assert torch_all_close(output_slice, expected_output_slice, atol=3e-3) |
| 192 | + |
| 193 | + @parameterized.expand( |
| 194 | + [ |
| 195 | + # fmt: off |
| 196 | + [13, [-0.0521, -0.2939, 0.1540, -0.1855, -0.5936, -0.3138, -0.4579, -0.2275]], |
| 197 | + [37, [-0.1820, -0.4345, -0.0455, -0.2923, -0.8035, -0.5089, -0.4795, -0.3106]], |
| 198 | + # fmt: on |
| 199 | + ] |
| 200 | + ) |
| 201 | + @require_torch_accelerator |
| 202 | + @skip_mps |
| 203 | + def test_stable_diffusion_decode(self, seed, expected_slice): |
| 204 | + model = self.get_sd_vae_model() |
| 205 | + encoding = self.get_sd_image(seed, shape=(3, 4, 64, 64)) |
| 206 | + |
| 207 | + with torch.no_grad(): |
| 208 | + sample = model.decode(encoding).sample |
| 209 | + |
| 210 | + assert list(sample.shape) == [3, 3, 512, 512] |
| 211 | + |
| 212 | + output_slice = sample[-1, -2:, :2, -2:].flatten().cpu() |
| 213 | + expected_output_slice = torch.tensor(expected_slice) |
| 214 | + |
| 215 | + assert torch_all_close(output_slice, expected_output_slice, atol=2e-3) |
| 216 | + |
| 217 | + @parameterized.expand([(13,), (16,), (37,)]) |
| 218 | + @require_torch_gpu |
| 219 | + @unittest.skipIf( |
| 220 | + not is_xformers_available(), |
| 221 | + reason="xformers is not required when using PyTorch 2.0.", |
| 222 | + ) |
| 223 | + def test_stable_diffusion_decode_xformers_vs_2_0(self, seed): |
| 224 | + model = self.get_sd_vae_model() |
| 225 | + encoding = self.get_sd_image(seed, shape=(3, 4, 64, 64)) |
| 226 | + |
| 227 | + with torch.no_grad(): |
| 228 | + sample = model.decode(encoding).sample |
| 229 | + |
| 230 | + model.enable_xformers_memory_efficient_attention() |
| 231 | + with torch.no_grad(): |
| 232 | + sample_2 = model.decode(encoding).sample |
| 233 | + |
| 234 | + assert list(sample.shape) == [3, 3, 512, 512] |
| 235 | + |
| 236 | + assert torch_all_close(sample, sample_2, atol=5e-2) |
| 237 | + |
| 238 | + @parameterized.expand( |
| 239 | + [ |
| 240 | + # fmt: off |
| 241 | + [33, [-0.3001, 0.0918, -2.6984, -3.9720, -3.2099, -5.0353, 1.7338, -0.2065, 3.4267]], |
| 242 | + [47, [-1.5030, -4.3871, -6.0355, -9.1157, -1.6661, -2.7853, 2.1607, -5.0823, 2.5633]], |
| 243 | + # fmt: on |
| 244 | + ] |
| 245 | + ) |
| 246 | + def test_stable_diffusion_encode_sample(self, seed, expected_slice): |
| 247 | + model = self.get_sd_vae_model() |
| 248 | + image = self.get_sd_image(seed) |
| 249 | + generator = self.get_generator(seed) |
| 250 | + |
| 251 | + with torch.no_grad(): |
| 252 | + dist = model.encode(image).latent_dist |
| 253 | + sample = dist.sample(generator=generator) |
| 254 | + |
| 255 | + assert list(sample.shape) == [image.shape[0], 4] + [i // 8 for i in image.shape[2:]] |
| 256 | + |
| 257 | + output_slice = sample[0, -1, -3:, -3:].flatten().cpu() |
| 258 | + expected_output_slice = torch.tensor(expected_slice) |
| 259 | + |
| 260 | + tolerance = 3e-3 if torch_device != "mps" else 1e-2 |
| 261 | + assert torch_all_close(output_slice, expected_output_slice, atol=tolerance) |
0 commit comments