diff --git a/src/diffusers/pipelines/ddim/pipeline_ddim.py b/src/diffusers/pipelines/ddim/pipeline_ddim.py index fdf6a00c72c2..19522c491063 100644 --- a/src/diffusers/pipelines/ddim/pipeline_ddim.py +++ b/src/diffusers/pipelines/ddim/pipeline_ddim.py @@ -90,7 +90,8 @@ def __call__( If `True` or `False`, see documentation for [`DDIMScheduler.step`]. If `None`, nothing is passed downstream to the scheduler (use `None` for schedulers which don't support this argument). output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generated image. Choose between `PIL.Image` or `np.array`. + The output format of the generated image. Choose between `"pil"` (`PIL.Image`), `"np"` (`np.array`) or + `"pt"` (`torch.Tensor`). return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. @@ -167,9 +168,13 @@ def __call__( xm.mark_step() image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).numpy() - if output_type == "pil": - image = self.numpy_to_pil(image) + if output_type != "pt": + image = image.cpu().permute(0, 2, 3, 1).numpy() + if output_type == "pil": + image = self.numpy_to_pil(image) + + # Offload all models + self.maybe_free_model_hooks() if not return_dict: return (image,) diff --git a/src/diffusers/pipelines/ddpm/pipeline_ddpm.py b/src/diffusers/pipelines/ddpm/pipeline_ddpm.py index c8c849aab0b2..8cde092684ad 100644 --- a/src/diffusers/pipelines/ddpm/pipeline_ddpm.py +++ b/src/diffusers/pipelines/ddpm/pipeline_ddpm.py @@ -73,7 +73,8 @@ def __call__( The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generated image. Choose between `PIL.Image` or `np.array`. + The output format of the generated image. Choose between `"pil"` (`PIL.Image`), `"np"` (`np.array`) or + `"pt"` (`torch.Tensor`). return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. @@ -108,12 +109,13 @@ def __call__( else: image_shape = (batch_size, self.unet.config.in_channels, *self.unet.config.sample_size) - if self.device.type == "mps": + device = self._execution_device + if device.type == "mps": # randn does not work reproducibly on mps image = randn_tensor(image_shape, generator=generator, dtype=self.unet.dtype) - image = image.to(self.device) + image = image.to(device) else: - image = randn_tensor(image_shape, generator=generator, device=self.device, dtype=self.unet.dtype) + image = randn_tensor(image_shape, generator=generator, device=device, dtype=self.unet.dtype) # set step values self.scheduler.set_timesteps(num_inference_steps) @@ -129,9 +131,13 @@ def __call__( xm.mark_step() image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).numpy() - if output_type == "pil": - image = self.numpy_to_pil(image) + if output_type != "pt": + image = image.cpu().permute(0, 2, 3, 1).numpy() + if output_type == "pil": + image = self.numpy_to_pil(image) + + # Offload all models + self.maybe_free_model_hooks() if not return_dict: return (image,) diff --git a/src/diffusers/pipelines/dit/pipeline_dit.py b/src/diffusers/pipelines/dit/pipeline_dit.py index dd73d364067e..c5e09280f37f 100644 --- a/src/diffusers/pipelines/dit/pipeline_dit.py +++ b/src/diffusers/pipelines/dit/pipeline_dit.py @@ -124,7 +124,8 @@ def __call__( The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generated image. Choose between `PIL.Image` or `np.array`. + The output format of the generated image. Choose between `"pil"` (`PIL.Image`), `"np"` (`np.array`) or + `"pt"` (`torch.Tensor`). return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`ImagePipelineOutput`] instead of a plain tuple. @@ -234,11 +235,12 @@ def __call__( samples = (samples / 2 + 0.5).clamp(0, 1) - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - samples = samples.cpu().permute(0, 2, 3, 1).float().numpy() + if output_type != "pt": + # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 + samples = samples.cpu().permute(0, 2, 3, 1).float().numpy() - if output_type == "pil": - samples = self.numpy_to_pil(samples) + if output_type == "pil": + samples = self.numpy_to_pil(samples) # Offload all models self.maybe_free_model_hooks() diff --git a/tests/pipelines/ddim/test_ddim.py b/tests/pipelines/ddim/test_ddim.py index 1b5237f7726b..a69ef4f6a05a 100644 --- a/tests/pipelines/ddim/test_ddim.py +++ b/tests/pipelines/ddim/test_ddim.py @@ -13,31 +13,33 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - import numpy as np import torch from diffusers import DDIMPipeline, DDIMScheduler, UNet2DModel -from ...testing_utils import enable_full_determinism, require_torch_accelerator, slow, torch_device +from ...testing_utils import ( + assert_tensors_close, + enable_full_determinism, + require_torch_accelerator, + slow, + torch_device, +) from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS -from ..test_pipelines_common import PipelineTesterMixin +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class DDIMPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class DDIMPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = DDIMPipeline - params = UNCONDITIONAL_IMAGE_GENERATION_PARAMS - required_optional_params = PipelineTesterMixin.required_optional_params - { - "num_images_per_prompt", - "latents", - "callback", - "callback_steps", - } - batch_params = UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS + required_input_params_in_call_signature = UNCONDITIONAL_IMAGE_GENERATION_PARAMS + batch_input_params = UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS + # DDIM is unconditional and samples its own noise: there is no prompt to repeat + # (`num_images_per_prompt`) and no user-suppliable `latents`. + optional_input_params = BasePipelineTesterConfig.optional_input_params - {"num_images_per_prompt", "latents"} + output_shape = (3, 8, 8) def get_dummy_components(self): torch.manual_seed(0) @@ -52,55 +54,53 @@ def get_dummy_components(self): up_block_types=("AttnUpBlock2D", "UpBlock2D"), ) scheduler = DDIMScheduler() - components = {"unet": unet, "scheduler": scheduler} - 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 = { + return {"unet": unet, "scheduler": scheduler} + + def get_dummy_inputs(self): + return { "batch_size": 1, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + "output_type": "pt", } - return inputs + +class TestDDIMPipeline(DDIMPipelineTesterConfig, PipelineTesterMixin): def test_inference(self): - device = "cpu" + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) + image = pipe(**self.get_dummy_inputs()).images + generated_image = image[0] + assert generated_image.shape == self.output_shape - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] + # fmt: off + expected_slice = torch.tensor([0.0, 9.979e-01, 0.0, 9.999e-01, 9.986e-01, 9.991e-01, 7.106e-04, 0.0, 0.0]) + # fmt: on - self.assertEqual(image.shape, (1, 8, 8, 3)) - expected_slice = np.array([0.0, 9.979e-01, 0.0, 9.999e-01, 9.986e-01, 9.991e-01, 7.106e-04, 0.0, 0.0]) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - self.assertLessEqual(max_diff, 1e-3) + # `"pt"` images are `(channels, height, width)`, so the trailing-channel corner slice of the old + # `"np"` layout is the last channel's bottom-right 3x3 block here. + generated_slice = generated_image[-1, -3:, -3:].flatten() + assert_tensors_close(generated_slice, expected_slice, atol=1e-3) def test_dict_tuple_outputs_equivalent(self): super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3) - def test_save_load_local(self): - super().test_save_load_local(expected_max_difference=3e-3) - - def test_save_load_optional_components(self): - super().test_save_load_optional_components(expected_max_difference=3e-3) + def test_save_load_local(self, tmp_path, base_pipe_output): + super().test_save_load_local(tmp_path, base_pipe_output, expected_max_difference=3e-3) def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=3e-3) +class TestDDIMPipelineMemory(DDIMPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the DDIM pipeline.""" + + @slow @require_torch_accelerator -class DDIMPipelineIntegrationTests(unittest.TestCase): +class TestDDIMPipelineIntegration: def test_inference_cifar10(self): model_id = "google/ddpm-cifar10-32" diff --git a/tests/pipelines/ddpm/test_ddpm.py b/tests/pipelines/ddpm/test_ddpm.py index 85676b56a1a3..885c04bde776 100644 --- a/tests/pipelines/ddpm/test_ddpm.py +++ b/tests/pipelines/ddpm/test_ddpm.py @@ -13,24 +13,37 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - import numpy as np import torch from diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel -from ...testing_utils import enable_full_determinism, require_torch_accelerator, slow, torch_device +from ...testing_utils import ( + assert_tensors_close, + enable_full_determinism, + require_torch_accelerator, + slow, + torch_device, +) +from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class DDPMPipelineFastTests(unittest.TestCase): - @property - def dummy_uncond_unet(self): +class DDPMPipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = DDPMPipeline + required_input_params_in_call_signature = UNCONDITIONAL_IMAGE_GENERATION_PARAMS + batch_input_params = UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS + # DDPM is unconditional and samples its own noise: there is no prompt to repeat + # (`num_images_per_prompt`) and no user-suppliable `latents`. + optional_input_params = BasePipelineTesterConfig.optional_input_params - {"num_images_per_prompt", "latents"} + output_shape = (3, 8, 8) + + def get_dummy_components(self): torch.manual_seed(0) - model = UNet2DModel( + unet = UNet2DModel( block_out_channels=(4, 8), layers_per_block=1, norm_num_groups=4, @@ -40,57 +53,61 @@ def dummy_uncond_unet(self): down_block_types=("DownBlock2D", "AttnDownBlock2D"), up_block_types=("AttnUpBlock2D", "UpBlock2D"), ) - return model - - def test_fast_inference(self): - device = "cpu" - unet = self.dummy_uncond_unet scheduler = DDPMScheduler() + return {"unet": unet, "scheduler": scheduler} - ddpm = DDPMPipeline(unet=unet, scheduler=scheduler) - ddpm.to(device) - ddpm.set_progress_bar_config(disable=None) + def get_dummy_inputs(self): + return { + "batch_size": 1, + "generator": self.get_generator(0), + "num_inference_steps": 2, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + "output_type": "pt", + } - generator = torch.Generator(device=device).manual_seed(0) - image = ddpm(generator=generator, num_inference_steps=2, output_type="np").images - generator = torch.Generator(device=device).manual_seed(0) - image_from_tuple = ddpm(generator=generator, num_inference_steps=2, output_type="np", return_dict=False)[0] +class TestDDPMPipeline(DDPMPipelineTesterConfig, PipelineTesterMixin): + def test_inference(self): + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - image_slice = image[0, -3:, -3:, -1] - image_from_tuple_slice = image_from_tuple[0, -3:, -3:, -1] + image = pipe(**self.get_dummy_inputs()).images + generated_image = image[0] + assert generated_image.shape == self.output_shape - assert image.shape == (1, 8, 8, 3) - expected_slice = np.array([0.0, 0.9996672, 0.00329116, 1.0, 0.9995991, 1.0, 0.0060907, 0.00115037, 0.0]) + # fmt: off + expected_slice = torch.tensor([0.0, 0.9996672, 0.00329116, 1.0, 0.9995991, 1.0, 0.0060907, 0.00115037, 0.0]) + # fmt: on - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 - assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2 + # `"pt"` images are `(channels, height, width)`, so the trailing-channel corner slice of the old + # `"np"` layout is the last channel's bottom-right 3x3 block here. + generated_slice = generated_image[-1, -3:, -3:].flatten() + assert_tensors_close(generated_slice, expected_slice, atol=1e-2) def test_inference_predict_sample(self): - unet = self.dummy_uncond_unet - scheduler = DDPMScheduler(prediction_type="sample") + # `prediction_type="sample"` makes the UNet output the denoised sample rather than the noise, so the + # scheduler consumes it differently and the pipeline must produce a different image than the default + # `epsilon` parameterization. + pipe = self.get_pipeline().to(torch_device) + output_epsilon = self.run_pipe(pipe) - ddpm = DDPMPipeline(unet=unet, scheduler=scheduler) - ddpm.to(torch_device) - ddpm.set_progress_bar_config(disable=None) + components = self.get_dummy_components() + components["scheduler"] = DDPMScheduler(prediction_type="sample") + pipe_sample = self.get_pipeline(**components).to(torch_device) + output_sample = self.run_pipe(pipe_sample) - generator = torch.manual_seed(0) - image = ddpm(generator=generator, num_inference_steps=2, output_type="np").images + assert output_sample.shape == output_epsilon.shape + assert not torch.isnan(output_sample).any() + assert not torch.allclose(output_sample, output_epsilon, atol=1e-3) - generator = torch.manual_seed(0) - image_eps = ddpm(generator=generator, num_inference_steps=2, output_type="np")[0] - - image_slice = image[0, -3:, -3:, -1] - image_eps_slice = image_eps[0, -3:, -3:, -1] - assert image.shape == (1, 8, 8, 3) - tolerance = 1e-2 if torch_device != "mps" else 3e-2 - assert np.abs(image_slice.flatten() - image_eps_slice.flatten()).max() < tolerance +class TestDDPMPipelineMemory(DDPMPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the DDPM pipeline.""" @slow @require_torch_accelerator -class DDPMPipelineIntegrationTests(unittest.TestCase): +class TestDDPMPipelineIntegration: def test_inference_cifar10(self): model_id = "google/ddpm-cifar10-32" diff --git a/tests/pipelines/diffusion_gemma/test_diffusion_gemma.py b/tests/pipelines/diffusion_gemma/test_diffusion_gemma.py index f1fbf5b633bb..b7ccb3e9d91d 100644 --- a/tests/pipelines/diffusion_gemma/test_diffusion_gemma.py +++ b/tests/pipelines/diffusion_gemma/test_diffusion_gemma.py @@ -1,17 +1,30 @@ -import unittest from types import SimpleNamespace +import numpy as np +import pytest import torch - -from diffusers import BlockRefinementScheduler, DiffusionGemmaPipeline, EntropyBoundScheduler +from PIL import Image + +from diffusers import ( + BlockRefinementScheduler, + DiffusionGemmaPipeline, + DiscreteDDIMScheduler, + EntropyBoundScheduler, +) from diffusers.utils.import_utils import is_peft_available -from diffusers.utils.testing_utils import require_peft_backend, require_peft_version_greater + +from ...testing_utils import require_peft_backend, require_peft_version_greater if is_peft_available(): from peft import LoraConfig +# `DiffusionGemmaPipeline` is a discrete *text* diffusion pipeline: it returns token sequences rather than images, +# so the image/video oriented `BasePipelineTesterConfig` + `PipelineTesterMixin` contract in `..testing_utils` +# does not apply here. These are plain pytest classes instead. + + # --- Lightweight stand-in for input-validation tests that never reach the model --- @@ -41,22 +54,22 @@ def _make_dummy_pipeline(processor=None, canvas_length: int = 8): return DiffusionGemmaPipeline(model=model, scheduler=BlockRefinementScheduler(), processor=processor) -class DiffusionGemmaPipelineInputTest(unittest.TestCase): +class TestDiffusionGemmaPipelineInput: """Input validation and prompt encoding, which short-circuit before the model is called.""" def test_no_inputs_raises(self): pipe = _make_dummy_pipeline() - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe(gen_length=8, num_inference_steps=2, output_type="seq") def test_output_type_invalid_raises(self): pipe = _make_dummy_pipeline() - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe(prompt="hi", gen_length=8, output_type="invalid") def test_prompt_and_messages_together_raises(self): pipe = _make_dummy_pipeline() - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe(prompt="hi", messages=[{"role": "user", "content": "hi"}], gen_length=8, output_type="seq") @@ -65,27 +78,28 @@ def test_prompt_and_messages_together_raises(self): _MODEL_ID = "trl-internal-testing/tiny-DiffusionGemmaForBlockDiffusion" -def _load_pipeline(test): +def _load_pipeline(): try: from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion except ImportError as e: - test.skipTest(f"transformers without DiffusionGemma: {e}") + pytest.skip(f"transformers without DiffusionGemma: {e}") try: model = DiffusionGemmaForBlockDiffusion.from_pretrained(_MODEL_ID, dtype=torch.float32).eval() processor = AutoProcessor.from_pretrained(_MODEL_ID) except Exception as e: # noqa: BLE001 - offline / hub errors should skip, not fail - test.skipTest(f"tiny DiffusionGemma checkpoint unavailable: {e}") + pytest.skip(f"tiny DiffusionGemma checkpoint unavailable: {e}") pipe = DiffusionGemmaPipeline(model=model, scheduler=BlockRefinementScheduler(), processor=processor) pipe.set_progress_bar_config(disable=True) return pipe, model.config.canvas_length -class DiffusionGemmaPipelineTest(unittest.TestCase): +class TestDiffusionGemmaPipeline: adaptive_stopping_vocab_size = 8 + prompt = "Name a color." - def setUp(self): - self.pipe, self.canvas_length = _load_pipeline(self) - self.prompt = "Name a color." + @pytest.fixture(autouse=True) + def pipeline(self): + self.pipe, self.canvas_length = _load_pipeline() def _run_adaptive_stopping(self, prompt): self.pipe.model.config.get_text_config(decoder=True).vocab_size = self.adaptive_stopping_vocab_size @@ -107,8 +121,8 @@ def test_generate(self): eos_early_stop=False, output_type="seq", ) - self.assertEqual(out.sequences.shape, (1, self.canvas_length * 2)) - self.assertIsNone(out.texts) + assert out.sequences.shape == (1, self.canvas_length * 2) + assert out.texts is None sequences, texts = self.pipe( prompt=self.prompt, @@ -119,8 +133,8 @@ def test_generate(self): output_type="text", return_dict=False, ) - self.assertEqual(sequences.shape, (1, self.canvas_length)) - self.assertEqual(len(texts), 1) + assert sequences.shape == (1, self.canvas_length) + assert len(texts) == 1 def test_adaptive_stopping_freezes_finished_rows(self): forward_calls = 0 @@ -143,9 +157,9 @@ def forward(decoder_input_ids, **kwargs): self.pipe.scheduler = BlockRefinementScheduler() output = self._run_adaptive_stopping(["Short prompt.", "A somewhat longer prompt for the second batch row."]) - self.assertEqual(forward_calls, 4) - self.assertTrue((output.sequences[0] == 1).all()) - self.assertTrue((output.sequences[1] == 5).all()) + assert forward_calls == 4 + assert (output.sequences[0] == 1).all() + assert (output.sequences[1] == 5).all() def test_adaptive_stopping_uses_scheduler_logits(self): forward_calls = 0 @@ -164,8 +178,8 @@ def forward(decoder_input_ids, **kwargs): self.pipe.scheduler = EntropyBoundScheduler(t_max=0.1, t_min=0.1) output = self._run_adaptive_stopping(self.prompt) - self.assertEqual(forward_calls, 2) - self.assertTrue((output.sequences == 0).all()) + assert forward_calls == 2 + assert (output.sequences == 0).all() def test_callback_receives_advertised_keys(self): observed: list[str] = [] @@ -185,12 +199,9 @@ def callback(pipe, step, timestep, callback_kwargs): callback_on_step_end=callback, callback_on_step_end_tensor_inputs=keys, ) - self.assertEqual(set(observed), set(keys)) + assert set(observed) == set(keys) def test_generate_with_image(self): - import numpy as np - from PIL import Image - image = Image.fromarray((np.random.rand(64, 64, 3) * 255).astype("uint8")) out = self.pipe( prompt="What?", @@ -201,11 +212,9 @@ def test_generate_with_image(self): eos_early_stop=False, output_type="seq", ) - self.assertEqual(out.sequences.shape, (1, self.canvas_length)) + assert out.sequences.shape == (1, self.canvas_length) def test_schedulers_are_interchangeable(self): - from diffusers import DiscreteDDIMScheduler, EntropyBoundScheduler - for scheduler in (DiscreteDDIMScheduler(), EntropyBoundScheduler(entropy_bound=0.1)): self.pipe.scheduler = scheduler out = self.pipe( @@ -216,11 +225,9 @@ def test_schedulers_are_interchangeable(self): eos_early_stop=False, output_type="seq", ) - self.assertEqual(out.sequences.shape, (1, self.canvas_length)) + assert out.sequences.shape == (1, self.canvas_length) def test_predictor_corrector_sampling(self): - from diffusers import DiscreteDDIMScheduler - self.pipe.scheduler = DiscreteDDIMScheduler(corrector_steps=2, corrector_k=2) out = self.pipe( prompt=self.prompt, @@ -230,7 +237,7 @@ def test_predictor_corrector_sampling(self): eos_early_stop=False, output_type="seq", ) - self.assertEqual(out.sequences.shape, (1, self.canvas_length)) + assert out.sequences.shape == (1, self.canvas_length) @require_peft_backend @require_peft_version_greater("0.18.9") @@ -242,7 +249,7 @@ def test_peft_adapter_api(self): adapter_name="test", ) self.pipe.model.set_adapter("test") - self.assertIn("test", self.pipe.model.active_adapters()) + assert "test" in self.pipe.model.active_adapters() out = self.pipe( prompt=self.prompt, @@ -252,7 +259,7 @@ def test_peft_adapter_api(self): eos_early_stop=False, output_type="seq", ) - self.assertEqual(out.sequences.shape, (1, self.canvas_length)) + assert out.sequences.shape == (1, self.canvas_length) self.pipe.model.disable_adapters() self.pipe.model.enable_adapters() @@ -274,8 +281,4 @@ def test_static_cache_matches_dynamic(self): generator=torch.Generator().manual_seed(0), cache_implementation="static", **kwargs ).sequences ndiff = (dynamic != static).sum().item() - self.assertEqual(ndiff, 0, f"static/dynamic agree on only ndiff={ndiff}/{dynamic.numel()} tokens") - - -if __name__ == "__main__": - unittest.main() + assert ndiff == 0, f"static/dynamic agree on only ndiff={ndiff}/{dynamic.numel()} tokens" diff --git a/tests/pipelines/dit/test_dit.py b/tests/pipelines/dit/test_dit.py index 406f3003c96a..11c2ac3488ed 100644 --- a/tests/pipelines/dit/test_dit.py +++ b/tests/pipelines/dit/test_dit.py @@ -14,15 +14,15 @@ # limitations under the License. import gc -import unittest import numpy as np +import pytest import torch from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, DiTTransformer2DModel, DPMSolverMultistepScheduler -from diffusers.utils import is_xformers_available from ...testing_utils import ( + assert_tensors_close, backend_empty_cache, enable_full_determinism, load_numpy, @@ -35,22 +35,20 @@ CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS, CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS, ) -from ..test_pipelines_common import PipelineTesterMixin +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class DiTPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class DiTPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = DiTPipeline - params = CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS - required_optional_params = PipelineTesterMixin.required_optional_params - { - "latents", - "num_images_per_prompt", - "callback", - "callback_steps", - } - batch_params = CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS + required_input_params_in_call_signature = CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS + batch_input_params = CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS + # DiT is class-conditioned and samples its own noise: there is no prompt to repeat + # (`num_images_per_prompt`) and no user-suppliable `latents`. + optional_input_params = BasePipelineTesterConfig.optional_input_params - {"num_images_per_prompt", "latents"} + output_shape = (3, 16, 16) def get_dummy_components(self): torch.manual_seed(0) @@ -70,60 +68,52 @@ def get_dummy_components(self): ) vae = AutoencoderKL() scheduler = DDIMScheduler() - components = {"transformer": transformer.eval(), "vae": vae.eval(), "scheduler": scheduler} - 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 = { + return {"transformer": transformer, "vae": vae, "scheduler": scheduler} + + def get_dummy_inputs(self): + return { "class_labels": [1], - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + "output_type": "pt", } - return inputs + +class TestDiTPipeline(DiTPipelineTesterConfig, PipelineTesterMixin): def test_inference(self): - device = "cpu" + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) + image = pipe(**self.get_dummy_inputs()).images + generated_image = image[0] + assert generated_image.shape == self.output_shape - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] + # fmt: off + expected_slice = torch.tensor([0.2946, 0.6601, 0.4329, 0.3296, 0.4144, 0.5319, 0.7273, 0.5013, 0.4457]) + # fmt: on - self.assertEqual(image.shape, (1, 16, 16, 3)) - expected_slice = np.array([0.2946, 0.6601, 0.4329, 0.3296, 0.4144, 0.5319, 0.7273, 0.5013, 0.4457]) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - self.assertLessEqual(max_diff, 1e-3) + # `"pt"` images are `(channels, height, width)`, so the trailing-channel corner slice of the old + # `"np"` layout is the last channel's bottom-right 3x3 block here. + generated_slice = generated_image[-1, -3:, -3:].flatten() + assert_tensors_close(generated_slice, expected_slice, atol=1e-3) def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=1e-3) + super().test_inference_batch_single_identical(expected_max_diff=1e-3) - @unittest.skipIf( - torch_device != "cuda" or not is_xformers_available(), - reason="XFormers attention is only available with CUDA and `xformers` installed", - ) - def test_xformers_attention_forwardGenerator_pass(self): - self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3) + +class TestDiTPipelineMemory(DiTPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the DiT pipeline.""" @nightly @require_torch_accelerator -class DiTPipelineIntegrationTests(unittest.TestCase): - def setUp(self): - super().setUp() +class TestDiTPipelineIntegration: + @pytest.fixture(autouse=True) + def cleanup(self): gc.collect() backend_empty_cache(torch_device) - - def tearDown(self): - super().tearDown() + yield gc.collect() backend_empty_cache(torch_device) diff --git a/tests/pipelines/dreamlite/test_pipeline_dreamlite.py b/tests/pipelines/dreamlite/test_pipeline_dreamlite.py index cec804727d00..51a564f7caed 100644 --- a/tests/pipelines/dreamlite/test_pipeline_dreamlite.py +++ b/tests/pipelines/dreamlite/test_pipeline_dreamlite.py @@ -20,7 +20,7 @@ tiny config (mirroring the NucleusMoE-Image fast tests), and load the matching processor / tokenizer from the public ``hf-internal-testing`` mirror, so that the standard ``PipelineTesterMixin`` save/load and dtype/device tests work -out of the box. +out of the box. The shared component set lives in ``testing_utils.py``. For end-to-end verification against the original repo, see the ``parity_run_*.py`` scripts shipped with the integration. @@ -28,86 +28,30 @@ import gc import os -import unittest import numpy as np +import pytest import torch from PIL import Image -from transformers import AutoTokenizer, Qwen3VLConfig, Qwen3VLForConditionalGeneration, Qwen3VLProcessor -from diffusers import ( - AutoencoderTiny, - DreamLitePipeline, - DreamLiteUNetModel, - FlowMatchEulerDiscreteScheduler, -) -from diffusers.utils.testing_utils import ( +from diffusers import DreamLitePipeline, DreamLiteUNetModel + +from ...testing_utils import ( enable_full_determinism, nightly, require_torch_gpu, torch_device, ) - -from ..test_pipelines_common import ( - PipelineTesterMixin, - to_np, -) +from ..testing_utils import MemoryTesterMixin, PipelineTesterMixin +from .testing_utils import CROSS_ATTN_DIM, DreamLiteBaseTesterConfig enable_full_determinism() -# Match the tiny text encoder hidden size below; the UNet's cross-attention -# dimension must match what ``encode_prompt`` returns. -_CROSS_ATTN_DIM = 16 - - -def _build_tiny_text_encoder() -> Qwen3VLForConditionalGeneration: - """Build a tiny but functional Qwen3-VL model for the fast test fixture. - - Mirrors the recipe used by ``tests/pipelines/nucleusmoe_image``: small text - + vision configs that still go through the real Qwen3-VL forward path, so - DreamLite's ``encode_prompt`` (chat template + tokenizer + multimodal - processor) is exercised for real. - """ - config = Qwen3VLConfig( - text_config={ - "hidden_size": _CROSS_ATTN_DIM, - "intermediate_size": _CROSS_ATTN_DIM, - "num_hidden_layers": 2, - "num_attention_heads": 2, - "num_key_value_heads": 2, - "rope_scaling": { - "mrope_section": [1, 1, 2], - "rope_type": "default", - "type": "default", - }, - "rope_theta": 1000000.0, - "vocab_size": 151936, - "head_dim": 8, - }, - vision_config={ - "depth": 2, - "hidden_size": _CROSS_ATTN_DIM, - "intermediate_size": _CROSS_ATTN_DIM, - "num_heads": 2, - "out_channels": _CROSS_ATTN_DIM, - # ``out_hidden_size`` is the dim that vision tokens are projected to before - # being merged into the text stream; it must match ``text_config.hidden_size``. - "out_hidden_size": _CROSS_ATTN_DIM, - # Match the cached ``hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration`` - # image processor (``patch_size=14``); otherwise the pixel_values - # produced by the processor cannot be reshaped to the model's - # vision patch embed. - "patch_size": 14, - }, - ) - return Qwen3VLForConditionalGeneration(config).eval() - - -class DreamLitePipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class DreamLitePipelineTesterConfig(DreamLiteBaseTesterConfig): pipeline_class = DreamLitePipeline - params = frozenset( + required_input_params_in_call_signature = frozenset( [ "prompt", "height", @@ -117,20 +61,61 @@ class DreamLitePipelineFastTests(PipelineTesterMixin, unittest.TestCase): "num_inference_steps", ] ) - batch_params = frozenset(["prompt", "negative_prompt"]) - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "output_type", - "return_dict", - ] + batch_input_params = frozenset(["prompt", "negative_prompt"]) + + def get_dummy_inputs(self): + return { + "prompt": "a small dog", + "negative_prompt": "", + "generator": self.get_generator(0), + "num_inference_steps": 2, + "guidance_scale": 3.5, + "height": 64, + "width": 64, + "max_sequence_length": 16, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + "output_type": "pt", + } + + def get_dummy_i2i_inputs(self, seed=0): + inputs = self.get_dummy_inputs() + # 64x64 RGB image -- will be processed by VaeImageProcessor. + inputs["image"] = Image.fromarray((np.random.RandomState(seed).rand(64, 64, 3) * 255).astype(np.uint8)) + inputs["image_guidance_scale"] = 1.5 + return inputs + + +class TestDreamLitePipeline(DreamLitePipelineTesterConfig, PipelineTesterMixin): + # ---- skips for mixin tests that genuinely don't apply ---------------- + # The remaining skips reflect intrinsic design choices of the DreamLite pipeline: + # * ``encode_prompt`` returns a ``(prompt_embeds, prompt_embeds_mask)`` + # tuple, while the mixin's ``test_encode_prompt_works_in_isolation`` + # assumes a single tensor return value; + # * the pipeline forces ``batch_size = 1`` internally, so the mixin's + # batch sweep cannot apply. + @pytest.mark.skip( + "DreamLite intentionally limits ``batch_size`` to 1 (CFG memory blow-up); " + "only ``num_images_per_prompt > 1`` is supported." + ) + def test_num_images_per_prompt(self): + pass + + @pytest.mark.skip( + "DreamLite encode_prompt returns (embeds, mask) tuple, not a single tensor; " + "the mixin's test_encode_prompt_works_in_isolation assumes single tensor return." ) - test_xformers_attention = False - test_attention_slicing = False - test_layerwise_casting = False - test_group_offloading = False + def test_encode_prompt_works_in_isolation(self): + pass + @pytest.mark.skip("DreamLite forces batch_size=1 internally.") + def test_inference_batch_consistent(self): + pass + + @pytest.mark.skip("DreamLite forces batch_size=1 internally.") + def test_inference_batch_single_identical(self): + pass + + # ---- actual tests ------------------------------------------------------ def test_legacy_block_type_aliases(self): unet = DreamLiteUNetModel( sample_size=8, @@ -148,30 +133,24 @@ def test_legacy_block_type_aliases(self): "UpBlock2D", ), block_out_channels=(16, 32, 64), - cross_attention_dim=_CROSS_ATTN_DIM, + cross_attention_dim=CROSS_ATTN_DIM, attention_head_dim=8, layers_per_block=1, norm_num_groups=8, transformer_layers_per_block=1, ) - self.assertEqual( - [block.__class__.__name__ for block in unet.down_blocks], - [ - "DreamLiteCrossAttnNoSelfAttnDownBlock2D", - "DreamLiteCrossAttnNoSelfAttnDownBlock2D", - "DreamLiteCrossAttnDownBlock2D", - ], - ) - self.assertEqual(unet.mid_block.__class__.__name__, "DreamLiteUNetMidBlock2DCrossAttn") - self.assertEqual( - [block.__class__.__name__ for block in unet.up_blocks], - [ - "DreamLiteCrossAttnUpBlock2D", - "DreamLiteCrossAttnNoSelfAttnUpBlock2D", - "DreamLiteUpBlock2D", - ], - ) + assert [block.__class__.__name__ for block in unet.down_blocks] == [ + "DreamLiteCrossAttnNoSelfAttnDownBlock2D", + "DreamLiteCrossAttnNoSelfAttnDownBlock2D", + "DreamLiteCrossAttnDownBlock2D", + ] + assert unet.mid_block.__class__.__name__ == "DreamLiteUNetMidBlock2DCrossAttn" + assert [block.__class__.__name__ for block in unet.up_blocks] == [ + "DreamLiteCrossAttnUpBlock2D", + "DreamLiteCrossAttnNoSelfAttnUpBlock2D", + "DreamLiteUpBlock2D", + ] unet_with_non_v1_up_alias = DreamLiteUNetModel( sample_size=8, @@ -189,154 +168,37 @@ def test_legacy_block_type_aliases(self): "UpBlock2D", ), block_out_channels=(16, 32, 64), - cross_attention_dim=_CROSS_ATTN_DIM, - attention_head_dim=8, - layers_per_block=1, - norm_num_groups=8, - transformer_layers_per_block=1, - ) - self.assertEqual( - [block.__class__.__name__ for block in unet_with_non_v1_up_alias.up_blocks], - [ - "DreamLiteCrossAttnUpBlock2D", - "DreamLiteCrossAttnNoSelfAttnUpBlock2D", - "DreamLiteUpBlock2D", - ], - ) - - def get_dummy_components(self): - torch.manual_seed(0) - unet = DreamLiteUNetModel( - sample_size=8, - in_channels=4, - out_channels=4, - down_block_types=( - "DreamLiteCrossAttnNoSelfAttnDownBlock2D", - "DreamLiteCrossAttnDownBlock2D", - ), - up_block_types=("DreamLiteCrossAttnUpBlock2D", "DreamLiteUpBlock2D"), - block_out_channels=(32, 64), - cross_attention_dim=_CROSS_ATTN_DIM, + cross_attention_dim=CROSS_ATTN_DIM, attention_head_dim=8, layers_per_block=1, norm_num_groups=8, transformer_layers_per_block=1, ) + assert [block.__class__.__name__ for block in unet_with_non_v1_up_alias.up_blocks] == [ + "DreamLiteCrossAttnUpBlock2D", + "DreamLiteCrossAttnNoSelfAttnUpBlock2D", + "DreamLiteUpBlock2D", + ] - torch.manual_seed(0) - vae = AutoencoderTiny( - in_channels=3, - out_channels=3, - encoder_block_out_channels=(32, 32), - decoder_block_out_channels=(32, 32), - num_encoder_blocks=(1, 1), - num_decoder_blocks=(1, 1), - latent_channels=4, - ) - - scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000) - - torch.manual_seed(0) - text_encoder = _build_tiny_text_encoder() - tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") - processor = Qwen3VLProcessor.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") - - return { - "text_encoder": text_encoder, - "tokenizer": tokenizer, - "processor": processor, - "vae": vae, - "unet": unet, - "scheduler": scheduler, - } - - 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) - return { - "prompt": "a small dog", - "negative_prompt": "", - "generator": generator, - "num_inference_steps": 2, - "guidance_scale": 3.5, - "height": 64, - "width": 64, - "max_sequence_length": 16, - "output_type": "np", - } - - def get_dummy_i2i_inputs(self, device, seed=0): - inputs = self.get_dummy_inputs(device, seed) - # 64x64 RGB image -- will be processed by VaeImageProcessor. - inputs["image"] = Image.fromarray((np.random.RandomState(seed).rand(64, 64, 3) * 255).astype(np.uint8)) - inputs["image_guidance_scale"] = 1.5 - return inputs - - # ---- skips for mixin tests that genuinely don't apply ---------------- - # The remaining skips reflect intrinsic design choices of the DreamLite pipeline: - # * ``encode_prompt`` returns a ``(prompt_embeds, prompt_embeds_mask)`` - # tuple, while the mixin's ``test_encode_prompt_works_in_isolation`` - # assumes a single tensor return value; - # * the pipeline forces ``batch_size = 1`` internally, so the mixin's - # batch sweep cannot apply. - @unittest.skip( - "DreamLite intentionally limits ``batch_size`` to 1 (CFG memory blow-up); " - "only ``num_images_per_prompt > 1`` is supported." - ) - def test_num_images_per_prompt(self): - pass - - @unittest.skip( - "DreamLite encode_prompt returns (embeds, mask) tuple, not a single tensor; " - "the mixin's test_encode_prompt_works_in_isolation assumes single tensor return." - ) - def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict=None, atol=1e-4, rtol=1e-4): - pass - - @unittest.skip("DreamLite forces batch_size=1 internally.") - def test_inference_batch_consistent(self): - pass - - @unittest.skip("DreamLite forces batch_size=1 internally.") - def test_inference_batch_single_identical(self): - pass - - # ---- actual tests ------------------------------------------------------ def test_dreamlite_t2i_default_case(self): - device = torch_device - components = self.get_dummy_components() - pipe = self.pipeline_class(**components).to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(device) - out = pipe(**inputs).images - out_np = to_np(out) + out = pipe(**self.get_dummy_inputs()).images - # shape: (B=1, H, W, C=3) - self.assertEqual(out_np.shape, (1, 64, 64, 3)) - self.assertFalse(np.isnan(out_np).any()) + assert out.shape == (1, *self.output_shape) + assert not torch.isnan(out).any() def test_dreamlite_i2i_default_case(self): - device = torch_device - components = self.get_dummy_components() - pipe = self.pipeline_class(**components).to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_i2i_inputs(device) - out = pipe(**inputs).images - out_np = to_np(out) + out = pipe(**self.get_dummy_i2i_inputs()).images - self.assertEqual(out_np.shape, (1, 64, 64, 3)) - self.assertFalse(np.isnan(out_np).any()) + assert out.shape == (1, *self.output_shape) + assert not torch.isnan(out).any() def test_dreamlite_cfg_branch_count(self): """In edit mode the pipeline must run a 3-way CFG concat (uncond/img/text).""" - device = torch_device - components = self.get_dummy_components() - pipe = self.pipeline_class(**components).to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline().to(torch_device) original_forward = pipe.unet.forward seen_batches = [] @@ -347,16 +209,20 @@ def spy_forward(*args, **kwargs): return original_forward(*args, **kwargs) pipe.unet.forward = spy_forward - inputs = self.get_dummy_i2i_inputs(device) + inputs = self.get_dummy_i2i_inputs() inputs["num_inference_steps"] = 1 pipe(**inputs) - self.assertTrue(all(b == 3 for b in seen_batches), f"expected all 3-way, got {seen_batches}") + assert all(b == 3 for b in seen_batches), f"expected all 3-way, got {seen_batches}" + + +class TestDreamLitePipelineMemory(DreamLitePipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the DreamLite pipeline.""" @nightly @require_torch_gpu -class DreamLitePipelineSlowTests(unittest.TestCase): +class TestDreamLitePipelineIntegration: """End-to-end test against the real DreamLite-base checkpoint on the Hub. By default this loads ``carlofkl/DreamLite-base`` (``diffusers`` branch) @@ -367,13 +233,11 @@ class DreamLitePipelineSlowTests(unittest.TestCase): repo_id = "carlofkl/DreamLite-base" revision = "diffusers" - def setUp(self): - super().setUp() + @pytest.fixture(autouse=True) + def cleanup(self): gc.collect() torch.cuda.empty_cache() - - def tearDown(self): - super().tearDown() + yield gc.collect() torch.cuda.empty_cache() @@ -397,8 +261,8 @@ def test_dreamlite_t2i_real_checkpoint(self): output_type="np", ).images - self.assertEqual(out.shape, (1, 1024, 1024, 3)) - self.assertFalse(np.isnan(out).any()) + assert out.shape == (1, 1024, 1024, 3) + assert not np.isnan(out).any() def test_dreamlite_i2i_real_checkpoint(self): pipe = DreamLitePipeline.from_pretrained(**self._from_pretrained_kwargs(), torch_dtype=torch.bfloat16).to( @@ -418,9 +282,5 @@ def test_dreamlite_i2i_real_checkpoint(self): output_type="np", ).images - self.assertEqual(out.shape, (1, 1024, 1024, 3)) - self.assertFalse(np.isnan(out).any()) - - -if __name__ == "__main__": - unittest.main() + assert out.shape == (1, 1024, 1024, 3) + assert not np.isnan(out).any() diff --git a/tests/pipelines/dreamlite/test_pipeline_dreamlite_mobile.py b/tests/pipelines/dreamlite/test_pipeline_dreamlite_mobile.py index 325bdc592177..268f2ca6ede8 100644 --- a/tests/pipelines/dreamlite/test_pipeline_dreamlite_mobile.py +++ b/tests/pipelines/dreamlite/test_pipeline_dreamlite_mobile.py @@ -16,92 +16,36 @@ The mobile pipeline is a distilled, no-CFG sibling of ``DreamLitePipeline``. It runs a single UNet forward per step (no 3-way concat) and ignores ``guidance_scale`` / ``image_guidance_scale``. Test layout mirrors -``test_pipeline_dreamlite.py``; see that file for the rationale behind the -tiny Qwen3-VL test fixture. +``test_pipeline_dreamlite.py``; the shared tiny Qwen3-VL test fixture lives in +``testing_utils.py``. """ import gc import os -import unittest import numpy as np +import pytest import torch from PIL import Image -from transformers import AutoTokenizer, Qwen3VLConfig, Qwen3VLForConditionalGeneration, Qwen3VLProcessor -from diffusers import ( - AutoencoderTiny, - DreamLiteMobilePipeline, - DreamLiteUNetModel, - FlowMatchEulerDiscreteScheduler, -) -from diffusers.utils.testing_utils import ( +from diffusers import DreamLiteMobilePipeline + +from ...testing_utils import ( enable_full_determinism, nightly, require_torch_gpu, torch_device, ) - -from ..test_pipelines_common import ( - PipelineTesterMixin, - to_np, -) +from ..testing_utils import MemoryTesterMixin, PipelineTesterMixin +from .testing_utils import DreamLiteBaseTesterConfig enable_full_determinism() -# Match the tiny text encoder hidden size below; the UNet's cross-attention -# dimension must match what ``encode_prompt`` returns. -_CROSS_ATTN_DIM = 16 - - -def _build_tiny_text_encoder() -> Qwen3VLForConditionalGeneration: - """Build a tiny but functional Qwen3-VL model for the fast test fixture. - - Mirrors the recipe used by ``tests/pipelines/nucleusmoe_image``: small text - + vision configs that still go through the real Qwen3-VL forward path, so - DreamLite's ``encode_prompt`` (chat template + tokenizer + multimodal - processor) is exercised for real. - """ - config = Qwen3VLConfig( - text_config={ - "hidden_size": _CROSS_ATTN_DIM, - "intermediate_size": _CROSS_ATTN_DIM, - "num_hidden_layers": 2, - "num_attention_heads": 2, - "num_key_value_heads": 2, - "rope_scaling": { - "mrope_section": [1, 1, 2], - "rope_type": "default", - "type": "default", - }, - "rope_theta": 1000000.0, - "vocab_size": 151936, - "head_dim": 8, - }, - vision_config={ - "depth": 2, - "hidden_size": _CROSS_ATTN_DIM, - "intermediate_size": _CROSS_ATTN_DIM, - "num_heads": 2, - "out_channels": _CROSS_ATTN_DIM, - # ``out_hidden_size`` is the dim that vision tokens are projected to before - # being merged into the text stream; it must match ``text_config.hidden_size``. - "out_hidden_size": _CROSS_ATTN_DIM, - # Match the cached ``hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration`` - # image processor (``patch_size=14``); otherwise the pixel_values - # produced by the processor cannot be reshaped to the model's - # vision patch embed. - "patch_size": 14, - }, - ) - return Qwen3VLForConditionalGeneration(config).eval() - - -class DreamLiteMobilePipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class DreamLiteMobilePipelineTesterConfig(DreamLiteBaseTesterConfig): pipeline_class = DreamLiteMobilePipeline - params = frozenset( + required_input_params_in_call_signature = frozenset( [ "prompt", "height", @@ -109,144 +53,72 @@ class DreamLiteMobilePipelineFastTests(PipelineTesterMixin, unittest.TestCase): "num_inference_steps", ] ) - batch_params = frozenset(["prompt"]) - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "output_type", - "return_dict", - ] - ) - test_xformers_attention = False - test_attention_slicing = False - test_layerwise_casting = False - test_group_offloading = False - - def get_dummy_components(self): - torch.manual_seed(0) - unet = DreamLiteUNetModel( - sample_size=8, - in_channels=4, - out_channels=4, - down_block_types=( - "DreamLiteCrossAttnNoSelfAttnDownBlock2D", - "DreamLiteCrossAttnDownBlock2D", - ), - up_block_types=("DreamLiteCrossAttnUpBlock2D", "DreamLiteUpBlock2D"), - block_out_channels=(32, 64), - cross_attention_dim=_CROSS_ATTN_DIM, - attention_head_dim=8, - layers_per_block=1, - norm_num_groups=8, - transformer_layers_per_block=1, - ) - - torch.manual_seed(0) - vae = AutoencoderTiny( - in_channels=3, - out_channels=3, - encoder_block_out_channels=(32, 32), - decoder_block_out_channels=(32, 32), - num_encoder_blocks=(1, 1), - num_decoder_blocks=(1, 1), - latent_channels=4, - ) - - scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000) - - torch.manual_seed(0) - text_encoder = _build_tiny_text_encoder() - tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") - processor = Qwen3VLProcessor.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") - - return { - "text_encoder": text_encoder, - "tokenizer": tokenizer, - "processor": processor, - "vae": vae, - "unet": unet, - "scheduler": scheduler, - } + batch_input_params = frozenset(["prompt"]) - 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) + def get_dummy_inputs(self): return { "prompt": "a small dog", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "height": 64, "width": 64, "max_sequence_length": 16, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + "output_type": "pt", } - def get_dummy_i2i_inputs(self, device, seed=0): - inputs = self.get_dummy_inputs(device, seed) + def get_dummy_i2i_inputs(self, seed=0): + inputs = self.get_dummy_inputs() inputs["image"] = Image.fromarray((np.random.RandomState(seed).rand(64, 64, 3) * 255).astype(np.uint8)) return inputs + +class TestDreamLiteMobilePipeline(DreamLiteMobilePipelineTesterConfig, PipelineTesterMixin): # ---- skips for mixin tests that genuinely don't apply ---------------- # The remaining skips are intrinsic to the mobile pipeline's design: # * ``encode_prompt`` returns ``(prompt_embeds, prompt_embeds_mask)``; # * the pipeline forces ``batch_size = 1`` internally. - @unittest.skip( + @pytest.mark.skip( "DreamLiteMobile encode_prompt returns (embeds, mask) tuple, not a single tensor; " "the mixin's test_encode_prompt_works_in_isolation assumes single tensor return." ) - def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict=None, atol=1e-4, rtol=1e-4): + def test_encode_prompt_works_in_isolation(self): pass - @unittest.skip( + @pytest.mark.skip( "DreamLiteMobile intentionally limits ``batch_size`` to 1; only ``num_images_per_prompt > 1`` is supported." ) def test_num_images_per_prompt(self): pass - @unittest.skip("DreamLiteMobile forces batch_size=1 internally.") + @pytest.mark.skip("DreamLiteMobile forces batch_size=1 internally.") def test_inference_batch_consistent(self): pass - @unittest.skip("DreamLiteMobile forces batch_size=1 internally.") + @pytest.mark.skip("DreamLiteMobile forces batch_size=1 internally.") def test_inference_batch_single_identical(self): pass # ---- actual tests ------------------------------------------------------ def test_mobile_t2i_default_case(self): - device = torch_device - components = self.get_dummy_components() - pipe = self.pipeline_class(**components).to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(device) - out = pipe(**inputs).images - out_np = to_np(out) + out = pipe(**self.get_dummy_inputs()).images - self.assertEqual(out_np.shape, (1, 64, 64, 3)) - self.assertFalse(np.isnan(out_np).any()) + assert out.shape == (1, *self.output_shape) + assert not torch.isnan(out).any() def test_mobile_i2i_default_case(self): - device = torch_device - components = self.get_dummy_components() - pipe = self.pipeline_class(**components).to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_i2i_inputs(device) - out = pipe(**inputs).images - out_np = to_np(out) + out = pipe(**self.get_dummy_i2i_inputs()).images - self.assertEqual(out_np.shape, (1, 64, 64, 3)) - self.assertFalse(np.isnan(out_np).any()) + assert out.shape == (1, *self.output_shape) + assert not torch.isnan(out).any() def test_mobile_single_forward_per_step(self): """Mobile pipeline must run exactly ONE UNet forward per step (no CFG concat).""" - device = torch_device - components = self.get_dummy_components() - pipe = self.pipeline_class(**components).to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline().to(torch_device) original_forward = pipe.unet.forward seen_batches = [] @@ -257,30 +129,32 @@ def spy_forward(*args, **kwargs): return original_forward(*args, **kwargs) pipe.unet.forward = spy_forward - inputs = self.get_dummy_i2i_inputs(device) + inputs = self.get_dummy_i2i_inputs() inputs["num_inference_steps"] = 2 pipe(**inputs) - self.assertTrue(all(b == 1 for b in seen_batches), f"expected all 1-way, got {seen_batches}") - self.assertEqual(len(seen_batches), 2, "expected exactly 2 unet calls (1 per step)") + assert all(b == 1 for b in seen_batches), f"expected all 1-way, got {seen_batches}" + assert len(seen_batches) == 2, "expected exactly 2 unet calls (1 per step)" def test_mobile_guidance_scale_ignored(self): """Passing guidance_scale to the mobile pipeline should be accepted but ignored (with warning).""" - device = torch_device - components = self.get_dummy_components() - pipe = self.pipeline_class(**components).to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["guidance_scale"] = 7.5 # should not raise inputs["image_guidance_scale"] = 1.5 # should not raise out = pipe(**inputs).images - self.assertEqual(to_np(out).shape, (1, 64, 64, 3)) + + assert out.shape == (1, *self.output_shape) + + +class TestDreamLiteMobilePipelineMemory(DreamLiteMobilePipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the mobile pipeline.""" @nightly @require_torch_gpu -class DreamLiteMobilePipelineSlowTests(unittest.TestCase): +class TestDreamLiteMobilePipelineIntegration: """End-to-end test against the real DreamLite-mobile checkpoint on the Hub. By default this loads ``carlofkl/DreamLite-mobile`` (``diffusers`` branch) @@ -291,13 +165,11 @@ class DreamLiteMobilePipelineSlowTests(unittest.TestCase): repo_id = "carlofkl/DreamLite-mobile" revision = "diffusers" - def setUp(self): - super().setUp() + @pytest.fixture(autouse=True) + def cleanup(self): gc.collect() torch.cuda.empty_cache() - - def tearDown(self): - super().tearDown() + yield gc.collect() torch.cuda.empty_cache() @@ -320,8 +192,8 @@ def test_mobile_t2i_real_checkpoint(self): output_type="np", ).images - self.assertEqual(out.shape, (1, 1024, 1024, 3)) - self.assertFalse(np.isnan(out).any()) + assert out.shape == (1, 1024, 1024, 3) + assert not np.isnan(out).any() def test_mobile_i2i_real_checkpoint(self): pipe = DreamLiteMobilePipeline.from_pretrained( @@ -339,9 +211,5 @@ def test_mobile_i2i_real_checkpoint(self): output_type="np", ).images - self.assertEqual(out.shape, (1, 1024, 1024, 3)) - self.assertFalse(np.isnan(out).any()) - - -if __name__ == "__main__": - unittest.main() + assert out.shape == (1, 1024, 1024, 3) + assert not np.isnan(out).any() diff --git a/tests/pipelines/dreamlite/testing_utils.py b/tests/pipelines/dreamlite/testing_utils.py new file mode 100644 index 000000000000..d4fea7d64323 --- /dev/null +++ b/tests/pipelines/dreamlite/testing_utils.py @@ -0,0 +1,131 @@ +# Copyright (c) 2026 ByteDance Ltd. and/or its affiliates. +# +# 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. +"""Shared test fixtures for the DreamLite pipelines. + +``DreamLitePipeline`` and its distilled sibling ``DreamLiteMobilePipeline`` take the exact same components, so the +tiny Qwen3-VL text encoder and the rest of the dummy component set live here and are shared by both test files. +""" + +import torch +from transformers import AutoTokenizer, Qwen3VLConfig, Qwen3VLForConditionalGeneration, Qwen3VLProcessor + +from diffusers import ( + AutoencoderTiny, + DreamLiteUNetModel, + FlowMatchEulerDiscreteScheduler, +) + +from ..testing_utils import BasePipelineTesterConfig + + +# Match the tiny text encoder hidden size below; the UNet's cross-attention +# dimension must match what ``encode_prompt`` returns. +CROSS_ATTN_DIM = 16 + + +def build_tiny_text_encoder() -> Qwen3VLForConditionalGeneration: + """Build a tiny but functional Qwen3-VL model for the fast test fixture. + + Mirrors the recipe used by ``tests/pipelines/nucleusmoe_image``: small text + + vision configs that still go through the real Qwen3-VL forward path, so + DreamLite's ``encode_prompt`` (chat template + tokenizer + multimodal + processor) is exercised for real. + """ + config = Qwen3VLConfig( + text_config={ + "hidden_size": CROSS_ATTN_DIM, + "intermediate_size": CROSS_ATTN_DIM, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "rope_scaling": { + "mrope_section": [1, 1, 2], + "rope_type": "default", + "type": "default", + }, + "rope_theta": 1000000.0, + "vocab_size": 151936, + "head_dim": 8, + }, + vision_config={ + "depth": 2, + "hidden_size": CROSS_ATTN_DIM, + "intermediate_size": CROSS_ATTN_DIM, + "num_heads": 2, + "out_channels": CROSS_ATTN_DIM, + # ``out_hidden_size`` is the dim that vision tokens are projected to before + # being merged into the text stream; it must match ``text_config.hidden_size``. + "out_hidden_size": CROSS_ATTN_DIM, + # Match the cached ``hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration`` + # image processor (``patch_size=14``); otherwise the pixel_values + # produced by the processor cannot be reshaped to the model's + # vision patch embed. + "patch_size": 14, + }, + ) + return Qwen3VLForConditionalGeneration(config).eval() + + +class DreamLiteBaseTesterConfig(BasePipelineTesterConfig): + """Component set shared by ``DreamLitePipeline`` and ``DreamLiteMobilePipeline``.""" + + # DreamLite samples its own noise; `latents` cannot be supplied by the caller. + optional_input_params = BasePipelineTesterConfig.optional_input_params - {"latents"} + output_shape = (3, 64, 64) + + def get_dummy_components(self): + torch.manual_seed(0) + unet = DreamLiteUNetModel( + sample_size=8, + in_channels=4, + out_channels=4, + down_block_types=( + "DreamLiteCrossAttnNoSelfAttnDownBlock2D", + "DreamLiteCrossAttnDownBlock2D", + ), + up_block_types=("DreamLiteCrossAttnUpBlock2D", "DreamLiteUpBlock2D"), + block_out_channels=(32, 64), + cross_attention_dim=CROSS_ATTN_DIM, + attention_head_dim=8, + layers_per_block=1, + norm_num_groups=8, + transformer_layers_per_block=1, + ) + + torch.manual_seed(0) + vae = AutoencoderTiny( + in_channels=3, + out_channels=3, + encoder_block_out_channels=(32, 32), + decoder_block_out_channels=(32, 32), + num_encoder_blocks=(1, 1), + num_decoder_blocks=(1, 1), + latent_channels=4, + ) + + scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000) + + torch.manual_seed(0) + text_encoder = build_tiny_text_encoder() + tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") + processor = Qwen3VLProcessor.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") + + return { + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "processor": processor, + "vae": vae, + "unet": unet, + "scheduler": scheduler, + }