mirror of
https://github.com/modelscope/DiffSynth-Studio.git
synced 2026-04-02 07:18:10 +00:00
16 lines
531 B
Python
16 lines
531 B
Python
def AceStepDiTStateDictConverter(state_dict):
|
|
"""
|
|
Convert ACE-Step DiT state dict to add 'model.' prefix for wrapper class.
|
|
|
|
The wrapper class has self.model = AceStepConditionGenerationModel(config),
|
|
so all keys need to be prefixed with 'model.'
|
|
"""
|
|
state_dict_ = {}
|
|
keys = state_dict.keys() if hasattr(state_dict, 'keys') else state_dict
|
|
for k in keys:
|
|
v = state_dict[k]
|
|
if not k.startswith("model."):
|
|
k = "model." + k
|
|
state_dict_[k] = v
|
|
return state_dict_
|