[tests] Pipeline test refactor

tests/pipelines/test_pipelines_common.py

@@ -4,25 +4,20 @@
 import os
 import tempfile
 import unittest
-import uuid
 from typing import Any, Callable, Dict
 
 import numpy as np
 import PIL.Image
 import pytest
 import torch
 import torch.nn as nn
-from huggingface_hub import ModelCard, delete_repo
-from huggingface_hub.utils import is_jinja_available
-from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
 
 import diffusers
 from diffusers import (
     AsymmetricAutoencoderKL,
     AutoencoderKL,
     AutoencoderTiny,
     ConsistencyDecoderVAE,
-    DDIMScheduler,
     DiffusionPipeline,
     FasterCacheConfig,
     KolorsPipeline,
@@ -63,7 +58,6 @@
     create_ip_adapter_faceid_state_dict,
     create_ip_adapter_state_dict,
 )
-from ..others.test_utils import TOKEN, USER, is_staging_test
 from ..testing_utils import (
     CaptureLogger,
     backend_empty_cache,
@@ -2507,146 +2501,6 @@ def test_pipeline_level_group_offloading_inference(self, expected_max_difference
         self.assertLess(max_diff, expected_max_difference)
 
 
-@is_staging_test
-class PipelinePushToHubTester(unittest.TestCase):
-    identifier = uuid.uuid4()
-    repo_id = f"test-pipeline-{identifier}"
-    org_repo_id = f"valid_org/{repo_id}-org"
-
-    def get_pipeline_components(self):
-        unet = UNet2DConditionModel(
-            block_out_channels=(32, 64),
-            layers_per_block=2,
-            sample_size=32,
-            in_channels=4,
-            out_channels=4,
-            down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
-            up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),
-            cross_attention_dim=32,
-        )
-
-        scheduler = DDIMScheduler(
-            beta_start=0.00085,
-            beta_end=0.012,
-            beta_schedule="scaled_linear",
-            clip_sample=False,
-            set_alpha_to_one=False,
-        )
-
-        vae = AutoencoderKL(
-            block_out_channels=[32, 64],
-            in_channels=3,
-            out_channels=3,
-            down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"],
-            up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],
-            latent_channels=4,
-        )
-
-        text_encoder_config = CLIPTextConfig(
-            bos_token_id=0,
-            eos_token_id=2,
-            hidden_size=32,
-            intermediate_size=37,
-            layer_norm_eps=1e-05,
-            num_attention_heads=4,
-            num_hidden_layers=5,
-            pad_token_id=1,
-            vocab_size=1000,
-        )
-        text_encoder = CLIPTextModel(text_encoder_config)
-
-        with tempfile.TemporaryDirectory() as tmpdir:
-            dummy_vocab = {"<|startoftext|>": 0, "<|endoftext|>": 1, "!": 2}
-            vocab_path = os.path.join(tmpdir, "vocab.json")
-            with open(vocab_path, "w") as f:
-                json.dump(dummy_vocab, f)
-
-            merges = "Ġ t\nĠt h"
-            merges_path = os.path.join(tmpdir, "merges.txt")
-            with open(merges_path, "w") as f:
-                f.writelines(merges)
-            tokenizer = CLIPTokenizer(vocab_file=vocab_path, merges_file=merges_path)
-
-        components = {
-            "unet": unet,
-            "scheduler": scheduler,
-            "vae": vae,
-            "text_encoder": text_encoder,
-            "tokenizer": tokenizer,
-            "safety_checker": None,
-            "feature_extractor": None,
-        }
-        return components
-
-    def test_push_to_hub(self):
-        components = self.get_pipeline_components()
-        pipeline = StableDiffusionPipeline(**components)
-        pipeline.push_to_hub(self.repo_id, token=TOKEN)
-
-        new_model = UNet2DConditionModel.from_pretrained(f"{USER}/{self.repo_id}", subfolder="unet")
-        unet = components["unet"]
-        for p1, p2 in zip(unet.parameters(), new_model.parameters()):
-            self.assertTrue(torch.equal(p1, p2))
-
-        # Push to hub via save_pretrained to a separate repo. Reusing `self.repo_id` after
-        # deleting it makes the staging server's LFS GC reject the next commit with
-        # "LFS pointer pointed to a file that does not exist" when the model bytes are identical.
-        save_repo_id = f"{self.repo_id}-saved"
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            pipeline.save_pretrained(tmp_dir, repo_id=save_repo_id, push_to_hub=True, token=TOKEN)
-
-        new_model = UNet2DConditionModel.from_pretrained(f"{USER}/{save_repo_id}", subfolder="unet")
-        for p1, p2 in zip(unet.parameters(), new_model.parameters()):
-            self.assertTrue(torch.equal(p1, p2))
-
-        # Reset repos
-        delete_repo(token=TOKEN, repo_id=self.repo_id)
-        delete_repo(save_repo_id, token=TOKEN)
-
-    def test_push_to_hub_in_organization(self):
-        components = self.get_pipeline_components()
-        pipeline = StableDiffusionPipeline(**components)
-        pipeline.push_to_hub(self.org_repo_id, token=TOKEN)
-
-        new_model = UNet2DConditionModel.from_pretrained(self.org_repo_id, subfolder="unet")
-        unet = components["unet"]
-        for p1, p2 in zip(unet.parameters(), new_model.parameters()):
-            self.assertTrue(torch.equal(p1, p2))
-
-        # Push to hub via save_pretrained to a separate repo. Reusing `self.org_repo_id` after
-        # deleting it makes the staging server's LFS GC reject the next commit with
-        # "LFS pointer pointed to a file that does not exist" when the model bytes are identical.
-        save_org_repo_id = f"{self.org_repo_id}-saved"
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            pipeline.save_pretrained(tmp_dir, push_to_hub=True, token=TOKEN, repo_id=save_org_repo_id)
-
-        new_model = UNet2DConditionModel.from_pretrained(save_org_repo_id, subfolder="unet")
-        for p1, p2 in zip(unet.parameters(), new_model.parameters()):
-            self.assertTrue(torch.equal(p1, p2))
-
-        # Reset repos
-        delete_repo(token=TOKEN, repo_id=self.org_repo_id)
-        delete_repo(save_org_repo_id, token=TOKEN)
-
-    @unittest.skipIf(
-        not is_jinja_available(),
-        reason="Model card tests cannot be performed without Jinja installed.",
-    )
-    def test_push_to_hub_library_name(self):
-        components = self.get_pipeline_components()
-        pipeline = StableDiffusionPipeline(**components)
-        # Use a method-unique repo to avoid recycling a name that `test_push_to_hub` just deleted,
-        # which the staging server rejects with an LFS pointer error.
-        repo_id = f"test-pipeline-library-name-{uuid.uuid4()}"
-        pipeline.push_to_hub(repo_id, token=TOKEN)
-
-        model_card = ModelCard.load(f"{USER}/{repo_id}", token=TOKEN).data
-        assert model_card.library_name == "diffusers"
-
-        # Reset repo
-        delete_repo(repo_id, token=TOKEN)
-
-
 class PyramidAttentionBroadcastTesterMixin:
     pab_config = PyramidAttentionBroadcastConfig(
         spatial_attention_block_skip_range=2,

tests/pipelines/testing_utils/__init__.py

@@ -0,0 +1,41 @@
+from .cache import (
+    CacheTesterMixin,
+    FasterCacheTesterMixin,
+    FirstBlockCacheTesterMixin,
+    MagCacheTesterMixin,
+    PyramidAttentionBroadcastTesterMixin,
+    TaylorSeerCacheTesterMixin,
+)
+from .common import BasePipelineTesterConfig, PipelineTesterMixin
+from .memory import (
+    GroupOffloadTesterMixin,
+    LayerwiseCastingTesterMixin,
+    MemoryTesterMixin,
+    PipelineOffloadTesterMixin,
+)
+from .utils import (
+    check_qkv_fused_layers_exist,
+    check_qkv_fusion_matches_attn_procs_length,
+    check_qkv_fusion_processors_exist,
+    check_same_shape,
+)
+
+
+__all__ = [
+    "BasePipelineTesterConfig",
+    "PipelineTesterMixin",
+    "MemoryTesterMixin",
+    "PipelineOffloadTesterMixin",
+    "GroupOffloadTesterMixin",
+    "LayerwiseCastingTesterMixin",
+    "CacheTesterMixin",
+    "PyramidAttentionBroadcastTesterMixin",
+    "FasterCacheTesterMixin",
+    "FirstBlockCacheTesterMixin",
+    "TaylorSeerCacheTesterMixin",
+    "MagCacheTesterMixin",
+    "check_qkv_fused_layers_exist",
+    "check_qkv_fusion_matches_attn_procs_length",
+    "check_qkv_fusion_processors_exist",
+    "check_same_shape",
+]