Updated docs and types for Lumina pipelines

Author: guiyrt

src/diffusers/pipelines/lumina/pipeline_lumina.py

@@ -20,7 +20,7 @@
 from typing import List, Optional, Tuple, Union
 
 import torch
-from transformers import PreTrainedModel, PreTrainedTokenizerBase
+from transformers import GemmaPreTrainedModel, GemmaTokenizer, GemmaTokenizerFast
 
 from ...image_processor import VaeImageProcessor
 from ...models import AutoencoderKL
@@ -143,13 +143,10 @@ class LuminaText2ImgPipeline(DiffusionPipeline):
     Args:
         vae ([`AutoencoderKL`]):
             Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
-        text_encoder ([`PreTrainedModel`]):
-            Frozen text-encoder. Lumina-T2I uses
-            [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.AutoModel), specifically the
-            [t5-v1_1-xxl](https://huggingface.co/Alpha-VLLM/tree/main/t5-v1_1-xxl) variant.
-        tokenizer (`AutoTokenizer`):
-            Tokenizer of class
-            [AutoTokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.AutoModel).
+        text_encoder ([`GemmaPreTrainedModel`]):
+            Frozen Gemma text-encoder.
+        tokenizer (`GemmaTokenizer` or `GemmaTokenizerFast`):
+            Gemma tokenizer.
         transformer ([`Transformer2DModel`]):
             A text conditioned `Transformer2DModel` to denoise the encoded image latents.
         scheduler ([`SchedulerMixin`]):
@@ -180,8 +177,8 @@ def __init__(
         transformer: LuminaNextDiT2DModel,
         scheduler: FlowMatchEulerDiscreteScheduler,
         vae: AutoencoderKL,
-        text_encoder: PreTrainedModel,
-        tokenizer: PreTrainedTokenizerBase,
+        text_encoder: GemmaPreTrainedModel,
+        tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
     ):
         super().__init__()
 

src/diffusers/pipelines/lumina2/pipeline_lumina2.py

@@ -17,7 +17,7 @@
 
 import numpy as np
 import torch
-from transformers import PreTrainedModel, PreTrainedTokenizerBase
+from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast
 
 from ...image_processor import VaeImageProcessor
 from ...models import AutoencoderKL
@@ -150,13 +150,10 @@ class Lumina2Text2ImgPipeline(DiffusionPipeline):
     Args:
         vae ([`AutoencoderKL`]):
             Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
-        text_encoder ([`PreTrainedModel`]):
-            Frozen text-encoder. Lumina-T2I uses
-            [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.AutoModel), specifically the
-            [t5-v1_1-xxl](https://huggingface.co/Alpha-VLLM/tree/main/t5-v1_1-xxl) variant.
-        tokenizer (`PreTrainedTokenizerBase`):
-            Tokenizer of class
-            [AutoModel](https://huggingface.co/docs/transformers/model_doc/t5#transformers.AutoModel).
+        text_encoder ([`Gemma2PreTrainedModel`]):
+            Frozen Gemma2 text-encoder.
+        tokenizer (`GemmaTokenizer` or `GemmaTokenizerFast`):
+            Gemma tokenizer.
         transformer ([`Transformer2DModel`]):
             A text conditioned `Transformer2DModel` to denoise the encoded image latents.
         scheduler ([`SchedulerMixin`]):
@@ -172,8 +169,8 @@ def __init__(
         transformer: Lumina2Transformer2DModel,
         scheduler: FlowMatchEulerDiscreteScheduler,
         vae: AutoencoderKL,
-        text_encoder: PreTrainedModel,
-        tokenizer: PreTrainedTokenizerBase,
+        text_encoder: Gemma2PreTrainedModel,
+        tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
     ):
         super().__init__()
 

tests/pipelines/lumina2/test_pipeline_lumina2.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 import torch
-from transformers import AutoTokenizer, GemmaConfig, GemmaForCausalLM
+from transformers import AutoTokenizer, Gemma2Config, Gemma2ForCausalLM
 
 from diffusers import (
     AutoencoderKL,
@@ -81,21 +81,21 @@ def get_dummy_components(self):
         tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/dummy-gemma")
 
         torch.manual_seed(0)
-        config = GemmaConfig(
+        config = Gemma2Config(
             head_dim=2,
             hidden_size=8,
             intermediate_size=37,
             num_attention_heads=4,
             num_hidden_layers=2,
             num_key_value_heads=4,
         )
-        text_encoder = GemmaForCausalLM(config)
+        text_encoder = Gemma2ForCausalLM(config)
 
         components = {
-            "transformer": transformer.eval(),
+            "transformer": transformer,
             "vae": vae.eval(),
             "scheduler": scheduler,
-            "text_encoder": text_encoder.eval(),
+            "text_encoder": text_encoder,
             "tokenizer": tokenizer,
         }
         return components