Repository: huggingface/diffusers
File: tests/pipelines/test_pipelines_common.py (Python)

Code before:
        inputs = self.get_dummy_inputs(torch_device)
        pipe_out_2 = full_pipe(**inputs)[0]

        if isinstance(pipe_out, np.ndarray) and isinstance(pipe_out_2, np.ndarray):
            self.assertTrue(np.allclose(pipe_out, pipe_out_2, atol=atol, rtol=rtol))
        elif isinstance(pipe_out, torch.Tensor) and isinstance(pipe_out_2, torch.Tensor):
            self.assertTrue(torch.allclose(pipe_out, pipe_out_2, atol=atol, rtol=rtol))

    def test_StableDiffusionMixin_component(self):
        """Any pipeline that have LDMFuncMixin should have vae and unet components."""
        if not issubclass(self.pipeline_class, StableDiffusionMixin):
            return
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        self.assertTrue(hasattr(pipe, "vae") and isinstance(pipe.vae, (AutoencoderKL, AutoencoderTiny)))
        self.assertTrue(
            hasattr(pipe, "unet")
            and isinstance(
                pipe.unet,
                (UNet2DConditionModel, UNet3DConditionModel, I2VGenXLUNet, UNetMotionModel, UNetControlNetXSModel),
            )
        )

    @require_hf_hub_version_greater("0.26.5")
    @require_transformers_version_greater("4.47.1")
    def test_save_load_dduf(self, atol=1e-3, rtol=1e-3):
        if not self.supports_dduf:
            return

        from huggingface_hub import export_folder_as_dduf

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

        inputs = self.get_dummy_inputs(device="cpu")
        inputs.pop("generator")
        inputs["generator"] = torch.manual_seed(0)

        pipeline_out = pipe(**inputs)[0]

        with tempfile.TemporaryDirectory() as tmpdir:
            dduf_filename = os.path.join(tmpdir, f"{pipe.__class__.__name__.lower()}.dduf")
            pipe.save_pretrained(tmpdir, safe_serialization=True)
            export_folder_as_dduf(dduf_filename, folder_path=tmpdir)
            loaded_pipe = self.pipeline_class.from_pretrained(tmpdir, dduf_file=dduf_filename).to(torch_device)

        inputs["generator"] = torch.manual_seed(0)
        loaded_pipeline_out = loaded_pipe(**inputs)[0]


Review comment (bug):
We shouldn't change it at the main test level. If a model test failing, we should override the method at the corresponding test class and relax tolerance.

Code after:
        inputs = self.get_dummy_inputs(torch_device)
        pipe_out_2 = full_pipe(**inputs)[0]

        if isinstance(pipe_out, np.ndarray) and isinstance(pipe_out_2, np.ndarray):
            self.assertTrue(np.allclose(pipe_out, pipe_out_2, atol=atol, rtol=rtol))
        elif isinstance(pipe_out, torch.Tensor) and isinstance(pipe_out_2, torch.Tensor):
            self.assertTrue(torch.allclose(pipe_out, pipe_out_2, atol=atol, rtol=rtol))

    def test_StableDiffusionMixin_component(self):
        """Any pipeline that have LDMFuncMixin should have vae and unet components."""
        if not issubclass(self.pipeline_class, StableDiffusionMixin):
            return
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        self.assertTrue(hasattr(pipe, "vae") and isinstance(pipe.vae, (AutoencoderKL, AutoencoderTiny)))
        self.assertTrue(
            hasattr(pipe, "unet")
            and isinstance(
                pipe.unet,
                (UNet2DConditionModel, UNet3DConditionModel, I2VGenXLUNet, UNetMotionModel, UNetControlNetXSModel),
            )
        )

    @require_hf_hub_version_greater("0.26.5")
    @require_transformers_version_greater("4.47.1")
    def test_save_load_dduf(self, atol=1e-4, rtol=1e-4):
        if not self.supports_dduf:
            return

        from huggingface_hub import export_folder_as_dduf

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

        inputs = self.get_dummy_inputs(device="cpu")
        inputs.pop("generator")
        inputs["generator"] = torch.manual_seed(0)

        pipeline_out = pipe(**inputs)[0]

        with tempfile.TemporaryDirectory() as tmpdir:
            dduf_filename = os.path.join(tmpdir, f"{pipe.__class__.__name__.lower()}.dduf")
            pipe.save_pretrained(tmpdir, safe_serialization=True)
            export_folder_as_dduf(dduf_filename, folder_path=tmpdir)
            loaded_pipe = self.pipeline_class.from_pretrained(tmpdir, dduf_file=dduf_filename).to(torch_device)

        inputs["generator"] = torch.manual_seed(0)
        loaded_pipeline_out = loaded_pipe(**inputs)[0]