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

Code before:
        model = torch.compile(model, fullgraph=True)

        with (
            torch._inductor.utils.fresh_inductor_cache(),
            torch._dynamo.config.patch(error_on_recompile=True),
            torch.no_grad(),
        ):
            _ = model(**inputs_dict)
            _ = model(**inputs_dict)

    def test_torch_compile_repeated_blocks(self):
        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()

        model = self.model_class(**init_dict).to(torch_device)
        model.compile_repeated_blocks(fullgraph=True)

        recompile_limit = 1
        if self.model_class.__name__ == "UNet2DConditionModel":
            recompile_limit = 2

        with (
            torch._inductor.utils.fresh_inductor_cache(),
            torch._dynamo.config.patch(recompile_limit=recompile_limit),
            torch.no_grad(),
        ):
            _ = model(**inputs_dict)
            _ = model(**inputs_dict)

    def test_compile_with_group_offloading(self):
        torch._dynamo.config.cache_size_limit = 10000

        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
        model = self.model_class(**init_dict)

        if not getattr(model, "_supports_group_offloading", True):
            return

        model.eval()
        # TODO: Can test for other group offloading kwargs later if needed.
        group_offload_kwargs = {
            "onload_device": "cuda",
            "offload_device": "cpu",
            "offload_type": "block_level",
            "num_blocks_per_group": 1,
            "use_stream": True,
            "non_blocking": True,
        }
        model.enable_group_offload(**group_offload_kwargs)
        model.compile()
        with torch.no_grad():
            _ = model(**inputs_dict)

Review comment (suggestion):
@sayakpaul For UNet, the shape changes between different instances of the transformer model, so we need 2 entries.

Code after:
    def test_torch_compile_recompilation_and_graph_break(self):
        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()

        model = self.model_class(**init_dict).to(torch_device)
        model = torch.compile(model, fullgraph=True)

        with (
            torch._inductor.utils.fresh_inductor_cache(),
            torch._dynamo.config.patch(error_on_recompile=True),
            torch.no_grad(),
        ):
            _ = model(**inputs_dict)
            _ = model(**inputs_dict)

    def test_torch_compile_repeated_blocks(self):
        if self.model_class._repeated_blocks is None:
            pytest.skip("Skipping test as the model class doesn't have `_repeated_blocks` set.")

        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()

        model = self.model_class(**init_dict).to(torch_device)
        model.compile_repeated_blocks(fullgraph=True)

        recompile_limit = 1
        if self.model_class.__name__ == "UNet2DConditionModel":
            recompile_limit = 2

        with (
            torch._inductor.utils.fresh_inductor_cache(),
            torch._dynamo.config.patch(recompile_limit=recompile_limit),
            torch.no_grad(),
        ):
            _ = model(**inputs_dict)
            _ = model(**inputs_dict)

    def test_compile_with_group_offloading(self):
        torch._dynamo.config.cache_size_limit = 10000

        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
        model = self.model_class(**init_dict)

        if not getattr(model, "_supports_group_offloading", True):
            return

        model.eval()
        # TODO: Can test for other group offloading kwargs later if needed.
        group_offload_kwargs = {
            "onload_device": "cuda",
            "offload_device": "cpu",
            "offload_type": "block_level",
            "num_blocks_per_group": 1,