test_pipeline_flux.py 13.3 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
    check_qkv_fused_layers_exist,
32
)
Sayak Paul's avatar
Sayak Paul committed
33
34


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

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

Aryan's avatar
Aryan committed
52
53
54
55
56
57
58
59
    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,
    )

60
    def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1):
Sayak Paul's avatar
Sayak Paul committed
61
62
63
64
        torch.manual_seed(0)
        transformer = FluxTransformer2DModel(
            patch_size=1,
            in_channels=4,
65
66
            num_layers=num_layers,
            num_single_layers=num_single_layers,
Sayak Paul's avatar
Sayak Paul committed
67
68
            attention_head_dim=16,
            num_attention_heads=2,
Sayak Paul's avatar
Sayak Paul committed
69
            joint_attention_dim=32,
Sayak Paul's avatar
Sayak Paul committed
70
71
            pooled_projection_dim=32,
            axes_dims_rope=[4, 4, 8],
Sayak Paul's avatar
Sayak Paul committed
72
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
        )
        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
103
            latent_channels=1,
Sayak Paul's avatar
Sayak Paul committed
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
            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
121
122
            "image_encoder": None,
            "feature_extractor": None,
Sayak Paul's avatar
Sayak Paul committed
123
124
125
126
127
128
129
130
131
132
133
134
135
        }

    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
136
137
138
            "height": 8,
            "width": 8,
            "max_sequence_length": 48,
Sayak Paul's avatar
Sayak Paul committed
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
            "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
156
        # For some reasons, they don't show large differences
Aryan's avatar
Aryan committed
157
        self.assertGreater(max_diff, 1e-6, "Outputs should be different for different prompts.")
Sayak Paul's avatar
Sayak Paul committed
158

159
160
161
162
163
164
165
166
167
168
169
170
171
172
    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()
173
174
175
        self.assertTrue(
            check_qkv_fused_layers_exist(pipe.transformer, ["to_qkv"]),
            ("Something wrong with the fused attention layers. Expected all the attention projections to be fused."),
176
        )
177
178
179
180
181
182
183
184
185
186

        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
187
188
189
        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."),
190
        )
Aryan's avatar
Aryan committed
191
192
193
        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."),
194
        )
Aryan's avatar
Aryan committed
195
196
197
        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."),
198
        )
199

Dhruv Nair's avatar
Dhruv Nair committed
200
201
202
203
204
205
206
207
208
209
210
211
    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
212
213
214
215
216
            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
217

218
219
220
221
222
223
224
225
226
    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
227
228
229
        self.assertFalse(
            np.allclose(no_true_cfg_out, true_cfg_out), "Outputs should be different when true_cfg_scale is set."
        )
230

Sayak Paul's avatar
Sayak Paul committed
231

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

    def setUp(self):
        super().setUp()
        gc.collect()
241
        backend_empty_cache(torch_device)
Sayak Paul's avatar
Sayak Paul committed
242
243
244
245

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

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

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

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

        inputs = self.get_inputs(torch_device)

        image = pipe(**inputs).images[0]
        image_slice = image[0, :10, :10]
Aryan's avatar
Aryan committed
278
        # fmt: off
Sayak Paul's avatar
Sayak Paul committed
279
        expected_slice = np.array(
Aryan's avatar
Aryan committed
280
            [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
281
282
            dtype=np.float32,
        )
Aryan's avatar
Aryan committed
283
        # fmt: on
Sayak Paul's avatar
Sayak Paul committed
284
285

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


@slow
292
@require_big_accelerator
hlky's avatar
hlky committed
293
294
295
296
297
298
299
300
301
302
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()
303
        backend_empty_cache(torch_device)
hlky's avatar
hlky committed
304
305
306
307

    def tearDown(self):
        super().tearDown()
        gc.collect()
308
        backend_empty_cache(torch_device)
hlky's avatar
hlky committed
309
310
311
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

    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
358
        # fmt: off
hlky's avatar
hlky committed
359
        expected_slice = np.array(
Aryan's avatar
Aryan committed
360
            [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
361
362
            dtype=np.float32,
        )
Aryan's avatar
Aryan committed
363
        # fmt: on
hlky's avatar
hlky committed
364
365

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