From 4a8432884167f2526d60ef36e985bdabebb9e1e0 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Sat, 23 Nov 2024 06:51:29 +0100 Subject: [PATCH] Add AuraFlow support --- nodes.py | 2 +- tools/README.md | 2 +- tools/convert.py | 23 ++++++++++++++---- tools/lcpp_sd3.patch | 57 ++++++++++++++++++++++++++++++++++---------- 4 files changed, 64 insertions(+), 20 deletions(-) diff --git a/nodes.py b/nodes.py index 1a655a5..d5fa5fa 100644 --- a/nodes.py +++ b/nodes.py @@ -62,7 +62,7 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model."): if len(arch_field.types) != 1 or arch_field.types[0] != gguf.GGUFValueType.STRING: raise TypeError(f"Bad type for GGUF general.architecture key: expected string, got {arch_field.types!r}") arch_str = str(arch_field.parts[arch_field.data[-1]], encoding="utf-8") - if arch_str not in {"flux", "sd1", "sdxl", "sd3", "t5", "t5encoder"}: + if arch_str not in {"flux", "sd1", "sdxl", "sd3", "aura", "t5", "t5encoder"}: raise ValueError(f"Unexpected architecture type in GGUF file, expected one of flux, sd1, sdxl, t5encoder but got {arch_str!r}") else: # stable-diffusion.cpp # import here to avoid changes to convert.py breaking regular models diff --git a/tools/README.md b/tools/README.md index 9c6fb15..fb962c3 100644 --- a/tools/README.md +++ b/tools/README.md @@ -19,7 +19,7 @@ git checkout tags/b3600 git apply ..\lcpp.patch ``` -If you wish to use the SD3 version (`lcpp_sd3.patch`), which has slightly modified logic for quantizing, you'll want to target `tags/b3962` instead. +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. diff --git a/tools/convert.py b/tools/convert.py index eb9298a..0af531e 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -33,6 +33,14 @@ class ModelSD3(ModelTemplate): ] keys_banned = ["transformer_blocks.0.attn.add_q_proj.weight",] +class ModelAura(ModelTemplate): + arch = "aura" + keys_detect = [ + ("double_layers.3.modX.1.weight",), + ("joint_transformer_blocks.3.ff_context.out_projection.weight",), + ] + keys_banned = ["joint_transformer_blocks.3.ff_context.out_projection.weight",] + class ModelSDXL(ModelTemplate): arch = "sdxl" shape_fix = True @@ -57,7 +65,7 @@ class ModelSD1(ModelTemplate): ] # The architectures are checked in order and the first successful match terminates the search. -arch_list = [ModelFlux, ModelSD3, ModelSDXL, ModelSD1] +arch_list = [ModelFlux, ModelSD3, ModelAura, ModelSDXL, ModelSD1] def is_model_arch(model, state_dict): # check if model is correct @@ -99,13 +107,18 @@ def load_state_dict(path): state_dict = load_file(path) # only keep unet with no prefix! + prefix = None + for pfx in ["model.diffusion_model.", "model."]: + if any([x.startswith(pfx) for x in state_dict.keys()]): + prefix = pfx + break + sd = {} - has_prefix = any(["model.diffusion_model." in x for x in state_dict.keys()]) for k, v in state_dict.items(): - if has_prefix and "model.diffusion_model." not in k: + if prefix and prefix not in k: continue - if has_prefix: - k = k.replace("model.diffusion_model.", "") + if prefix: + k = k.replace(prefix, "") sd[k] = v return sd diff --git a/tools/lcpp_sd3.patch b/tools/lcpp_sd3.patch index 2c51b90..2367944 100644 --- a/tools/lcpp_sd3.patch +++ b/tools/lcpp_sd3.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..a54fd6a2 100644 +index 24e1f1f0..aeccc173 100644 --- a/src/llama.cpp +++ b/src/llama.cpp -@@ -205,6 +205,10 @@ enum llm_arch { +@@ -205,6 +205,11 @@ enum llm_arch { LLM_ARCH_GRANITE, LLM_ARCH_GRANITE_MOE, LLM_ARCH_CHAMELEON, @@ -50,10 +50,11 @@ index 24e1f1f0..a54fd6a2 100644 + LLM_ARCH_SD1, + LLM_ARCH_SDXL, + LLM_ARCH_SD3, ++ LLM_ARCH_AURA, LLM_ARCH_UNKNOWN, }; -@@ -258,6 +262,10 @@ static const std::map LLM_ARCH_NAMES = { +@@ -258,6 +263,11 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GRANITE, "granite" }, { LLM_ARCH_GRANITE_MOE, "granitemoe" }, { LLM_ARCH_CHAMELEON, "chameleon" }, @@ -61,10 +62,11 @@ index 24e1f1f0..a54fd6a2 100644 + { LLM_ARCH_SD1, "sd1" }, + { LLM_ARCH_SDXL, "sdxl" }, + { LLM_ARCH_SD3, "sd3" }, ++ { LLM_ARCH_AURA, "aura" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; -@@ -1531,6 +1539,10 @@ static const std::map> LLM_TENSOR_N +@@ -1531,6 +1541,11 @@ static const std::map> LLM_TENSOR_N { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, }, }, @@ -72,15 +74,16 @@ index 24e1f1f0..a54fd6a2 100644 + { LLM_ARCH_SD1, {}}, + { LLM_ARCH_SDXL, {}}, + { LLM_ARCH_SD3, {}}, ++ { LLM_ARCH_AURA, {}}, { LLM_ARCH_UNKNOWN, { -@@ -5403,6 +5415,12 @@ static void llm_load_hparams( +@@ -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) { ++ 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; + } @@ -88,7 +91,7 @@ index 24e1f1f0..a54fd6a2 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 +18034,122 @@ static void llama_tensor_dequantize_internal( +@@ -18016,6 +18037,125 @@ static void llama_tensor_dequantize_internal( workers.clear(); } @@ -101,7 +104,8 @@ index 24e1f1f0..a54fd6a2 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"); + } @@ -119,7 +123,9 @@ index 24e1f1f0..a54fd6a2 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(".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; @@ -211,7 +217,7 @@ index 24e1f1f0..a54fd6a2 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); -@@ -18547,6 +18681,15 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -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); @@ -223,11 +229,25 @@ index 24e1f1f0..a54fd6a2 100644 + 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 +18790,45 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -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; @@ -265,6 +285,17 @@ index 24e1f1f0..a54fd6a2 100644 + 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; @@ -273,7 +304,7 @@ index 24e1f1f0..a54fd6a2 100644 enum ggml_type new_type; void * new_data; size_t new_size; -@@ -18655,6 +18837,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -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. @@ -283,7 +314,7 @@ index 24e1f1f0..a54fd6a2 100644 if (!params->pure && ggml_is_quantized(default_type)) { new_type = llama_tensor_get_type(qs, new_type, tensor, ftype); } -@@ -18664,6 +18849,7 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -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; }