diff --git a/toolkit/models/v2/PLANNING.md b/toolkit/models/v2/PLANNING.md index 0048e19..7b17bef 100644 --- a/toolkit/models/v2/PLANNING.md +++ b/toolkit/models/v2/PLANNING.md @@ -382,8 +382,12 @@ Decisions: raw `quantize()` call sites and silently dropped by `quantize_model()`; the ARA path inside `quantize_model` hardcodes `uint8`. Decide the unified behavior when consolidating. -- [ ] `toolkit/models/loaders/umt5.py` dead `comfy_files` param (wan21 comfy-TE - no-op) — fix when wan migrates. +- [x] wan comfy-TE resolved for real: UMT5TextEncoder carries comfy + candidates (fp8_e4m3fn_scaled via the legacy importer, fp16), files + already in transformers key layout (spiece blob dropped, tied + embed_tokens materialized). Verified: wan21 samples with the local + comfy fp8 TE. The loaders/umt5.py `comfy_files` param stays as a + no-op shim for old callers. - [ ] Registry hardening: error (don't fall back to SD1) on unknown arch; lazy per-arch imports; single source of truth shared with the UI's `options.tsx` model list. diff --git a/toolkit/models/v2/text_encoders/umt5.py b/toolkit/models/v2/text_encoders/umt5.py index d0b2e99..8bd02e0 100644 --- a/toolkit/models/v2/text_encoders/umt5.py +++ b/toolkit/models/v2/text_encoders/umt5.py @@ -33,11 +33,36 @@ class UMT5TextEncoder(UMT5EncoderModel, OstrisTransformersMixin): aitk_subfolder = "text_encoder" aitk_tokenizer_subfolder = "tokenizer" + aitk_config_repo = "ai-toolkit/umt5_xxl_encoder" + + aitk_comfy_repo = "Comfy-Org/Wan_2.1_ComfyUI_repackaged" + # comfy umt5 files already use the transformers key layout; the + # fp8_e4m3fn_scaled variant is the legacy scaled-fp8 format (handled by + # the importer's float8 backend) + aitk_comfy_weight_names = { + "ai-toolkit/umt5_xxl_encoder": [ + "split_files/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors", + "split_files/text_encoders/umt5_xxl_fp16.safetensors", + ], + } @classmethod def get_transformer_block_names(cls): return ["encoder.block"] + @classmethod + def convert_state_dict_on_load(cls, state_dict): + # drop the embedded sentencepiece blob and materialize the tied + # embed_tokens reference + state_dict = dict(state_dict) + state_dict.pop("spiece_model", None) + if ( + "encoder.embed_tokens.weight" not in state_dict + and "shared.weight" in state_dict + ): + state_dict["encoder.embed_tokens.weight"] = state_dict["shared.weight"] + return state_dict + @classmethod def load_tokenizer(cls, name_or_path=None, subfolder=None, **kwargs): # T5's tokenizer needs the _spm_precompiled_charsmap patch