From 8cb9649382b34b1a81c2a4ecd4e0238e5e68a5d5 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Fri, 17 Apr 2026 12:01:09 -0600
Subject: [PATCH] Add decode latent to ernie pipe

---
 .../ernie_image/ernie_image.py                | 22 +++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/extensions_built_in/diffusion_models/ernie_image/ernie_image.py b/extensions_built_in/diffusion_models/ernie_image/ernie_image.py
index 665cd6f..d7177ad 100644
--- a/extensions_built_in/diffusion_models/ernie_image/ernie_image.py
+++ b/extensions_built_in/diffusion_models/ernie_image/ernie_image.py
@@ -26,8 +26,6 @@ except ImportError:
     )
 
 
-
-
 scheduler_config = {
     "base_image_seq_len": 256,
     "base_shift": 0.5,
@@ -239,6 +237,26 @@ class ErnieImageModel(BaseModel):
 
         return latents
 
+    def decode_latents(self, latents: torch.Tensor, device=None, dtype=None):
+        if self.vae.device == torch.device("cpu"):
+            self.vae.to(self.device_torch)
+        if device is None:
+            device = self.vae_device_torch
+        if dtype is None:
+            dtype = self.vae_torch_dtype
+
+        latents = latents.to(device, dtype=dtype)
+        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)
+
+        # Decode
+        images = self.vae.decode(latents, return_dict=False)[0]
+        return images
+
     def generate_single_image(
         self,
         pipeline: ErnieImagePipeline,