[diffusion] server: use meta to avoid Linear init for TextEncoder (#13564)
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
@@ -46,6 +46,20 @@ from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
|
|||||||
logger = init_logger(__name__)
|
logger = init_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class skip_init_modules:
|
||||||
|
def __enter__(self):
|
||||||
|
# Save originals
|
||||||
|
self._orig_reset = {}
|
||||||
|
for cls in (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d):
|
||||||
|
self._orig_reset[cls] = cls.reset_parameters
|
||||||
|
cls.reset_parameters = lambda self: None # skip init
|
||||||
|
|
||||||
|
def __exit__(self, exc_type, exc_value, traceback):
|
||||||
|
# Restore originals
|
||||||
|
for cls, orig in self._orig_reset.items():
|
||||||
|
cls.reset_parameters = orig
|
||||||
|
|
||||||
|
|
||||||
class ComponentLoader(ABC):
|
class ComponentLoader(ABC):
|
||||||
"""Base class for loading a specific type of model component."""
|
"""Base class for loading a specific type of model component."""
|
||||||
|
|
||||||
@@ -287,7 +301,7 @@ class TextEncoderLoader(ComponentLoader):
|
|||||||
)
|
)
|
||||||
|
|
||||||
with set_default_torch_dtype(PRECISION_TO_TYPE[dtype]):
|
with set_default_torch_dtype(PRECISION_TO_TYPE[dtype]):
|
||||||
with target_device:
|
with target_device, skip_init_modules():
|
||||||
architectures = getattr(model_config, "architectures", [])
|
architectures = getattr(model_config, "architectures", [])
|
||||||
model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
|
model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
|
||||||
model = model_cls(model_config)
|
model = model_cls(model_config)
|
||||||
@@ -454,7 +468,7 @@ class VAELoader(ComponentLoader):
|
|||||||
|
|
||||||
with set_default_torch_dtype(
|
with set_default_torch_dtype(
|
||||||
PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
|
PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
|
||||||
):
|
), skip_init_modules():
|
||||||
vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
|
vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
|
||||||
vae = vae_cls(vae_config).to(target_device)
|
vae = vae_cls(vae_config).to(target_device)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user