From 0f99fce0042089a727f069aaedd2bed0743345f5 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 16 Apr 2025 09:24:23 -0600 Subject: [PATCH] Adjust hidream lora names to work with comfy --- .../diffusion_models/hidream/hidream_model.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/extensions_built_in/diffusion_models/hidream/hidream_model.py b/extensions_built_in/diffusion_models/hidream/hidream_model.py index 3322a111..35b4a192 100644 --- a/extensions_built_in/diffusion_models/hidream/hidream_model.py +++ b/extensions_built_in/diffusion_models/hidream/hidream_model.py @@ -426,5 +426,20 @@ class HidreamModel(BaseModel): def get_transformer_block_names(self) -> Optional[List[str]]: return ['double_stream_blocks', 'single_stream_blocks'] + + def convert_lora_weights_before_save(self, state_dict): + # currently starte with transformer. but needs to start with diffusion_model. for comfyui + new_sd = {} + for key, value in state_dict.items(): + new_key = key.replace("transformer.", "diffusion_model.") + new_sd[new_key] = value + return new_sd + + def convert_lora_weights_before_load(self, state_dict): + # saved as diffusion_model. but needs to be transformer. for ai-toolkit + new_sd = {} + for key, value in state_dict.items(): + new_key = key.replace("diffusion_model.", "transformer.") + new_sd[new_key] = value + return new_sd - \ No newline at end of file