Fix RNA conditioning, MUSK encoder init, and text/inference mismatch

pipeline.py

@@ -7,29 +7,17 @@ import math
 import numpy as np
 from PIL import Image
 from diffusers import DiffusionPipeline, AutoencoderKL
-from dictdot import dictdot
 import gc
 
-from models.autoencoder import vae_models
-from samplers import euler_sampler, euler_maruyama_sampler
+# NOTE: the repo's own modules (samplers.py, models/, musk/) are deliberately NOT
+# imported at module scope. When loaded via custom_pipeline, diffusers copies only
+# this file into its modules cache, so those imports would fail for anyone whose
+# working directory doesn't happen to contain them. They are imported lazily once
+# _snapshot_root() has put the downloaded snapshot on sys.path.
 
 def denormalize_latents(latents, latents_scale, latents_bias):
     return latents / latents_scale + latents_bias
 
-try:
-    from models.sit import SiT_models
-except ImportError:
-    # Fallback for published structure where sit is in transformer/
-    from transformer.sit import SiT_models
-
-try:
-    from models.projection_loss import CosineProjectionLoss
-except ImportError:
-    try:
-        from transformer.projection_loss import CosineProjectionLoss
-    except ImportError:
-        pass # Might not be needed for inference
-
 # Must match the training-time order (train.py / generate.py); the checkpoint's
 # modality_embedding has one row per entry.
 MODALITY_ORDER = ("text", "image", "st", "st_meta")
@@ -69,7 +57,19 @@ class SiTPipeline(DiffusionPipeline):
         else:
             # Legacy/Experiment loading logic
             device = kwargs.get("device", torch.device("cuda" if torch.cuda.is_available() else "cpu"))
-            
+
+            # This layout builds the model from source, so make the bundled modules importable.
+            legacy_root = os.path.abspath(pretrained_model_path)
+            if legacy_root not in sys.path:
+                sys.path.append(legacy_root)
+            from dictdot import dictdot
+            from models.autoencoder import vae_models
+            try:
+                from models.sit import SiT_models
+            except ImportError:
+                # Published structure keeps sit in transformer/
+                from transformer.sit import SiT_models
+
             # 1. Determine Layout and Load Config
             # Check for published config first (sit_config.json)
             published_config_path = os.path.join(pretrained_model_path, "sit_config.json")
@@ -306,32 +306,56 @@ class SiTPipeline(DiffusionPipeline):
         ])
         print("MUSK encoder loaded successfully.")
 
-    def _ensure_musk_encoder(self):
+    def _snapshot_root(self):
         """
-        Load MUSK on first use when the pipeline was built by the standard
-        DiffusionPipeline.from_pretrained route, which bypasses our from_pretrained.
+        Local directory holding the bundled code and assets (samplers.py, models/,
+        musk/, musk_weights/), with that directory added to sys.path.
+
+        Under custom_pipeline, only pipeline.py lands in the diffusers modules cache,
+        so the rest has to be located explicitly rather than imported as siblings.
         """
-        if self.musk_encoder is not None:
-            return
+        cached = getattr(self, "_resolved_root", None)
+        if cached:
+            return cached
 
         # diffusers records where the pipeline came from after instantiation.
-        model_root = getattr(self.config, "_name_or_path", None)
-        if not model_root:
+        root = getattr(self.config, "_name_or_path", None)
+        if not root:
             raise RuntimeError(
-                "MUSK encoder not initialized and the model path is unknown. "
-                "Call pipeline.load_musk_encoder('/path/to/snapshot') explicitly."
+                "Cannot locate the model snapshot (no _name_or_path). Load the pipeline "
+                "with SiTPipeline.from_pretrained('/path/to/snapshot')."
             )
 
-        if not os.path.isdir(model_root):
-            # It's a Hub repo id; fetch the assets the encoder needs.
+        if not os.path.isdir(root):
+            # It's a Hub repo id; fetch the code/assets (weights are already cached
+            # by from_pretrained, and this resolves to the same snapshot directory).
             from huggingface_hub import snapshot_download
-            model_root = snapshot_download(
-                repo_id=model_root,
+            root = snapshot_download(
+                repo_id=root,
                 repo_type="model",
-                allow_patterns=["musk/**", "musk_weights/**"],
+                allow_patterns=["*.py", "models/**", "musk/**", "musk_weights/**"],
             )
 
-        self.load_musk_encoder(model_root, device=self.device)
+        root = os.path.abspath(root)
+        if root not in sys.path:
+            sys.path.append(root)
+        self._resolved_root = root
+        return root
+
+    def _load_samplers(self):
+        self._snapshot_root()
+        from samplers import euler_sampler, euler_maruyama_sampler
+        return euler_sampler, euler_maruyama_sampler
+
+    def _ensure_musk_encoder(self):
+        """
+        Load MUSK on first use when the pipeline was built by the standard
+        DiffusionPipeline.from_pretrained route, which bypasses our from_pretrained.
+        """
+        if self.musk_encoder is not None:
+            return
+
+        self.load_musk_encoder(self._snapshot_root(), device=self.device)
 
     def to(self, *args, **kwargs):
         super().to(*args, **kwargs)
@@ -437,6 +461,8 @@ class SiTPipeline(DiffusionPipeline):
             cls_latents=cls_latents,
         )
         
+        euler_sampler, euler_maruyama_sampler = self._load_samplers()
+
         if mode == "sde":
             latents = euler_maruyama_sampler(**sampling_kwargs).to(torch.float32)
         elif mode == "ode":