Repository: huggingface/diffusers
File: tests/models/unets/test_models_unet_2d_condition.py (Python)

Code before:
        assert loaded_model
        assert new_output.sample.shape == (4, 4, 16, 16)

    @require_torch_accelerator
    def test_load_sharded_checkpoint_device_map_from_hub_local(self):
        _, inputs_dict = self.prepare_init_args_and_inputs_for_common()
        ckpt_path = snapshot_download("hf-internal-testing/unet2d-sharded-dummy")
        loaded_model = self.model_class.from_pretrained(ckpt_path, local_files_only=True, device_map="auto")
        new_output = loaded_model(**inputs_dict)

        assert loaded_model
        assert new_output.sample.shape == (4, 4, 16, 16)

    @require_torch_accelerator
    def test_load_sharded_checkpoint_device_map_from_hub_local_subfolder(self):
        _, inputs_dict = self.prepare_init_args_and_inputs_for_common()
        ckpt_path = snapshot_download("hf-internal-testing/unet2d-sharded-dummy-subfolder")
        loaded_model = self.model_class.from_pretrained(
            ckpt_path, local_files_only=True, subfolder="unet", device_map="auto"
        )
        new_output = loaded_model(**inputs_dict)

        assert loaded_model
        assert new_output.sample.shape == (4, 4, 16, 16)

    def test_wrong_device_map_raises_error(self):
        with self.assertRaises(ValueError) as err_ctx:
            _ = self.model_class.from_pretrained(
                "hf-internal-testing/unet2d-sharded-dummy-subfolder", subfolder="unet", device_map=-1
            )

        msg_substring = "You can't pass device_map as a negative int"
        assert msg_substring in str(err_ctx.exception)

    @parameterized.expand([0, "cuda", torch.device("cuda"), torch.device("cuda:0")])
    @require_torch_gpu
    def test_passing_non_dict_device_map_works(self, device_map):
        _, inputs_dict = self.prepare_init_args_and_inputs_for_common()
        loaded_model = self.model_class.from_pretrained(
            "hf-internal-testing/unet2d-sharded-dummy-subfolder", subfolder="unet", device_map=device_map
        )
        output = loaded_model(**inputs_dict)
        assert output.sample.shape == (4, 4, 16, 16)

    @require_peft_backend
    def test_load_attn_procs_raise_warning(self):
        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
        model = self.model_class(**init_dict)
        model.to(torch_device)

        # forward pass without LoRA

Review comment (bug):
Testing more of the failure code paths (if there are any; I am not fully aware of the relevant parts of the codebase) could be nice

Code after:
        assert loaded_model
        assert new_output.sample.shape == (4, 4, 16, 16)

    @require_torch_accelerator
    def test_load_sharded_checkpoint_device_map_from_hub_local(self):
        _, inputs_dict = self.prepare_init_args_and_inputs_for_common()
        ckpt_path = snapshot_download("hf-internal-testing/unet2d-sharded-dummy")
        loaded_model = self.model_class.from_pretrained(ckpt_path, local_files_only=True, device_map="auto")
        new_output = loaded_model(**inputs_dict)

        assert loaded_model
        assert new_output.sample.shape == (4, 4, 16, 16)

    @require_torch_accelerator
    def test_load_sharded_checkpoint_device_map_from_hub_local_subfolder(self):
        _, inputs_dict = self.prepare_init_args_and_inputs_for_common()
        ckpt_path = snapshot_download("hf-internal-testing/unet2d-sharded-dummy-subfolder")
        loaded_model = self.model_class.from_pretrained(
            ckpt_path, local_files_only=True, subfolder="unet", device_map="auto"
        )
        new_output = loaded_model(**inputs_dict)

        assert loaded_model
        assert new_output.sample.shape == (4, 4, 16, 16)

    @parameterized.expand(
        [
            (-1, "You can't pass device_map as a negative int"),
            ("foo", "When passing device_map as a string, the value needs to be a device name"),
        ]
    )
    def test_wrong_device_map_raises_error(self, device_map, msg_substring):
        with self.assertRaises(ValueError) as err_ctx:
            _ = self.model_class.from_pretrained(
                "hf-internal-testing/unet2d-sharded-dummy-subfolder", subfolder="unet", device_map=device_map
            )

        assert msg_substring in str(err_ctx.exception)

    @parameterized.expand([0, "cuda", torch.device("cuda"), torch.device("cuda:0")])
    @require_torch_gpu
    def test_passing_non_dict_device_map_works(self, device_map):
        _, inputs_dict = self.prepare_init_args_and_inputs_for_common()
        loaded_model = self.model_class.from_pretrained(
            "hf-internal-testing/unet2d-sharded-dummy-subfolder", subfolder="unet", device_map=device_map
        )
        output = loaded_model(**inputs_dict)
        assert output.sample.shape == (4, 4, 16, 16)

    @parameterized.expand([("", "cuda"), ("", torch.device("cuda"))])
    @require_torch_gpu