diff --git a/modules/sdnq/loader.py b/modules/sdnq/loader.py index 9af93227f..d808d67ff 100644 --- a/modules/sdnq/loader.py +++ b/modules/sdnq/loader.py @@ -130,13 +130,18 @@ def load_sdnq_model(model_path: str, model_cls: ModelMixin = None, file_name: st state_dict = load_files(files, key_mapping=key_mapping, device=device, method=load_method) - if model.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"} and "encoder.embed_tokens.weight" not in state_dict.keys(): - state_dict["encoder.embed_tokens.weight"] = state_dict["shared.weight"] - - # Handle Qwen models with tied lm_head weights - if model.__class__.__name__ in {"Qwen3ForCausalLM"} and "lm_head.weight" not in state_dict.keys(): - if "model.embed_tokens.weight" in state_dict.keys(): - state_dict["lm_head.weight"] = state_dict["model.embed_tokens.weight"] + if isinstance(getattr(model, "_tied_weights_keys", None), dict): + for key, value in model._tied_weights_keys.items(): + print(key, value) + if value in state_dict.keys() and key not in state_dict.keys(): + state_dict[key] = state_dict[value] + else: + # older transformers case, handle known models manually + if model.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"} and "encoder.embed_tokens.weight" not in state_dict.keys(): + state_dict["encoder.embed_tokens.weight"] = state_dict["shared.weight"] + elif model.__class__.__name__ in {"Qwen3ForCausalLM"} and "lm_head.weight" not in state_dict.keys(): + if "model.embed_tokens.weight" in state_dict.keys(): + state_dict["lm_head.weight"] = state_dict["model.embed_tokens.weight"] model.load_state_dict(state_dict, assign=True) del state_dict