add acestep models

This commit is contained in:
Artiprocher
2026-04-02 10:58:45 +08:00
parent 166e6d2d38
commit 59f512b574
8 changed files with 3169 additions and 1 deletions

View File

@@ -0,0 +1,15 @@
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_