test_pixart.py 12.9 KB
Newer Older
Sayak Paul's avatar
Sayak Paul committed
1
# coding=utf-8
2
# Copyright 2025 HuggingFace Inc.
Sayak Paul's avatar
Sayak Paul committed
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import gc
import tempfile
import unittest

import numpy as np
import torch
from transformers import AutoTokenizer, T5EncoderModel

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    PixArtAlphaPipeline,
28
    PixArtTransformer2DModel,
Sayak Paul's avatar
Sayak Paul committed
29
)
30
31

from ...testing_utils import (
32
    backend_empty_cache,
Dhruv Nair's avatar
Dhruv Nair committed
33
34
    enable_full_determinism,
    numpy_cosine_similarity_distance,
35
    require_torch_accelerator,
Dhruv Nair's avatar
Dhruv Nair committed
36
37
38
    slow,
    torch_device,
)
Sayak Paul's avatar
Sayak Paul committed
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, to_np


enable_full_determinism()


class PixArtAlphaPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
    pipeline_class = PixArtAlphaPipeline
    params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
    batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
    image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
    image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS

    required_optional_params = PipelineTesterMixin.required_optional_params
Aryan's avatar
Aryan committed
54
    test_layerwise_casting = True
Aryan's avatar
Aryan committed
55
    test_group_offloading = True
Sayak Paul's avatar
Sayak Paul committed
56
57
58

    def get_dummy_components(self):
        torch.manual_seed(0)
59
        transformer = PixArtTransformer2DModel(
Sayak Paul's avatar
Sayak Paul committed
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
            sample_size=8,
            num_layers=2,
            patch_size=2,
            attention_head_dim=8,
            num_attention_heads=3,
            caption_channels=32,
            in_channels=4,
            cross_attention_dim=24,
            out_channels=8,
            attention_bias=True,
            activation_fn="gelu-approximate",
            num_embeds_ada_norm=1000,
            norm_type="ada_norm_single",
            norm_elementwise_affine=False,
            norm_eps=1e-6,
        )
76
        torch.manual_seed(0)
Sayak Paul's avatar
Sayak Paul committed
77
        vae = AutoencoderKL()
78

Sayak Paul's avatar
Sayak Paul committed
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
        scheduler = DDIMScheduler()
        text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")

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

        components = {
            "transformer": transformer.eval(),
            "vae": vae.eval(),
            "scheduler": scheduler,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
        }
        return components

    def get_dummy_inputs(self, device, seed=0):
        if str(device).startswith("mps"):
            generator = torch.manual_seed(seed)
        else:
            generator = torch.Generator(device=device).manual_seed(seed)
        inputs = {
            "prompt": "A painting of a squirrel eating a burger",
            "generator": generator,
            "num_inference_steps": 2,
            "guidance_scale": 5.0,
103
104
            "use_resolution_binning": False,
            "output_type": "np",
Sayak Paul's avatar
Sayak Paul committed
105
106
107
        }
        return inputs

108
    @unittest.skip("Not supported.")
Sayak Paul's avatar
Sayak Paul committed
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
    def test_sequential_cpu_offload_forward_pass(self):
        # TODO(PVP, Sayak) need to fix later
        return

    def test_inference(self):
        device = "cpu"

        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe.to(device)
        pipe.set_progress_bar_config(disable=None)

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

        self.assertEqual(image.shape, (1, 8, 8, 3))
126
        expected_slice = np.array([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.483, 0.2583, 0.5331, 0.4852])
Sayak Paul's avatar
Sayak Paul committed
127
128
129
        max_diff = np.abs(image_slice.flatten() - expected_slice).max()
        self.assertLessEqual(max_diff, 1e-3)

130
131
132
133
134
135
136
137
138
139
140
141
    def test_inference_non_square_images(self):
        device = "cpu"

        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe.to(device)
        pipe.set_progress_bar_config(disable=None)

        inputs = self.get_dummy_inputs(device)
        image = pipe(**inputs, height=32, width=48).images
        image_slice = image[0, -3:, -3:, -1]
        self.assertEqual(image.shape, (1, 32, 48, 3))
142

143
        expected_slice = np.array([0.6493, 0.537, 0.4081, 0.4762, 0.3695, 0.4711, 0.3026, 0.5218, 0.5263])
144
145
146
        max_diff = np.abs(image_slice.flatten() - expected_slice).max()
        self.assertLessEqual(max_diff, 1e-3)

147
148
149
150
    @unittest.skip("Test is already covered through encode_prompt isolation.")
    def test_save_load_optional_components(self):
        pass

151
152
153
154
155
156
157
158
159
160
161
162
163
    def test_inference_with_embeddings_and_multiple_images(self):
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

        inputs = self.get_dummy_inputs(torch_device)

        prompt = inputs["prompt"]
        generator = inputs["generator"]
        num_inference_steps = inputs["num_inference_steps"]
        output_type = inputs["output_type"]

164
        prompt_embeds, prompt_attn_mask, negative_prompt_embeds, neg_prompt_attn_mask = pipe.encode_prompt(prompt)
165
166
167
168

        # inputs with prompt converted to embeddings
        inputs = {
            "prompt_embeds": prompt_embeds,
169
            "prompt_attention_mask": prompt_attn_mask,
170
171
            "negative_prompt": None,
            "negative_prompt_embeds": negative_prompt_embeds,
172
            "negative_prompt_attention_mask": neg_prompt_attn_mask,
173
174
175
176
            "generator": generator,
            "num_inference_steps": num_inference_steps,
            "output_type": output_type,
            "num_images_per_prompt": 2,
177
            "use_resolution_binning": False,
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
        }

        # set all optional components to None
        for optional_component in pipe._optional_components:
            setattr(pipe, optional_component, None)

        output = pipe(**inputs)[0]

        with tempfile.TemporaryDirectory() as tmpdir:
            pipe.save_pretrained(tmpdir)
            pipe_loaded = self.pipeline_class.from_pretrained(tmpdir)
            pipe_loaded.to(torch_device)
            pipe_loaded.set_progress_bar_config(disable=None)

        for optional_component in pipe._optional_components:
            self.assertTrue(
                getattr(pipe_loaded, optional_component) is None,
                f"`{optional_component}` did not stay set to None after loading.",
            )

        inputs = self.get_dummy_inputs(torch_device)

        generator = inputs["generator"]
        num_inference_steps = inputs["num_inference_steps"]
        output_type = inputs["output_type"]

        # inputs with prompt converted to embeddings
        inputs = {
            "prompt_embeds": prompt_embeds,
207
            "prompt_attention_mask": prompt_attn_mask,
208
209
            "negative_prompt": None,
            "negative_prompt_embeds": negative_prompt_embeds,
210
            "negative_prompt_attention_mask": neg_prompt_attn_mask,
211
212
213
214
            "generator": generator,
            "num_inference_steps": num_inference_steps,
            "output_type": output_type,
            "num_images_per_prompt": 2,
215
            "use_resolution_binning": False,
216
217
218
219
220
221
222
        }

        output_loaded = pipe_loaded(**inputs)[0]

        max_diff = np.abs(to_np(output) - to_np(output_loaded)).max()
        self.assertLess(max_diff, 1e-4)

223
224
225
226
227
228
229
230
231
232
233
234
235
236
    def test_inference_with_multiple_images_per_prompt(self):
        device = "cpu"

        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe.to(device)
        pipe.set_progress_bar_config(disable=None)

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

        self.assertEqual(image.shape, (2, 8, 8, 3))
237
        expected_slice = np.array([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.483, 0.2583, 0.5331, 0.4852])
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
        max_diff = np.abs(image_slice.flatten() - expected_slice).max()
        self.assertLessEqual(max_diff, 1e-3)

    def test_raises_warning_for_mask_feature(self):
        device = "cpu"

        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe.to(device)
        pipe.set_progress_bar_config(disable=None)

        inputs = self.get_dummy_inputs(device)
        inputs.update({"mask_feature": True})

        with self.assertWarns(FutureWarning) as warning_ctx:
            _ = pipe(**inputs).images

        assert "mask_feature" in str(warning_ctx.warning)

Sayak Paul's avatar
Sayak Paul committed
257
258
259
260
261
    def test_inference_batch_single_identical(self):
        self._test_inference_batch_single_identical(expected_max_diff=1e-3)


@slow
262
@require_torch_accelerator
Sayak Paul's avatar
Sayak Paul committed
263
class PixArtAlphaPipelineIntegrationTests(unittest.TestCase):
264
265
266
267
    ckpt_id_1024 = "PixArt-alpha/PixArt-XL-2-1024-MS"
    ckpt_id_512 = "PixArt-alpha/PixArt-XL-2-512x512"
    prompt = "A small cactus with a happy face in the Sahara desert."

268
269
270
    def setUp(self):
        super().setUp()
        gc.collect()
271
        backend_empty_cache(torch_device)
272

Sayak Paul's avatar
Sayak Paul committed
273
274
275
    def tearDown(self):
        super().tearDown()
        gc.collect()
276
        backend_empty_cache(torch_device)
Sayak Paul's avatar
Sayak Paul committed
277
278

    def test_pixart_1024(self):
Dhruv Nair's avatar
Dhruv Nair committed
279
        generator = torch.Generator("cpu").manual_seed(0)
Sayak Paul's avatar
Sayak Paul committed
280

281
        pipe = PixArtAlphaPipeline.from_pretrained(self.ckpt_id_1024, torch_dtype=torch.float16)
282
        pipe.enable_model_cpu_offload(device=torch_device)
283
        prompt = self.prompt
Sayak Paul's avatar
Sayak Paul committed
284

Dhruv Nair's avatar
Dhruv Nair committed
285
        image = pipe(prompt, generator=generator, num_inference_steps=2, output_type="np").images
Sayak Paul's avatar
Sayak Paul committed
286
287

        image_slice = image[0, -3:, -3:, -1]
Dhruv Nair's avatar
Dhruv Nair committed
288
        expected_slice = np.array([0.0742, 0.0835, 0.2114, 0.0295, 0.0784, 0.2361, 0.1738, 0.2251, 0.3589])
Sayak Paul's avatar
Sayak Paul committed
289

Dhruv Nair's avatar
Dhruv Nair committed
290
291
        max_diff = numpy_cosine_similarity_distance(image_slice.flatten(), expected_slice)
        self.assertLessEqual(max_diff, 1e-4)
Sayak Paul's avatar
Sayak Paul committed
292
293

    def test_pixart_512(self):
Dhruv Nair's avatar
Dhruv Nair committed
294
        generator = torch.Generator("cpu").manual_seed(0)
Sayak Paul's avatar
Sayak Paul committed
295

296
        pipe = PixArtAlphaPipeline.from_pretrained(self.ckpt_id_512, torch_dtype=torch.float16)
297
        pipe.enable_model_cpu_offload(device=torch_device)
Sayak Paul's avatar
Sayak Paul committed
298

299
        prompt = self.prompt
Sayak Paul's avatar
Sayak Paul committed
300

Dhruv Nair's avatar
Dhruv Nair committed
301
        image = pipe(prompt, generator=generator, num_inference_steps=2, output_type="np").images
Sayak Paul's avatar
Sayak Paul committed
302
303

        image_slice = image[0, -3:, -3:, -1]
Dhruv Nair's avatar
Dhruv Nair committed
304
        expected_slice = np.array([0.3477, 0.3882, 0.4541, 0.3413, 0.3821, 0.4463, 0.4001, 0.4409, 0.4958])
Sayak Paul's avatar
Sayak Paul committed
305

Dhruv Nair's avatar
Dhruv Nair committed
306
307
        max_diff = numpy_cosine_similarity_distance(image_slice.flatten(), expected_slice)
        self.assertLessEqual(max_diff, 1e-4)
308
309
310
311

    def test_pixart_1024_without_resolution_binning(self):
        generator = torch.manual_seed(0)

312
        pipe = PixArtAlphaPipeline.from_pretrained(self.ckpt_id_1024, torch_dtype=torch.float16)
313
        pipe.enable_model_cpu_offload(device=torch_device)
314

315
316
        prompt = self.prompt
        height, width = 1024, 768
Dhruv Nair's avatar
Dhruv Nair committed
317
        num_inference_steps = 2
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339

        image = pipe(
            prompt,
            height=height,
            width=width,
            generator=generator,
            num_inference_steps=num_inference_steps,
            output_type="np",
        ).images
        image_slice = image[0, -3:, -3:, -1]

        generator = torch.manual_seed(0)
        no_res_bin_image = pipe(
            prompt,
            height=height,
            width=width,
            generator=generator,
            num_inference_steps=num_inference_steps,
            output_type="np",
            use_resolution_binning=False,
        ).images
        no_res_bin_image_slice = no_res_bin_image[0, -3:, -3:, -1]
340

341
342
343
344
345
346
        assert not np.allclose(image_slice, no_res_bin_image_slice, atol=1e-4, rtol=1e-4)

    def test_pixart_512_without_resolution_binning(self):
        generator = torch.manual_seed(0)

        pipe = PixArtAlphaPipeline.from_pretrained(self.ckpt_id_512, torch_dtype=torch.float16)
347
        pipe.enable_model_cpu_offload(device=torch_device)
348
349
350

        prompt = self.prompt
        height, width = 512, 768
Dhruv Nair's avatar
Dhruv Nair committed
351
        num_inference_steps = 2
352
353
354
355
356
357
358
359
360

        image = pipe(
            prompt,
            height=height,
            width=width,
            generator=generator,
            num_inference_steps=num_inference_steps,
            output_type="np",
        ).images
361
362
363
364
        image_slice = image[0, -3:, -3:, -1]

        generator = torch.manual_seed(0)
        no_res_bin_image = pipe(
365
366
367
368
369
370
371
            prompt,
            height=height,
            width=width,
            generator=generator,
            num_inference_steps=num_inference_steps,
            output_type="np",
            use_resolution_binning=False,
372
373
374
375
        ).images
        no_res_bin_image_slice = no_res_bin_image[0, -3:, -3:, -1]

        assert not np.allclose(image_slice, no_res_bin_image_slice, atol=1e-4, rtol=1e-4)