Update pipeline/pipeline.py

pipeline/pipeline.py

@@ -1,5 +1,3 @@
-# pipeline/pipeline.py
-
 import torch
 from diffusers import StableDiffusionPipeline
 
@@ -26,28 +24,38 @@ class OmegaDiffusionPipeline(StableDiffusionPipeline):
             feature_extractor=feature_extractor,
             **kwargs,
         )
-        # register your custom bridge *as a module*
         self.register_modules(bridge=bridge)
 
+        # --- Monkey-patch the VAE.decode method ---
+        orig_decode = self.vae.decode
+
+        def decode_with_bridge(z_scaled, *args, return_dict=False, generator=None, **decode_kwargs):
+            # z_scaled = z4 / scaling_factor
+            sc = self.vae.config.scaling_factor
+            z4 = z_scaled * sc                  # recover 4-channel latent
+            z16 = self.bridge.dec(z4)           # map 4→16 channels
+            z16_scaled = z16 / sc               # rescale for VAE
+            # call original decoder
+            return orig_decode(z16_scaled, *args, return_dict=return_dict, generator=generator, **decode_kwargs)
+
+        # Replace VAE.decode at runtime
+        self.vae.decode = decode_with_bridge
+
     @property
     def components(self):
-        # Return exactly the expected components (no extras!)
         return {
-            "scheduler":        self.scheduler,
-            "tokenizer":        self.tokenizer,
-            "vae":              self.vae,
-            "unet":             self.unet,
-            "text_encoder":     self.text_encoder,
-            "bridge":           self.bridge,
-            "feature_extractor":self.feature_extractor,
-            "safety_checker":   self.safety_checker,
-            # the loader also expects a 'kwargs' slot
-            "kwargs":           {},
+            "scheduler":         self.scheduler,
+            "tokenizer":         self.tokenizer,
+            "vae":               self.vae,
+            "unet":              self.unet,
+            "text_encoder":      self.text_encoder,
+            "bridge":            self.bridge,
+            "feature_extractor": self.feature_extractor,
+            "safety_checker":    self.safety_checker,
+            "kwargs":            {},
         }
 
     @torch.no_grad()
-    def decode_latents(self, latents):
-        # your custom decode logic
-        lat16  = self.bridge.dec(latents)
-        images = self.vae.decode(lat16 / self.vae.config.scaling_factor).sample
-        return (images.clamp(-1, 1) + 1) / 2
+    def __call__(self, *args, **kwargs):
+        # let the parent handle everything, including calling vae.decode (now patched)
+        return super().__call__(*args, **kwargs)