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.15 16import unittest17 18import torch19 20from diffusers import DDIMScheduler, DDPMScheduler, UNet2DModel21from diffusers.training_utils import set_seed22from diffusers.utils.testing_utils import slow23 24 25torch.backends.cuda.matmul.allow_tf32 = False26 27 28class TrainingTests(unittest.TestCase):29    def get_model_optimizer(self, resolution=32):30        set_seed(0)31        model = UNet2DModel(sample_size=resolution, in_channels=3, out_channels=3)32        optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)33        return model, optimizer34 35    @slow36    def test_training_step_equality(self):37        device = "cpu"  # ensure full determinism without setting the CUBLAS_WORKSPACE_CONFIG env variable38        ddpm_scheduler = DDPMScheduler(39            num_train_timesteps=1000,40            beta_start=0.0001,41            beta_end=0.02,42            beta_schedule="linear",43            clip_sample=True,44        )45        ddim_scheduler = DDIMScheduler(46            num_train_timesteps=1000,47            beta_start=0.0001,48            beta_end=0.02,49            beta_schedule="linear",50            clip_sample=True,51        )52 53        assert ddpm_scheduler.config.num_train_timesteps == ddim_scheduler.config.num_train_timesteps54 55        # shared batches for DDPM and DDIM56        set_seed(0)57        clean_images = [torch.randn((4, 3, 32, 32)).clip(-1, 1).to(device) for _ in range(4)]58        noise = [torch.randn((4, 3, 32, 32)).to(device) for _ in range(4)]59        timesteps = [torch.randint(0, 1000, (4,)).long().to(device) for _ in range(4)]60 61        # train with a DDPM scheduler62        model, optimizer = self.get_model_optimizer(resolution=32)63        model.train().to(device)64        for i in range(4):65            optimizer.zero_grad()66            ddpm_noisy_images = ddpm_scheduler.add_noise(clean_images[i], noise[i], timesteps[i])67            ddpm_noise_pred = model(ddpm_noisy_images, timesteps[i]).sample68            loss = torch.nn.functional.mse_loss(ddpm_noise_pred, noise[i])69            loss.backward()70            optimizer.step()71        del model, optimizer72 73        # recreate the model and optimizer, and retry with DDIM74        model, optimizer = self.get_model_optimizer(resolution=32)75        model.train().to(device)76        for i in range(4):77            optimizer.zero_grad()78            ddim_noisy_images = ddim_scheduler.add_noise(clean_images[i], noise[i], timesteps[i])79            ddim_noise_pred = model(ddim_noisy_images, timesteps[i]).sample80            loss = torch.nn.functional.mse_loss(ddim_noise_pred, noise[i])81            loss.backward()82            optimizer.step()83        del model, optimizer84 85        self.assertTrue(torch.allclose(ddpm_noisy_images, ddim_noisy_images, atol=1e-5))86        self.assertTrue(torch.allclose(ddpm_noise_pred, ddim_noise_pred, atol=1e-5))87