Skip to content

[tests] Speed up AnimateDiff test#7707

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
a-r-r-o-w wants to merge 6 commits into from
Closed

[tests] Speed up AnimateDiff test #7707

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
a8e9135
speed up animatediff test
a-r-r-o-w Apr 18, 2024
b21ba5a
make style
a-r-r-o-w Apr 18, 2024
7b39a62
Merge branch 'main' into animatediff/speedup-test
a-r-r-o-w Apr 18, 2024
2bb1dfc
Merge branch 'main' into animatediff/speedup-test
a-r-r-o-w Apr 22, 2024
8fb71e8
Merge branch 'main' into animatediff/speedup-test
sayakpaul Apr 23, 2024
68b0208
Merge branch 'main' into animatediff/speedup-test
DN6 Apr 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions src/diffusers/models/unets/unet_3d_blocks.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -1461,7 +1461,6 @@ def __init__(
resnet_pre_norm: bool = True,
output_scale_factor: float = 1.0,
add_upsample: bool = True,
temporal_norm_num_groups: int = 32,
temporal_cross_attention_dim: Optional[int] = None,
temporal_num_attention_heads: int = 8,
temporal_max_seq_length: int = 32,
Expand DownExpand Up@@ -1493,7 +1492,7 @@ def __init__(
TransformerTemporalModel(
num_attention_heads=temporal_num_attention_heads,
in_channels=out_channels,
norm_num_groups=temporal_norm_num_groups,
norm_num_groups=resnet_groups,

@a-r-r-o-wa-r-r-o-wApr 18, 2024•
edited
Loading

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is because in DownBlockMotion, we make use of resnet_groups too and temporal_norm_num_groups does not seem to be present there. I think this is correct but not too sure. Not doing this causes multiple errors with the group norm in_channels being indivisible by num_groups. There is no control over temporal_norm_num_groups parameter because it is not present in get_up_block(). Everywhere else that passes norm_num_groups into TemporalTransformerModel does what is present here.

@sayakpaulsayakpaulApr 23, 2024

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's not touch the default values here please especially when we're not too sure. Will defer to @DN6 to verify.

@DN6DN6Apr 25, 2024

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah this change should be safe. I think it was a mistake here passing in temporal_norm_num_groups

@a-r-r-o-wa-r-r-o-wApr 25, 2024•
edited
Loading

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

awesome! What are your thoughts on the first failing test after these changes? If I make any changes there, multiple other tests are most likely going to break. We are trying to load hf-internal-testing/tiny-stable-diffusion-pipe weights into a different model definition due to the changes in unet causing weight mismatch...

@DN6DN6Apr 30, 2024•
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's okay. We can make a change to that model repo or change the test specifically for animatediff.

cross_attention_dim=temporal_cross_attention_dim,
attention_bias=False,
activation_fn="geglu",
Expand Down
22 changes: 13 additions & 9 deletions tests/pipelines/animatediff/test_animatediff.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -51,16 +51,19 @@ class AnimateDiffPipelineFastTests(
)

def get_dummy_components(self):
block_out_channels = (8, 8)
cross_attention_dim = 8

torch.manual_seed(0)
unet = UNet2DConditionModel(
block_out_channels=(32, 64),
layers_per_block=2,
sample_size=32,
block_out_channels=block_out_channels,
layers_per_block=1,
sample_size=8,
in_channels=4,
out_channels=4,
down_block_types=("CrossAttnDownBlock2D", "DownBlock2D"),
up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),
cross_attention_dim=32,
cross_attention_dim=cross_attention_dim,
norm_num_groups=2,
)
scheduler = DDIMScheduler(
Expand All@@ -71,18 +74,19 @@ def get_dummy_components(self):
)
torch.manual_seed(0)
vae = AutoencoderKL(
block_out_channels=[32, 64],
block_out_channels=block_out_channels,
in_channels=3,
out_channels=3,
down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"],
up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],
latent_channels=4,
norm_num_groups=2,
)
torch.manual_seed(0)
text_encoder_config = CLIPTextConfig(
bos_token_id=0,
eos_token_id=2,
hidden_size=32,
hidden_size=cross_attention_dim,
intermediate_size=37,
layer_norm_eps=1e-05,
num_attention_heads=4,
Expand All@@ -93,8 +97,8 @@ def get_dummy_components(self):
text_encoder = CLIPTextModel(text_encoder_config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
motion_adapter = MotionAdapter(
block_out_channels=(32, 64),
motion_layers_per_block=2,
block_out_channels=block_out_channels,
motion_layers_per_block=1,
motion_norm_num_groups=2,
motion_num_attention_heads=4,
)
Expand DownExpand Up@@ -279,7 +283,7 @@ def test_prompt_embeds(self):

inputs = self.get_dummy_inputs(torch_device)
inputs.pop("prompt")
inputs["prompt_embeds"] = torch.randn((1, 4, 32), device=torch_device)
inputs["prompt_embeds"] = torch.randn((1, 4, pipe.text_encoder.config.hidden_size), device=torch_device)
Comment thread
a-r-r-o-w marked this conversation as resolved.
pipe(**inputs)

def test_free_init(self):
Expand Down