update

example/export_samples.py

@@ -112,6 +112,8 @@ def main():
     int_like = stats.get("int_like", {})
     max_decimals = stats.get("max_decimals", {})
     transforms = stats.get("transform", {})
+    cont_target = str(cfg.get("cont_target", "eps"))
+    cont_clamp_x0 = float(cfg.get("cont_clamp_x0", 0.0))
 
     vocab_json = json.load(open(args.vocab_path, "r", encoding="utf-8"))
     vocab = vocab_json["vocab"]
@@ -191,6 +193,12 @@ def main():
         t_batch = torch.full((args.batch_size,), t, device=device, dtype=torch.long)
         eps_pred, logits = model(x_cont, x_disc, t_batch, cond)
 
+        if cont_target == "x0":
+            x0_pred = eps_pred
+            if cont_clamp_x0 > 0:
+                x0_pred = torch.clamp(x0_pred, -cont_clamp_x0, cont_clamp_x0)
+            eps_pred = (x_cont - torch.sqrt(a_bar_t) * x0_pred) / torch.sqrt(1.0 - a_bar_t)
+
         a_t = alphas[t]
         a_bar_t = alphas_cumprod[t]
         coef1 = 1.0 / torch.sqrt(a_t)