Update pipeline/pipeline.py

pipeline/pipeline.py

@@ -1,61 +1,58 @@
+# pipeline/pipeline.py
+
 import torch
 from diffusers import StableDiffusionPipeline
 
 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)
         self.register_modules(bridge=bridge)
 
-        # --- Monkey-patch the VAE.decode method ---
-        orig_decode = self.vae.decode
+        # Capture the original decode
+        _orig_decode = self.vae.decode
+        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:
+                z = self.bridge.dec(z)
+
+            # Rescale for the VAE’s conv_in
+            z_scaled2 = z / sc
 
-        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)
+            # Now call the real VAE.decode on exactly 16 channels
+            return _orig_decode(z_scaled2, *args,
+                                 return_dict=return_dict,
+                                 generator=generator,
+                                 **decode_kwargs)
 
-        # Replace VAE.decode at runtime
-        self.vae.decode = decode_with_bridge
+        # Override the instance’s decode method
+        self.vae.decode = _decode_with_bridge
 
     @property
     def components(self):
         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 __call__(self, *args, **kwargs):
-        # let the parent handle everything, including calling vae.decode (now patched)
+        # Defer entirely to the parent, which will call our patched decode
         return super().__call__(*args, **kwargs)