declare-lab/tango2

1# coding=utf-82# Copyright 2023 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 at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# 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 and14# limitations under the License.15import unittest16from typing import Tuple17 18import torch19 20from diffusers.utils import floats_tensor, randn_tensor, torch_all_close, torch_device21from diffusers.utils.testing_utils import require_torch22 23 24@require_torch25class UNetBlockTesterMixin:26    @property27    def dummy_input(self):28        return self.get_dummy_input()29 30    @property31    def output_shape(self):32        if self.block_type == "down":33            return (4, 32, 16, 16)34        elif self.block_type == "mid":35            return (4, 32, 32, 32)36        elif self.block_type == "up":37            return (4, 32, 64, 64)38 39        raise ValueError(f"'{self.block_type}' is not a supported block_type. Set it to 'up', 'mid', or 'down'.")40 41    def get_dummy_input(42        self,43        include_temb=True,44        include_res_hidden_states_tuple=False,45        include_encoder_hidden_states=False,46        include_skip_sample=False,47    ):48        batch_size = 449        num_channels = 3250        sizes = (32, 32)51 52        generator = torch.manual_seed(0)53        device = torch.device(torch_device)54        shape = (batch_size, num_channels) + sizes55        hidden_states = randn_tensor(shape, generator=generator, device=device)56        dummy_input = {"hidden_states": hidden_states}57 58        if include_temb:59            temb_channels = 12860            dummy_input["temb"] = randn_tensor((batch_size, temb_channels), generator=generator, device=device)61 62        if include_res_hidden_states_tuple:63            generator_1 = torch.manual_seed(1)64            dummy_input["res_hidden_states_tuple"] = (randn_tensor(shape, generator=generator_1, device=device),)65 66        if include_encoder_hidden_states:67            dummy_input["encoder_hidden_states"] = floats_tensor((batch_size, 32, 32)).to(torch_device)68 69        if include_skip_sample:70            dummy_input["skip_sample"] = randn_tensor(((batch_size, 3) + sizes), generator=generator, device=device)71 72        return dummy_input73 74    def prepare_init_args_and_inputs_for_common(self):75        init_dict = {76            "in_channels": 32,77            "out_channels": 32,78            "temb_channels": 128,79        }80        if self.block_type == "up":81            init_dict["prev_output_channel"] = 3282 83        if self.block_type == "mid":84            init_dict.pop("out_channels")85 86        inputs_dict = self.dummy_input87        return init_dict, inputs_dict88 89    def test_output(self, expected_slice):90        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()91        unet_block = self.block_class(**init_dict)92        unet_block.to(torch_device)93        unet_block.eval()94 95        with torch.no_grad():96            output = unet_block(**inputs_dict)97 98        if isinstance(output, Tuple):99            output = output[0]100 101        self.assertEqual(output.shape, self.output_shape)102 103        output_slice = output[0, -1, -3:, -3:]104        expected_slice = torch.tensor(expected_slice).to(torch_device)105        assert torch_all_close(output_slice.flatten(), expected_slice, atol=5e-3)106 107    @unittest.skipIf(torch_device == "mps", "Training is not supported in mps")108    def test_training(self):109        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()110        model = self.block_class(**init_dict)111        model.to(torch_device)112        model.train()113        output = model(**inputs_dict)114 115        if isinstance(output, Tuple):116            output = output[0]117 118        device = torch.device(torch_device)119        noise = randn_tensor(output.shape, device=device)120        loss = torch.nn.functional.mse_loss(output, noise)121        loss.backward()122