new model, new sampling method.

inference.py

@@ -73,7 +73,7 @@ def mel_to_wav(net, netcpp, mel, config):
     )
     ft = wavegru_inference(net, mel)
     ft = jax.device_get(ft[0])
-    wav = netcpp.inference(ft, 0.9)
+    wav = netcpp.inference(ft, 1.0)
     wav = np.array(wav)
     wav = librosa.mu_expand(wav - 127, mu=255)
     wav = librosa.effects.deemphasis(wav, coef=0.86)

wavegru.ckpt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c06310d989fd524359d5f3ecf8ea1dc146980bb594b7b90553d0d42a64c512d8
+oid sha256:64b1ce6558cfe09b95bb29d48c34900ed0c3490d17c81fd6190969b226f4617a
 size 58039876

wavegru_mod.cc

@@ -39,6 +39,7 @@ struct WaveGRU {
   vec o1b, o2b;
   vec t;
   vec h;
+  vec logits;
   mat o1, o2;
   std::vector<vec> embed;
 
@@ -55,7 +56,8 @@ struct WaveGRU {
         fco2(256),
         h(hidden_dim),
         o1b(hidden_dim),
-        o2b(256) {
+        o2b(256),
+        logits(256) {
     m = create_mat(hidden_dim, 3*hidden_dim);
     o1 = create_mat(hidden_dim, hidden_dim);
     o2 = create_mat(hidden_dim, 256);
@@ -120,6 +122,18 @@ struct WaveGRU {
       }
       o1.SpMM_bias(h, o1b, &fco1, true);
       o2.SpMM_bias(fco1, o2b, &fco2, false);
+      auto max_logit = fco2[0];
+      for (int i = 1; i <= 255; ++i) {
+        max_logit = max(max_logit, fco2[i]);
+      }
+      float total = 0.0;
+      for (int i = 0; i <= 255; ++i) {
+        logits[i] = csrblocksparse::fast_exp(fco2[i] - max_logit);
+        total += logits[i];
+      }
+      for (int i = 0; i <= 255; ++i) {
+        if (logits[i] < total / 256.0) fco2[i] = -1e9;
+      }
       value = fco2.Sample(temperature);
       signal[index] = value;
     }

wavegru_mod.so

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:12c27f0ea07f8da3a3ab48bc01bb0f68971ce7d57b19ada87669eab138623a9c
+oid sha256:b65c5312f24f8ab9cfa51e8340a24ac1165b247046a331386d636fba9036c19c
 size 525536