Skip to content

GitLab

Unverified Commit 5b448a5e authored by p1kit's avatar p1kit Committed by GitHub
Browse files

[Tests] Optimize test configurations for faster execution (#5535) 

Optimize test configurations for faster execution
parent a69ebe55
Show whitespace changes
Inline Side-by-side
Showing with 14 additions and 7 deletions
tests/pipelines/controlnet/test_controlnet_img2img.py
View file @ 5b448a5e
...@@ -72,7 +72,7 @@ class ControlNetImg2ImgPipelineFastTests( ...@@ -72,7 +72,7 @@ class ControlNetImg2ImgPipelineFastTests(
def get_dummy_components(self): def get_dummy_components(self):
torch.manual_seed(0) torch.manual_seed(0)
unet = UNet2DConditionModel( unet = UNet2DConditionModel(
block_out_channels=(32, 64), block_out_channels=(4, 8),
layers_per_block=2, layers_per_block=2,
sample_size=32, sample_size=32,
in_channels=4, in_channels=4,
...@@ -80,15 +80,17 @@ class ControlNetImg2ImgPipelineFastTests( ...@@ -80,15 +80,17 @@ class ControlNetImg2ImgPipelineFastTests(
down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"), up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),
cross_attention_dim=32, cross_attention_dim=32,
norm_num_groups=1,
) )
torch.manual_seed(0) torch.manual_seed(0)
controlnet = ControlNetModel( controlnet = ControlNetModel(
block_out_channels=(32, 64), block_out_channels=(4, 8),
layers_per_block=2, layers_per_block=2,
in_channels=4, in_channels=4,
down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
cross_attention_dim=32, cross_attention_dim=32,
conditioning_embedding_out_channels=(16, 32), conditioning_embedding_out_channels=(16, 32),
norm_num_groups=1,
) )
torch.manual_seed(0) torch.manual_seed(0)
scheduler = DDIMScheduler( scheduler = DDIMScheduler(
...@@ -100,12 +102,13 @@ class ControlNetImg2ImgPipelineFastTests( ...@@ -100,12 +102,13 @@ class ControlNetImg2ImgPipelineFastTests(
) )
torch.manual_seed(0) torch.manual_seed(0)
vae = AutoencoderKL( vae = AutoencoderKL(
block_out_channels=[32, 64], block_out_channels=[4, 8],
in_channels=3, in_channels=3,
out_channels=3, out_channels=3,
down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"], down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"],
up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"], up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],
latent_channels=4, latent_channels=4,
norm_num_groups=2,
) )
torch.manual_seed(0) torch.manual_seed(0)
text_encoder_config = CLIPTextConfig( text_encoder_config = CLIPTextConfig(
...@@ -186,7 +189,7 @@ class StableDiffusionMultiControlNetPipelineFastTests( ...@@ -186,7 +189,7 @@ class StableDiffusionMultiControlNetPipelineFastTests(
def get_dummy_components(self): def get_dummy_components(self):
torch.manual_seed(0) torch.manual_seed(0)
unet = UNet2DConditionModel( unet = UNet2DConditionModel(
block_out_channels=(32, 64), block_out_channels=(4, 8),
layers_per_block=2, layers_per_block=2,
sample_size=32, sample_size=32,
in_channels=4, in_channels=4,
...@@ -194,6 +197,7 @@ class StableDiffusionMultiControlNetPipelineFastTests( ...@@ -194,6 +197,7 @@ class StableDiffusionMultiControlNetPipelineFastTests(
down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"), up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),
cross_attention_dim=32, cross_attention_dim=32,
norm_num_groups=1,
) )
torch.manual_seed(0) torch.manual_seed(0)
...@@ -203,23 +207,25 @@ class StableDiffusionMultiControlNetPipelineFastTests( ...@@ -203,23 +207,25 @@ class StableDiffusionMultiControlNetPipelineFastTests(
m.bias.data.fill_(1.0) m.bias.data.fill_(1.0)
controlnet1 = ControlNetModel( controlnet1 = ControlNetModel(
block_out_channels=(32, 64), block_out_channels=(4, 8),
layers_per_block=2, layers_per_block=2,
in_channels=4, in_channels=4,
down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
cross_attention_dim=32, cross_attention_dim=32,
conditioning_embedding_out_channels=(16, 32), conditioning_embedding_out_channels=(16, 32),
norm_num_groups=1,
) )
controlnet1.controlnet_down_blocks.apply(init_weights) controlnet1.controlnet_down_blocks.apply(init_weights)
torch.manual_seed(0) torch.manual_seed(0)
controlnet2 = ControlNetModel( controlnet2 = ControlNetModel(
block_out_channels=(32, 64), block_out_channels=(4, 8),
layers_per_block=2, layers_per_block=2,
in_channels=4, in_channels=4,
down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
cross_attention_dim=32, cross_attention_dim=32,
conditioning_embedding_out_channels=(16, 32), conditioning_embedding_out_channels=(16, 32),
norm_num_groups=1,
) )
controlnet2.controlnet_down_blocks.apply(init_weights) controlnet2.controlnet_down_blocks.apply(init_weights)
...@@ -233,12 +239,13 @@ class StableDiffusionMultiControlNetPipelineFastTests( ...@@ -233,12 +239,13 @@ class StableDiffusionMultiControlNetPipelineFastTests(
) )
torch.manual_seed(0) torch.manual_seed(0)
vae = AutoencoderKL( vae = AutoencoderKL(
block_out_channels=[32, 64], block_out_channels=[4, 8],
in_channels=3, in_channels=3,
out_channels=3, out_channels=3,
down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"], down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"],
up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"], up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],
latent_channels=4, latent_channels=4,
norm_num_groups=2,
) )
torch.manual_seed(0) torch.manual_seed(0)
text_encoder_config = CLIPTextConfig( text_encoder_config = CLIPTextConfig(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment