From abdfec54c4b929eeca5cff3bf99545969f0d3026 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Thu, 20 Feb 2025 05:10:27 +0100 Subject: [PATCH] Hunyuan Jank --- tools/README.md | 7 +- tools/convert.py | 20 ++- tools/fix_5d_tensors.py | 46 ++++++ tools/lcpp.patch | 298 +++++++++++++++++++++++++++--------- tools/lcpp_sd3.patch | 324 ---------------------------------------- 5 files changed, 294 insertions(+), 401 deletions(-) create mode 100644 tools/fix_5d_tensors.py delete mode 100644 tools/lcpp_sd3.patch diff --git a/tools/README.md b/tools/README.md index fb962c3..5b040ba 100644 --- a/tools/README.md +++ b/tools/README.md @@ -15,12 +15,10 @@ python convert.py --src E:\models\unet\flux1-dev.safetensors To quantize the model, first apply the provided patch to the llama.cpp repo you've just cloned. If you get a "corrupt patch" error, you may have to [change the line endings in the patch file](https://github.com/city96/ComfyUI-GGUF/issues/90#issuecomment-2323011648). ``` cd llama.cpp -git checkout tags/b3600 +git checkout tags/b3962 git apply ..\lcpp.patch ``` -If you wish to quantize **SD3** or **AuraFlow** models, you should use the patch named `lcpp_sd3.patch`, which has slightly modified logic for quantizing. For this you'll want to target `tags/b3962` instead. - Then, compile the llama-quantize binary. This example uses cmake, on linux you can just use make. ``` @@ -41,6 +39,9 @@ llama.cpp\build\bin\Debug\llama-quantize.exe E:\models\unet\flux1-dev-BF16.gguf You can extract the patch again with `git diff src\llama.cpp > lcpp.patch` if you wish to change something and contribute back. +> [!WARNING] +>For hunyuan video, you will have to uncomment the block in convert.py that deals with 5D tensors. This will save a **non functional** model to disk first, that you can quantize. After quantization, run `fix_5d_tensor.py` to add back the missing key that was saved by the conversion code. You will have to edit this file to set the correct paths/architecture. This may change in the future. + > [!WARNING] > Do not use the diffusers UNET for flux, it won't work, use the default/reference checkpoint format. This is due to q/k/v being merged into one qkv key. You can convert it by loading it in ComfyUI and saving it using the built-in "ModelSave" node. diff --git a/tools/convert.py b/tools/convert.py index 274f01d..61e1592 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -45,6 +45,24 @@ class ModelAura(ModelTemplate): ] keys_banned = ["joint_transformer_blocks.3.ff_context.out_projection.weight",] +class ModelHyVid(ModelTemplate): + arch = "hyvid" + keys_detect = [ + ( + "double_blocks.0.img_attn_proj.weight", + "txt_in.individual_token_refiner.blocks.1.self_attn_qkv.weight", + ) + ] + + # def handle_nd_tensor(self, key, data): + # # hacky but this is the only arch that uses it + # path = f"./fix_5d_tensors_{self.arch}.pt" + # if os.path.isfile(path): + # raise RuntimeError(f"5D tensor fix file already exists! {path}") + # fsd = {key: data} + # tqdm.write(f"5D key found in state dict! Manual fix required! - {key} {data.shape}") + # torch.save(fsd, path) + class ModelLTXV(ModelTemplate): arch = "ltxv" keys_detect = [ @@ -79,7 +97,7 @@ class ModelSD1(ModelTemplate): ] # The architectures are checked in order and the first successful match terminates the search. -arch_list = [ModelFlux, ModelSD3, ModelAura, ModelLTXV, ModelSDXL, ModelSD1] +arch_list = [ModelFlux, ModelSD3, ModelAura, ModelLTXV, ModelHyVid, ModelSDXL, ModelSD1] def is_model_arch(model, state_dict): # check if model is correct diff --git a/tools/fix_5d_tensors.py b/tools/fix_5d_tensors.py new file mode 100644 index 0000000..7465cb4 --- /dev/null +++ b/tools/fix_5d_tensors.py @@ -0,0 +1,46 @@ +import gguf +import torch +from tqdm import tqdm + +arch = "hyvid" # TODO: Really should autodetect this +prec = "Q4_K_S" # edit manually for each step +src = fr"raw/hunyuan-video-t2v-720p-{prec}.gguf" +dst = fr"hunyuan-video-t2v-720p-{prec}.gguf" + +sd5d = torch.load(f"./fix_5d_tensors_{arch}.pt") +print("5D:", sd5d.keys()) + +reader = gguf.GGUFReader(src) +writer = gguf.GGUFWriter(path=None, arch=arch) + +writer.add_quantization_version(gguf.GGML_QUANT_VERSION) +writer.add_file_type(getattr(gguf.LlamaFileType, f"MOSTLY_{prec}")) # TODO: also autodetect + +added = [] +def add_extra_key(writer, key, data): + old_dtype = data.dtype + data_qtype = gguf.GGMLQuantizationType.F32 + n_dims = len(data.shape) + data_shape = data.shape + data = gguf.quants.quantize(data, data_qtype) + tqdm.write(f"Adding key {key} ({data_shape})") + writer.add_tensor(key, data, raw_dtype=data_qtype) + global added + added.append(key) + +# main loop to add missing +for tensor in tqdm(reader.tensors): + writer.add_tensor(tensor.name, tensor.data, raw_dtype=tensor.tensor_type) + key5d = tensor.name.replace(".bias", ".weight") + if key5d in sd5d.keys(): + add_extra_key(writer, key5d, sd5d[key5d]) + +# brute force for any missed +for key, data in sd5d.items(): + if key not in added: + add_extra_key(writer, key, data) + +writer.write_header_to_file(path=dst) +writer.write_kv_data_to_file() +writer.write_tensors_to_file(progress=True) +writer.close() diff --git a/tools/lcpp.patch b/tools/lcpp.patch index e7bfbff..2706a51 100644 --- a/tools/lcpp.patch +++ b/tools/lcpp.patch @@ -1,67 +1,115 @@ diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h -index 1d2a3540..b1a9ee96 100644 +index de3c706f..0267c1fa 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h -@@ -230,7 +230,7 @@ - #define GGML_MAX_CONTEXTS 64 - #define GGML_MAX_SRC 10 - #ifndef GGML_MAX_NAME --#define GGML_MAX_NAME 64 -+#define GGML_MAX_NAME 128 - #endif +@@ -223,7 +223,7 @@ #define GGML_MAX_OP_PARAMS 64 + + #ifndef GGML_MAX_NAME +-# define GGML_MAX_NAME 64 ++# define GGML_MAX_NAME 128 + #endif + #define GGML_DEFAULT_N_THREADS 4 +@@ -2449,6 +2449,7 @@ extern "C" { + + // manage tensor info + GGML_API void gguf_add_tensor(struct gguf_context * ctx, const struct ggml_tensor * tensor); ++ GGML_API void gguf_set_tensor_ndim(struct gguf_context * ctx, const char * name, int n_dim); + GGML_API void gguf_set_tensor_type(struct gguf_context * ctx, const char * name, enum ggml_type type); + GGML_API void gguf_set_tensor_data(struct gguf_context * ctx, const char * name, const void * data, size_t size); + +diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c +index b16c462f..6d1568f1 100644 +--- a/ggml/src/ggml.c ++++ b/ggml/src/ggml.c +@@ -22960,6 +22960,14 @@ void gguf_add_tensor( + ctx->header.n_tensors++; + } + ++void gguf_set_tensor_ndim(struct gguf_context * ctx, const char * name, const int n_dim) { ++ const int idx = gguf_find_tensor(ctx, name); ++ if (idx < 0) { ++ GGML_ABORT("tensor not found"); ++ } ++ ctx->infos[idx].n_dims = n_dim; ++} ++ + void gguf_set_tensor_type(struct gguf_context * ctx, const char * name, enum ggml_type type) { + const int idx = gguf_find_tensor(ctx, name); + if (idx < 0) { diff --git a/src/llama.cpp b/src/llama.cpp -index 5ab65ea9..35580d9d 100644 +index 24e1f1f0..c2142c13 100644 --- a/src/llama.cpp +++ b/src/llama.cpp -@@ -212,6 +212,9 @@ enum llm_arch { - LLM_ARCH_JAIS, - LLM_ARCH_NEMOTRON, - LLM_ARCH_EXAONE, +@@ -205,6 +205,13 @@ enum llm_arch { + LLM_ARCH_GRANITE, + LLM_ARCH_GRANITE_MOE, + LLM_ARCH_CHAMELEON, + LLM_ARCH_FLUX, + LLM_ARCH_SD1, + LLM_ARCH_SDXL, ++ LLM_ARCH_SD3, ++ LLM_ARCH_AURA, ++ LLM_ARCH_LTXV, ++ LLM_ARCH_HYVID, LLM_ARCH_UNKNOWN, }; - -@@ -259,6 +262,9 @@ static const std::map LLM_ARCH_NAMES = { - { LLM_ARCH_JAIS, "jais" }, - { LLM_ARCH_NEMOTRON, "nemotron" }, - { LLM_ARCH_EXAONE, "exaone" }, + +@@ -258,6 +265,13 @@ static const std::map LLM_ARCH_NAMES = { + { LLM_ARCH_GRANITE, "granite" }, + { LLM_ARCH_GRANITE_MOE, "granitemoe" }, + { LLM_ARCH_CHAMELEON, "chameleon" }, + { LLM_ARCH_FLUX, "flux" }, + { LLM_ARCH_SD1, "sd1" }, + { LLM_ARCH_SDXL, "sdxl" }, ++ { LLM_ARCH_SD3, "sd3" }, ++ { LLM_ARCH_AURA, "aura" }, ++ { LLM_ARCH_LTXV, "ltxv" }, ++ { LLM_ARCH_HYVID, "hyvid" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; - -@@ -1337,6 +1343,9 @@ static const std::map> LLM_TENSOR_NA - { LLM_TENSOR_FFN_UP, "blk.%d.ffn_up" }, + +@@ -1531,6 +1545,13 @@ static const std::map> LLM_TENSOR_N + { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, }, }, + { LLM_ARCH_FLUX, {}}, + { LLM_ARCH_SD1, {}}, + { LLM_ARCH_SDXL, {}}, ++ { LLM_ARCH_SD3, {}}, ++ { LLM_ARCH_AURA, {}}, ++ { LLM_ARCH_LTXV, {}}, ++ { LLM_ARCH_HYVID,{}}, { LLM_ARCH_UNKNOWN, { -@@ -4629,6 +4638,12 @@ static void llm_load_hparams( +@@ -5403,6 +5424,21 @@ static void llm_load_hparams( // get general kv ml.get_key(LLM_KV_GENERAL_NAME, model.name, false); - + + // Disable LLM metadata for image models -+ if (model.arch == LLM_ARCH_FLUX || model.arch == LLM_ARCH_SD1 || model.arch == LLM_ARCH_SDXL) { -+ model.ftype = ml.ftype; -+ return; ++ switch (model.arch) { ++ case LLM_ARCH_FLUX: ++ case LLM_ARCH_SD1: ++ case LLM_ARCH_SDXL: ++ case LLM_ARCH_SD3: ++ case LLM_ARCH_AURA: ++ case LLM_ARCH_LTXV: ++ case LLM_ARCH_HYVID: ++ model.ftype = ml.ftype; ++ return; ++ default: ++ break; + } + // 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); - -@@ -15827,11 +15842,162 @@ static void llama_tensor_dequantize_internal( + +@@ -18016,6 +18052,126 @@ static void llama_tensor_dequantize_internal( workers.clear(); } - + +static ggml_type img_tensor_get_type(quantize_state_internal & qs, ggml_type new_type, const ggml_tensor * tensor, llama_ftype ftype) { + // Special function for quantizing image model tensors + const std::string name = ggml_get_name(tensor); @@ -71,7 +119,8 @@ index 5ab65ea9..35580d9d 100644 + if ( + (name.find("model.diffusion_model.") != std::string::npos) || + (name.find("first_stage_model.") != std::string::npos) || -+ (name.find("single_transformer_blocks.") != std::string::npos) ++ (name.find("single_transformer_blocks.") != std::string::npos) || ++ (name.find("joint_transformer_blocks.") != std::string::npos) + ) { + throw std::runtime_error("Invalid input GGUF file. This is not a supported UNET model"); + } @@ -87,41 +136,11 @@ index 5ab65ea9..35580d9d 100644 + throw std::runtime_error("Invalid quantization type for image model (Not supported)"); + } + -+ if ( // Tensors to keep in FP32 precision -+ (arch == LLM_ARCH_FLUX) && ( -+ (name.find("img_in.") != std::string::npos) || -+ (name.find("time_in.in_layer.") != std::string::npos) || -+ (name.find("vector_in.in_layer.") != std::string::npos) || -+ (name.find("guidance_in.in_layer.") != std::string::npos) || -+ (name.find("final_layer.linear.") != std::string::npos) -+ ) || (arch == LLM_ARCH_SD1 || arch == LLM_ARCH_SDXL) && ( -+ (name.find("conv_in.") != std::string::npos) || -+ (name.find("conv_out.") != std::string::npos) || -+ (name == "input_blocks.0.0.weight") || -+ (name == "out.2.weight") -+ )) { -+ new_type = GGML_TYPE_F32; -+ } else if ( // Tensors to keep in FP16 precision -+ (arch == LLM_ARCH_FLUX) && ( -+ (name.find("txt_in.") != std::string::npos) || -+ (name.find("time_in.") != std::string::npos) || -+ (name.find("vector_in.") != std::string::npos) || -+ (name.find("guidance_in.") != std::string::npos) || -+ (name.find("final_layer.") != std::string::npos) -+ ) || (arch == LLM_ARCH_SD1 || arch == LLM_ARCH_SDXL) && ( -+ (name.find("class_embedding.") != std::string::npos) || -+ (name.find("time_embedding.") != std::string::npos) || -+ (name.find("add_embedding.") != std::string::npos) || -+ (name.find("time_embed.") != std::string::npos) || -+ (name.find("label_emb.") != std::string::npos) || -+ (name.find("proj_in.") != std::string::npos) || -+ (name.find("proj_out.") != std::string::npos) -+ // (name.find("conv_shortcut.") != std::string::npos) // marginal improvement -+ )) { -+ new_type = GGML_TYPE_F16; -+ } else if ( // Rules for to_v attention ++ if ( // Rules for to_v attention + (name.find("attn_v.weight") != std::string::npos) || -+ (name.find(".to_v.weight") != std::string::npos) ++ (name.find(".to_v.weight") != std::string::npos) || ++ (name.find(".attn.w1v.weight") != std::string::npos) || ++ (name.find(".attn.w2v.weight") != std::string::npos) + ){ + if (ftype == LLAMA_FTYPE_MOSTLY_Q2_K) { + new_type = GGML_TYPE_Q3_K; @@ -154,7 +173,7 @@ index 5ab65ea9..35580d9d 100644 + } + } else if ( // Rules for ffn + (name.find("ffn_down") != std::string::npos) || -+ (name.find("DenseReluDense.wo") != std::string::npos) ++ (name.find(".ff.net.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) { @@ -210,14 +229,147 @@ index 5ab65ea9..35580d9d 100644 + } + return new_type; +} -+ + 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); - - // TODO: avoid hardcoded tensor names - use the TN_* constants - const llm_arch arch = qs.model.arch; -+ if (arch == LLM_ARCH_FLUX || arch == LLM_ARCH_SD1 || arch == LLM_ARCH_SDXL) { return img_tensor_get_type(qs, new_type, tensor, ftype); }; - const auto tn = LLM_TN(arch); - - auto use_more_bits = [](int i_layer, int n_layers) -> bool { + +@@ -18513,7 +18669,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s + if (llama_model_has_encoder(&model)) { + n_attn_layer *= 3; + } +- GGML_ASSERT((qs.n_attention_wv == n_attn_layer) && "n_attention_wv is unexpected"); ++ if (model.arch != LLM_ARCH_HYVID) { // TODO: Check why this fails ++ GGML_ASSERT((qs.n_attention_wv == n_attn_layer) && "n_attention_wv is unexpected"); ++ } + } + + size_t total_size_org = 0; +@@ -18547,6 +18705,36 @@ 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); ++ // SD3 pos_embed needs special fix as first dim is 1, which gets truncated here ++ if (model.arch == LLM_ARCH_SD3) { ++ const std::string name = ggml_get_name(tensor); ++ if (name == "pos_embed" && tensor->ne[2] == 1) { ++ const int n_dim = 3; ++ gguf_set_tensor_ndim(ctx_outs[i_split], "pos_embed", n_dim); ++ LLAMA_LOG_INFO("\n%s: Correcting pos_embed shape for SD3: [key:%s]\n", __func__, tensor->name); ++ } ++ } ++ // same goes for auraflow ++ if (model.arch == LLM_ARCH_AURA) { ++ const std::string name = ggml_get_name(tensor); ++ if (name == "positional_encoding" && tensor->ne[2] == 1) { ++ const int n_dim = 3; ++ gguf_set_tensor_ndim(ctx_outs[i_split], "positional_encoding", n_dim); ++ LLAMA_LOG_INFO("\n%s: Correcting positional_encoding shape for AuraFlow: [key:%s]\n", __func__, tensor->name); ++ } ++ if (name == "register_tokens" && tensor->ne[2] == 1) { ++ const int n_dim = 3; ++ gguf_set_tensor_ndim(ctx_outs[i_split], "register_tokens", n_dim); ++ LLAMA_LOG_INFO("\n%s: Correcting register_tokens shape for AuraFlow: [key:%s]\n", __func__, tensor->name); ++ } ++ } ++ // conv3d fails due to max dims - unsure what to do here as we never even reach this check ++ if (model.arch == LLM_ARCH_HYVID) { ++ const std::string name = ggml_get_name(tensor); ++ if (name == "img_in.proj.weight" && tensor->ne[5] != 1 ) { ++ throw std::runtime_error("img_in.proj.weight size failed for HyVid"); ++ } ++ } + } + + // Set split info if needed +@@ -18647,6 +18835,73 @@ 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; + ++ // rules for image models ++ bool image_model = false; ++ if (model.arch == LLM_ARCH_FLUX) { ++ image_model = true; ++ quantize &= name.find("txt_in.") == std::string::npos; ++ quantize &= name.find("img_in.") == std::string::npos; ++ quantize &= name.find("time_in.") == std::string::npos; ++ quantize &= name.find("vector_in.") == std::string::npos; ++ quantize &= name.find("guidance_in.") == std::string::npos; ++ quantize &= name.find("final_layer.") == std::string::npos; ++ } ++ if (model.arch == LLM_ARCH_SD1 || model.arch == LLM_ARCH_SDXL) { ++ image_model = true; ++ quantize &= name.find("class_embedding.") == std::string::npos; ++ quantize &= name.find("time_embedding.") == std::string::npos; ++ quantize &= name.find("add_embedding.") == std::string::npos; ++ quantize &= name.find("time_embed.") == std::string::npos; ++ quantize &= name.find("label_emb.") == std::string::npos; ++ quantize &= name.find("conv_in.") == std::string::npos; ++ quantize &= name.find("conv_out.") == std::string::npos; ++ quantize &= name != "input_blocks.0.0.weight"; ++ quantize &= name != "out.2.weight"; ++ } ++ if (model.arch == LLM_ARCH_SD3) { ++ image_model = true; ++ quantize &= name.find("final_layer.") == std::string::npos; ++ quantize &= name.find("time_text_embed.") == std::string::npos; ++ quantize &= name.find("context_embedder.") == std::string::npos; ++ quantize &= name.find("t_embedder.") == std::string::npos; ++ quantize &= name.find("y_embedder.") == std::string::npos; ++ quantize &= name.find("x_embedder.") == std::string::npos; ++ quantize &= name != "proj_out.weight"; ++ quantize &= name != "pos_embed"; ++ } ++ if (model.arch == LLM_ARCH_AURA) { ++ image_model = true; ++ quantize &= name.find("t_embedder.") == std::string::npos; ++ quantize &= name.find("init_x_linear.") == std::string::npos; ++ quantize &= name != "modF.1.weight"; ++ quantize &= name != "cond_seq_linear.weight"; ++ quantize &= name != "final_linear.weight"; ++ quantize &= name != "final_linear.weight"; ++ quantize &= name != "positional_encoding"; ++ quantize &= name != "register_tokens"; ++ } ++ if (model.arch == LLM_ARCH_LTXV) { ++ image_model = true; ++ quantize &= name.find("adaln_single.") == std::string::npos; ++ quantize &= name.find("caption_projection.") == std::string::npos; ++ quantize &= name.find("patchify_proj.") == std::string::npos; ++ quantize &= name.find("proj_out.") == std::string::npos; ++ quantize &= name.find("scale_shift_table") == std::string::npos; // last block too ++ } ++ if (model.arch == LLM_ARCH_HYVID) { ++ image_model = true; ++ quantize &= name.find("txt_in.") == std::string::npos; ++ quantize &= name.find("img_in.") == std::string::npos; ++ quantize &= name.find("time_in.") == std::string::npos; ++ quantize &= name.find("vector_in.") == std::string::npos; ++ quantize &= name.find("guidance_in.") == 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; ++ } ++ + 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 + new_type = default_type; + + // get more optimal quantization type based on the tensor shape, layer, etc. ++ if (image_model) { ++ new_type = img_tensor_get_type(qs, new_type, tensor, ftype); ++ } else { + 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 + if (params->output_tensor_type < GGML_TYPE_COUNT && strcmp(tensor->name, "output.weight") == 0) { + new_type = params->output_tensor_type; + } ++ } + + // If we've decided to quantize to the same type the tensor is already + // in then there's nothing to do. diff --git a/tools/lcpp_sd3.patch b/tools/lcpp_sd3.patch deleted file mode 100644 index 2367944..0000000 --- a/tools/lcpp_sd3.patch +++ /dev/null @@ -1,324 +0,0 @@ -diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h -index de3c706f..0267c1fa 100644 ---- a/ggml/include/ggml.h -+++ b/ggml/include/ggml.h -@@ -223,7 +223,7 @@ - #define GGML_MAX_OP_PARAMS 64 - - #ifndef GGML_MAX_NAME --# define GGML_MAX_NAME 64 -+# define GGML_MAX_NAME 128 - #endif - - #define GGML_DEFAULT_N_THREADS 4 -@@ -2449,6 +2449,7 @@ extern "C" { - - // manage tensor info - GGML_API void gguf_add_tensor(struct gguf_context * ctx, const struct ggml_tensor * tensor); -+ GGML_API void gguf_set_tensor_ndim(struct gguf_context * ctx, const char * name, int n_dim); - GGML_API void gguf_set_tensor_type(struct gguf_context * ctx, const char * name, enum ggml_type type); - GGML_API void gguf_set_tensor_data(struct gguf_context * ctx, const char * name, const void * data, size_t size); - -diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c -index b16c462f..6d1568f1 100644 ---- a/ggml/src/ggml.c -+++ b/ggml/src/ggml.c -@@ -22960,6 +22960,14 @@ void gguf_add_tensor( - ctx->header.n_tensors++; - } - -+void gguf_set_tensor_ndim(struct gguf_context * ctx, const char * name, const int n_dim) { -+ const int idx = gguf_find_tensor(ctx, name); -+ if (idx < 0) { -+ GGML_ABORT("tensor not found"); -+ } -+ ctx->infos[idx].n_dims = n_dim; -+} -+ - void gguf_set_tensor_type(struct gguf_context * ctx, const char * name, enum ggml_type type) { - const int idx = gguf_find_tensor(ctx, name); - if (idx < 0) { -diff --git a/src/llama.cpp b/src/llama.cpp -index 24e1f1f0..aeccc173 100644 ---- a/src/llama.cpp -+++ b/src/llama.cpp -@@ -205,6 +205,11 @@ enum llm_arch { - LLM_ARCH_GRANITE, - LLM_ARCH_GRANITE_MOE, - LLM_ARCH_CHAMELEON, -+ LLM_ARCH_FLUX, -+ LLM_ARCH_SD1, -+ LLM_ARCH_SDXL, -+ LLM_ARCH_SD3, -+ LLM_ARCH_AURA, - LLM_ARCH_UNKNOWN, - }; - -@@ -258,6 +263,11 @@ static const std::map LLM_ARCH_NAMES = { - { LLM_ARCH_GRANITE, "granite" }, - { LLM_ARCH_GRANITE_MOE, "granitemoe" }, - { LLM_ARCH_CHAMELEON, "chameleon" }, -+ { LLM_ARCH_FLUX, "flux" }, -+ { LLM_ARCH_SD1, "sd1" }, -+ { LLM_ARCH_SDXL, "sdxl" }, -+ { LLM_ARCH_SD3, "sd3" }, -+ { LLM_ARCH_AURA, "aura" }, - { LLM_ARCH_UNKNOWN, "(unknown)" }, - }; - -@@ -1531,6 +1541,11 @@ static const std::map> LLM_TENSOR_N - { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, - }, - }, -+ { LLM_ARCH_FLUX, {}}, -+ { LLM_ARCH_SD1, {}}, -+ { LLM_ARCH_SDXL, {}}, -+ { LLM_ARCH_SD3, {}}, -+ { LLM_ARCH_AURA, {}}, - { - LLM_ARCH_UNKNOWN, - { -@@ -5403,6 +5418,12 @@ static void llm_load_hparams( - // get general kv - ml.get_key(LLM_KV_GENERAL_NAME, model.name, false); - -+ // Disable LLM metadata for image models -+ if (model.arch == LLM_ARCH_FLUX || model.arch == LLM_ARCH_SD1 || model.arch == LLM_ARCH_SDXL || model.arch == LLM_ARCH_SD3 || model.arch == LLM_ARCH_AURA) { -+ model.ftype = ml.ftype; -+ return; -+ } -+ - // 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 +18037,125 @@ static void llama_tensor_dequantize_internal( - workers.clear(); - } - -+static ggml_type img_tensor_get_type(quantize_state_internal & qs, ggml_type new_type, const ggml_tensor * tensor, llama_ftype ftype) { -+ // Special function for quantizing image model tensors -+ const std::string name = ggml_get_name(tensor); -+ const llm_arch arch = qs.model.arch; -+ -+ // Sanity check -+ if ( -+ (name.find("model.diffusion_model.") != std::string::npos) || -+ (name.find("first_stage_model.") != std::string::npos) || -+ (name.find("single_transformer_blocks.") != std::string::npos) || -+ (name.find("joint_transformer_blocks.") != std::string::npos) -+ ) { -+ throw std::runtime_error("Invalid input GGUF file. This is not a supported UNET model"); -+ } -+ -+ // Unsupported quant types - exclude all IQ quants for now -+ if (ftype == LLAMA_FTYPE_MOSTLY_IQ2_XXS || ftype == LLAMA_FTYPE_MOSTLY_IQ2_XS || -+ ftype == LLAMA_FTYPE_MOSTLY_IQ2_S || ftype == LLAMA_FTYPE_MOSTLY_IQ2_M || -+ ftype == LLAMA_FTYPE_MOSTLY_IQ3_XXS || ftype == LLAMA_FTYPE_MOSTLY_IQ1_S || -+ ftype == LLAMA_FTYPE_MOSTLY_IQ1_M || ftype == LLAMA_FTYPE_MOSTLY_IQ4_NL || -+ ftype == LLAMA_FTYPE_MOSTLY_IQ4_XS || ftype == LLAMA_FTYPE_MOSTLY_IQ3_S || -+ ftype == LLAMA_FTYPE_MOSTLY_IQ3_M || ftype == LLAMA_FTYPE_MOSTLY_Q4_0_4_4 || -+ ftype == LLAMA_FTYPE_MOSTLY_Q4_0_4_8 || ftype == LLAMA_FTYPE_MOSTLY_Q4_0_8_8) { -+ throw std::runtime_error("Invalid quantization type for image model (Not supported)"); -+ } -+ -+ if ( // Rules for to_v attention -+ (name.find("attn_v.weight") != std::string::npos) || -+ (name.find(".to_v.weight") != std::string::npos) || -+ (name.find(".attn.w1v.weight") != std::string::npos) || -+ (name.find(".attn.w2v.weight") != std::string::npos) -+ ){ -+ if (ftype == LLAMA_FTYPE_MOSTLY_Q2_K) { -+ new_type = GGML_TYPE_Q3_K; -+ } -+ else if (ftype == LLAMA_FTYPE_MOSTLY_Q3_K_M) { -+ new_type = qs.i_attention_wv < 2 ? GGML_TYPE_Q5_K : GGML_TYPE_Q4_K; -+ } -+ else if (ftype == LLAMA_FTYPE_MOSTLY_Q3_K_L) { -+ new_type = GGML_TYPE_Q5_K; -+ } -+ else if (ftype == LLAMA_FTYPE_MOSTLY_Q4_K_M || ftype == LLAMA_FTYPE_MOSTLY_Q5_K_M) { -+ new_type = GGML_TYPE_Q6_K; -+ } -+ else if (ftype == LLAMA_FTYPE_MOSTLY_Q4_K_S && qs.i_attention_wv < 4) { -+ new_type = GGML_TYPE_Q5_K; -+ } -+ ++qs.i_attention_wv; -+ } else if ( // Rules for fused qkv attention -+ (name.find("attn_qkv.weight") != std::string::npos) || -+ (name.find("attn.qkv.weight") != std::string::npos) -+ ) { -+ if (ftype == LLAMA_FTYPE_MOSTLY_Q3_K_M || ftype == LLAMA_FTYPE_MOSTLY_Q3_K_L) { -+ new_type = GGML_TYPE_Q4_K; -+ } -+ else if (ftype == LLAMA_FTYPE_MOSTLY_Q4_K_M) { -+ new_type = GGML_TYPE_Q5_K; -+ } -+ else if (ftype == LLAMA_FTYPE_MOSTLY_Q5_K_M) { -+ new_type = GGML_TYPE_Q6_K; -+ } -+ } else if ( // Rules for ffn -+ (name.find("ffn_down") != std::string::npos) -+ ) { -+ // TODO: add back `layer_info` with some model specific logic + logic further down -+ if (ftype == LLAMA_FTYPE_MOSTLY_Q3_K_M) { -+ new_type = GGML_TYPE_Q4_K; -+ } -+ else if (ftype == LLAMA_FTYPE_MOSTLY_Q3_K_L) { -+ new_type = GGML_TYPE_Q5_K; -+ } -+ else if (ftype == LLAMA_FTYPE_MOSTLY_Q4_K_S) { -+ new_type = GGML_TYPE_Q5_K; -+ } -+ else if (ftype == LLAMA_FTYPE_MOSTLY_Q4_K_M) { -+ new_type = GGML_TYPE_Q6_K; -+ } -+ else if (ftype == LLAMA_FTYPE_MOSTLY_Q5_K_M) { -+ new_type = GGML_TYPE_Q6_K; -+ } -+ else if (ftype == LLAMA_FTYPE_MOSTLY_Q4_0) { -+ new_type = GGML_TYPE_Q4_1; -+ } -+ else if (ftype == LLAMA_FTYPE_MOSTLY_Q5_0) { -+ new_type = GGML_TYPE_Q5_1; -+ } -+ ++qs.i_ffn_down; -+ } -+ -+ // Sanity check for row shape -+ bool convert_incompatible_tensor = false; -+ if (new_type == GGML_TYPE_Q2_K || new_type == GGML_TYPE_Q3_K || new_type == GGML_TYPE_Q4_K || -+ new_type == GGML_TYPE_Q5_K || new_type == GGML_TYPE_Q6_K) { -+ int nx = tensor->ne[0]; -+ int ny = tensor->ne[1]; -+ if (nx % QK_K != 0) { -+ LLAMA_LOG_WARN("\n\n%s : tensor cols %d x %d are not divisible by %d, required for %s", __func__, nx, ny, QK_K, ggml_type_name(new_type)); -+ convert_incompatible_tensor = true; -+ } else { -+ ++qs.n_k_quantized; -+ } -+ } -+ if (convert_incompatible_tensor) { -+ // TODO: Possibly reenable this in the future -+ // switch (new_type) { -+ // case GGML_TYPE_Q2_K: -+ // case GGML_TYPE_Q3_K: -+ // case GGML_TYPE_Q4_K: new_type = GGML_TYPE_Q5_0; break; -+ // case GGML_TYPE_Q5_K: new_type = GGML_TYPE_Q5_1; break; -+ // case GGML_TYPE_Q6_K: new_type = GGML_TYPE_Q8_0; break; -+ // default: throw std::runtime_error("\nUnsupported tensor size encountered\n"); -+ // } -+ new_type = GGML_TYPE_F16; -+ LLAMA_LOG_WARN(" - using fallback quantization %s\n", ggml_type_name(new_type)); -+ ++qs.n_fallback; -+ } -+ return new_type; -+} -+ - 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); - -@@ -18547,6 +18687,29 @@ 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); -+ // SD3 pos_embed needs special fix as first dim is 1, which gets truncated here -+ if (model.arch == LLM_ARCH_SD3) { -+ const std::string name = ggml_get_name(tensor); -+ if (name == "pos_embed" && tensor->ne[2] == 1) { -+ const int n_dim = 3; -+ gguf_set_tensor_ndim(ctx_outs[i_split], "pos_embed", n_dim); -+ LLAMA_LOG_INFO("\n%s: Correcting pos_embed shape for SD3: [key:%s]\n", __func__, tensor->name); -+ } -+ } -+ // same goes for auraflow -+ if (model.arch == LLM_ARCH_AURA) { -+ const std::string name = ggml_get_name(tensor); -+ if (name == "positional_encoding" && tensor->ne[2] == 1) { -+ const int n_dim = 3; -+ gguf_set_tensor_ndim(ctx_outs[i_split], "positional_encoding", n_dim); -+ LLAMA_LOG_INFO("\n%s: Correcting positional_encoding shape for AuraFlow: [key:%s]\n", __func__, tensor->name); -+ } -+ if (name == "register_tokens" && tensor->ne[2] == 1) { -+ const int n_dim = 3; -+ gguf_set_tensor_ndim(ctx_outs[i_split], "register_tokens", n_dim); -+ LLAMA_LOG_INFO("\n%s: Correcting register_tokens shape for AuraFlow: [key:%s]\n", __func__, tensor->name); -+ } -+ } - } - - // Set split info if needed -@@ -18647,6 +18810,56 @@ 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; - -+ // rules for image models -+ bool image_model = false; -+ if (model.arch == LLM_ARCH_FLUX) { -+ image_model = true; -+ quantize &= name.find("txt_in.") == std::string::npos; -+ quantize &= name.find("img_in.") == std::string::npos; -+ quantize &= name.find("time_in.") == std::string::npos; -+ quantize &= name.find("vector_in.") == std::string::npos; -+ quantize &= name.find("guidance_in.") == std::string::npos; -+ quantize &= name.find("final_layer.") == std::string::npos; -+ } -+ if (model.arch == LLM_ARCH_SD1 || model.arch == LLM_ARCH_SDXL) { -+ image_model = true; -+ quantize &= name.find("class_embedding.") == std::string::npos; -+ quantize &= name.find("time_embedding.") == std::string::npos; -+ quantize &= name.find("add_embedding.") == std::string::npos; -+ quantize &= name.find("time_embed.") == std::string::npos; -+ quantize &= name.find("label_emb.") == std::string::npos; -+ quantize &= name.find("conv_in.") == std::string::npos; -+ quantize &= name.find("conv_out.") == std::string::npos; -+ quantize &= name != "input_blocks.0.0.weight"; -+ quantize &= name != "out.2.weight"; -+ } -+ if (model.arch == LLM_ARCH_SD3) { -+ image_model = true; -+ quantize &= name.find("final_layer.") == std::string::npos; -+ quantize &= name.find("time_text_embed.") == std::string::npos; -+ quantize &= name.find("context_embedder.") == std::string::npos; -+ quantize &= name.find("t_embedder.") == std::string::npos; -+ quantize &= name.find("y_embedder.") == std::string::npos; -+ quantize &= name.find("x_embedder.") == std::string::npos; -+ quantize &= name != "proj_out.weight"; -+ quantize &= name != "pos_embed"; -+ } -+ if (model.arch == LLM_ARCH_AURA) { -+ image_model = true; -+ quantize &= name.find("t_embedder.") == std::string::npos; -+ quantize &= name.find("init_x_linear.") == std::string::npos; -+ quantize &= name != "modF.1.weight"; -+ quantize &= name != "cond_seq_linear.weight"; -+ quantize &= name != "final_linear.weight"; -+ quantize &= name != "final_linear.weight"; -+ quantize &= name != "positional_encoding"; -+ quantize &= name != "register_tokens"; -+ } -+ // 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; -+ } -+ - enum ggml_type new_type; - void * new_data; - size_t new_size; -@@ -18655,6 +18868,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. -+ if (image_model) { -+ new_type = img_tensor_get_type(qs, new_type, tensor, ftype); -+ } else { - if (!params->pure && ggml_is_quantized(default_type)) { - new_type = llama_tensor_get_type(qs, new_type, tensor, ftype); - } -@@ -18664,6 +18880,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; - } -+ } - - // If we've decided to quantize to the same type the tensor is already - // in then there's nothing to do.