diff --git a/loader.py b/loader.py index 8e7dd81..898fce6 100644 --- a/loader.py +++ b/loader.py @@ -7,7 +7,7 @@ import gguf from .ops import GGMLTensor from .dequant import is_quantized, dequantize_tensor -IMG_ARCH_LIST = {"flux", "sd1", "sdxl", "sd3", "aura", "hidream", "ltxv", "hyvid", "wan"} +IMG_ARCH_LIST = {"flux", "sd1", "sdxl", "sd3", "aura", "hidream", "cosmos", "ltxv", "hyvid", "wan"} TXT_ARCH_LIST = {"t5", "t5encoder", "llama"} def get_orig_shape(reader, tensor_name): diff --git a/tools/convert.py b/tools/convert.py index fa12655..0078622 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -18,6 +18,7 @@ class ModelTemplate: keys_detect = [] # list of lists to match in state dict keys_banned = [] # list of keys that should mark model as invalid for conversion keys_hiprec = [] # list of keys that need to be kept in fp32 for some reason + keys_ignore = [] # list of strings to ignore keys by when found def handle_nd_tensor(self, key, data): raise NotImplementedError(f"Tensor detected that exceeds dims supported by C++ code! ({key} @ {data.shape})") @@ -60,6 +61,17 @@ class ModelHiDream(ModelTemplate): "img_emb.emb_pos" ] +class CosmosPredict2(ModelTemplate): + arch = "cosmos" + keys_detect = [ + ( + "blocks.0.mlp.layer1.weight", + "blocks.0.adaln_modulation_cross_attn.1.weight", + ) + ] + keys_hiprec = ["pos_embedder"] + keys_ignore = ["_extra_state", "accum_"] + class ModelHyVid(ModelTemplate): arch = "hyvid" keys_detect = [ @@ -128,7 +140,7 @@ class ModelSD1(ModelTemplate): ] # The architectures are checked in order and the first successful match terminates the search. -arch_list = [ModelFlux, ModelSD3, ModelAura, ModelHiDream, ModelLTXV, ModelHyVid, ModelWan, ModelSDXL, ModelSD1] +arch_list = [ModelFlux, ModelSD3, ModelAura, ModelHiDream, CosmosPredict2, ModelLTXV, ModelHyVid, ModelWan, ModelSDXL, ModelSD1] def is_model_arch(model, state_dict): # check if model is correct @@ -163,20 +175,32 @@ def parse_args(): return args def strip_prefix(state_dict): - # only keep unet with no prefix! + # prefix for mixed state dict prefix = None for pfx in ["model.diffusion_model.", "model."]: if any([x.startswith(pfx) for x in state_dict.keys()]): prefix = pfx break - sd = {} - for k, v in state_dict.items(): - if prefix and prefix not in k: - continue - if prefix: + # prefix for uniform state dict + if prefix is None: + for pfx in ["net."]: + if all([x.startswith(pfx) for x in state_dict.keys()]): + prefix = pfx + break + + # strip prefix if found + if prefix is not None: + logging.info(f"State dict prefix found: '{prefix}'") + sd = {} + for k, v in state_dict.items(): + if prefix not in k: + continue k = k.replace(prefix, "") - sd[k] = v + sd[k] = v + else: + logging.debug("State dict has no prefix") + sd = state_dict return sd @@ -209,6 +233,10 @@ def handle_tensors(writer, state_dict, model_arch): for key, data in tqdm(state_dict.items()): old_dtype = data.dtype + if any(x in key for x in model_arch.keys_ignore): + tqdm.write(f"Filtering ignored key: '{key}'") + continue + if data.dtype == torch.bfloat16: data = data.to(torch.float32).numpy() # this is so we don't break torch 2.0.X diff --git a/tools/lcpp.patch b/tools/lcpp.patch index cecb849..f5f4174 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..39045ca5 100644 +index 24e1f1f0..9957ea30 100644 --- a/src/llama.cpp +++ b/src/llama.cpp -@@ -205,6 +205,15 @@ enum llm_arch { +@@ -205,6 +205,16 @@ enum llm_arch { LLM_ARCH_GRANITE, LLM_ARCH_GRANITE_MOE, LLM_ARCH_CHAMELEON, @@ -55,10 +55,11 @@ index 24e1f1f0..39045ca5 100644 + LLM_ARCH_HYVID, + LLM_ARCH_WAN, + LLM_ARCH_HIDREAM, ++ LLM_ARCH_COSMOS, LLM_ARCH_UNKNOWN, }; -@@ -258,6 +267,15 @@ static const std::map LLM_ARCH_NAMES = { +@@ -258,6 +268,16 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GRANITE, "granite" }, { LLM_ARCH_GRANITE_MOE, "granitemoe" }, { LLM_ARCH_CHAMELEON, "chameleon" }, @@ -71,10 +72,11 @@ index 24e1f1f0..39045ca5 100644 + { LLM_ARCH_HYVID, "hyvid" }, + { LLM_ARCH_WAN, "wan" }, + { LLM_ARCH_HIDREAM, "hidream" }, ++ { LLM_ARCH_COSMOS, "cosmos" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; -@@ -1531,6 +1549,15 @@ static const std::map> LLM_TENSOR_N +@@ -1531,6 +1551,16 @@ static const std::map> LLM_TENSOR_N { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, }, }, @@ -87,10 +89,11 @@ index 24e1f1f0..39045ca5 100644 + { LLM_ARCH_HYVID, {}}, + { LLM_ARCH_WAN, {}}, + { LLM_ARCH_HIDREAM, {}}, ++ { LLM_ARCH_COSMOS, {}}, { LLM_ARCH_UNKNOWN, { -@@ -5403,6 +5430,23 @@ static void llm_load_hparams( +@@ -5403,6 +5433,24 @@ static void llm_load_hparams( // get general kv ml.get_key(LLM_KV_GENERAL_NAME, model.name, false); @@ -105,6 +108,7 @@ index 24e1f1f0..39045ca5 100644 + case LLM_ARCH_HYVID: + case LLM_ARCH_WAN: + case LLM_ARCH_HIDREAM: ++ case LLM_ARCH_COSMOS: + model.ftype = ml.ftype; + return; + default: @@ -114,7 +118,7 @@ index 24e1f1f0..39045ca5 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 +18060,129 @@ static void llama_tensor_dequantize_internal( +@@ -18016,6 +18064,132 @@ static void llama_tensor_dequantize_internal( workers.clear(); } @@ -149,7 +153,8 @@ index 24e1f1f0..39045ca5 100644 + (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) ++ (name.find(".attn.w2v.weight") != std::string::npos) || ++ (name.find("_attn.v_proj.weight") != std::string::npos) + ){ + if (ftype == LLAMA_FTYPE_MOSTLY_Q2_K) { + new_type = GGML_TYPE_Q3_K; @@ -184,7 +189,9 @@ index 24e1f1f0..39045ca5 100644 + (name.find("ffn_down") != std::string::npos) || + ((name.find("experts.") != std::string::npos) && (name.find(".w2.weight") != 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) ++ (name.find(".ff.net.2.weight") != std::string::npos) || ++ (name.find(".mlp.layer2.weight") != std::string::npos) || ++ (name.find(".adaln_modulation_mlp.2.weight") != std::string::npos) + ) { + // TODO: add back `layer_info` with some model specific logic + logic further down + if (ftype == LLAMA_FTYPE_MOSTLY_Q3_K_M) { @@ -244,7 +251,7 @@ index 24e1f1f0..39045ca5 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 +18680,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18513,7 +18687,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s if (llama_model_has_encoder(&model)) { n_attn_layer *= 3; } @@ -255,7 +262,7 @@ index 24e1f1f0..39045ca5 100644 } size_t total_size_org = 0; -@@ -18547,6 +18716,51 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18547,6 +18723,51 @@ 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); @@ -307,7 +314,7 @@ index 24e1f1f0..39045ca5 100644 } // Set split info if needed -@@ -18647,6 +18861,92 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18647,6 +18868,101 @@ 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; @@ -392,6 +399,15 @@ index 24e1f1f0..39045ca5 100644 + quantize &= name.find(".ff_i.gate.weight") == std::string::npos; + quantize &= name.find("caption_projection.") == std::string::npos; + } ++ if (model.arch == LLM_ARCH_COSMOS) { ++ image_model = true; ++ quantize &= name.find("p_embedder.") == std::string::npos; ++ quantize &= name.find("t_embedder.") == std::string::npos; ++ quantize &= name.find("t_embedding_norm.") == std::string::npos; ++ quantize &= name.find("x_embedder.") == std::string::npos; ++ quantize &= name.find("pos_embedder.") == std::string::npos; ++ quantize &= name.find("final_layer.") == 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; @@ -400,7 +416,7 @@ index 24e1f1f0..39045ca5 100644 enum ggml_type new_type; void * new_data; size_t new_size; -@@ -18655,6 +18955,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18655,6 +18971,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. @@ -410,7 +426,7 @@ index 24e1f1f0..39045ca5 100644 if (!params->pure && ggml_is_quantized(default_type)) { new_type = llama_tensor_get_type(qs, new_type, tensor, ftype); } -@@ -18664,6 +18967,7 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18664,6 +18983,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; }