incline path finding

python/sglang/multimodal_gen/runtime/loader/utils.py

@@ -9,7 +9,6 @@
 import re
 from collections import defaultdict
 from collections.abc import Callable, Iterator
-from pathlib import Path
 from typing import Any, Dict, Type
 
 import torch
@@ -149,18 +148,6 @@ def _list_safetensors_files(model_path: str) -> list[str]:
     return sorted(glob.glob(os.path.join(str(model_path), "*.safetensors")))
 
 
-def find_weights_dir(local_path: str, module_name: str) -> Path | None:
-    """Locate the safetensors directory for module_name under local_path.
-
-    Diffusion models store weights in per-module subdirectories (e.g.
-    transformer/, vae/, text_encoder/).
-    """
-    dir_path = Path(local_path) / module_name
-    if dir_path.exists():
-        return dir_path
-    return None
-
-
 def get_memory_usage_of_component(module) -> float | None:
     """
     returned value is in GB, rounded to 2 decimal digits

python/sglang/multimodal_gen/runtime/loader/weights_updater.py

@@ -50,14 +50,14 @@
 from __future__ import annotations
 
 import gc
+from pathlib import Path
 
 import torch
 from torch.distributed.tensor import DTensor, distribute_tensor
 
 from sglang.multimodal_gen.runtime.cache.teacache import TeaCacheMixin
 from sglang.multimodal_gen.runtime.loader.utils import (
     _list_safetensors_files,
-    find_weights_dir,
 )
 from sglang.multimodal_gen.runtime.loader.weight_utils import (
     safetensors_weights_iterator,
@@ -108,9 +108,9 @@ def _validate_weight_files(
     weights_map: dict[str, str] = {}
     missing: list[str] = []
     for module_name, _ in modules_to_update:
-        weights_dir = find_weights_dir(local_model_path, module_name)
-        if weights_dir and _list_safetensors_files(weights_dir):
-            weights_map[module_name] = weights_dir
+        weights_dir = Path(local_model_path) / module_name
+        if weights_dir.exists() and _list_safetensors_files(str(weights_dir)):
+            weights_map[module_name] = str(weights_dir)
         else:
             missing.append(module_name)
     return weights_map, missing
@@ -306,8 +306,8 @@ def _rollback(self, updated_modules: list[str]) -> None:
             module = self.pipeline.get_module(name)
             if module is None:
                 continue
-            weights_dir = find_weights_dir(original_path, name)
-            if weights_dir is None:
+            weights_dir = Path(original_path) / name
+            if not weights_dir.exists():
                 continue
-            weights_iter = _get_weights_iter(weights_dir)
+            weights_iter = _get_weights_iter(str(weights_dir))
             _load_weights_into_module(module, weights_iter)

python/sglang/multimodal_gen/test/server/test_update_weights_from_disk.py

@@ -149,7 +149,6 @@
 
 from sglang.multimodal_gen.runtime.loader.utils import (
     _list_safetensors_files,
-    find_weights_dir,
 )
 from sglang.multimodal_gen.runtime.loader.weight_utils import (
     compute_weights_checksum,
@@ -223,8 +222,10 @@ def _compute_checksum_from_disk(model_path: str, module_name: str) -> str:
     same disk checksum is requested multiple times across tests.
     """
     local_path = maybe_download_model(model_path)
-    weights_dir = find_weights_dir(local_path, module_name)
-    assert weights_dir is not None, f"No weights dir for {module_name} in {local_path}"
+    weights_dir = os.path.join(local_path, module_name)
+    assert os.path.exists(
+        weights_dir
+    ), f"No weights dir for {module_name} in {local_path}"
 
     safetensors_files = _list_safetensors_files(weights_dir)
     assert safetensors_files, f"No safetensors files in {weights_dir}"