Fix tests

Author: Virginia Fernandez

tests/inferers/test_diffusion_inferer.py

@@ -69,7 +69,7 @@ def test_call(self, model_params, input_shape):
         inferer = DiffusionInferer(scheduler=scheduler)
         scheduler.set_timesteps(num_inference_steps=10)
         timesteps = torch.randint(0, scheduler.num_train_timesteps, (input_shape[0],), device=input.device).long()
-        sample = inferer(inputs=input, noise=noise, diffusion_model=model, timesteps=timesteps)
+        sample = inferer(inputs=input, noise=noise, network=model, timesteps=timesteps)
         self.assertEqual(sample.shape, input_shape)
 
     @parameterized.expand(TEST_CASES)
@@ -84,7 +84,7 @@ def test_sample_intermediates(self, model_params, input_shape):
         inferer = DiffusionInferer(scheduler=scheduler)
         scheduler.set_timesteps(num_inference_steps=10)
         sample, intermediates = inferer.sample(
-            input_noise=noise, diffusion_model=model, scheduler=scheduler, save_intermediates=True, intermediate_steps=1
+            input_noise=noise, network=model, scheduler=scheduler, save_intermediates=True, intermediate_steps=1
         )
         self.assertEqual(len(intermediates), 10)
 
@@ -101,7 +101,7 @@ def test_sample_cfg(self, model_params, input_shape):
         scheduler.set_timesteps(num_inference_steps=10)
         sample, intermediates = inferer.sample(
             input_noise=noise,
-            diffusion_model=model,
+            network=model,
             scheduler=scheduler,
             save_intermediates=True,
             intermediate_steps=1,
@@ -121,7 +121,7 @@ def test_ddpm_sampler(self, model_params, input_shape):
         inferer = DiffusionInferer(scheduler=scheduler)
         scheduler.set_timesteps(num_inference_steps=10)
         sample, intermediates = inferer.sample(
-            input_noise=noise, diffusion_model=model, scheduler=scheduler, save_intermediates=True, intermediate_steps=1
+            input_noise=noise, network=model, scheduler=scheduler, save_intermediates=True, intermediate_steps=1
         )
         self.assertEqual(len(intermediates), 10)
 
@@ -137,7 +137,7 @@ def test_ddim_sampler(self, model_params, input_shape):
         inferer = DiffusionInferer(scheduler=scheduler)
         scheduler.set_timesteps(num_inference_steps=10)
         sample, intermediates = inferer.sample(
-            input_noise=noise, diffusion_model=model, scheduler=scheduler, save_intermediates=True, intermediate_steps=1
+            input_noise=noise, network=model, scheduler=scheduler, save_intermediates=True, intermediate_steps=1
         )
         self.assertEqual(len(intermediates), 10)
 
@@ -153,7 +153,7 @@ def test_rflow_sampler(self, model_params, input_shape):
         inferer = DiffusionInferer(scheduler=scheduler)
         scheduler.set_timesteps(num_inference_steps=10)
         sample, intermediates = inferer.sample(
-            input_noise=noise, diffusion_model=model, scheduler=scheduler, save_intermediates=True, intermediate_steps=1
+            input_noise=noise, network=model, scheduler=scheduler, save_intermediates=True, intermediate_steps=1
         )
         self.assertEqual(len(intermediates), 10)
 

tests/inferers/test_latent_diffusion_inferer.py

@@ -355,14 +355,14 @@ def test_prediction_shape(
                 prediction = inferer(
                     inputs=input,
                     autoencoder_model=stage_1,
-                    diffusion_model=stage_2,
+                    network=stage_2,
                     seg=input_seg,
                     noise=noise,
                     timesteps=timesteps,
                 )
             else:
                 prediction = inferer(
-                    inputs=input, autoencoder_model=stage_1, diffusion_model=stage_2, noise=noise, timesteps=timesteps
+                    inputs=input, autoencoder_model=stage_1, network=stage_2, noise=noise, timesteps=timesteps
                 )
             self.assertEqual(prediction.shape, latent_shape)
 
@@ -404,13 +404,13 @@ def test_sample_shape(
                 sample = inferer.sample(
                     input_noise=noise,
                     autoencoder_model=stage_1,
-                    diffusion_model=stage_2,
+                    network=stage_2,
                     scheduler=scheduler,
                     seg=input_seg,
                 )
             else:
                 sample = inferer.sample(
-                    input_noise=noise, autoencoder_model=stage_1, diffusion_model=stage_2, scheduler=scheduler
+                    input_noise=noise, autoencoder_model=stage_1, network=stage_2, scheduler=scheduler
                 )
             self.assertEqual(sample.shape, input_shape)
 
@@ -452,14 +452,14 @@ def test_sample_shape_with_cfg(
                 sample = inferer.sample(
                     input_noise=noise,
                     autoencoder_model=stage_1,
-                    diffusion_model=stage_2,
+                    network=stage_2,
                     scheduler=scheduler,
                     seg=input_seg,
                     cfg=5,
                 )
             else:
                 sample = inferer.sample(
-                    input_noise=noise, autoencoder_model=stage_1, diffusion_model=stage_2, scheduler=scheduler, cfg=5
+                    input_noise=noise, autoencoder_model=stage_1, network=stage_2, scheduler=scheduler, cfg=5
                 )
             self.assertEqual(sample.shape, input_shape)
 
@@ -503,7 +503,7 @@ def test_sample_intermediates(
                 sample, intermediates = inferer.sample(
                     input_noise=noise,
                     autoencoder_model=stage_1,
-                    diffusion_model=stage_2,
+                    network=stage_2,
                     scheduler=scheduler,
                     seg=input_seg,
                     save_intermediates=True,
@@ -513,7 +513,7 @@ def test_sample_intermediates(
                 sample, intermediates = inferer.sample(
                     input_noise=noise,
                     autoencoder_model=stage_1,
-                    diffusion_model=stage_2,
+                    network=stage_2,
                     scheduler=scheduler,
                     save_intermediates=True,
                     intermediate_steps=1,
@@ -560,7 +560,7 @@ def test_get_likelihoods(
             sample, intermediates = inferer.get_likelihood(
                 inputs=input,
                 autoencoder_model=stage_1,
-                diffusion_model=stage_2,
+                network=stage_2,
                 scheduler=scheduler,
                 save_intermediates=True,
                 seg=input_seg,
@@ -569,7 +569,7 @@ def test_get_likelihoods(
             sample, intermediates = inferer.get_likelihood(
                 inputs=input,
                 autoencoder_model=stage_1,
-                diffusion_model=stage_2,
+                network=stage_2,
                 scheduler=scheduler,
                 save_intermediates=True,
             )
@@ -615,7 +615,7 @@ def test_resample_likelihoods(
             sample, intermediates = inferer.get_likelihood(
                 inputs=input,
                 autoencoder_model=stage_1,
-                diffusion_model=stage_2,
+                network=stage_2,
                 scheduler=scheduler,
                 save_intermediates=True,
                 resample_latent_likelihoods=True,
@@ -625,7 +625,7 @@ def test_resample_likelihoods(
             sample, intermediates = inferer.get_likelihood(
                 inputs=input,
                 autoencoder_model=stage_1,
-                diffusion_model=stage_2,
+                network=stage_2,
                 scheduler=scheduler,
                 save_intermediates=True,
                 resample_latent_likelihoods=True,
@@ -682,7 +682,7 @@ def test_prediction_shape_conditioned_concat(
                 prediction = inferer(
                     inputs=input,
                     autoencoder_model=stage_1,
-                    diffusion_model=stage_2,
+                    network=stage_2,
                     noise=noise,
                     timesteps=timesteps,
                     condition=conditioning,
@@ -693,7 +693,7 @@ def test_prediction_shape_conditioned_concat(
                 prediction = inferer(
                     inputs=input,
                     autoencoder_model=stage_1,
-                    diffusion_model=stage_2,
+                    network=stage_2,
                     noise=noise,
                     timesteps=timesteps,
                     condition=conditioning,
@@ -747,7 +747,7 @@ def test_sample_shape_conditioned_concat(
                 sample = inferer.sample(
                     input_noise=noise,
                     autoencoder_model=stage_1,
-                    diffusion_model=stage_2,
+                    network=stage_2,
                     scheduler=scheduler,
                     conditioning=conditioning,
                     mode="concat",
@@ -757,7 +757,7 @@ def test_sample_shape_conditioned_concat(
                 sample = inferer.sample(
                     input_noise=noise,
                     autoencoder_model=stage_1,
-                    diffusion_model=stage_2,
+                    network=stage_2,
                     scheduler=scheduler,
                     conditioning=conditioning,
                     mode="concat",
@@ -813,14 +813,14 @@ def test_shape_different_latents(
                 prediction = inferer(
                     inputs=input,
                     autoencoder_model=stage_1,
-                    diffusion_model=stage_2,
+                    network=stage_2,
                     noise=noise,
                     timesteps=timesteps,
                     seg=input_seg,
                 )
             else:
                 prediction = inferer(
-                    inputs=input, autoencoder_model=stage_1, diffusion_model=stage_2, noise=noise, timesteps=timesteps
+                    inputs=input, autoencoder_model=stage_1, network=stage_2, noise=noise, timesteps=timesteps
                 )
             self.assertEqual(prediction.shape, latent_shape)
 
@@ -875,14 +875,14 @@ def test_sample_shape_different_latents(
                 input_seg = torch.randn(input_shape_seg).to(device)
                 prediction, _ = inferer.sample(
                     autoencoder_model=stage_1,
-                    diffusion_model=stage_2,
+                    network=stage_2,
                     input_noise=noise,
                     save_intermediates=True,
                     seg=input_seg,
                 )
             else:
                 prediction = inferer.sample(
-                    autoencoder_model=stage_1, diffusion_model=stage_2, input_noise=noise, save_intermediates=False
+                    autoencoder_model=stage_1, network=stage_2, input_noise=noise, save_intermediates=False
                 )
             self.assertEqual(prediction.shape, input_shape)
 
@@ -929,7 +929,7 @@ def test_incompatible_spade_setup(self):
                 _ = inferer.sample(
                     input_noise=noise,
                     autoencoder_model=stage_1,
-                    diffusion_model=stage_2,
+                    network=stage_2,
                     scheduler=scheduler,
                     seg=input_seg,
                 )