From a8c16b2e5575790f9d6c3db7e44ec1edf1ef71f1 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 22 Aug 2026 19:12:48 +0800 Subject: [PATCH] [diffusion] feature: use the directory for the vae mapping gate (#35946) Co-authored-by: Claude Opus 5 --- .../loader/component_loaders/vae_loader.py | 7 ++++- .../test/unit/test_vae_loader.py | 29 ++++++++++++++++++- 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index 94cda9011..bafb45f2e 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/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)", ) ) diff --git a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py index bd0dbc828..02927db63 100644 --- a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py +++ b/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."""