Use VaeImageProcessor.postprocess in standard and modular ernie

Author: akshan-main

src/diffusers/modular_pipelines/ernie_image/decoders.py

@@ -12,11 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import numpy as np
 import torch
-from PIL import Image
 
 from ...configuration_utils import FrozenDict
+from ...image_processor import VaeImageProcessor
 from ...models import AutoencoderKLFlux2
 from ...utils import logging
 from ..modular_pipeline import ModularPipelineBlocks, PipelineState
@@ -44,6 +43,12 @@ def expected_components(self) -> list[ComponentSpec]:
                 config=FrozenDict({"patch_size": 2}),
                 default_creation_method="from_config",
             ),
+            ComponentSpec(
+                "image_processor",
+                VaeImageProcessor,
+                config=FrozenDict({"vae_scale_factor": 16}),
+                default_creation_method="from_config",
+            ),
         ]
 
     @property
@@ -81,18 +86,7 @@ def __call__(self, components: ErnieImageModularPipeline, state: PipelineState)
         latents = components.pachifier.unpack_latents(latents)
 
         images = vae.decode(latents.to(vae.dtype), return_dict=False)[0]
-        images = (images.clamp(-1, 1) + 1) / 2
-
-        output_type = block_state.output_type
-        if output_type == "pt":
-            block_state.images = images
-        elif output_type == "np":
-            block_state.images = images.cpu().permute(0, 2, 3, 1).float().numpy()
-        elif output_type == "pil":
-            images_np = images.cpu().permute(0, 2, 3, 1).float().numpy()
-            block_state.images = [Image.fromarray((img * 255).astype(np.uint8)) for img in images_np]
-        else:
-            raise ValueError(f"Unsupported `output_type`: {output_type!r}. Expected one of 'pil', 'np', 'pt'.")
+        block_state.images = components.image_processor.postprocess(images, output_type=block_state.output_type)
 
         self.set_block_state(state, block_state)
         return components, state

src/diffusers/pipelines/ernie_image/pipeline_ernie_image.py

@@ -20,9 +20,9 @@
 from typing import Callable, List, Optional, Union
 
 import torch
-from PIL import Image
 from transformers import AutoTokenizer, Ministral3ForCausalLM, Mistral3Model
 
+from ...image_processor import VaeImageProcessor
 from ...models import AutoencoderKLFlux2
 from ...models.transformers import ErnieImageTransformer2DModel
 from ...pipelines.pipeline_utils import DiffusionPipeline
@@ -68,6 +68,7 @@ def __init__(
             pe_tokenizer=pe_tokenizer,
         )
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels)) if getattr(self, "vae", None) else 16
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
 
     @property
     def guidance_scale(self):
@@ -379,11 +380,7 @@ def __call__(
             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()
-
-            if output_type == "pil":
-                images = [Image.fromarray((img * 255).astype("uint8")) for img in images]
+            images = self.image_processor.postprocess(images, output_type=output_type)
 
         # Offload all models
         self.maybe_free_model_hooks()