replace 'diffusion_model' with 'transformer'

This commit is contained in:
rickard 2024-12-25 18:04:56 +01:00
parent 21f20638bd
commit ceee378d86

View File

@ -401,9 +401,9 @@ def model_lora_keys_unet(model, key_map={}):
if isinstance(model, comfy.model_base.LTXV):
for k in sdk:
if k.startswith("transformer.") and k.endswith(".weight"): #Official Mochi lora format
key_lora = k[len("transformer."):-len(".weight")]
key_map["{}".format(key_lora)] = k
if k.startswith("diffusion_model.") and k.endswith(".weight"):
key_lora = k[len("diffusion_model."):-len(".weight")]
key_map["transformer.{}".format(key_lora)] = k
return key_map