|
| 1 | +import gc |
| 2 | +import os |
| 3 | +import unittest |
| 4 | +from pathlib import Path |
| 5 | + |
| 6 | +import torch |
| 7 | +from diffusers import AutoencoderKL, EulerDiscreteScheduler, StableDiffusionPipeline, StableDiffusionXLPipeline, UNet2DConditionModel |
| 8 | +from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection |
| 9 | + |
| 10 | +from compel import CompelForSD, CompelForSDXL |
| 11 | + |
| 12 | +try: |
| 13 | + from prompting_test_utils import DummyTokenizer |
| 14 | +except ModuleNotFoundError: |
| 15 | + from test.prompting_test_utils import DummyTokenizer |
| 16 | + |
| 17 | + |
| 18 | +LOCAL_SDXL_CHECKPOINTS = { |
| 19 | + "sd_xl_turbo": Path("/mnt/s/Code/Models/ImageGen/CUI-Archived/checkpoints/sd_xl_turbo_1.0_fp16.safetensors"), |
| 20 | + "cyberrealisticXL_v80": Path("/mnt/s/Code/Models/ImageGen/CUI-Archived/checkpoints/cyberrealisticXL_v80_fp16.safetensors"), |
| 21 | +} |
| 22 | +RUN_LOCAL_CHECKPOINT_TESTS = os.getenv("COMPEL_RUN_LOCAL_CHECKPOINT_TESTS") == "1" |
| 23 | + |
| 24 | + |
| 25 | +def make_clip_text_config(hidden_size: int = 32, projection_dim: int = 32, max_position_embeddings: int = 16) -> CLIPTextConfig: |
| 26 | + return CLIPTextConfig( |
| 27 | + vocab_size=32, |
| 28 | + hidden_size=hidden_size, |
| 29 | + intermediate_size=37, |
| 30 | + projection_dim=projection_dim, |
| 31 | + num_hidden_layers=2, |
| 32 | + num_attention_heads=4, |
| 33 | + max_position_embeddings=max_position_embeddings, |
| 34 | + bos_token_id=10, |
| 35 | + pad_token_id=11, |
| 36 | + eos_token_id=12, |
| 37 | + ) |
| 38 | + |
| 39 | + |
| 40 | +def make_tiny_vae() -> AutoencoderKL: |
| 41 | + return AutoencoderKL( |
| 42 | + in_channels=3, |
| 43 | + out_channels=3, |
| 44 | + down_block_types=["DownEncoderBlock2D"], |
| 45 | + up_block_types=["UpDecoderBlock2D"], |
| 46 | + block_out_channels=[32], |
| 47 | + layers_per_block=1, |
| 48 | + latent_channels=4, |
| 49 | + norm_num_groups=8, |
| 50 | + sample_size=32, |
| 51 | + ) |
| 52 | + |
| 53 | + |
| 54 | +def make_tiny_sd_unet(cross_attention_dim: int) -> UNet2DConditionModel: |
| 55 | + return UNet2DConditionModel( |
| 56 | + sample_size=32, |
| 57 | + in_channels=4, |
| 58 | + out_channels=4, |
| 59 | + layers_per_block=1, |
| 60 | + block_out_channels=(32, 64), |
| 61 | + down_block_types=("CrossAttnDownBlock2D", "DownBlock2D"), |
| 62 | + up_block_types=("UpBlock2D", "CrossAttnUpBlock2D"), |
| 63 | + cross_attention_dim=cross_attention_dim, |
| 64 | + attention_head_dim=(4, 8), |
| 65 | + norm_num_groups=8, |
| 66 | + ) |
| 67 | + |
| 68 | + |
| 69 | +def make_tiny_sdxl_unet(cross_attention_dim: int, projection_dim: int) -> UNet2DConditionModel: |
| 70 | + return UNet2DConditionModel( |
| 71 | + sample_size=32, |
| 72 | + in_channels=4, |
| 73 | + out_channels=4, |
| 74 | + layers_per_block=1, |
| 75 | + block_out_channels=(32, 64), |
| 76 | + down_block_types=("CrossAttnDownBlock2D", "DownBlock2D"), |
| 77 | + up_block_types=("UpBlock2D", "CrossAttnUpBlock2D"), |
| 78 | + cross_attention_dim=cross_attention_dim, |
| 79 | + attention_head_dim=(4, 8), |
| 80 | + norm_num_groups=8, |
| 81 | + addition_embed_type="text_time", |
| 82 | + addition_time_embed_dim=8, |
| 83 | + projection_class_embeddings_input_dim=(8 * 6) + projection_dim, |
| 84 | + ) |
| 85 | + |
| 86 | + |
| 87 | +class DiffusersSmokeTestCase(unittest.TestCase): |
| 88 | + def test_stable_diffusion_pipeline_accepts_compel_prompt_embeds(self): |
| 89 | + tokenizer = DummyTokenizer(model_max_length=16) |
| 90 | + text_encoder = CLIPTextModel(make_clip_text_config(hidden_size=32, projection_dim=32, max_position_embeddings=16)) |
| 91 | + pipe = StableDiffusionPipeline( |
| 92 | + vae=make_tiny_vae(), |
| 93 | + text_encoder=text_encoder, |
| 94 | + tokenizer=tokenizer, |
| 95 | + unet=make_tiny_sd_unet(cross_attention_dim=32), |
| 96 | + scheduler=EulerDiscreteScheduler(num_train_timesteps=10, steps_offset=1), |
| 97 | + safety_checker=None, |
| 98 | + feature_extractor=None, |
| 99 | + requires_safety_checker=False, |
| 100 | + ).to("cpu") |
| 101 | + pipe.set_progress_bar_config(disable=True) |
| 102 | + |
| 103 | + conditioning = CompelForSD(pipe)("a b c", negative_prompt="a") |
| 104 | + result = pipe( |
| 105 | + prompt_embeds=conditioning.embeds, |
| 106 | + negative_prompt_embeds=conditioning.negative_embeds, |
| 107 | + num_inference_steps=1, |
| 108 | + guidance_scale=2.0, |
| 109 | + output_type="np", |
| 110 | + ) |
| 111 | + |
| 112 | + self.assertEqual(conditioning.embeds.shape, (1, 16, 32)) |
| 113 | + self.assertEqual(conditioning.negative_embeds.shape, (1, 16, 32)) |
| 114 | + self.assertEqual(len(result.images), 1) |
| 115 | + self.assertEqual(result.images[0].shape, (32, 32, 3)) |
| 116 | + |
| 117 | + def test_sdxl_pipeline_accepts_compel_prompt_and_pooled_embeds(self): |
| 118 | + tokenizer = DummyTokenizer(model_max_length=16) |
| 119 | + text_encoder = CLIPTextModel(make_clip_text_config(hidden_size=32, projection_dim=16, max_position_embeddings=16)) |
| 120 | + text_encoder_2 = CLIPTextModelWithProjection( |
| 121 | + make_clip_text_config(hidden_size=32, projection_dim=16, max_position_embeddings=16) |
| 122 | + ) |
| 123 | + pipe = StableDiffusionXLPipeline( |
| 124 | + vae=make_tiny_vae(), |
| 125 | + text_encoder=text_encoder, |
| 126 | + text_encoder_2=text_encoder_2, |
| 127 | + tokenizer=tokenizer, |
| 128 | + tokenizer_2=DummyTokenizer(model_max_length=16), |
| 129 | + unet=make_tiny_sdxl_unet(cross_attention_dim=64, projection_dim=16), |
| 130 | + scheduler=EulerDiscreteScheduler(num_train_timesteps=10, steps_offset=1), |
| 131 | + image_encoder=None, |
| 132 | + feature_extractor=None, |
| 133 | + ).to("cpu") |
| 134 | + pipe.set_progress_bar_config(disable=True) |
| 135 | + |
| 136 | + conditioning = CompelForSDXL(pipe)( |
| 137 | + "a b c", |
| 138 | + style_prompt="b c", |
| 139 | + negative_prompt="a", |
| 140 | + negative_style_prompt="a", |
| 141 | + ) |
| 142 | + result = pipe( |
| 143 | + prompt_embeds=conditioning.embeds, |
| 144 | + pooled_prompt_embeds=conditioning.pooled_embeds, |
| 145 | + negative_prompt_embeds=conditioning.negative_embeds, |
| 146 | + negative_pooled_prompt_embeds=conditioning.negative_pooled_embeds, |
| 147 | + num_inference_steps=1, |
| 148 | + guidance_scale=2.0, |
| 149 | + output_type="np", |
| 150 | + ) |
| 151 | + |
| 152 | + self.assertEqual(conditioning.embeds.shape, (1, 16, 64)) |
| 153 | + self.assertEqual(conditioning.pooled_embeds.shape, (1, 16)) |
| 154 | + self.assertEqual(conditioning.negative_embeds.shape, (1, 16, 64)) |
| 155 | + self.assertEqual(conditioning.negative_pooled_embeds.shape, (1, 16)) |
| 156 | + self.assertEqual(len(result.images), 1) |
| 157 | + self.assertEqual(result.images[0].shape, (32, 32, 3)) |
| 158 | + |
| 159 | + |
| 160 | +@unittest.skipUnless(RUN_LOCAL_CHECKPOINT_TESTS, "set COMPEL_RUN_LOCAL_CHECKPOINT_TESTS=1 to run local checkpoint smoke tests") |
| 161 | +class LocalCheckpointSmokeTestCase(unittest.TestCase): |
| 162 | + def test_sdxl_single_file_checkpoints_load_locally(self): |
| 163 | + dtype = torch.float16 if torch.cuda.is_available() else torch.float32 |
| 164 | + |
| 165 | + for checkpoint_name, checkpoint_path in LOCAL_SDXL_CHECKPOINTS.items(): |
| 166 | + with self.subTest(checkpoint=checkpoint_name): |
| 167 | + self.assertTrue(checkpoint_path.exists(), f"missing local checkpoint: {checkpoint_path}") |
| 168 | + |
| 169 | + pipe = StableDiffusionXLPipeline.from_single_file( |
| 170 | + str(checkpoint_path), |
| 171 | + local_files_only=True, |
| 172 | + torch_dtype=dtype, |
| 173 | + ) |
| 174 | + pipe.set_progress_bar_config(disable=True) |
| 175 | + |
| 176 | + self.assertIsInstance(pipe, StableDiffusionXLPipeline) |
| 177 | + self.assertIsNotNone(pipe.text_encoder) |
| 178 | + self.assertIsNotNone(pipe.text_encoder_2) |
| 179 | + self.assertGreater(pipe.unet.config.cross_attention_dim, 0) |
| 180 | + |
| 181 | + del pipe |
| 182 | + gc.collect() |
| 183 | + |
| 184 | + |
| 185 | +if __name__ == "__main__": |
| 186 | + unittest.main() |
0 commit comments