Skip to content

Commit 1d13041

Browse files
leejetleejet
committed
fix: resolve precision issues in SDXL VAE under fp16
1 parent 2e9242e commit 1d13041

File tree

6 files changed

+55
-44
lines changed

6 files changed

+55
-44
lines changed

‎README.md‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ API and command-line option may change frequently.***
1717
- Image Models
1818
- SD1.x, SD2.x, [SD-Turbo](https://huggingface.co/stabilityai/sd-turbo)
1919
- SDXL, [SDXL-Turbo](https://huggingface.co/stabilityai/sdxl-turbo)
20-
- !!!The VAE in SDXL encounters NaN issues under FP16, but unfortunately, the ggml_conv_2d only operates under FP16. Hence, a parameter is needed to specify the VAE that has fixed the FP16 NaN issue. You can find it here: [SDXL VAE FP16 Fix](https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/blob/main/sdxl_vae.safetensors).
2120
- [SD3/SD3.5](./docs/sd3.md)
2221
- [Flux-dev/Flux-schnell](./docs/flux.md)
2322
- [Chroma](./docs/chroma.md)

‎conditioner.hpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1457,7 +1457,7 @@ struct Qwen2_5_VLCLIPEmbedder : public Conditioner {
14571457
const ConditionerParams& conditioner_params) {
14581458
std::string prompt;
14591459
std::vector<std::pair<int, ggml_tensor*>> image_embeds;
1460-
size_t system_prompt_length = 0;
1460+
size_t system_prompt_length = 0;
14611461
int prompt_template_encode_start_idx = 34;
14621462
if (qwenvl->enable_vision && conditioner_params.ref_images.size() > 0) {
14631463
LOG_INFO("QwenImageEditPlusPipeline");

‎ggml_extend.hpp‎

Lines changed: 36 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -975,38 +975,28 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_conv_2d(struct ggml_context* ctx,
975975
struct ggml_tensor* x,
976976
struct ggml_tensor* w,
977977
struct ggml_tensor* b,
978-
int s0 = 1,
979-
int s1 = 1,
980-
int p0 = 0,
981-
int p1 = 0,
982-
int d0 = 1,
983-
int d1 = 1) {
984-
x = ggml_conv_2d(ctx, w, x, s0, s1, p0, p1, d0, d1);
985-
if (b != NULL) {
986-
b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
987-
// b = ggml_repeat(ctx, b, x);
988-
x = ggml_add_inplace(ctx, x, b);
978+
int s0 = 1,
979+
int s1 = 1,
980+
int p0 = 0,
981+
int p1 = 0,
982+
int d0 = 1,
983+
int d1 = 1,
984+
bool direct = false,
985+
float scale = 1.f) {
986+
if (scale != 1.f) {
987+
x = ggml_scale(ctx, x, scale);
988+
}
989+
if (direct) {
990+
x = ggml_conv_2d_direct(ctx, w, x, s0, s1, p0, p1, d0, d1);
991+
} else {
992+
x = ggml_conv_2d(ctx, w, x, s0, s1, p0, p1, d0, d1);
993+
}
994+
if (scale != 1.f) {
995+
x = ggml_scale(ctx, x, 1.f / scale);
989996
}
990-
return x;
991-
}
992-
993-
// w: [OC*IC, KD, KH, KW]
994-
// x: [N*IC, ID, IH, IW]
995-
__STATIC_INLINE__ struct ggml_tensor* ggml_nn_conv_2d_direct(struct ggml_context* ctx,
996-
struct ggml_tensor* x,
997-
struct ggml_tensor* w,
998-
struct ggml_tensor* b,
999-
int s0 = 1,
1000-
int s1 = 1,
1001-
int p0 = 0,
1002-
int p1 = 0,
1003-
int d0 = 1,
1004-
int d1 = 1) {
1005-
x = ggml_conv_2d_direct(ctx, w, x, s0, s1, p0, p1, d0, d1);
1006997
if (b != NULL) {
1007998
b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
1008-
// b = ggml_repeat(ctx, b, x);
1009-
x = ggml_add(ctx, x, b);
999+
x = ggml_add_inplace(ctx, x, b);
10101000
}
10111001
return x;
10121002
}
@@ -2067,6 +2057,7 @@ class Conv2d : public UnaryBlock {
20672057
std::pair<int, int> dilation;
20682058
bool bias;
20692059
bool direct = false;
2060+
float scale = 1.f;
20702061

20712062
void init_params(struct ggml_context* ctx, const String2GGMLType& tensor_types, const std::string prefix = "") {
20722063
enum ggml_type wtype = GGML_TYPE_F16;
@@ -2097,6 +2088,10 @@ class Conv2d : public UnaryBlock {
20972088
direct = true;
20982089
}
20992090

2091+
void set_scale(float scale_value) {
2092+
scale = scale_value;
2093+
}
2094+
21002095
std::string get_desc() {
21012096
return "Conv2d";
21022097
}
@@ -2107,11 +2102,18 @@ class Conv2d : public UnaryBlock {
21072102
if (bias) {
21082103
b = params["bias"];
21092104
}
2110-
if (direct) {
2111-
return ggml_nn_conv_2d_direct(ctx, x, w, b, stride.second, stride.first, padding.second, padding.first, dilation.second, dilation.first);
2112-
} else {
2113-
return ggml_nn_conv_2d(ctx, x, w, b, stride.second, stride.first, padding.second, padding.first, dilation.second, dilation.first);
2114-
}
2105+
return ggml_nn_conv_2d(ctx,
2106+
x,
2107+
w,
2108+
b,
2109+
stride.second,
2110+
stride.first,
2111+
padding.second,
2112+
padding.first,
2113+
dilation.second,
2114+
dilation.first,
2115+
direct,
2116+
scale);
21152117
}
21162118
};
21172119

‎qwen_image.hpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -535,7 +535,7 @@ namespace Qwen {
535535
}
536536
}
537537
LOG_ERROR("qwen_image_params.num_layers: %ld", qwen_image_params.num_layers);
538-
qwen_image = QwenImageModel(qwen_image_params);
538+
qwen_image = QwenImageModel(qwen_image_params);
539539
qwen_image.init(params_ctx, tensor_types, prefix);
540540
}
541541

‎stable-diffusion.cpp‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -330,13 +330,6 @@ class StableDiffusionGGML {
330330

331331
if (sd_version_is_sdxl(version)) {
332332
scale_factor = 0.13025f;
333-
if (strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 && strlen(SAFE_STR(sd_ctx_params->taesd_path)) == 0) {
334-
LOG_WARN(
335-
"!!!It looks like you are using SDXL model. "
336-
"If you find that the generated images are completely black, "
337-
"try specifying SDXL VAE FP16 Fix with the --vae parameter. "
338-
"You can find it here: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/blob/main/sdxl_vae.safetensors");
339-
}
340333
} else if (sd_version_is_sd3(version)) {
341334
scale_factor = 1.5305f;
342335
} else if (sd_version_is_flux(version)) {
@@ -517,6 +510,11 @@ class StableDiffusionGGML {
517510
LOG_INFO("Using Conv2d direct in the vae model");
518511
first_stage_model->enable_conv2d_direct();
519512
}
513+
if (version == VERSION_SDXL && strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0) {
514+
float vae_conv_2d_scale = 1.f / 32.f;
515+
LOG_WARN("No VAE specified with --vae, using Conv2D scale %.3f", vae_conv_2d_scale);
516+
first_stage_model->set_conv2d_scale(vae_conv_2d_scale);
517+
}
520518
first_stage_model->alloc_params_buffer();
521519
first_stage_model->get_param_tensors(tensors, "first_stage_model");
522520
} else {

‎vae.hpp‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -530,6 +530,7 @@ struct VAE : public GGMLRunner {
530530
struct ggml_context* output_ctx) = 0;
531531
virtual void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) = 0;
532532
virtual void enable_conv2d_direct(){};
533+
virtual void set_conv2d_scale(float scale) { SD_UNUSED(scale); };
533534
};
534535

535536
struct AutoEncoderKL : public VAE {
@@ -558,6 +559,17 @@ struct AutoEncoderKL : public VAE {
558559
}
559560
}
560561

562+
void set_conv2d_scale(float scale) {
563+
std::vector<GGMLBlock*> blocks;
564+
ae.get_all_blocks(blocks);
565+
for (auto block : blocks) {
566+
if (block->get_desc() == "Conv2d") {
567+
auto conv_block = (Conv2d*)block;
568+
conv_block->set_scale(scale);
569+
}
570+
}
571+
}
572+
561573
std::string get_desc() {
562574
return "vae";
563575
}

0 commit comments

Comments
 (0)