fix: correct the handling of weight loading

models/convert.py

@@ -31,7 +31,7 @@ ggml_ttype_str_to_int = {
 
 QK4_0 = 32
 def quantize_q4_0(x):
-    assert x.shape[-1] % QK4_0 == 0
+    assert x.shape[-1] % QK4_0 == 0 and x.shape[-1] > QK4_0
     x = x.reshape(-1, QK4_0)
     max = np.take_along_axis(x, np.argmax(np.abs(x), axis=-1)[:, np.newaxis], axis=-1)
     d = max / -8
@@ -44,7 +44,7 @@ def quantize_q4_0(x):
 
 QK4_1 = 32
 def quantize_q4_1(x):
-    assert x.shape[-1] % QK4_1 == 0
+    assert x.shape[-1] % QK4_1 == 0 and x.shape[-1] > QK4_1
     x = x.reshape(-1, QK4_1)
     min = np.min(x, axis=-1, keepdims=True)
     max = np.max(x, axis=-1, keepdims=True)
@@ -59,7 +59,7 @@ def quantize_q4_1(x):
 
 QK5_0 = 32
 def quantize_q5_0(x):
-    assert x.shape[1] % QK5_0 == 0
+    assert x.shape[-1] % QK5_0 == 0 and x.shape[-1] > QK5_0
     x = x.reshape(-1, QK5_0)
     max = np.take_along_axis(x, np.argmax(np.abs(x), axis=-1)[:, np.newaxis], axis=-1)
     d = max / -16
@@ -76,7 +76,7 @@ def quantize_q5_0(x):
 
 QK5_1 = 32
 def quantize_q5_1(x):
-    assert x.shape[-1] % QK5_1 == 0
+    assert x.shape[-1] % QK5_1 == 0 and x.shape[-1] > QK5_1
     x = x.reshape(-1, QK5_1)
     min = np.min(x, axis=-1, keepdims=True)
     max = np.max(x, axis=-1, keepdims=True)
@@ -95,7 +95,7 @@ def quantize_q5_1(x):
 
 QK8_0 = 32
 def quantize_q8_0(x):
-    assert x.shape[-1] % QK8_0 == 0
+    assert x.shape[-1] % QK8_0 == 0 and x.shape[-1] > QK8_0
     x = x.reshape(-1, QK8_0)
     amax = np.max(np.abs(x), axis=-1, keepdims=True) 
     d = amax / ((1 << 7) - 1)
@@ -156,7 +156,10 @@ unused_tensors = [
     "posterior_mean_coef2",
     "cond_stage_model.transformer.text_model.embeddings.position_ids",
     "model_ema.decay",
-    "model_ema.num_updates"
+    "model_ema.num_updates",
+    "control_model",
+    "lora_te_text_model",
+    "embedding_manager"
 ]
 
 def convert(model_path, out_type = None, out_file=None):
@@ -182,6 +185,10 @@ def convert(model_path, out_type = None, out_file=None):
             out_type = "f32"
         elif weight.dtype == np.float16:
             out_type = "f16"
+        elif weight.dtype == np.float64:
+            out_type = "f32"
+        else:
+            raise Exception("unsupported weight type %s" % weight.dtype)
     if out_file == None:
         out_file = os.path.splitext(os.path.basename(model_path))[0] + f"-ggml-model-{out_type}.bin"
         out_file = os.path.join(os.getcwd(), out_file)
@@ -207,6 +214,13 @@ def convert(model_path, out_type = None, out_file=None):
         for name in state_dict.keys():
             if not isinstance(state_dict[name], torch.Tensor):
                 continue
+            skip = False
+            for unused_tensor in unused_tensors:
+                if name.startswith(unused_tensor):
+                    skip = True
+                    break
+            if skip:
+                continue
             if name in unused_tensors:
                 continue
             data = state_dict[name].numpy()

stable-diffusion.cpp

@@ -26,12 +26,16 @@ static SDLogLevel log_level = SDLogLevel::INFO;
         }                                                                                             \
         if (level == SDLogLevel::DEBUG) {                                                             \
             printf("[DEBUG] %s:%-4d - " format "\n", __FILENAME__, __LINE__, ##__VA_ARGS__);          \
+            fflush(stdout);                                                                           \
         } else if (level == SDLogLevel::INFO) {                                                       \
             printf("[INFO]  %s:%-4d - " format "\n", __FILENAME__, __LINE__, ##__VA_ARGS__);          \
+            fflush(stdout);                                                                           \
         } else if (level == SDLogLevel::WARN) {                                                       \
             fprintf(stderr, "[WARN]  %s:%-4d - " format "\n", __FILENAME__, __LINE__, ##__VA_ARGS__); \
+            fflush(stdout);                                                                           \
         } else if (level == SDLogLevel::ERROR) {                                                      \
             fprintf(stderr, "[ERROR] %s:%-4d - " format "\n", __FILENAME__, __LINE__, ##__VA_ARGS__); \
+            fflush(stdout);                                                                           \
         }                                                                                             \
     } while (0)
 
@@ -135,6 +139,7 @@ float ggml_tensor_get_f32(const ggml_tensor* tensor, int l, int k = 0, int j = 0
 
 void print_ggml_tensor(struct ggml_tensor* tensor, bool shape_only = false) {
     printf("shape(%zu, %zu, %zu, %zu)\n", tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]);
+    fflush(stdout);
     if (shape_only) {
         return;
     }
@@ -156,6 +161,7 @@ void print_ggml_tensor(struct ggml_tensor* tensor, bool shape_only = false) {
                         continue;
                     }
                     printf("  [%d, %d, %d, %d] = %f\n", i, j, k, l, ggml_tensor_get_f32(tensor, l, k, j, i));
+                    fflush(stdout);
                 }
             }
         }
@@ -251,6 +257,11 @@ void image_vec_to_ggml(const std::vector<uint8_t>& vec,
     }
 }
 
+struct ggml_tensor * ggml_group_norm_32(struct ggml_context * ctx,
+                                        struct ggml_tensor * a) {
+    return ggml_group_norm(ctx, a, 32);
+}
+
 /*================================================== CLIPTokenizer ===================================================*/
 
 const std::string UNK_TOKEN = "<|endoftext|>";
@@ -899,7 +910,7 @@ struct ResBlock {
 
         // in_layers
         // group norm 32
-        auto h = ggml_group_norm(ctx, x);
+        auto h = ggml_group_norm_32(ctx, x);
         h = ggml_add(ctx,
                      ggml_mul(ctx,
                               ggml_repeat(ctx,
@@ -929,7 +940,7 @@ struct ResBlock {
         // out_layers
         h = ggml_add(ctx, h, emb_out);
         // group norm 32
-        h = ggml_group_norm_inplace(ctx, h);
+        h = ggml_group_norm_inplace(ctx, h, 32);
         h = ggml_add(ctx,
                      ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, out_layer_0_w, 1, 1, out_layer_0_w->ne[0], 1), h), h),
                      ggml_repeat(ctx, ggml_reshape_4d(ctx, out_layer_0_b, 1, 1, out_layer_0_b->ne[0], 1), h));
@@ -1122,7 +1133,7 @@ struct SpatialTransformer {
 
         auto x_in = x;
         // group norm 32
-        x = ggml_group_norm(ctx, x);
+        x = ggml_group_norm_32(ctx, x);
         x = ggml_add(ctx,
                      ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_w, 1, 1, norm_w->ne[0], 1), x), x),
                      ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_b, 1, 1, norm_b->ne[0], 1), x));
@@ -1424,7 +1435,7 @@ struct UpSample {
 
     struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {
         // x: [N, channels, h, w]
-        x = ggml_upscale(ctx, x);  // [N, channels, h*2, w*2]
+        x = ggml_upscale(ctx, x, 2);  // [N, channels, h*2, w*2]
         x = ggml_conv_2d(ctx, conv_w, x, 1, 1, 1, 1, 1, 1);
 
         x = ggml_add(ctx,
@@ -1815,7 +1826,7 @@ struct UNetModel {
 
         // out
         // group norm 32
-        h = ggml_group_norm(ctx, h);
+        h = ggml_group_norm_32(ctx, h);
         h = ggml_add(ctx,
                      ggml_mul(ctx,
                               ggml_repeat(ctx,
@@ -1919,7 +1930,7 @@ struct ResnetBlock {
         // z: [N, in_channels, h, w]
 
         // group norm 32
-        auto h = ggml_group_norm(ctx, z);
+        auto h = ggml_group_norm_32(ctx, z);
         h = ggml_mul(ctx,
                      ggml_repeat(ctx,
                                  ggml_reshape_4d(ctx, norm1_w, 1, 1, norm1_w->ne[0], 1),
@@ -1941,7 +1952,7 @@ struct ResnetBlock {
                                  h));  // [N, out_channels, h, w]
 
         // group norm 32
-        h = ggml_group_norm(ctx, h);
+        h = ggml_group_norm_32(ctx, h);
         h = ggml_add(ctx,
                      ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, norm2_w, 1, 1, norm2_w->ne[0], 1), h), h),
                      ggml_repeat(ctx, ggml_reshape_4d(ctx, norm2_b, 1, 1, norm2_b->ne[0], 1), h));
@@ -2028,7 +2039,7 @@ struct AttnBlock {
         // x: [N, in_channels, h, w]
 
         // group norm 32
-        auto h_ = ggml_group_norm(ctx, x);
+        auto h_ = ggml_group_norm_32(ctx, x);
         h_ = ggml_add(ctx,
                       ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_w, 1, 1, norm_w->ne[0], 1), h_), h_),
                       ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_b, 1, 1, norm_b->ne[0], 1), h_));
@@ -2253,7 +2264,7 @@ struct Encoder {
         h = mid.block_2.forward(ctx, h);  // [N, block_in, h, w]
 
         // group norm 32
-        h = ggml_group_norm(ctx, h);
+        h = ggml_group_norm_32(ctx, h);
         h = ggml_add(ctx,
                      ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_out_w, 1, 1, norm_out_w->ne[0], 1), h), h),
                      ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_out_b, 1, 1, norm_out_b->ne[0], 1), h));
@@ -2435,7 +2446,7 @@ struct Decoder {
         }
 
         // group norm 32
-        h = ggml_group_norm(ctx, h);
+        h = ggml_group_norm_32(ctx, h);
         h = ggml_add(ctx,
                      ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_out_w, 1, 1, norm_out_w->ne[0], 1), h), h),
                      ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_out_b, 1, 1, norm_out_b->ne[0], 1), h));
@@ -2853,6 +2864,8 @@ class StableDiffusionGGML {
                     nelements *= ne[i];
                 }
 
+                const size_t num_bytes = nelements / ggml_blck_size(ggml_type(ttype)) * ggml_type_size(ggml_type(ttype));
+
                 std::string name(length, 0);
                 file.read(&name[0], length);
 
@@ -2880,7 +2893,7 @@ class StableDiffusionGGML {
                             return false;
                         }
                     }
-                    file.ignore(nelements * ggml_type_size((ggml_type)ttype));
+                    file.ignore(num_bytes);
                     continue;
                 }
 
@@ -2908,8 +2921,6 @@ class StableDiffusionGGML {
                     return false;
                 }
 
-                const size_t num_bytes = nelements / ggml_blck_size(ggml_type(ttype)) * ggml_type_size(ggml_type(ttype));
-
                 file.read(reinterpret_cast<char*>(tensor->data), num_bytes);
 
                 total_size += ggml_nbytes(tensor);