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

Code before:
        model = self.model_class(**init_dict)

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

        model.to(torch_device)
        model.eval()
        _ = model(**inputs_dict)[0]

        torch.manual_seed(0)
        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
        inputs_dict = cast_maybe_tensor_dtype(inputs_dict, torch.float32, compute_dtype)
        model = self.model_class(**init_dict)
        model.eval()
        model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=1)
        model.enable_layerwise_casting(storage_dtype=storage_dtype, compute_dtype=compute_dtype)
        _ = model(**inputs_dict)[0]

        torch.manual_seed(0)
        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
        inputs_dict = cast_maybe_tensor_dtype(inputs_dict, torch.float32, compute_dtype)
        model = self.model_class(**init_dict)
        model.eval()
        model.enable_group_offload(
            torch_device,
            offload_type="block_level",
            num_blocks_per_group=1,
            use_stream=True,
            non_blocking=True,
            record_stream=record_stream,
        )
        _ = model(**inputs_dict)[0]

    def test_auto_model(self, expected_max_diff=5e-5):
        if self.forward_requires_fresh_args:
            model = self.model_class(**self.init_dict)
        else:
            init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
            model = self.model_class(**init_dict)

        model = model.eval()
        model = model.to(torch_device)

        if hasattr(model, "set_default_attn_processor"):
            model.set_default_attn_processor()

        with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname:
            model.save_pretrained(tmpdirname, safe_serialization=False)

            auto_model = AutoModel.from_pretrained(tmpdirname)
            if hasattr(auto_model, "set_default_attn_processor"):

Review comment (suggestion):
I think it would make more sense to parameterize test on offload_type instead of record_stream

Code after:
        model = self.model_class(**init_dict)

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

        model.to(torch_device)
        model.eval()
        _ = model(**inputs_dict)[0]

        torch.manual_seed(0)
        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
        storage_dtype, compute_dtype = torch.float16, torch.float32
        inputs_dict = cast_maybe_tensor_dtype(inputs_dict, torch.float32, compute_dtype)
        model = self.model_class(**init_dict)
        model.eval()
        additional_kwargs = {} if offload_type == "leaf_level" else {"num_blocks_per_group": 1}
        model.enable_group_offload(
            torch_device, offload_type=offload_type, use_stream=True, record_stream=record_stream, **additional_kwargs
        )
        model.enable_layerwise_casting(storage_dtype=storage_dtype, compute_dtype=compute_dtype)
        _ = model(**inputs_dict)[0]

    def test_auto_model(self, expected_max_diff=5e-5):
        if self.forward_requires_fresh_args:
            model = self.model_class(**self.init_dict)
        else:
            init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
            model = self.model_class(**init_dict)

        model = model.eval()
        model = model.to(torch_device)

        if hasattr(model, "set_default_attn_processor"):
            model.set_default_attn_processor()

        with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname:
            model.save_pretrained(tmpdirname, safe_serialization=False)

            auto_model = AutoModel.from_pretrained(tmpdirname)
            if hasattr(auto_model, "set_default_attn_processor"):
                auto_model.set_default_attn_processor()

        auto_model = auto_model.eval()
        auto_model = auto_model.to(torch_device)

        with torch.no_grad():
            if self.forward_requires_fresh_args:
                output_original = model(**self.inputs_dict(0))
                output_auto = auto_model(**self.inputs_dict(0))
            else:
                output_original = model(**inputs_dict)