diff --git a/tools/convert.py b/tools/convert.py index 61e1592..001cd86 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -55,7 +55,7 @@ class ModelHyVid(ModelTemplate): ] # def handle_nd_tensor(self, key, data): - # # hacky but this is the only arch that uses it + # # hacky but don't have any better ideas # path = f"./fix_5d_tensors_{self.arch}.pt" # if os.path.isfile(path): # raise RuntimeError(f"5D tensor fix file already exists! {path}") @@ -63,6 +63,16 @@ class ModelHyVid(ModelTemplate): # tqdm.write(f"5D key found in state dict! Manual fix required! - {key} {data.shape}") # torch.save(fsd, path) +class ModelWan(ModelHyVid): + arch = "wan" + keys_detect = [ + ( + "blocks.0.self_attn.norm_q.weight", + "text_embedding.2.weight", + "head.modulation", + ) + ] + class ModelLTXV(ModelTemplate): arch = "ltxv" keys_detect = [ @@ -97,7 +107,7 @@ class ModelSD1(ModelTemplate): ] # The architectures are checked in order and the first successful match terminates the search. -arch_list = [ModelFlux, ModelSD3, ModelAura, ModelLTXV, ModelHyVid, ModelSDXL, ModelSD1] +arch_list = [ModelFlux, ModelSD3, ModelAura, ModelLTXV, ModelHyVid, ModelWan, ModelSDXL, ModelSD1] def is_model_arch(model, state_dict): # check if model is correct diff --git a/tools/lcpp.patch b/tools/lcpp.patch index 2706a51..0ad92d1 100644 --- a/tools/lcpp.patch +++ b/tools/lcpp.patch @@ -39,10 +39,10 @@ index b16c462f..6d1568f1 100644 const int idx = gguf_find_tensor(ctx, name); if (idx < 0) { diff --git a/src/llama.cpp b/src/llama.cpp -index 24e1f1f0..c2142c13 100644 +index 24e1f1f0..f90738cf 100644 --- a/src/llama.cpp +++ b/src/llama.cpp -@@ -205,6 +205,13 @@ enum llm_arch { +@@ -205,6 +205,14 @@ enum llm_arch { LLM_ARCH_GRANITE, LLM_ARCH_GRANITE_MOE, LLM_ARCH_CHAMELEON, @@ -53,10 +53,11 @@ index 24e1f1f0..c2142c13 100644 + LLM_ARCH_AURA, + LLM_ARCH_LTXV, + LLM_ARCH_HYVID, ++ LLM_ARCH_WAN, LLM_ARCH_UNKNOWN, }; -@@ -258,6 +265,13 @@ static const std::map LLM_ARCH_NAMES = { +@@ -258,6 +266,14 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GRANITE, "granite" }, { LLM_ARCH_GRANITE_MOE, "granitemoe" }, { LLM_ARCH_CHAMELEON, "chameleon" }, @@ -67,10 +68,11 @@ index 24e1f1f0..c2142c13 100644 + { LLM_ARCH_AURA, "aura" }, + { LLM_ARCH_LTXV, "ltxv" }, + { LLM_ARCH_HYVID, "hyvid" }, ++ { LLM_ARCH_WAN, "wan" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; -@@ -1531,6 +1545,13 @@ static const std::map> LLM_TENSOR_N +@@ -1531,6 +1547,14 @@ static const std::map> LLM_TENSOR_N { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, }, }, @@ -81,10 +83,11 @@ index 24e1f1f0..c2142c13 100644 + { LLM_ARCH_AURA, {}}, + { LLM_ARCH_LTXV, {}}, + { LLM_ARCH_HYVID,{}}, ++ { LLM_ARCH_WAN,{}}, { LLM_ARCH_UNKNOWN, { -@@ -5403,6 +5424,21 @@ static void llm_load_hparams( +@@ -5403,6 +5427,22 @@ static void llm_load_hparams( // get general kv ml.get_key(LLM_KV_GENERAL_NAME, model.name, false); @@ -97,6 +100,7 @@ index 24e1f1f0..c2142c13 100644 + case LLM_ARCH_AURA: + case LLM_ARCH_LTXV: + case LLM_ARCH_HYVID: ++ case LLM_ARCH_WAN: + model.ftype = ml.ftype; + return; + default: @@ -106,7 +110,7 @@ index 24e1f1f0..c2142c13 100644 // get hparams kv ml.get_key(LLM_KV_VOCAB_SIZE, hparams.n_vocab, false) || ml.get_arr_n(LLM_KV_TOKENIZER_LIST, hparams.n_vocab); -@@ -18016,6 +18052,126 @@ static void llama_tensor_dequantize_internal( +@@ -18016,6 +18056,128 @@ static void llama_tensor_dequantize_internal( workers.clear(); } @@ -139,6 +143,7 @@ index 24e1f1f0..c2142c13 100644 + if ( // Rules for to_v attention + (name.find("attn_v.weight") != std::string::npos) || + (name.find(".to_v.weight") != std::string::npos) || ++ (name.find(".v.weight") != std::string::npos) || + (name.find(".attn.w1v.weight") != std::string::npos) || + (name.find(".attn.w2v.weight") != std::string::npos) + ){ @@ -173,6 +178,7 @@ index 24e1f1f0..c2142c13 100644 + } + } else if ( // Rules for ffn + (name.find("ffn_down") != std::string::npos) || ++ (name.find(".ffn.2.weight") != std::string::npos) || // is this even the right way around? + (name.find(".ff.net.2.weight") != std::string::npos) + ) { + // TODO: add back `layer_info` with some model specific logic + logic further down @@ -233,7 +239,7 @@ index 24e1f1f0..c2142c13 100644 static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type new_type, const ggml_tensor * tensor, llama_ftype ftype) { const std::string name = ggml_get_name(tensor); -@@ -18513,7 +18669,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18513,7 +18675,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s if (llama_model_has_encoder(&model)) { n_attn_layer *= 3; } @@ -244,7 +250,7 @@ index 24e1f1f0..c2142c13 100644 } size_t total_size_org = 0; -@@ -18547,6 +18705,36 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18547,6 +18711,45 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s ctx_outs[i_split] = gguf_init_empty(); } gguf_add_tensor(ctx_outs[i_split], tensor); @@ -277,11 +283,20 @@ index 24e1f1f0..c2142c13 100644 + if (name == "img_in.proj.weight" && tensor->ne[5] != 1 ) { + throw std::runtime_error("img_in.proj.weight size failed for HyVid"); + } ++ } ++ // All the modulation layers also have dim1, and I think conv3d fails here too but we segfaul way before that... ++ if (model.arch == LLM_ARCH_WAN) { ++ const std::string name = ggml_get_name(tensor); ++ if (name.find(".modulation") != std::string::npos && tensor->ne[2] == 1) { ++ const int n_dim = 3; ++ gguf_set_tensor_ndim(ctx_outs[i_split], tensor->name, n_dim); ++ LLAMA_LOG_INFO("\n%s: Correcting shape for Wan: [key:%s]\n", __func__, tensor->name); ++ } + } } // Set split info if needed -@@ -18647,6 +18835,73 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18647,6 +18850,82 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s // do not quantize relative position bias (T5) quantize &= name.find("attn_rel_b.weight") == std::string::npos; @@ -347,6 +362,15 @@ index 24e1f1f0..c2142c13 100644 + quantize &= name.find("guidance_in.") == std::string::npos; + quantize &= name.find("final_layer.") == std::string::npos; + } ++ if (model.arch == LLM_ARCH_WAN) { ++ image_model = true; ++ quantize &= name.find("modulation.") == std::string::npos; ++ quantize &= name.find("patch_embedding.") == std::string::npos; ++ quantize &= name.find("text_embedding.") == std::string::npos; ++ quantize &= name.find("time_projection.") == std::string::npos; ++ quantize &= name.find("time_embedding.") == std::string::npos; ++ quantize &= name.find("head.") == std::string::npos; ++ } + // ignore 3D/4D tensors for image models as the code was never meant to handle these + if (image_model) { + quantize &= ggml_n_dims(tensor) == 2; @@ -355,7 +379,7 @@ index 24e1f1f0..c2142c13 100644 enum ggml_type new_type; void * new_data; size_t new_size; -@@ -18655,6 +18910,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18655,6 +18934,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s new_type = default_type; // get more optimal quantization type based on the tensor shape, layer, etc. @@ -365,7 +389,7 @@ index 24e1f1f0..c2142c13 100644 if (!params->pure && ggml_is_quantized(default_type)) { new_type = llama_tensor_get_type(qs, new_type, tensor, ftype); } -@@ -18664,6 +18922,7 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18664,6 +18946,7 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s if (params->output_tensor_type < GGML_TYPE_COUNT && strcmp(tensor->name, "output.weight") == 0) { new_type = params->output_tensor_type; }