[diffusion] feature: use the directory for the vae mapping gate (#35946)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py

@@ -276,7 +276,12 @@ class VAELoader(ComponentLoader):
         keep_mapping = component_starts_on_cpu and (
             current_platform.is_mps()
             or keep_checkpoint_mapped(
-                weight_bytes=checkpoint_bytes(server_args.model_path),
+                # server_args.model_path can be a hub repo id, which is not a
+                # directory anywhere; the component path is always local, and
+                # its parent holds the rest of the variant being deployed.
+                weight_bytes=checkpoint_bytes(
+                    os.path.dirname(str(component_model_path))
+                ),
                 component=f"{component_name or 'vae'} (VAE)",
             )
         )

python/sglang/multimodal_gen/test/unit/test_vae_loader.py

@@ -1,3 +1,4 @@
+import pathlib
 import unittest
 from tempfile import TemporaryDirectory
 from unittest.mock import patch
@@ -24,7 +25,10 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.vae_loader import (
     _require_native_loader_for_quantized_vae,
     _should_use_channels_last_3d,
 )
-from sglang.multimodal_gen.runtime.loader.utils import keep_checkpoint_mapped
+from sglang.multimodal_gen.runtime.loader.utils import (
+    checkpoint_bytes,
+    keep_checkpoint_mapped,
+)
 from sglang.multimodal_gen.runtime.managers.memory_managers import (
     host_memory_budget,
 )
@@ -50,6 +54,29 @@ class _FakeServerArgs:
         return component_name in self.layerwise_components
 
 
+class TestDeploymentBytesRoot(unittest.TestCase):
+    """A hub repo id is not a directory; the component path always is."""
+
+    def test_the_component_parent_carries_the_variant_weight(self):
+        with TemporaryDirectory() as root:
+            variant = pathlib.Path(root) / "FL2VA"
+            (variant / "video_vae").mkdir(parents=True)
+            (variant / "transformer").mkdir()
+            (variant / "video_vae" / "w.safetensors").write_bytes(b"x" * 128)
+            (variant / "transformer" / "w.safetensors").write_bytes(b"x" * 512)
+            self.assertEqual(
+                checkpoint_bytes(str(variant)),
+                640,
+                "the parent of a component dir sums every sibling's shards",
+            )
+            self.assertEqual(
+                checkpoint_bytes("MiniMaxAI/MiniMax-H3"),
+                0,
+                "a repo id globs nothing -- which is why the gate must never "
+                "be fed one",
+            )
+
+
 class TestKeepCheckpointMapped(unittest.TestCase):
     """The mapping is for hosts that cannot afford the whole deployment."""