Add support for SD1x and SDXL models
This commit is contained in:
@@ -34,9 +34,7 @@ def gguf_sd_loader(path):
|
||||
sd[str(tensor.name)] = GGMLTensor(
|
||||
torch.from_numpy(tensor.data), # mmap
|
||||
tensor_type = tensor.tensor_type,
|
||||
tensor_shape = torch.Size(
|
||||
np.flip(list(tensor.shape))
|
||||
)
|
||||
tensor_shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape)))
|
||||
)
|
||||
dt[str(tensor.tensor_type)] = dt.get(str(tensor.tensor_type), 0) + 1
|
||||
|
||||
|
||||
+7
-5
@@ -13,6 +13,7 @@ def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Generate F16 GGUF files from single UNET")
|
||||
parser.add_argument("--src", required=True, help="Source model ckpt file.")
|
||||
parser.add_argument("--dst", help="Output unet gguf file.")
|
||||
parser.add_argument("--force-arch", type=str, choices=("flux", "sd1", "sdxl"), default=None, help="Advanced option, allows forcing the model architecture. May be necessary for non-Flux models.")
|
||||
args = parser.parse_args()
|
||||
|
||||
if not os.path.isfile(args.src):
|
||||
@@ -39,11 +40,12 @@ def load_state_dict(path):
|
||||
|
||||
return sd
|
||||
|
||||
def load_model(path):
|
||||
def load_model(path, force_arch=None):
|
||||
state_dict = load_state_dict(path)
|
||||
|
||||
# from ComfyUI model detection
|
||||
if "transformer_blocks.0.attn.norm_added_k.weight" in state_dict:
|
||||
if force_arch is not None:
|
||||
arch = force_arch
|
||||
elif "transformer_blocks.0.attn.norm_added_k.weight" in state_dict:
|
||||
arch = "flux"
|
||||
raise ValueError(f"The Diffusers UNET can not be used for this!")
|
||||
elif "double_blocks.0.img_attn.proj.weight" in state_dict:
|
||||
@@ -54,7 +56,7 @@ def load_model(path):
|
||||
if "add_embedding.linear_1.weight" in state_dict:
|
||||
arch = "sdxl"
|
||||
else:
|
||||
arch = "sd1"
|
||||
arch = "sd1"
|
||||
else:
|
||||
breakpoint()
|
||||
raise ValueError(f"Unknown model architecture!")
|
||||
@@ -141,7 +143,7 @@ def handle_tensors(args, writer, state_dict):
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
path = args.src
|
||||
writer, state_dict = load_model(path)
|
||||
writer, state_dict = load_model(path, force_arch=args.force_arch)
|
||||
|
||||
writer.add_quantization_version(gguf.GGML_QUANT_VERSION)
|
||||
if next(iter(state_dict.values())).dtype == torch.bfloat16:
|
||||
|
||||
+18
-14
@@ -1,40 +1,43 @@
|
||||
diff --git a/src/llama.cpp b/src/llama.cpp
|
||||
index 5ab65ea9..44fe3fe6 100644
|
||||
index 5ab65ea9..ef63e28f 100644
|
||||
--- a/src/llama.cpp
|
||||
+++ b/src/llama.cpp
|
||||
@@ -212,6 +212,8 @@ enum llm_arch {
|
||||
@@ -212,6 +212,9 @@ enum llm_arch {
|
||||
LLM_ARCH_JAIS,
|
||||
LLM_ARCH_NEMOTRON,
|
||||
LLM_ARCH_EXAONE,
|
||||
+ LLM_ARCH_FLUX,
|
||||
+ LLM_ARCH_SD1,
|
||||
+ LLM_ARCH_SDXL,
|
||||
LLM_ARCH_UNKNOWN,
|
||||
};
|
||||
|
||||
@@ -259,6 +261,8 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
|
||||
@@ -259,6 +262,9 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
|
||||
{ LLM_ARCH_JAIS, "jais" },
|
||||
{ LLM_ARCH_NEMOTRON, "nemotron" },
|
||||
{ LLM_ARCH_EXAONE, "exaone" },
|
||||
+ { LLM_ARCH_FLUX, "flux" },
|
||||
+ { LLM_ARCH_SD1, "sd1" },
|
||||
+ { LLM_ARCH_SDXL, "sdxl" },
|
||||
{ LLM_ARCH_UNKNOWN, "(unknown)" },
|
||||
};
|
||||
|
||||
@@ -1337,6 +1341,8 @@ static const std::map<llm_arch, std::map<llm_tensor, std::string>> LLM_TENSOR_NA
|
||||
@@ -1337,6 +1343,9 @@ static const std::map<llm_arch, std::map<llm_tensor, std::string>> LLM_TENSOR_NA
|
||||
{ LLM_TENSOR_FFN_UP, "blk.%d.ffn_up" },
|
||||
},
|
||||
},
|
||||
+ { LLM_ARCH_FLUX, { {} }, },
|
||||
+ { LLM_ARCH_SD1, { {} }, },
|
||||
+ { LLM_ARCH_SDXL, { {} }, },
|
||||
{
|
||||
LLM_ARCH_UNKNOWN,
|
||||
{
|
||||
@@ -4629,6 +4635,13 @@ static void llm_load_hparams(
|
||||
@@ -4629,6 +4638,13 @@ 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) {
|
||||
+ if (model.arch == LLM_ARCH_FLUX || model.arch == LLM_ARCH_SD1 || model.arch == LLM_ARCH_SDXL) {
|
||||
+ model.ftype = ml.ftype;
|
||||
+ hparams.rope_type = llama_rope_type(&model);
|
||||
+ return;
|
||||
@@ -43,7 +46,7 @@ index 5ab65ea9..44fe3fe6 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);
|
||||
|
||||
@@ -15854,26 +15867,40 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
|
||||
@@ -15854,26 +15870,40 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
|
||||
return std::make_pair(i_layer, n_layer);
|
||||
};
|
||||
|
||||
@@ -102,7 +105,7 @@ index 5ab65ea9..44fe3fe6 100644
|
||||
if (qs.params->token_embedding_type < GGML_TYPE_COUNT) {
|
||||
new_type = qs.params->token_embedding_type;
|
||||
} else {
|
||||
@@ -15891,10 +15918,13 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
|
||||
@@ -15891,10 +15921,13 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
|
||||
new_type == GGML_TYPE_Q4_0_8_8) {
|
||||
new_type = GGML_TYPE_Q4_0;
|
||||
}
|
||||
@@ -117,7 +120,7 @@ index 5ab65ea9..44fe3fe6 100644
|
||||
if (qs.model.hparams.n_gqa() >= 4 || qs.model.hparams.n_expert >= 4) new_type = GGML_TYPE_Q4_K;
|
||||
else new_type = ftype == LLAMA_FTYPE_MOSTLY_IQ2_S || ftype == LLAMA_FTYPE_MOSTLY_IQ2_M ? GGML_TYPE_IQ3_S : GGML_TYPE_Q2_K;
|
||||
++qs.i_attention_wv;
|
||||
@@ -15916,7 +15946,7 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
|
||||
@@ -15916,7 +15949,7 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
|
||||
else if (ftype == LLAMA_FTYPE_MOSTLY_IQ2_S || ftype == LLAMA_FTYPE_MOSTLY_IQ2_M) new_type = GGML_TYPE_IQ3_S;
|
||||
}
|
||||
}
|
||||
@@ -126,7 +129,7 @@ index 5ab65ea9..44fe3fe6 100644
|
||||
if (ftype == LLAMA_FTYPE_MOSTLY_Q2_K) {
|
||||
new_type = qs.model.hparams.n_gqa() >= 4 ? GGML_TYPE_Q4_K : GGML_TYPE_Q3_K;
|
||||
}
|
||||
@@ -15954,7 +15984,7 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
|
||||
@@ -15954,7 +15987,7 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
|
||||
new_type = GGML_TYPE_Q8_0;
|
||||
}
|
||||
++qs.i_attention_wv;
|
||||
@@ -135,7 +138,7 @@ index 5ab65ea9..44fe3fe6 100644
|
||||
if (qs.model.hparams.n_expert == 8) {
|
||||
// for the 8-expert model, bumping this to Q8_0 trades just ~128MB
|
||||
// TODO: explore better strategies
|
||||
@@ -15966,7 +15996,7 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
|
||||
@@ -15966,7 +15999,7 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
|
||||
else if (ftype == LLAMA_FTYPE_MOSTLY_IQ3_XXS) {
|
||||
new_type = GGML_TYPE_IQ2_S;
|
||||
}
|
||||
@@ -144,7 +147,7 @@ index 5ab65ea9..44fe3fe6 100644
|
||||
if (ftype == LLAMA_FTYPE_MOSTLY_IQ3_XS) {
|
||||
new_type = GGML_TYPE_IQ3_XXS;
|
||||
}
|
||||
@@ -16038,7 +16068,7 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
|
||||
@@ -16038,7 +16071,7 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
|
||||
if (ftype == LLAMA_FTYPE_MOSTLY_Q3_K_L) new_type = GGML_TYPE_Q4_K;
|
||||
}
|
||||
}
|
||||
@@ -153,7 +156,7 @@ index 5ab65ea9..44fe3fe6 100644
|
||||
if (ftype == LLAMA_FTYPE_MOSTLY_Q3_K_M || ftype == LLAMA_FTYPE_MOSTLY_Q3_K_L || ftype == LLAMA_FTYPE_MOSTLY_IQ3_M) {
|
||||
new_type = GGML_TYPE_Q4_K;
|
||||
}
|
||||
@@ -16107,6 +16137,8 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
|
||||
@@ -16107,6 +16140,8 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
|
||||
}
|
||||
LLAMA_LOG_WARN(" - using fallback quantization %s\n", ggml_type_name(new_type));
|
||||
++qs.n_fallback;
|
||||
@@ -162,12 +165,13 @@ index 5ab65ea9..44fe3fe6 100644
|
||||
}
|
||||
|
||||
return new_type;
|
||||
@@ -17432,6 +17464,8 @@ enum llama_rope_type llama_rope_type(const struct llama_model * model) {
|
||||
@@ -17432,6 +17467,9 @@ enum llama_rope_type llama_rope_type(const struct llama_model * model) {
|
||||
case LLM_ARCH_T5:
|
||||
case LLM_ARCH_T5ENCODER:
|
||||
case LLM_ARCH_JAIS:
|
||||
+ case LLM_ARCH_FLUX:
|
||||
+ case LLM_ARCH_SD1:
|
||||
+ case LLM_ARCH_SDXL:
|
||||
return LLAMA_ROPE_TYPE_NONE;
|
||||
|
||||
// use what we call a normal RoPE, operating on pairs of consecutive head values
|
||||
|
||||
Reference in New Issue
Block a user