[diffusion] model: support ltx-2.5 (#34471)

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

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

@@ -5,6 +5,7 @@
 
 import contextlib
 import glob
+import json
 import os
 import re
 from collections import defaultdict
@@ -12,6 +13,7 @@ from collections.abc import Callable, Iterator
 from typing import Any, Dict, Type
 
 import torch
+from safetensors.torch import load_file as safetensors_load_file
 from torch import nn
 
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -260,8 +262,6 @@ def _list_safetensors_files(model_path: str) -> list[str]:
         str(model_path), "diffusion_pytorch_model.safetensors.index.json"
     )
     if os.path.exists(index_path):
-        import json
-
         with open(index_path) as f:
             index = json.load(f)
         expected_shards = sorted(set(index.get("weight_map", {}).values()))
@@ -284,6 +284,33 @@ def _list_safetensors_files(model_path: str) -> list[str]:
     return found
 
 
+def load_safetensors_state_dict(model_path: str) -> dict[str, torch.Tensor]:
+    """Load one safetensors checkpoint, including an indexed sharded set."""
+    index_path = os.path.join(
+        str(model_path), "diffusion_pytorch_model.safetensors.index.json"
+    )
+    safetensors_files = _list_safetensors_files(model_path)
+    if os.path.exists(index_path):
+        with open(index_path) as f:
+            index = json.load(f)
+        shard_names = sorted(set(index.get("weight_map", {}).values()))
+        state_dict: dict[str, torch.Tensor] = {}
+        for shard_name in shard_names:
+            state_dict.update(
+                safetensors_load_file(os.path.join(str(model_path), shard_name))
+            )
+        return state_dict
+
+    if not safetensors_files:
+        raise ValueError(f"No safetensors files found in {model_path}")
+    if len(safetensors_files) != 1:
+        raise ValueError(
+            f"Found {len(safetensors_files)} safetensors files in {model_path} "
+            "and no index to disambiguate them."
+        )
+    return safetensors_load_file(safetensors_files[0])
+
+
 BYTES_PER_GB = 1024**3