Address ernie-image review findings #13577

src/diffusers/modular_pipelines/ernie_image/decoders.py

@@ -75,9 +75,7 @@ def __call__(self, components: ErnieImageModularPipeline, state: PipelineState)
 
         latents = block_state.latents
         bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(device=device, dtype=latents.dtype)
-        bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to(
-            device=device, dtype=latents.dtype
-        )
+        bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + 1e-5).to(device=device, dtype=latents.dtype)
         latents = latents * bn_std + bn_mean
 
         latents = components.pachifier.unpack_latents(latents)

src/diffusers/modular_pipelines/ernie_image/encoders.py

@@ -15,7 +15,7 @@
 import json
 
 import torch
-from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, Ministral3ForCausalLM, Mistral3Model
 
 from ...configuration_utils import FrozenDict
 from ...guiders import ClassifierFreeGuidance
@@ -38,7 +38,7 @@ def description(self) -> str:
     @property
     def expected_components(self) -> list[ComponentSpec]:
         return [
-            ComponentSpec("pe", AutoModelForCausalLM),
+            ComponentSpec("pe", Ministral3ForCausalLM),
             ComponentSpec("pe_tokenizer", AutoTokenizer),
         ]
 
@@ -83,7 +83,7 @@ def intermediate_outputs(self) -> list[OutputParam]:
 
     @staticmethod
     def _enhance_prompt(
-        pe: AutoModelForCausalLM,
+        pe: Ministral3ForCausalLM,
         pe_tokenizer: AutoTokenizer,
         prompt: str,
         device: torch.device,
@@ -160,7 +160,7 @@ def description(self) -> str:
     @property
     def expected_components(self) -> list[ComponentSpec]:
         return [
-            ComponentSpec("text_encoder", AutoModel),
+            ComponentSpec("text_encoder", Mistral3Model),
             ComponentSpec("tokenizer", AutoTokenizer),
             ComponentSpec(
                 "guider",
@@ -200,7 +200,7 @@ def intermediate_outputs(self) -> list[OutputParam]:
 
     @staticmethod
     def _encode(
-        text_encoder: AutoModel,
+        text_encoder: Mistral3Model,
         tokenizer: AutoTokenizer,
         prompt: list[str],
         device: torch.device,

src/diffusers/modular_pipelines/ernie_image/modular_blocks_ernie_image.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from ...utils import logging
-from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline import ConditionalPipelineBlocks, SequentialPipelineBlocks
 from ..modular_pipeline_utils import OutputParam
 from .before_denoise import (
     ErnieImagePrepareLatentsStep,
@@ -29,11 +29,11 @@
 
 
 # auto_docstring
-class ErnieImageAutoPromptEnhancerStep(AutoPipelineBlocks):
+class ErnieImageAutoPromptEnhancerStep(ConditionalPipelineBlocks):
     """
-    Auto block that runs the optional prompt enhancer when `use_pe` is provided.
-       - `ErnieImagePromptEnhancerStep` is used when `use_pe` is set.
-       - If `use_pe` is not provided, the step is skipped.
+    Conditional block that runs the optional prompt enhancer when `use_pe` is truthy.
+       - `ErnieImagePromptEnhancerStep` is used when `use_pe=True`.
+       - If `use_pe` is `None` or `False`, the step is skipped.
 
       Components:
           pe (`AutoModelForCausalLM`) pe_tokenizer (`AutoTokenizer`)
@@ -66,12 +66,17 @@ class ErnieImageAutoPromptEnhancerStep(AutoPipelineBlocks):
     block_names = ["prompt_enhancer"]
     block_trigger_inputs = ["use_pe"]
 
+    def select_block(self, use_pe=None) -> str | None:
+        if use_pe:
+            return "prompt_enhancer"
+        return None
+
     @property
     def description(self):
         return (
-            "Auto block that runs the optional prompt enhancer when `use_pe` is provided.\n"
-            " - `ErnieImagePromptEnhancerStep` is used when `use_pe` is set.\n"
-            " - If `use_pe` is not provided, the step is skipped."
+            "Conditional block that runs the optional prompt enhancer when `use_pe` is truthy.\n"
+            " - `ErnieImagePromptEnhancerStep` is used when `use_pe=True`.\n"
+            " - If `use_pe` is `None` or `False`, the step is skipped."
         )
 
 

src/diffusers/pipelines/ernie_image/pipeline_ernie_image.py

@@ -21,7 +21,7 @@
 
 import torch
 from PIL import Image
-from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, Ministral3ForCausalLM, Mistral3Model
 
 from ...models import AutoencoderKLFlux2
 from ...models.transformers import ErnieImageTransformer2DModel
@@ -51,10 +51,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__()
@@ -361,26 +361,26 @@ def __call__(
                 progress_bar.update()
 
         if output_type == "latent":
-            return latents
-
-        # Decode latents to images
-        # Unnormalize latents using VAE's BN stats
-        bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(device)
-        bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + 1e-5).to(device)
-        latents = latents * bn_std + bn_mean
+            images = latents
+        else:
+            # Decode latents to images
+            # Unnormalize latents using VAE's BN stats
+            bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(device)
+            bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + 1e-5).to(device)
+            latents = latents * bn_std + bn_mean
 
-        # Unpatchify
-        latents = self._unpatchify_latents(latents)
+            # Unpatchify
+            latents = self._unpatchify_latents(latents)
 
-        # Decode
-        images = self.vae.decode(latents, return_dict=False)[0]
+            # Decode
+            images = self.vae.decode(latents, return_dict=False)[0]
 
-        # Post-process
-        images = (images.clamp(-1, 1) + 1) / 2
-        images = images.cpu().permute(0, 2, 3, 1).float().numpy()
+            # Post-process
+            images = (images.clamp(-1, 1) + 1) / 2
+            images = images.cpu().permute(0, 2, 3, 1).float().numpy()
 
-        if output_type == "pil":
-            images = [Image.fromarray((img * 255).astype("uint8")) for img in images]
+            if output_type == "pil":
+                images = [Image.fromarray((img * 255).astype("uint8")) for img in images]
 
         # Offload all models
         self.maybe_free_model_hooks()