chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[diffusion] feat: support distilled vae generic (#14195)

Browse Source 
Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Dongjie Zou
2025-12-02 21:27:31 -05:00
committed by GitHub
parent 922054079c
commit f764c6910d
10 changed files with 302 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
python/sglang/multimodal_gen/runtime/loader/component_loader.py
View File

@@ -4,10 +4,11 @@
import dataclasses
import glob
import importlib.util
import json
import os
import time
from abc import ABC, abstractmethod
from abc import ABC
from collections.abc import Generator, Iterable
from copy import deepcopy
from typing import Any, cast
@@ -101,17 +102,13 @@ def load_native(library, component_module_path: str, server_args: ServerArgs):
revision=server_args.revision,
)
elif library == "diffusers":
import diffusers
from diffusers import AutoModel
config = get_diffusers_component_config(model_path=component_module_path)
class_name = config.pop("_class_name", None)
if class_name:
cls = getattr(diffusers, class_name)
return cls.from_pretrained(
component_module_path, revision=server_args.revision, **config
)
else:
raise ValueError("Cannot determine class name for generic diffusers loader")
return AutoModel.from_pretrained(
component_module_path,
revision=server_args.revision,
trust_remote_code=server_args.trust_remote_code,
)
else:
raise ValueError(f"Unsupported library: {library}")
@@ -202,13 +199,6 @@ class ComponentLoader(ABC):
f"load_customized not implemented for {self.__class__.__name__}"
)
@abstractmethod
def load_customized(
self, model_path: str, server_args: ServerArgs, module_name: str
) -> Any:
"""Implement the minimal core load logic in subclasses."""
raise NotImplementedError
@classmethod
def for_module_type(
cls, module_type: str, transformers_or_diffusers: str
@@ -503,7 +493,6 @@ class ImageEncoderLoader(TextEncoderLoader):
encoder_config.update_model_arch(model_config)
# Always start with local device; load_model will adjust for offload if needed
should_offload = self.should_offload(server_args)
# TODO(will): add support for other dtypes
return self.load_model(
component_model_path,
@@ -555,7 +544,7 @@ class VAELoader(ComponentLoader):
):
"""Load the VAE based on the model path, and inference args."""
config = get_diffusers_component_config(model_path=component_model_path)
class_name = config.pop("_class_name")
class_name = config.pop("_class_name", None)
assert (
class_name is not None
), "Model config does not contain a _class_name attribute. Only diffusers format is supported."
@@ -571,23 +560,42 @@ class VAELoader(ComponentLoader):
target_device = self.target_device(server_args.vae_cpu_offload)
with set_default_torch_dtype(
PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
), skip_init_modules():
# Check for auto_map first (custom VAE classes)
auto_map = config.get("auto_map", {})
auto_model_map = auto_map.get("AutoModel")
if auto_model_map:
module_path, cls_name = auto_model_map.rsplit(".", 1)
custom_module_file = os.path.join(component_model_path, f"{module_path}.py")
spec = importlib.util.spec_from_file_location("_custom", custom_module_file)
custom_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(custom_module)
vae_cls = getattr(custom_module, cls_name)
vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
with set_default_torch_dtype(vae_dtype):
vae = vae_cls.from_pretrained(
component_model_path,
revision=server_args.revision,
trust_remote_code=server_args.trust_remote_code,
)
vae = vae.to(device=target_device, dtype=vae_dtype)
return vae.eval()
# Load from ModelRegistry (standard VAE classes)
with (
set_default_torch_dtype(
PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
),
skip_init_modules(),
):
vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
vae = vae_cls(vae_config).to(target_device)
# Find all safetensors files
safetensors_list = _list_safetensors_files(component_model_path)
# TODO(PY)
assert (
len(safetensors_list) == 1
), f"Found {len(safetensors_list)} safetensors files in {component_model_path}"
loaded = safetensors_load_file(safetensors_list[0])
vae.load_state_dict(
loaded, strict=False
) # We might only load encoder or decoder
vae.load_state_dict(loaded, strict=False)
return vae.eval()