Skip to content

Commit cff933c

Browse files
authored
[tests] Pipeline test refactor (#14113)
* up * up * quantization level refactors. * style * up * up * use mixins and remove stale attributes * fixture for the deprecated pipelines * remove qkv * remove lergacy for attention processor. * use pytorch tensors exclusively. * rejig dodgy test. * improve how base pipeline output fixture is used. * rename params. * up * use pipeline load api for ip adapter. * move to flux specific. * redesign config for caching * remove eval. * up * up * add a todo note on check_qkv_fusion_matches_attn_procs_length * nan test in layerwise casting * proper skipping * skip properly. * up * up
1 parent 48fadcc commit cff933c

16 files changed

Lines changed: 2399 additions & 1263 deletions

tests/pipelines/flux/test_pipeline_flux.py

Lines changed: 204 additions & 110 deletions
Large diffs are not rendered by default.

tests/pipelines/test_pipeline_utils.py

Lines changed: 210 additions & 76 deletions
Large diffs are not rendered by default.

tests/pipelines/test_pipelines_common.py

Lines changed: 0 additions & 146 deletions
Original file line numberDiff line numberDiff line change
@@ -4,25 +4,20 @@
44
import os
55
import tempfile
66
import unittest
7-
import uuid
87
from typing import Any, Callable, Dict
98

109
import numpy as np
1110
import PIL.Image
1211
import pytest
1312
import torch
1413
import torch.nn as nn
15-
from huggingface_hub import ModelCard, delete_repo
16-
from huggingface_hub.utils import is_jinja_available
17-
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
1814

1915
import diffusers
2016
from diffusers import (
2117
AsymmetricAutoencoderKL,
2218
AutoencoderKL,
2319
AutoencoderTiny,
2420
ConsistencyDecoderVAE,
25-
DDIMScheduler,
2621
DiffusionPipeline,
2722
FasterCacheConfig,
2823
KolorsPipeline,
@@ -63,7 +58,6 @@
6358
create_ip_adapter_faceid_state_dict,
6459
create_ip_adapter_state_dict,
6560
)
66-
from ..others.test_utils import TOKEN, USER, is_staging_test
6761
from ..testing_utils import (
6862
CaptureLogger,
6963
backend_empty_cache,
@@ -2485,146 +2479,6 @@ def test_pipeline_level_group_offloading_inference(self, expected_max_difference
24852479
self.assertLess(max_diff, expected_max_difference)
24862480

24872481

2488-
@is_staging_test
2489-
class PipelinePushToHubTester(unittest.TestCase):
2490-
identifier = uuid.uuid4()
2491-
repo_id = f"test-pipeline-{identifier}"
2492-
org_repo_id = f"valid_org/{repo_id}-org"
2493-
2494-
def get_pipeline_components(self):
2495-
unet = UNet2DConditionModel(
2496-
block_out_channels=(32, 64),
2497-
layers_per_block=2,
2498-
sample_size=32,
2499-
in_channels=4,
2500-
out_channels=4,
2501-
down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
2502-
up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),
2503-
cross_attention_dim=32,
2504-
)
2505-
2506-
scheduler = DDIMScheduler(
2507-
beta_start=0.00085,
2508-
beta_end=0.012,
2509-
beta_schedule="scaled_linear",
2510-
clip_sample=False,
2511-
set_alpha_to_one=False,
2512-
)
2513-
2514-
vae = AutoencoderKL(
2515-
block_out_channels=[32, 64],
2516-
in_channels=3,
2517-
out_channels=3,
2518-
down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"],
2519-
up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],
2520-
latent_channels=4,
2521-
)
2522-
2523-
text_encoder_config = CLIPTextConfig(
2524-
bos_token_id=0,
2525-
eos_token_id=2,
2526-
hidden_size=32,
2527-
intermediate_size=37,
2528-
layer_norm_eps=1e-05,
2529-
num_attention_heads=4,
2530-
num_hidden_layers=5,
2531-
pad_token_id=1,
2532-
vocab_size=1000,
2533-
)
2534-
text_encoder = CLIPTextModel(text_encoder_config)
2535-
2536-
with tempfile.TemporaryDirectory() as tmpdir:
2537-
dummy_vocab = {"<|startoftext|>": 0, "<|endoftext|>": 1, "!": 2}
2538-
vocab_path = os.path.join(tmpdir, "vocab.json")
2539-
with open(vocab_path, "w") as f:
2540-
json.dump(dummy_vocab, f)
2541-
2542-
merges = "Ġ t\nĠt h"
2543-
merges_path = os.path.join(tmpdir, "merges.txt")
2544-
with open(merges_path, "w") as f:
2545-
f.writelines(merges)
2546-
tokenizer = CLIPTokenizer(vocab_file=vocab_path, merges_file=merges_path)
2547-
2548-
components = {
2549-
"unet": unet,
2550-
"scheduler": scheduler,
2551-
"vae": vae,
2552-
"text_encoder": text_encoder,
2553-
"tokenizer": tokenizer,
2554-
"safety_checker": None,
2555-
"feature_extractor": None,
2556-
}
2557-
return components
2558-
2559-
def test_push_to_hub(self):
2560-
components = self.get_pipeline_components()
2561-
pipeline = StableDiffusionPipeline(**components)
2562-
pipeline.push_to_hub(self.repo_id, token=TOKEN)
2563-
2564-
new_model = UNet2DConditionModel.from_pretrained(f"{USER}/{self.repo_id}", subfolder="unet")
2565-
unet = components["unet"]
2566-
for p1, p2 in zip(unet.parameters(), new_model.parameters()):
2567-
self.assertTrue(torch.equal(p1, p2))
2568-
2569-
# Push to hub via save_pretrained to a separate repo. Reusing `self.repo_id` after
2570-
# deleting it makes the staging server's LFS GC reject the next commit with
2571-
# "LFS pointer pointed to a file that does not exist" when the model bytes are identical.
2572-
save_repo_id = f"{self.repo_id}-saved"
2573-
with tempfile.TemporaryDirectory() as tmp_dir:
2574-
pipeline.save_pretrained(tmp_dir, repo_id=save_repo_id, push_to_hub=True, token=TOKEN)
2575-
2576-
new_model = UNet2DConditionModel.from_pretrained(f"{USER}/{save_repo_id}", subfolder="unet")
2577-
for p1, p2 in zip(unet.parameters(), new_model.parameters()):
2578-
self.assertTrue(torch.equal(p1, p2))
2579-
2580-
# Reset repos
2581-
delete_repo(token=TOKEN, repo_id=self.repo_id)
2582-
delete_repo(save_repo_id, token=TOKEN)
2583-
2584-
def test_push_to_hub_in_organization(self):
2585-
components = self.get_pipeline_components()
2586-
pipeline = StableDiffusionPipeline(**components)
2587-
pipeline.push_to_hub(self.org_repo_id, token=TOKEN)
2588-
2589-
new_model = UNet2DConditionModel.from_pretrained(self.org_repo_id, subfolder="unet")
2590-
unet = components["unet"]
2591-
for p1, p2 in zip(unet.parameters(), new_model.parameters()):
2592-
self.assertTrue(torch.equal(p1, p2))
2593-
2594-
# Push to hub via save_pretrained to a separate repo. Reusing `self.org_repo_id` after
2595-
# deleting it makes the staging server's LFS GC reject the next commit with
2596-
# "LFS pointer pointed to a file that does not exist" when the model bytes are identical.
2597-
save_org_repo_id = f"{self.org_repo_id}-saved"
2598-
with tempfile.TemporaryDirectory() as tmp_dir:
2599-
pipeline.save_pretrained(tmp_dir, push_to_hub=True, token=TOKEN, repo_id=save_org_repo_id)
2600-
2601-
new_model = UNet2DConditionModel.from_pretrained(save_org_repo_id, subfolder="unet")
2602-
for p1, p2 in zip(unet.parameters(), new_model.parameters()):
2603-
self.assertTrue(torch.equal(p1, p2))
2604-
2605-
# Reset repos
2606-
delete_repo(token=TOKEN, repo_id=self.org_repo_id)
2607-
delete_repo(save_org_repo_id, token=TOKEN)
2608-
2609-
@unittest.skipIf(
2610-
not is_jinja_available(),
2611-
reason="Model card tests cannot be performed without Jinja installed.",
2612-
)
2613-
def test_push_to_hub_library_name(self):
2614-
components = self.get_pipeline_components()
2615-
pipeline = StableDiffusionPipeline(**components)
2616-
# Use a method-unique repo to avoid recycling a name that `test_push_to_hub` just deleted,
2617-
# which the staging server rejects with an LFS pointer error.
2618-
repo_id = f"test-pipeline-library-name-{uuid.uuid4()}"
2619-
pipeline.push_to_hub(repo_id, token=TOKEN)
2620-
2621-
model_card = ModelCard.load(f"{USER}/{repo_id}", token=TOKEN).data
2622-
assert model_card.library_name == "diffusers"
2623-
2624-
# Reset repo
2625-
delete_repo(repo_id, token=TOKEN)
2626-
2627-
26282482
class PyramidAttentionBroadcastTesterMixin:
26292483
pab_config = PyramidAttentionBroadcastConfig(
26302484
spatial_attention_block_skip_range=2,
Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
from .cache import (
2+
CacheTesterMixin,
3+
FasterCacheTesterMixin,
4+
FirstBlockCacheTesterMixin,
5+
MagCacheTesterMixin,
6+
PyramidAttentionBroadcastTesterMixin,
7+
TaylorSeerCacheTesterMixin,
8+
)
9+
from .common import BasePipelineTesterConfig, PipelineTesterMixin
10+
from .memory import (
11+
GroupOffloadTesterMixin,
12+
LayerwiseCastingTesterMixin,
13+
MemoryTesterMixin,
14+
PipelineOffloadTesterMixin,
15+
)
16+
from .utils import (
17+
check_qkv_fused_layers_exist,
18+
check_qkv_fusion_matches_attn_procs_length,
19+
check_qkv_fusion_processors_exist,
20+
check_same_shape,
21+
)
22+
23+
24+
__all__ = [
25+
"BasePipelineTesterConfig",
26+
"PipelineTesterMixin",
27+
"MemoryTesterMixin",
28+
"PipelineOffloadTesterMixin",
29+
"GroupOffloadTesterMixin",
30+
"LayerwiseCastingTesterMixin",
31+
"CacheTesterMixin",
32+
"PyramidAttentionBroadcastTesterMixin",
33+
"FasterCacheTesterMixin",
34+
"FirstBlockCacheTesterMixin",
35+
"TaylorSeerCacheTesterMixin",
36+
"MagCacheTesterMixin",
37+
"check_qkv_fused_layers_exist",
38+
"check_qkv_fusion_matches_attn_procs_length",
39+
"check_qkv_fusion_processors_exist",
40+
"check_same_shape",
41+
]

0 commit comments

Comments
 (0)