test_pipeline_flux.py 13.5 KB
Newer Older
Sayak Paul's avatar
Sayak Paul committed
1
2
3
4
5
import gc
import unittest

import numpy as np
import torch
6
from huggingface_hub import hf_hub_download
Sayak Paul's avatar
Sayak Paul committed
7
8
from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel

Aryan's avatar
Aryan committed
9
10
11
12
13
14
15
from diffusers import (
    AutoencoderKL,
    FasterCacheConfig,
    FlowMatchEulerDiscreteScheduler,
    FluxPipeline,
    FluxTransformer2DModel,
)
Sayak Paul's avatar
Sayak Paul committed
16
from diffusers.utils.testing_utils import (
17
    backend_empty_cache,
18
    nightly,
Sayak Paul's avatar
Sayak Paul committed
19
    numpy_cosine_similarity_distance,
20
    require_big_accelerator,
Sayak Paul's avatar
Sayak Paul committed
21
22
23
24
    slow,
    torch_device,
)

25
from ..test_pipelines_common import (
Aryan's avatar
Aryan committed
26
    FasterCacheTesterMixin,
Aryan's avatar
Aryan committed
27
    FirstBlockCacheTesterMixin,
hlky's avatar
hlky committed
28
    FluxIPAdapterTesterMixin,
29
    PipelineTesterMixin,
30
    PyramidAttentionBroadcastTesterMixin,
31
32
33
    check_qkv_fusion_matches_attn_procs_length,
    check_qkv_fusion_processors_exist,
)
Sayak Paul's avatar
Sayak Paul committed
34
35


36
class FluxPipelineFastTests(
Aryan's avatar
Aryan committed
37
38
39
40
    PipelineTesterMixin,
    FluxIPAdapterTesterMixin,
    PyramidAttentionBroadcastTesterMixin,
    FasterCacheTesterMixin,
Aryan's avatar
Aryan committed
41
42
    FirstBlockCacheTesterMixin,
    unittest.TestCase,
43
):
Sayak Paul's avatar
Sayak Paul committed
44
    pipeline_class = FluxPipeline
Sayak Paul's avatar
Sayak Paul committed
45
46
    params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"])
    batch_params = frozenset(["prompt"])
Sayak Paul's avatar
Sayak Paul committed
47

48
49
    # there is no xformers processor for Flux
    test_xformers_attention = False
Aryan's avatar
Aryan committed
50
    test_layerwise_casting = True
Aryan's avatar
Aryan committed
51
    test_group_offloading = True
52

Aryan's avatar
Aryan committed
53
54
55
56
57
58
59
60
    faster_cache_config = FasterCacheConfig(
        spatial_attention_block_skip_range=2,
        spatial_attention_timestep_skip_range=(-1, 901),
        unconditional_batch_skip_range=2,
        attention_weight_callback=lambda _: 0.5,
        is_guidance_distilled=True,
    )

61
    def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1):
Sayak Paul's avatar
Sayak Paul committed
62
63
64
65
        torch.manual_seed(0)
        transformer = FluxTransformer2DModel(
            patch_size=1,
            in_channels=4,
66
67
            num_layers=num_layers,
            num_single_layers=num_single_layers,
Sayak Paul's avatar
Sayak Paul committed
68
69
            attention_head_dim=16,
            num_attention_heads=2,
Sayak Paul's avatar
Sayak Paul committed
70
            joint_attention_dim=32,
Sayak Paul's avatar
Sayak Paul committed
71
72
            pooled_projection_dim=32,
            axes_dims_rope=[4, 4, 8],
Sayak Paul's avatar
Sayak Paul committed
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
        )
        clip_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,
            hidden_act="gelu",
            projection_dim=32,
        )

        torch.manual_seed(0)
        text_encoder = CLIPTextModel(clip_text_encoder_config)

        torch.manual_seed(0)
        text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")

        tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
        tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")

        torch.manual_seed(0)
        vae = AutoencoderKL(
            sample_size=32,
            in_channels=3,
            out_channels=3,
            block_out_channels=(4,),
            layers_per_block=1,
Sayak Paul's avatar
Sayak Paul committed
104
            latent_channels=1,
Sayak Paul's avatar
Sayak Paul committed
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
            norm_num_groups=1,
            use_quant_conv=False,
            use_post_quant_conv=False,
            shift_factor=0.0609,
            scaling_factor=1.5035,
        )

        scheduler = FlowMatchEulerDiscreteScheduler()

        return {
            "scheduler": scheduler,
            "text_encoder": text_encoder,
            "text_encoder_2": text_encoder_2,
            "tokenizer": tokenizer,
            "tokenizer_2": tokenizer_2,
            "transformer": transformer,
            "vae": vae,
hlky's avatar
hlky committed
122
123
            "image_encoder": None,
            "feature_extractor": None,
Sayak Paul's avatar
Sayak Paul committed
124
125
126
127
128
129
130
131
132
133
134
135
136
        }

    def get_dummy_inputs(self, device, seed=0):
        if str(device).startswith("mps"):
            generator = torch.manual_seed(seed)
        else:
            generator = torch.Generator(device="cpu").manual_seed(seed)

        inputs = {
            "prompt": "A painting of a squirrel eating a burger",
            "generator": generator,
            "num_inference_steps": 2,
            "guidance_scale": 5.0,
Sayak Paul's avatar
Sayak Paul committed
137
138
139
            "height": 8,
            "width": 8,
            "max_sequence_length": 48,
Sayak Paul's avatar
Sayak Paul committed
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
            "output_type": "np",
        }
        return inputs

    def test_flux_different_prompts(self):
        pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)

        inputs = self.get_dummy_inputs(torch_device)
        output_same_prompt = pipe(**inputs).images[0]

        inputs = self.get_dummy_inputs(torch_device)
        inputs["prompt_2"] = "a different prompt"
        output_different_prompts = pipe(**inputs).images[0]

        max_diff = np.abs(output_same_prompt - output_different_prompts).max()

        # Outputs should be different here
Sayak Paul's avatar
Sayak Paul committed
157
        # For some reasons, they don't show large differences
Aryan's avatar
Aryan committed
158
        self.assertGreater(max_diff, 1e-6, "Outputs should be different for different prompts.")
Sayak Paul's avatar
Sayak Paul committed
159

160
161
162
163
164
165
166
167
168
169
170
171
172
173
    def test_fused_qkv_projections(self):
        device = "cpu"  # ensure determinism for the device-dependent torch.Generator
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe = pipe.to(device)
        pipe.set_progress_bar_config(disable=None)

        inputs = self.get_dummy_inputs(device)
        image = pipe(**inputs).images
        original_image_slice = image[0, -3:, -3:, -1]

        # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added
        # to the pipeline level.
        pipe.transformer.fuse_qkv_projections()
174
175
176
        assert check_qkv_fusion_processors_exist(pipe.transformer), (
            "Something wrong with the fused attention processors. Expected all the attention processors to be fused."
        )
177
178
179
180
181
182
183
184
185
186
187
188
189
        assert check_qkv_fusion_matches_attn_procs_length(
            pipe.transformer, pipe.transformer.original_attn_processors
        ), "Something wrong with the attention processors concerning the fused QKV projections."

        inputs = self.get_dummy_inputs(device)
        image = pipe(**inputs).images
        image_slice_fused = image[0, -3:, -3:, -1]

        pipe.transformer.unfuse_qkv_projections()
        inputs = self.get_dummy_inputs(device)
        image = pipe(**inputs).images
        image_slice_disabled = image[0, -3:, -3:, -1]

Aryan's avatar
Aryan committed
190
191
192
        self.assertTrue(
            np.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3),
            ("Fusion of QKV projections shouldn't affect the outputs."),
193
        )
Aryan's avatar
Aryan committed
194
195
196
        self.assertTrue(
            np.allclose(image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3),
            ("Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled."),
197
        )
Aryan's avatar
Aryan committed
198
199
200
        self.assertTrue(
            np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2),
            ("Original outputs should match when fused QKV projections are disabled."),
201
        )
202

Dhruv Nair's avatar
Dhruv Nair committed
203
204
205
206
207
208
209
210
211
212
213
214
    def test_flux_image_output_shape(self):
        pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
        inputs = self.get_dummy_inputs(torch_device)

        height_width_pairs = [(32, 32), (72, 57)]
        for height, width in height_width_pairs:
            expected_height = height - height % (pipe.vae_scale_factor * 2)
            expected_width = width - width % (pipe.vae_scale_factor * 2)

            inputs.update({"height": height, "width": width})
            image = pipe(**inputs).images[0]
            output_height, output_width, _ = image.shape
Aryan's avatar
Aryan committed
215
216
217
218
219
            self.assertEqual(
                (output_height, output_width),
                (expected_height, expected_width),
                f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}",
            )
Dhruv Nair's avatar
Dhruv Nair committed
220

221
222
223
224
225
226
227
228
229
    def test_flux_true_cfg(self):
        pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
        inputs = self.get_dummy_inputs(torch_device)
        inputs.pop("generator")

        no_true_cfg_out = pipe(**inputs, generator=torch.manual_seed(0)).images[0]
        inputs["negative_prompt"] = "bad quality"
        inputs["true_cfg_scale"] = 2.0
        true_cfg_out = pipe(**inputs, generator=torch.manual_seed(0)).images[0]
Aryan's avatar
Aryan committed
230
231
232
        self.assertFalse(
            np.allclose(no_true_cfg_out, true_cfg_out), "Outputs should be different when true_cfg_scale is set."
        )
233

Sayak Paul's avatar
Sayak Paul committed
234

235
@nightly
236
@require_big_accelerator
Sayak Paul's avatar
Sayak Paul committed
237
238
239
240
241
242
243
class FluxPipelineSlowTests(unittest.TestCase):
    pipeline_class = FluxPipeline
    repo_id = "black-forest-labs/FLUX.1-schnell"

    def setUp(self):
        super().setUp()
        gc.collect()
244
        backend_empty_cache(torch_device)
Sayak Paul's avatar
Sayak Paul committed
245
246
247
248

    def tearDown(self):
        super().tearDown()
        gc.collect()
249
        backend_empty_cache(torch_device)
Sayak Paul's avatar
Sayak Paul committed
250
251

    def get_inputs(self, device, seed=0):
252
        generator = torch.Generator(device="cpu").manual_seed(seed)
Sayak Paul's avatar
Sayak Paul committed
253

254
255
        prompt_embeds = torch.load(
            hf_hub_download(repo_id="diffusers/test-slices", repo_type="dataset", filename="flux/prompt_embeds.pt")
256
        ).to(torch_device)
257
258
259
260
        pooled_prompt_embeds = torch.load(
            hf_hub_download(
                repo_id="diffusers/test-slices", repo_type="dataset", filename="flux/pooled_prompt_embeds.pt"
            )
261
        ).to(torch_device)
Sayak Paul's avatar
Sayak Paul committed
262
        return {
263
264
            "prompt_embeds": prompt_embeds,
            "pooled_prompt_embeds": pooled_prompt_embeds,
Sayak Paul's avatar
Sayak Paul committed
265
            "num_inference_steps": 2,
266
267
            "guidance_scale": 0.0,
            "max_sequence_length": 256,
Sayak Paul's avatar
Sayak Paul committed
268
269
270
271
272
            "output_type": "np",
            "generator": generator,
        }

    def test_flux_inference(self):
273
274
        pipe = self.pipeline_class.from_pretrained(
            self.repo_id, torch_dtype=torch.bfloat16, text_encoder=None, text_encoder_2=None
275
        ).to(torch_device)
Sayak Paul's avatar
Sayak Paul committed
276
277
278
279
280

        inputs = self.get_inputs(torch_device)

        image = pipe(**inputs).images[0]
        image_slice = image[0, :10, :10]
Aryan's avatar
Aryan committed
281
        # fmt: off
Sayak Paul's avatar
Sayak Paul committed
282
        expected_slice = np.array(
Aryan's avatar
Aryan committed
283
            [0.3242, 0.3203, 0.3164, 0.3164, 0.3125, 0.3125, 0.3281, 0.3242, 0.3203, 0.3301, 0.3262, 0.3242, 0.3281, 0.3242, 0.3203, 0.3262, 0.3262, 0.3164, 0.3262, 0.3281, 0.3184, 0.3281, 0.3281, 0.3203, 0.3281, 0.3281, 0.3164, 0.3320, 0.3320, 0.3203],
Sayak Paul's avatar
Sayak Paul committed
284
285
            dtype=np.float32,
        )
Aryan's avatar
Aryan committed
286
        # fmt: on
Sayak Paul's avatar
Sayak Paul committed
287
288

        max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), image_slice.flatten())
Aryan's avatar
Aryan committed
289
290
291
        self.assertLess(
            max_diff, 1e-4, f"Image slice is different from expected slice: {image_slice} != {expected_slice}"
        )
hlky's avatar
hlky committed
292
293
294


@slow
295
@require_big_accelerator
hlky's avatar
hlky committed
296
297
298
299
300
301
302
303
304
305
class FluxIPAdapterPipelineSlowTests(unittest.TestCase):
    pipeline_class = FluxPipeline
    repo_id = "black-forest-labs/FLUX.1-dev"
    image_encoder_pretrained_model_name_or_path = "openai/clip-vit-large-patch14"
    weight_name = "ip_adapter.safetensors"
    ip_adapter_repo_id = "XLabs-AI/flux-ip-adapter"

    def setUp(self):
        super().setUp()
        gc.collect()
306
        backend_empty_cache(torch_device)
hlky's avatar
hlky committed
307
308
309
310

    def tearDown(self):
        super().tearDown()
        gc.collect()
311
        backend_empty_cache(torch_device)
hlky's avatar
hlky committed
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360

    def get_inputs(self, device, seed=0):
        if str(device).startswith("mps"):
            generator = torch.manual_seed(seed)
        else:
            generator = torch.Generator(device="cpu").manual_seed(seed)

        prompt_embeds = torch.load(
            hf_hub_download(repo_id="diffusers/test-slices", repo_type="dataset", filename="flux/prompt_embeds.pt")
        )
        pooled_prompt_embeds = torch.load(
            hf_hub_download(
                repo_id="diffusers/test-slices", repo_type="dataset", filename="flux/pooled_prompt_embeds.pt"
            )
        )
        negative_prompt_embeds = torch.zeros_like(prompt_embeds)
        negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
        ip_adapter_image = np.zeros((1024, 1024, 3), dtype=np.uint8)
        return {
            "prompt_embeds": prompt_embeds,
            "pooled_prompt_embeds": pooled_prompt_embeds,
            "negative_prompt_embeds": negative_prompt_embeds,
            "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds,
            "ip_adapter_image": ip_adapter_image,
            "num_inference_steps": 2,
            "guidance_scale": 3.5,
            "true_cfg_scale": 4.0,
            "max_sequence_length": 256,
            "output_type": "np",
            "generator": generator,
        }

    def test_flux_ip_adapter_inference(self):
        pipe = self.pipeline_class.from_pretrained(
            self.repo_id, torch_dtype=torch.bfloat16, text_encoder=None, text_encoder_2=None
        )
        pipe.load_ip_adapter(
            self.ip_adapter_repo_id,
            weight_name=self.weight_name,
            image_encoder_pretrained_model_name_or_path=self.image_encoder_pretrained_model_name_or_path,
        )
        pipe.set_ip_adapter_scale(1.0)
        pipe.enable_model_cpu_offload()

        inputs = self.get_inputs(torch_device)

        image = pipe(**inputs).images[0]
        image_slice = image[0, :10, :10]

Aryan's avatar
Aryan committed
361
        # fmt: off
hlky's avatar
hlky committed
362
        expected_slice = np.array(
Aryan's avatar
Aryan committed
363
            [0.1855, 0.1680, 0.1406, 0.1953, 0.1699, 0.1465, 0.2012, 0.1738, 0.1484, 0.2051, 0.1797, 0.1523, 0.2012, 0.1719, 0.1445, 0.2070, 0.1777, 0.1465, 0.2090, 0.1836, 0.1484, 0.2129, 0.1875, 0.1523, 0.2090, 0.1816, 0.1484, 0.2110, 0.1836, 0.1543],
hlky's avatar
hlky committed
364
365
            dtype=np.float32,
        )
Aryan's avatar
Aryan committed
366
        # fmt: on
hlky's avatar
hlky committed
367
368

        max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), image_slice.flatten())
Aryan's avatar
Aryan committed
369
370
371
        self.assertLess(
            max_diff, 1e-4, f"Image slice is different from expected slice: {image_slice} != {expected_slice}"
        )