Update pipeline/pipeline.py

pipeline/pipeline.py

@@ -2,57 +2,100 @@
 
 import torch
 from diffusers import StableDiffusionPipeline
+from diffusers.utils import BaseOutput
 
 class OmegaDiffusionPipeline(StableDiffusionPipeline):
-    def __init__(self, vae, text_encoder, tokenizer, unet, scheduler, bridge,
-                 safety_checker=None, feature_extractor=None, **kwargs):
-        super().__init__(vae=vae, text_encoder=text_encoder, tokenizer=tokenizer,
-                         unet=unet, scheduler=scheduler,
-                         safety_checker=safety_checker,
-                         feature_extractor=feature_extractor, **kwargs)
+    def __init__(
+        self,
+        vae,
+        text_encoder,
+        tokenizer,
+        unet,
+        scheduler,
+        bridge,
+        safety_checker=None,
+        feature_extractor=None,
+        **kwargs,
+    ):
+        super().__init__(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            unet=unet,
+            scheduler=scheduler,
+            safety_checker=safety_checker,
+            feature_extractor=feature_extractor,
+            **kwargs,
+        )
+        # register your bridge so Diffusers knows about it
         self.register_modules(bridge=bridge)
 
-        # Capture the original decode
+        # ─── Monkey-patch the VAE.decode to insert your bridge ───
         _orig_decode = self.vae.decode
-        in_ch = self.unet.config.in_channels
-        sc   = self.vae.config.scaling_factor
+        in_ch        = self.unet.config.in_channels
+        sc           = self.vae.config.scaling_factor
 
-        def _decode_with_bridge(z_scaled, *args, return_dict=False, generator=None, **decode_kwargs):
-            # z_scaled is latents/scaling_factor
-            # Reconstruct the raw latent
-            z = z_scaled * sc
-
-            # If it’s 4-ch, run the bridge → 16-ch
-            if z.shape[1] == in_ch:
+        def _decode_with_bridge(
+            z_scaled, *args, return_dict=False, generator=None, **decode_kwargs
+        ):
+            # z_scaled = latents / scaling_factor
+            z = z_scaled * sc                      # back to raw latent
+            if z.shape[1] == in_ch:               # 4→16 only when needed
                 z = self.bridge.dec(z)
+            z = z / sc                             # scale again
+            # call the real decode
+            out = _orig_decode(
+                z,
+                *args,
+                return_dict=return_dict,
+                generator=generator,
+                **decode_kwargs
+            )
+            return out
 
-            # Rescale for the VAE’s conv_in
-            z_scaled2 = z / sc
-
-            # Now call the real VAE.decode on exactly 16 channels
-            return _orig_decode(z_scaled2, *args,
-                                 return_dict=return_dict,
-                                 generator=generator,
-                                 **decode_kwargs)
-
-        # Override the instance’s decode method
+        # override it in place
         self.vae.decode = _decode_with_bridge
 
     @property
     def components(self):
+        # ensure Diffusers sees exactly the modules you expect
         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,
-            "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):
+        """
+        The single hook that StableDiffusionPipeline.__call__
+        uses to turn final latents → images.
+        """
+        # calling self.vae.decode here actually invokes your _decode_with_bridge
+        decoded = self.vae.decode(latents, return_dict=False)
+        images  = decoded[0] if isinstance(decoded, (tuple, list)) else decoded
+        # normalize to [0,1]
+        return (images.clamp(-1, 1) + 1) / 2
+
     @torch.no_grad()
     def __call__(self, *args, **kwargs):
-        # Defer entirely to the parent, which will call our patched decode
+        # defer everything to the parent implementation, which
+        # will in turn call our patched _decode_latents
         return super().__call__(*args, **kwargs)
+
+    @torch.no_grad()
+    def decode_latents(self, latents, return_dict=True):
+        """
+        If you ever call pipe.decode_latents(...) manually,
+        this will route through the same bridge logic.
+        """
+        imgs = self._decode_latents(latents)
+        if return_dict:
+            return BaseOutput(images=imgs)
+        return imgs