Use concrete Mistral3Model / Ministral3ForCausalLM types (guarded)

Author: akshan-main

src/diffusers/modular_pipelines/ernie_image/encoders.py

@@ -17,6 +17,17 @@
 import torch
 from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
 
+
+try:
+    from transformers import Mistral3Model as _TextEncoderClass
+except ImportError:
+    _TextEncoderClass = AutoModel
+
+try:
+    from transformers import Ministral3ForCausalLM as _PromptEnhancerClass
+except ImportError:
+    _PromptEnhancerClass = AutoModelForCausalLM
+
 from ...configuration_utils import FrozenDict
 from ...guiders import ClassifierFreeGuidance
 from ...utils import logging
@@ -38,7 +49,7 @@ def description(self) -> str:
     @property
     def expected_components(self) -> list[ComponentSpec]:
         return [
-            ComponentSpec("pe", AutoModelForCausalLM),
+            ComponentSpec("pe", _PromptEnhancerClass),
             ComponentSpec("pe_tokenizer", AutoTokenizer),
         ]
 
@@ -160,7 +171,7 @@ def description(self) -> str:
     @property
     def expected_components(self) -> list[ComponentSpec]:
         return [
-            ComponentSpec("text_encoder", AutoModel),
+            ComponentSpec("text_encoder", _TextEncoderClass),
             ComponentSpec("tokenizer", AutoTokenizer),
             ComponentSpec(
                 "guider",

src/diffusers/pipelines/ernie_image/pipeline_ernie_image.py

@@ -17,10 +17,10 @@
 """
 
 import json
-from typing import Callable, List, Optional, Union
+from typing import TYPE_CHECKING, Callable, List, Optional, Union
 
 import torch
-from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer
 
 from ...image_processor import VaeImageProcessor
 from ...loaders import ErnieImageLoraLoaderMixin
@@ -32,6 +32,10 @@
 from .pipeline_output import ErnieImagePipelineOutput
 
 
+if TYPE_CHECKING:
+    from transformers import Ministral3ForCausalLM, Mistral3Model
+
+
 class ErnieImagePipeline(DiffusionPipeline, ErnieImageLoraLoaderMixin):
     """
     Pipeline for text-to-image generation using ErnieImageTransformer2DModel.
@@ -52,10 +56,10 @@ def __init__(
         self,
         transformer: ErnieImageTransformer2DModel,
         vae: AutoencoderKLFlux2,
-        text_encoder: AutoModel,
+        text_encoder: "Mistral3Model",
         tokenizer: AutoTokenizer,
         scheduler: FlowMatchEulerDiscreteScheduler,
-        pe: Optional[AutoModelForCausalLM] = None,
+        pe: Optional["Ministral3ForCausalLM"] = None,
         pe_tokenizer: Optional[AutoTokenizer] = None,
     ):
         super().__init__()