[tests] use proper gemma class and config in lumina2 tests.

[sayakpaul]

No description provided.

[sayakpaul]

Failing tests are unrelated.

[a-r-r-o-w]

Thanks, LGTM!

[guiyrt]

Similar changes are proposed in #10758 here, which also update docstrings and type annotations for Lumina pipelines. I changed to Gemma2ForCausalLM instead of Gemma2Model, as I didn't know the reason for using ForCausalLM variant. In the pipelines, the text_encoder type annotations were set to Gemma2PreTrainedModel because the tests use the ForCausalLM variant, but if we end up using Gemma2Model in tests, we can annotate simply as Gemma2Model in the pipelines (same for sana pipelines, proposed new type annotation as Gemma2PreTrainedModel, might be worth changing tests to use Gemma2Model?).

To avoid merge conflicts, I could either remove my changes on tests/pipelines/lumina2/test_pipeline_lumina2.py or incorporate these into that PR, as you prefer :)

Also, on the topic of tests changes, one point identified but not yet addressed in #10758 (comment) is regarding LDMTextToImagePipeline, which docstrings' refer to BERT tokenizer and text_encoder, but tests use CLIP, which forces us to use generic PreTrainedModel and PreTrainedTokenizer annotations. Is it relevant to change the tests to use BERT and correct the type annotations?

[sayakpaul]

@guiyrt thanks for your hard work :) Also sorry for not seeing your awesome PR first!

To avoid merge conflicts, I could either remove my changes on tests/pipelines/lumina2/test_pipeline_lumina2.py or incorporate these into that PR, as you prefer :)

Let's go with the first option.

Also, on the topic of tests changes, one point identified but not yet addressed in #10758 (comment) is regarding LDMTextToImagePipeline, which docstrings' refer to BERT tokenizer and text_encoder, but tests use CLIP, which forces us to use generic PreTrainedModel and PreTrainedTokenizer annotations. Is it relevant to change the tests to use BERT and correct the type annotations?

I would prefer:

Is it relevant to change the tests to use BERT and correct the type annotations?

In general, we should try to emulate the actual model architectures from a given pipeline as much as possible.

[guiyrt]

@guiyrt thanks for your hard work :) Also sorry for not seeing your awesome PR first!

Thanks! 🤗 One thing that I also changed in my PR is removing.eval() in the return values of get_dummy_components, otherwise the tests would fail for sequential offload. SanaPipelineFastTests also uses Gemma2Model, and none of the components are using eval(). I'm not sure why eval() is breaking sequential offload for Gemma2Model, but it works without it (is it worth taking a deeper look?). You might want to double check these tests locally, as the PR tests are failing for different reasons.

With "text_encoder": text_encoder.eval(), this is the output from pytest tests/pipelines/lumina2/test_pipeline_lumina2.py:

FAILED tests/pipelines/lumina2/test_pipeline_lumina2.py::Lumina2Text2ImgPipelinePipelineFastTests::test_sequential_cpu_offload_forward_pass
- RuntimeError: Tensor.item() cannot be called on meta tensors
FAILED tests/pipelines/lumina2/test_pipeline_lumina2.py::Lumina2Text2ImgPipelinePipelineFastTests::test_sequential_offload_forward_pass_twice
- RuntimeError: Tensor.item() cannot be called on meta tensors
=============================================== 2 failed, 26 passed, 1 skipped in 25.49s ================================================

Test traceback

============================= test session starts ==============================
platform linux -- Python 3.10.16, pytest-8.3.4, pluggy-1.5.0
rootdir: /home/guiyrt/diffusers
configfile: pyproject.toml
plugins: timeout-2.3.1, xdist-3.6.1, requests-mock-1.10.0
collected 1 item

tests/pipelines/lumina2/test_pipeline_lumina2.py The 'batch_size' argument of HybridCache is deprecated and will be removed in v4.49. Use the more precisely named 'max_batch_size' argument instead.
The 'batch_size' attribute of HybridCache is deprecated and will be removed in v4.49. Use the more precisely named 'self.max_batch_size' attribute instead.
F

=================================== FAILURES ===================================
_ Lumina2Text2ImgPipelinePipelineFastTests.test_sequential_cpu_offload_forward_pass _

self = <tests.pipelines.lumina2.test_pipeline_lumina2.Lumina2Text2ImgPipelinePipelineFastTests testMethod=test_sequential_cpu_offload_forward_pass>
expected_max_diff = 0.0001

    @require_accelerator
    @require_accelerate_version_greater("0.14.0")
    def test_sequential_cpu_offload_forward_pass(self, expected_max_diff=1e-4):
        import accelerate
    
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        for component in pipe.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()
        pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)
    
        generator_device = "cpu"
        inputs = self.get_dummy_inputs(generator_device)
        output_without_offload = pipe(**inputs)[0]
    
        pipe.enable_sequential_cpu_offload(device=torch_device)
        assert pipe._execution_device.type == torch_device
    
        inputs = self.get_dummy_inputs(generator_device)
>       output_with_offload = pipe(**inputs)[0]

tests/pipelines/test_pipelines_common.py:1518: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
../anaconda3/envs/diffusers/lib/python3.10/site-packages/torch/utils/_contextlib.py:116: in decorate_context
    return func(*args, **kwargs)
src/diffusers/pipelines/lumina2/pipeline_lumina2.py:645: in __call__
    ) = self.encode_prompt(
src/diffusers/pipelines/lumina2/pipeline_lumina2.py:290: in encode_prompt
    prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds(
src/diffusers/pipelines/lumina2/pipeline_lumina2.py:218: in _get_gemma_prompt_embeds
    prompt_embeds = self.text_encoder(
../anaconda3/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/module.py:1739: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
../anaconda3/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/module.py:1750: in _call_impl
    return forward_call(*args, **kwargs)
../anaconda3/envs/diffusers/lib/python3.10/site-packages/accelerate/hooks.py:170: in new_forward
    output = module._old_forward(*args, **kwargs)
../anaconda3/envs/diffusers/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py:600: in forward
    past_key_values = HybridCache(
../anaconda3/envs/diffusers/lib/python3.10/site-packages/transformers/cache_utils.py:1657: in __init__
    cache_shape = global_cache_shape if not self.is_sliding[i] else sliding_cache_shape
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = tensor(..., device='meta', size=(), dtype=torch.bool)

    @register_meta(aten._local_scalar_dense)
    def meta_local_scalar_dense(self: Tensor):
>       raise RuntimeError("Tensor.item() cannot be called on meta tensors")
E       RuntimeError: Tensor.item() cannot be called on meta tensors

../anaconda3/envs/diffusers/lib/python3.10/site-packages/torch/_meta_registrations.py:6585: RuntimeError
=========================== short test summary info ============================
FAILED tests/pipelines/lumina2/test_pipeline_lumina2.py::Lumina2Text2ImgPipelinePipelineFastTests::test_sequential_cpu_offload_forward_pass
============================== 1 failed in 5.68s ===============================

I would prefer:

Is it relevant to change the tests to use BERT and correct the type annotations?

In general, we should try to emulate the actual model architectures from a given pipeline as much as possible.

I will try to update the tests for LDMTextToImagePipeline then :)

[sayakpaul]

We should consider having eval(). We remove it?

[guiyrt]

There is no eval() in sana tests.

diffusers/tests/pipelines/sana/test_sana.py

Lines 120 to 127 in a4c1aac

	components = {
	"transformer": transformer,
	"vae": vae,
	"scheduler": scheduler,
	"text_encoder": text_encoder,
	"tokenizer": tokenizer,
	}
	return components

Adding eval() in text_encoder results in:

================================================= short test summary info ==============================================
FAILED tests/pipelines/sana/test_sana.py::SanaPipelineFastTests::test_group_offloading_inference
- RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)
FAILED tests/pipelines/sana/test_sana.py::SanaPipelineFastTests::test_sequential_cpu_offload_forward_pass
- RuntimeError: Tensor.item() cannot be called on meta tensors
FAILED tests/pipelines/sana/test_sana.py::SanaPipelineFastTests::test_sequential_offload_forward_pass_twice
- RuntimeError: Tensor.item() cannot be called on meta tensors
======================================= 3 failed, 25 passed, 5 skipped in 29.88s =======================================