Add Boomer FLA checkpoint (step 055000, ema weights, bf16)

modeling_boomer_fla.py

@@ -1261,7 +1261,8 @@ class BoomerFLADiT(nn.Module):
         sd    = load_file(str(path / "diffusion_pytorch_model.safetensors"))
         model.load_state_dict(sd, strict=False)
 
-        # Attach inference metadata (latent stats, component repos, etc.)
-        # so BoomerPipeline.__init__ can read them without a separate config file.
+        # Attach inference metadata so BoomerPipeline.__init__ can read without a separate file.
         model._boomer_cfg = {k: v for k, v in cfg_raw.items() if k.startswith("_")}
+        # Store snapshot root so the pipeline can add it to sys.path for sibling imports.
+        model._snapshot_dir = str(path.parent if subfolder else path)
         return model

pipeline_boomer.py

@@ -112,6 +112,8 @@ class BoomerPipeline(DiffusionPipeline):
         self._vae_repo       = cfg.get("_vae_repo",  "mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers")
         self._te_repo        = cfg.get("_te_repo",   "google/gemma-4-E2B-it")
         self._hf_token       = None
+        # Snapshot root — added to sys.path so sibling .py files (scheduler, etc.) are importable
+        self._snapshot_dir   = getattr(transformer, "_snapshot_dir", None)
 
     # ── lazy component loading ─────────────────────────────────────────────────
     def _ensure_vae(self):
@@ -164,6 +166,7 @@ class BoomerPipeline(DiffusionPipeline):
 
         transformer = BoomerFLADiT.from_pretrained(str(local), subfolder="transformer")
         transformer = transformer.to(dtype=dtype)
+        transformer._snapshot_dir = str(local)   # carry snapshot path for sibling imports
 
         pipe = cls(transformer=transformer)
         pipe._hf_token = token
@@ -193,6 +196,10 @@ class BoomerPipeline(DiffusionPipeline):
         substeps     : STORK-2 internal RK micro-steps
         offload_text_encoder : unload text encoder after encoding to free VRAM
         """
+        # Add snapshot dir to sys.path so scheduling_boomer_stork (and STORKScheduler)
+        # are findable — diffusers only caches pipeline_boomer.py itself, not sibling files.
+        if self._snapshot_dir and self._snapshot_dir not in sys.path:
+            sys.path.insert(0, self._snapshot_dir)
         from scheduling_boomer_stork import make_stork_scheduler  # noqa: PLC0415
 
         prompts = [prompt] if isinstance(prompt, str) else list(prompt)

transformer/modeling_boomer_fla.py

@@ -1261,7 +1261,8 @@ class BoomerFLADiT(nn.Module):
         sd    = load_file(str(path / "diffusion_pytorch_model.safetensors"))
         model.load_state_dict(sd, strict=False)
 
-        # Attach inference metadata (latent stats, component repos, etc.)
-        # so BoomerPipeline.__init__ can read them without a separate config file.
+        # Attach inference metadata so BoomerPipeline.__init__ can read without a separate file.
         model._boomer_cfg = {k: v for k, v in cfg_raw.items() if k.startswith("_")}
+        # Store snapshot root so the pipeline can add it to sys.path for sibling imports.
+        model._snapshot_dir = str(path.parent if subfolder else path)
         return model