     1→Device: cuda
     2→LeWM TTS model: 24.00M parameters
     3→Dataset: 6959 samples loaded
     4→Training: 6959 samples, 27 batches/epoch
     5→Epoch 1/200 | avg_loss=13.6511 pred=1.4613 kl=0.2462 recon=2.3551 spec=4.9864 | 27.3s
     6→  New best model (loss=13.6511)
     7→  [2/200] step 50 | loss=13.3484 pred=1.1684 kl=0.1995 recon=2.3801 spec=4.9465 | grad=2.33 | 22.7s
     8→Epoch 2/200 | avg_loss=13.2451 pred=1.2534 kl=0.2111 recon=2.3552 spec=4.8541 | 26.6s
     9→  New best model (loss=13.2451)
    10→Epoch 3/200 | avg_loss=12.5440 pred=1.0315 kl=0.2030 recon=2.3546 spec=4.5354 | 26.7s
    11→  New best model (loss=12.5440)
    12→  [4/200] step 100 | loss=10.7409 pred=0.9306 kl=0.3007 recon=2.3797 spec=3.3671 | grad=3.37 | 18.9s
    13→Epoch 4/200 | avg_loss=11.0376 pred=0.9384 kl=0.2805 recon=2.3558 spec=3.5916 | 26.7s
    14→  New best model (loss=11.0376)
    15→Epoch 5/200 | avg_loss=8.9693 pred=0.8248 kl=0.2909 recon=2.2691 spec=2.4041 | 26.7s
    16→  New best model (loss=8.9693)
    17→  [6/200] step 150 | loss=6.7562 pred=0.7736 kl=0.2844 recon=1.6614 spec=1.7729 | grad=4.86 | 15.0s
    18→Epoch 6/200 | avg_loss=6.9467 pred=0.7647 kl=0.2749 recon=1.6966 spec=1.8590 | 26.6s
    19→  New best model (loss=6.9467)
    20→Epoch 7/200 | avg_loss=5.4291 pred=0.6201 kl=0.3056 recon=1.4078 spec=1.3287 | 26.8s
    21→  New best model (loss=5.4291)
    22→  [8/200] step 200 | loss=4.5682 pred=0.5125 kl=0.3461 recon=1.2266 spec=1.0678 | grad=4.27 | 11.3s
    23→Epoch 8/200 | avg_loss=4.5015 pred=0.5005 kl=0.3539 recon=1.2074 spec=1.0571 | 26.8s
    24→  New best model (loss=4.5015)
    25→Epoch 9/200 | avg_loss=3.8949 pred=0.4052 kl=0.4210 recon=1.0379 spec=0.9421 | 26.8s
    26→  New best model (loss=3.8949)
    27→  [10/200] step 250 | loss=3.6904 pred=0.3596 kl=0.4723 recon=0.9872 spec=0.9036 | grad=12.19 | 7.1s
    28→Epoch 10/200 | avg_loss=3.6212 pred=0.3419 kl=0.4878 recon=0.9624 spec=0.9022 | 26.6s
    29→  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0010.pt
    30→  New best model (loss=3.6212)
    31→Epoch 11/200 | avg_loss=3.4785 pred=0.2867 kl=0.5559 recon=0.9338 spec=0.8818 | 26.7s
    32→  New best model (loss=3.4785)
    33→  [12/200] step 300 | loss=3.4295 pred=0.2622 kl=0.5906 recon=0.9275 spec=0.8737 | grad=4.26 | 3.4s
    34→Epoch 12/200 | avg_loss=3.3684 pred=0.2458 kl=0.6160 recon=0.9131 spec=0.8631 | 26.7s
    35→  New best model (loss=3.3684)
    36→  [13/200] step 350 | loss=3.2924 pred=0.1971 kl=0.7111 recon=0.8917 spec=0.8731 | grad=40.29 | 25.6s
    37→Epoch 13/200 | avg_loss=3.2882 pred=0.2122 kl=0.6808 recon=0.8986 spec=0.8510 | 26.6s
    38→  New best model (loss=3.2882)
    39→Epoch 14/200 | avg_loss=3.4840 pred=0.1903 kl=0.7435 recon=0.9788 spec=0.8891 | 26.7s
    40→  [15/200] step 400 | loss=3.3437 pred=0.1641 kl=0.8119 recon=0.9571 spec=0.8416 | grad=23.07 | 21.7s
    41→Epoch 15/200 | avg_loss=3.3427 pred=0.1730 kl=0.7889 recon=0.9350 spec=0.8645 | 26.6s
    42→Epoch 16/200 | avg_loss=3.2980 pred=0.1490 kl=0.8601 recon=0.9295 spec=0.8578 | 26.4s
    43→  [17/200] step 450 | loss=3.1454 pred=0.1204 kl=0.9734 recon=0.8944 spec=0.8214 | grad=13.68 | 18.0s
    44→Epoch 17/200 | avg_loss=3.1228 pred=0.1237 kl=0.9565 recon=0.8765 spec=0.8282 | 26.8s
    45→  New best model (loss=3.1228)
    46→Epoch 18/200 | avg_loss=3.0995 pred=0.0969 kl=1.0964 recon=0.8774 spec=0.8288 | 26.8s
    47→  New best model (loss=3.0995)
    48→  [19/200] step 500 | loss=3.1561 pred=0.0811 kl=1.2092 recon=0.9116 spec=0.8308 | grad=21.60 | 14.1s
    49→Epoch 19/200 | avg_loss=3.0741 pred=0.0815 kl=1.2060 recon=0.8744 spec=0.8255 | 26.7s
    50→  New best model (loss=3.0741)
    51→Epoch 20/200 | avg_loss=3.0775 pred=0.0656 kl=1.3424 recon=0.8822 spec=0.8273 | 26.8s
    52→  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0020.pt
    53→  [21/200] step 550 | loss=3.2473 pred=0.0602 kl=1.4247 recon=0.9552 spec=0.8463 | grad=22.38 | 10.2s
    54→Epoch 21/200 | avg_loss=3.0235 pred=0.0595 kl=1.4346 recon=0.8644 spec=0.8185 | 26.9s
    55→  New best model (loss=3.0235)
    56→Epoch 22/200 | avg_loss=2.9452 pred=0.0495 kl=1.5580 recon=0.8416 spec=0.8027 | 26.8s
    57→  New best model (loss=2.9452)
    58→  [23/200] step 600 | loss=2.8896 pred=0.0462 kl=1.6198 recon=0.8079 spec=0.8124 | grad=2.82 | 6.4s
    59→Epoch 23/200 | avg_loss=2.9057 pred=0.0442 kl=1.6588 recon=0.8278 spec=0.7978 | 26.7s
    60→  New best model (loss=2.9057)
    61→Epoch 24/200 | avg_loss=2.8642 pred=0.0400 kl=1.7860 recon=0.8170 spec=0.7866 | 26.7s
    62→  New best model (loss=2.8642)
    63→  [25/200] step 650 | loss=2.8204 pred=0.0366 kl=1.8447 recon=0.8009 spec=0.7807 | grad=4.15 | 2.5s
    64→Epoch 25/200 | avg_loss=2.8957 pred=0.0356 kl=1.8836 recon=0.8275 spec=0.7957 | 26.6s
    65→  [26/200] step 700 | loss=2.8124 pred=0.0335 kl=1.9722 recon=0.8053 spec=0.7704 | grad=2.36 | 24.9s
    66→Epoch 26/200 | avg_loss=2.8820 pred=0.0343 kl=1.9471 recon=0.8247 spec=0.7907 | 26.9s
    67→Epoch 27/200 | avg_loss=2.8777 pred=0.0320 kl=2.0381 recon=0.8212 spec=0.7933 | 26.6s
    68→  [28/200] step 750 | loss=2.8433 pred=0.0302 kl=2.1263 recon=0.8163 spec=0.7772 | grad=6.28 | 21.0s
    69→Epoch 28/200 | avg_loss=2.9005 pred=0.0306 kl=2.1059 recon=0.8323 spec=0.7939 | 26.8s
    70→Epoch 29/200 | avg_loss=2.9652 pred=0.0316 kl=2.1214 recon=0.8583 spec=0.8013 | 26.7s
    71→  [30/200] step 800 | loss=2.8278 pred=0.0308 kl=2.1655 recon=0.8064 spec=0.7788 | grad=7.38 | 16.8s
    72→Epoch 30/200 | avg_loss=2.8287 pred=0.0321 kl=2.1512 recon=0.8065 spec=0.7786 | 26.5s
    73→  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0030.pt
    74→  New best model (loss=2.8287)
    75→Epoch 31/200 | avg_loss=2.7880 pred=0.0300 kl=2.2180 recon=0.7923 spec=0.7710 | 26.7s
    76→  New best model (loss=2.7880)
    77→  [32/200] step 850 | loss=2.7250 pred=0.0279 kl=2.2650 recon=0.7754 spec=0.7523 | grad=5.50 | 13.0s
    78→Epoch 32/200 | avg_loss=2.7648 pred=0.0283 kl=2.2636 recon=0.7842 spec=0.7668 | 26.6s
    79→  New best model (loss=2.7648)
    80→Epoch 33/200 | avg_loss=2.7526 pred=0.0278 kl=2.2914 recon=0.7802 spec=0.7638 | 26.9s
    81→  New best model (loss=2.7526)
    82→  [34/200] step 900 | loss=2.7781 pred=0.0286 kl=2.3051 recon=0.7798 spec=0.7806 | grad=6.41 | 9.2s
    83→Epoch 34/200 | avg_loss=2.7445 pred=0.0293 kl=2.3163 recon=0.7752 spec=0.7636 | 26.7s
    84→  New best model (loss=2.7445)
    85→Epoch 35/200 | avg_loss=2.7538 pred=0.0283 kl=2.3394 recon=0.7791 spec=0.7648 | 26.6s
    86→  [36/200] step 950 | loss=2.7073 pred=0.0278 kl=2.3435 recon=0.7654 spec=0.7521 | grad=3.65 | 5.3s
    87→Epoch 36/200 | avg_loss=2.7155 pred=0.0279 kl=2.3379 recon=0.7666 spec=0.7558 | 26.8s
    88→  New best model (loss=2.7155)
    89→Epoch 37/200 | avg_loss=2.6678 pred=0.0273 kl=2.3513 recon=0.7473 spec=0.7496 | 26.8s
    90→  New best model (loss=2.6678)
    91→  [38/200] step 1000 | loss=2.6472 pred=0.0271 kl=2.3522 recon=0.7342 spec=0.7534 | grad=5.13 | 1.5s
    92→Epoch 38/200 | avg_loss=2.6343 pred=0.0280 kl=2.3500 recon=0.7348 spec=0.7431 | 27.5s
    93→  New best model (loss=2.6343)
    94→  [39/200] step 1050 | loss=2.5785 pred=0.0287 kl=2.3464 recon=0.7147 spec=0.7317 | grad=5.51 | 23.7s
    95→Epoch 39/200 | avg_loss=2.6180 pred=0.0287 kl=2.3480 recon=0.7280 spec=0.7405 | 26.6s
    96→  New best model (loss=2.6180)
    97→Epoch 40/200 | avg_loss=2.6203 pred=0.0283 kl=2.3584 recon=0.7266 spec=0.7436 | 26.6s
    98→  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0040.pt
    99→  [41/200] step 1100 | loss=2.7463 pred=0.0283 kl=2.3608 recon=0.7838 spec=0.7512 | grad=18.19 | 20.0s
   100→Epoch 41/200 | avg_loss=2.6081 pred=0.0271 kl=2.3644 recon=0.7251 spec=0.7380 | 26.8s
   101→  New best model (loss=2.6081)
   102→Epoch 42/200 | avg_loss=2.5948 pred=0.0279 kl=2.3493 recon=0.7197 spec=0.7360 | 26.7s
   103→  New best model (loss=2.5948)
   104→  [43/200] step 1150 | loss=2.6223 pred=0.0281 kl=2.3435 recon=0.7132 spec=0.7629 | grad=10.47 | 16.1s
   105→Epoch 43/200 | avg_loss=2.5626 pred=0.0275 kl=2.3467 recon=0.7075 spec=0.7311 | 26.9s
   106→  New best model (loss=2.5626)
   107→Epoch 44/200 | avg_loss=2.5482 pred=0.0262 kl=2.3584 recon=0.7020 spec=0.7296 | 26.8s
   108→  New best model (loss=2.5482)
   109→  [45/200] step 1200 | loss=2.5354 pred=0.0254 kl=2.3629 recon=0.6983 spec=0.7266 | grad=8.17 | 12.3s
   110→Epoch 45/200 | avg_loss=2.5314 pred=0.0258 kl=2.3601 recon=0.6974 spec=0.7248 | 26.8s
   111→  New best model (loss=2.5314)
   112→Epoch 46/200 | avg_loss=2.5207 pred=0.0255 kl=2.3618 recon=0.6923 spec=0.7246 | 26.7s
   113→  New best model (loss=2.5207)
   114→  [47/200] step 1250 | loss=2.5220 pred=0.0252 kl=2.3601 recon=0.6899 spec=0.7290 | grad=4.47 | 8.3s
   115→Epoch 47/200 | avg_loss=2.5180 pred=0.0253 kl=2.3646 recon=0.6911 spec=0.7246 | 26.6s
   116→  New best model (loss=2.5180)
   117→Epoch 48/200 | avg_loss=2.5174 pred=0.0252 kl=2.3674 recon=0.6896 spec=0.7262 | 26.5s
   118→  New best model (loss=2.5174)
   119→  [49/200] step 1300 | loss=2.5020 pred=0.0237 kl=2.3808 recon=0.6847 spec=0.7234 | grad=5.57 | 4.4s
   120→Epoch 49/200 | avg_loss=2.5067 pred=0.0252 kl=2.3666 recon=0.6874 spec=0.7220 | 26.6s
   121→  New best model (loss=2.5067)
   122→  [50/200] step 1350 | loss=2.6388 pred=0.0287 kl=2.3445 recon=0.7289 spec=0.7525 | grad=14.39 | 26.7s
   123→Epoch 50/200 | avg_loss=2.5203 pred=0.0259 kl=2.3627 recon=0.6931 spec=0.7231 | 26.8s
   124→  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0050.pt
   125→Epoch 51/200 | avg_loss=2.5382 pred=0.0275 kl=2.3595 recon=0.6992 spec=0.7258 | 26.7s
   126→  [52/200] step 1400 | loss=2.5137 pred=0.0257 kl=2.3582 recon=0.6845 spec=0.7303 | grad=7.45 | 22.8s
   127→Epoch 52/200 | avg_loss=2.5127 pred=0.0260 kl=2.3622 recon=0.6897 spec=0.7224 | 26.6s
   128→Epoch 53/200 | avg_loss=2.4939 pred=0.0257 kl=2.3670 recon=0.6829 spec=0.7192 | 26.7s
   129→  New best model (loss=2.4939)
   130→  [54/200] step 1450 | loss=2.4921 pred=0.0246 kl=2.3704 recon=0.6866 spec=0.7137 | grad=5.12 | 18.9s
   131→Epoch 54/200 | avg_loss=2.4788 pred=0.0255 kl=2.3646 recon=0.6781 spec=0.7156 | 26.8s
   132→  New best model (loss=2.4788)
   133→Epoch 55/200 | avg_loss=2.4712 pred=0.0239 kl=2.3831 recon=0.6744 spec=0.7165 | 26.7s
   134→  New best model (loss=2.4712)
   135→  [56/200] step 1500 | loss=2.4887 pred=0.0232 kl=2.3902 recon=0.6736 spec=0.7297 | grad=7.95 | 15.0s
   136→Epoch 56/200 | avg_loss=2.4544 pred=0.0236 kl=2.3850 recon=0.6701 spec=0.7112 | 26.7s
   137→  New best model (loss=2.4544)
   138→Epoch 57/200 | avg_loss=2.4547 pred=0.0233 kl=2.3901 recon=0.6701 spec=0.7115 | 26.6s
   139→  [58/200] step 1550 | loss=2.4111 pred=0.0234 kl=2.3922 recon=0.6599 spec=0.6960 | grad=4.94 | 11.0s
   140→Epoch 58/200 | avg_loss=2.4372 pred=0.0232 kl=2.3893 recon=0.6639 spec=0.7081 | 26.7s
   141→  New best model (loss=2.4372)
   142→Epoch 59/200 | avg_loss=2.4653 pred=0.0232 kl=2.3853 recon=0.6725 spec=0.7156 | 27.4s
   143→  [60/200] step 1600 | loss=2.4717 pred=0.0227 kl=2.3877 recon=0.6644 spec=0.7309 | grad=7.61 | 7.2s
   144→Epoch 60/200 | avg_loss=2.4280 pred=0.0229 kl=2.3884 recon=0.6599 spec=0.7076 | 26.8s
   145→  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0060.pt
   146→  New best model (loss=2.4280)
   147→Epoch 61/200 | avg_loss=2.4268 pred=0.0222 kl=2.3978 recon=0.6590 spec=0.7084 | 26.5s
   148→  New best model (loss=2.4268)
   149→  [62/200] step 1650 | loss=2.3963 pred=0.0226 kl=2.3918 recon=0.6497 spec=0.7002 | grad=3.71 | 3.4s
   150→Epoch 62/200 | avg_loss=2.4323 pred=0.0228 kl=2.3914 recon=0.6612 spec=0.7088 | 26.7s
   151→  [63/200] step 1700 | loss=2.4068 pred=0.0227 kl=2.3939 recon=0.6540 spec=0.7015 | grad=4.37 | 25.7s
   152→Epoch 63/200 | avg_loss=2.4369 pred=0.0234 kl=2.3862 recon=0.6625 spec=0.7098 | 26.7s
   153→Epoch 64/200 | avg_loss=2.3966 pred=0.0227 kl=2.3951 recon=0.6489 spec=0.7014 | 26.7s
   154→  New best model (loss=2.3966)
   155→  [65/200] step 1750 | loss=2.3868 pred=0.0220 kl=2.4025 recon=0.6484 spec=0.6960 | grad=4.48 | 22.0s
   156→Epoch 65/200 | avg_loss=2.3947 pred=0.0225 kl=2.4031 recon=0.6473 spec=0.7024 | 26.8s
   157→  New best model (loss=2.3947)
   158→Epoch 66/200 | avg_loss=2.4162 pred=0.0219 kl=2.4036 recon=0.6549 spec=0.7071 | 26.6s
   159→  [67/200] step 1800 | loss=2.3823 pred=0.0224 kl=2.4010 recon=0.6396 spec=0.7045 | grad=3.25 | 18.1s
   160→Epoch 67/200 | avg_loss=2.3911 pred=0.0222 kl=2.3991 recon=0.6470 spec=0.7005 | 26.8s
   161→  New best model (loss=2.3911)
   162→Epoch 68/200 | avg_loss=2.3972 pred=0.0223 kl=2.3984 recon=0.6486 spec=0.7025 | 27.5s
   163→  [69/200] step 1850 | loss=2.3773 pred=0.0217 kl=2.4010 recon=0.6350 spec=0.7077 | grad=5.08 | 14.0s
   164→Epoch 69/200 | avg_loss=2.3685 pred=0.0217 kl=2.4027 recon=0.6373 spec=0.6988 | 26.6s
   165→  New best model (loss=2.3685)
   166→Epoch 70/200 | avg_loss=2.3709 pred=0.0212 kl=2.4103 recon=0.6396 spec=0.6976 | 26.6s
   167→  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0070.pt
   168→  [71/200] step 1900 | loss=2.3780 pred=0.0212 kl=2.4116 recon=0.6415 spec=0.6998 | grad=5.27 | 10.2s
   169→Epoch 71/200 | avg_loss=2.3903 pred=0.0215 kl=2.4097 recon=0.6475 spec=0.6998 | 26.7s
   170→Epoch 72/200 | avg_loss=2.3743 pred=0.0214 kl=2.4124 recon=0.6401 spec=0.6991 | 26.6s
   171→  [73/200] step 1950 | loss=2.3665 pred=0.0216 kl=2.4074 recon=0.6433 spec=0.6895 | grad=6.87 | 6.3s
   172→Epoch 73/200 | avg_loss=2.3476 pred=0.0208 kl=2.4145 recon=0.6311 spec=0.6936 | 26.7s
   173→  New best model (loss=2.3476)
   174→Epoch 74/200 | avg_loss=2.3751 pred=0.0206 kl=2.4173 recon=0.6415 spec=0.6983 | 26.7s
   175→  [75/200] step 2000 | loss=2.3530 pred=0.0215 kl=2.4050 recon=0.6368 spec=0.6892 | grad=4.39 | 2.5s
   176→Epoch 75/200 | avg_loss=2.3483 pred=0.0208 kl=2.4150 recon=0.6328 spec=0.6918 | 26.8s
   177→  [76/200] step 2050 | loss=2.3388 pred=0.0200 kl=2.4332 recon=0.6292 spec=0.6907 | grad=3.30 | 24.7s
   178→Epoch 76/200 | avg_loss=2.3433 pred=0.0203 kl=2.4257 recon=0.6302 spec=0.6923 | 26.7s
   179→  New best model (loss=2.3433)
   180→Epoch 77/200 | avg_loss=2.3291 pred=0.0201 kl=2.4292 recon=0.6261 spec=0.6883 | 27.2s
   181→  New best model (loss=2.3291)
   182→  [78/200] step 2100 | loss=2.3289 pred=0.0195 kl=2.4343 recon=0.6279 spec=0.6862 | grad=5.88 | 21.0s
   183→Epoch 78/200 | avg_loss=2.3364 pred=0.0193 kl=2.4395 recon=0.6274 spec=0.6919 | 27.0s
   184→Epoch 79/200 | avg_loss=2.3418 pred=0.0195 kl=2.4350 recon=0.6287 spec=0.6937 | 26.7s
   185→  [80/200] step 2150 | loss=2.3271 pred=0.0202 kl=2.4211 recon=0.6260 spec=0.6871 | grad=5.95 | 17.1s
   186→Epoch 80/200 | avg_loss=2.3267 pred=0.0199 kl=2.4251 recon=0.6241 spec=0.6896 | 26.8s
   187→  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0080.pt
   188→  New best model (loss=2.3267)
   189→Epoch 81/200 | avg_loss=2.3267 pred=0.0196 kl=2.4266 recon=0.6243 spec=0.6895 | 26.6s
   190→  New best model (loss=2.3267)
   191→  [82/200] step 2200 | loss=2.3135 pred=0.0189 kl=2.4331 recon=0.6199 spec=0.6870 | grad=3.65 | 13.0s
   192→Epoch 82/200 | avg_loss=2.3166 pred=0.0193 kl=2.4277 recon=0.6217 spec=0.6864 | 26.8s
   193→  New best model (loss=2.3166)
   194→Epoch 83/200 | avg_loss=2.3199 pred=0.0192 kl=2.4266 recon=0.6214 spec=0.6891 | 26.8s
   195→  [84/200] step 2250 | loss=2.3533 pred=0.0191 kl=2.4279 recon=0.6260 spec=0.7052 | grad=4.00 | 9.2s
   196→Epoch 84/200 | avg_loss=2.3257 pred=0.0191 kl=2.4266 recon=0.6228 spec=0.6912 | 26.6s
   197→Epoch 85/200 | avg_loss=2.3389 pred=0.0193 kl=2.4254 recon=0.6291 spec=0.6913 | 26.8s
   198→  [86/200] step 2300 | loss=2.3102 pred=0.0194 kl=2.4256 recon=0.6223 spec=0.6813 | grad=6.62 | 5.4s
   199→Epoch 86/200 | avg_loss=2.3226 pred=0.0198 kl=2.4241 recon=0.6243 spec=0.6866 | 26.6s
   200→Epoch 87/200 | avg_loss=2.3069 pred=0.0200 kl=2.4188 recon=0.6192 spec=0.6829 | 26.7s
   201→  New best model (loss=2.3069)
   202→  [88/200] step 2350 | loss=2.3216 pred=0.0198 kl=2.4191 recon=0.6240 spec=0.6863 | grad=7.57 | 1.5s
   203→Epoch 88/200 | avg_loss=2.3111 pred=0.0203 kl=2.4208 recon=0.6184 spec=0.6865 | 26.8s
   204→  [89/200] step 2400 | loss=2.3275 pred=0.0199 kl=2.4210 recon=0.6096 spec=0.7095 | grad=3.11 | 23.7s
   205→Epoch 89/200 | avg_loss=2.2879 pred=0.0199 kl=2.4262 recon=0.6106 spec=0.6817 | 26.6s
   206→  New best model (loss=2.2879)
   207→Epoch 90/200 | avg_loss=2.2803 pred=0.0199 kl=2.4287 recon=0.6082 spec=0.6798 | 26.8s
   208→  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0090.pt
   209→  New best model (loss=2.2803)
   210→  [91/200] step 2450 | loss=2.2793 pred=0.0197 kl=2.4401 recon=0.6099 spec=0.6770 | grad=3.34 | 19.9s
   211→Epoch 91/200 | avg_loss=2.2796 pred=0.0199 kl=2.4380 recon=0.6070 spec=0.6808 | 26.8s
   212→  New best model (loss=2.2796)
   213→Epoch 92/200 | avg_loss=2.2626 pred=0.0199 kl=2.4405 recon=0.6016 spec=0.6766 | 27.6s
   214→  New best model (loss=2.2626)
   215→  [93/200] step 2500 | loss=2.2485 pred=0.0191 kl=2.4473 recon=0.5994 spec=0.6708 | grad=2.98 | 15.7s
   216→Epoch 93/200 | avg_loss=2.2596 pred=0.0194 kl=2.4440 recon=0.6001 spec=0.6770 | 26.6s
   217→  New best model (loss=2.2596)
   218→Epoch 94/200 | avg_loss=2.2549 pred=0.0192 kl=2.4421 recon=0.5986 spec=0.6761 | 26.6s
   219→  New best model (loss=2.2549)
   220→  [95/200] step 2550 | loss=2.2863 pred=0.0187 kl=2.4448 recon=0.5999 spec=0.6956 | grad=6.09 | 12.2s
   221→Epoch 95/200 | avg_loss=2.2549 pred=0.0190 kl=2.4427 recon=0.5989 spec=0.6758 | 26.6s
   222→

Device: cuda
LeWM TTS model: 24.00M parameters
Dataset: 6959 samples loaded
Training: 6959 samples, 27 batches/epoch
Epoch 1/200 | avg_loss=13.6511 pred=1.4613 kl=0.2462 recon=2.3551 spec=4.9864 | 27.3s
  New best model (loss=13.6511)
  [2/200] step 50 | loss=13.3484 pred=1.1684 kl=0.1995 recon=2.3801 spec=4.9465 | grad=2.33 | 22.7s
Epoch 2/200 | avg_loss=13.2451 pred=1.2534 kl=0.2111 recon=2.3552 spec=4.8541 | 26.6s
  New best model (loss=13.2451)
Epoch 3/200 | avg_loss=12.5440 pred=1.0315 kl=0.2030 recon=2.3546 spec=4.5354 | 26.7s
  New best model (loss=12.5440)
  [4/200] step 100 | loss=10.7409 pred=0.9306 kl=0.3007 recon=2.3797 spec=3.3671 | grad=3.37 | 18.9s
Epoch 4/200 | avg_loss=11.0376 pred=0.9384 kl=0.2805 recon=2.3558 spec=3.5916 | 26.7s
  New best model (loss=11.0376)
Epoch 5/200 | avg_loss=8.9693 pred=0.8248 kl=0.2909 recon=2.2691 spec=2.4041 | 26.7s
  New best model (loss=8.9693)
  [6/200] step 150 | loss=6.7562 pred=0.7736 kl=0.2844 recon=1.6614 spec=1.7729 | grad=4.86 | 15.0s
Epoch 6/200 | avg_loss=6.9467 pred=0.7647 kl=0.2749 recon=1.6966 spec=1.8590 | 26.6s
  New best model (loss=6.9467)
Epoch 7/200 | avg_loss=5.4291 pred=0.6201 kl=0.3056 recon=1.4078 spec=1.3287 | 26.8s
  New best model (loss=5.4291)
  [8/200] step 200 | loss=4.5682 pred=0.5125 kl=0.3461 recon=1.2266 spec=1.0678 | grad=4.27 | 11.3s
Epoch 8/200 | avg_loss=4.5015 pred=0.5005 kl=0.3539 recon=1.2074 spec=1.0571 | 26.8s
  New best model (loss=4.5015)
Epoch 9/200 | avg_loss=3.8949 pred=0.4052 kl=0.4210 recon=1.0379 spec=0.9421 | 26.8s
  New best model (loss=3.8949)
  [10/200] step 250 | loss=3.6904 pred=0.3596 kl=0.4723 recon=0.9872 spec=0.9036 | grad=12.19 | 7.1s
Epoch 10/200 | avg_loss=3.6212 pred=0.3419 kl=0.4878 recon=0.9624 spec=0.9022 | 26.6s
  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0010.pt
  New best model (loss=3.6212)
Epoch 11/200 | avg_loss=3.4785 pred=0.2867 kl=0.5559 recon=0.9338 spec=0.8818 | 26.7s
  New best model (loss=3.4785)
  [12/200] step 300 | loss=3.4295 pred=0.2622 kl=0.5906 recon=0.9275 spec=0.8737 | grad=4.26 | 3.4s
Epoch 12/200 | avg_loss=3.3684 pred=0.2458 kl=0.6160 recon=0.9131 spec=0.8631 | 26.7s
  New best model (loss=3.3684)
  [13/200] step 350 | loss=3.2924 pred=0.1971 kl=0.7111 recon=0.8917 spec=0.8731 | grad=40.29 | 25.6s
Epoch 13/200 | avg_loss=3.2882 pred=0.2122 kl=0.6808 recon=0.8986 spec=0.8510 | 26.6s
  New best model (loss=3.2882)
Epoch 14/200 | avg_loss=3.4840 pred=0.1903 kl=0.7435 recon=0.9788 spec=0.8891 | 26.7s
  [15/200] step 400 | loss=3.3437 pred=0.1641 kl=0.8119 recon=0.9571 spec=0.8416 | grad=23.07 | 21.7s
Epoch 15/200 | avg_loss=3.3427 pred=0.1730 kl=0.7889 recon=0.9350 spec=0.8645 | 26.6s
Epoch 16/200 | avg_loss=3.2980 pred=0.1490 kl=0.8601 recon=0.9295 spec=0.8578 | 26.4s
  [17/200] step 450 | loss=3.1454 pred=0.1204 kl=0.9734 recon=0.8944 spec=0.8214 | grad=13.68 | 18.0s
Epoch 17/200 | avg_loss=3.1228 pred=0.1237 kl=0.9565 recon=0.8765 spec=0.8282 | 26.8s
  New best model (loss=3.1228)
Epoch 18/200 | avg_loss=3.0995 pred=0.0969 kl=1.0964 recon=0.8774 spec=0.8288 | 26.8s
  New best model (loss=3.0995)
  [19/200] step 500 | loss=3.1561 pred=0.0811 kl=1.2092 recon=0.9116 spec=0.8308 | grad=21.60 | 14.1s
Epoch 19/200 | avg_loss=3.0741 pred=0.0815 kl=1.2060 recon=0.8744 spec=0.8255 | 26.7s
  New best model (loss=3.0741)
Epoch 20/200 | avg_loss=3.0775 pred=0.0656 kl=1.3424 recon=0.8822 spec=0.8273 | 26.8s
  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0020.pt
  [21/200] step 550 | loss=3.2473 pred=0.0602 kl=1.4247 recon=0.9552 spec=0.8463 | grad=22.38 | 10.2s
Epoch 21/200 | avg_loss=3.0235 pred=0.0595 kl=1.4346 recon=0.8644 spec=0.8185 | 26.9s
  New best model (loss=3.0235)
Epoch 22/200 | avg_loss=2.9452 pred=0.0495 kl=1.5580 recon=0.8416 spec=0.8027 | 26.8s
  New best model (loss=2.9452)
  [23/200] step 600 | loss=2.8896 pred=0.0462 kl=1.6198 recon=0.8079 spec=0.8124 | grad=2.82 | 6.4s
Epoch 23/200 | avg_loss=2.9057 pred=0.0442 kl=1.6588 recon=0.8278 spec=0.7978 | 26.7s
  New best model (loss=2.9057)
Epoch 24/200 | avg_loss=2.8642 pred=0.0400 kl=1.7860 recon=0.8170 spec=0.7866 | 26.7s
  New best model (loss=2.8642)
  [25/200] step 650 | loss=2.8204 pred=0.0366 kl=1.8447 recon=0.8009 spec=0.7807 | grad=4.15 | 2.5s
Epoch 25/200 | avg_loss=2.8957 pred=0.0356 kl=1.8836 recon=0.8275 spec=0.7957 | 26.6s
  [26/200] step 700 | loss=2.8124 pred=0.0335 kl=1.9722 recon=0.8053 spec=0.7704 | grad=2.36 | 24.9s
Epoch 26/200 | avg_loss=2.8820 pred=0.0343 kl=1.9471 recon=0.8247 spec=0.7907 | 26.9s
Epoch 27/200 | avg_loss=2.8777 pred=0.0320 kl=2.0381 recon=0.8212 spec=0.7933 | 26.6s
  [28/200] step 750 | loss=2.8433 pred=0.0302 kl=2.1263 recon=0.8163 spec=0.7772 | grad=6.28 | 21.0s
Epoch 28/200 | avg_loss=2.9005 pred=0.0306 kl=2.1059 recon=0.8323 spec=0.7939 | 26.8s
Epoch 29/200 | avg_loss=2.9652 pred=0.0316 kl=2.1214 recon=0.8583 spec=0.8013 | 26.7s
  [30/200] step 800 | loss=2.8278 pred=0.0308 kl=2.1655 recon=0.8064 spec=0.7788 | grad=7.38 | 16.8s
Epoch 30/200 | avg_loss=2.8287 pred=0.0321 kl=2.1512 recon=0.8065 spec=0.7786 | 26.5s
  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0030.pt
  New best model (loss=2.8287)
Epoch 31/200 | avg_loss=2.7880 pred=0.0300 kl=2.2180 recon=0.7923 spec=0.7710 | 26.7s
  New best model (loss=2.7880)
  [32/200] step 850 | loss=2.7250 pred=0.0279 kl=2.2650 recon=0.7754 spec=0.7523 | grad=5.50 | 13.0s
Epoch 32/200 | avg_loss=2.7648 pred=0.0283 kl=2.2636 recon=0.7842 spec=0.7668 | 26.6s
  New best model (loss=2.7648)
Epoch 33/200 | avg_loss=2.7526 pred=0.0278 kl=2.2914 recon=0.7802 spec=0.7638 | 26.9s
  New best model (loss=2.7526)
  [34/200] step 900 | loss=2.7781 pred=0.0286 kl=2.3051 recon=0.7798 spec=0.7806 | grad=6.41 | 9.2s
Epoch 34/200 | avg_loss=2.7445 pred=0.0293 kl=2.3163 recon=0.7752 spec=0.7636 | 26.7s
  New best model (loss=2.7445)
Epoch 35/200 | avg_loss=2.7538 pred=0.0283 kl=2.3394 recon=0.7791 spec=0.7648 | 26.6s
  [36/200] step 950 | loss=2.7073 pred=0.0278 kl=2.3435 recon=0.7654 spec=0.7521 | grad=3.65 | 5.3s
Epoch 36/200 | avg_loss=2.7155 pred=0.0279 kl=2.3379 recon=0.7666 spec=0.7558 | 26.8s
  New best model (loss=2.7155)
Epoch 37/200 | avg_loss=2.6678 pred=0.0273 kl=2.3513 recon=0.7473 spec=0.7496 | 26.8s
  New best model (loss=2.6678)
  [38/200] step 1000 | loss=2.6472 pred=0.0271 kl=2.3522 recon=0.7342 spec=0.7534 | grad=5.13 | 1.5s
Epoch 38/200 | avg_loss=2.6343 pred=0.0280 kl=2.3500 recon=0.7348 spec=0.7431 | 27.5s
  New best model (loss=2.6343)
  [39/200] step 1050 | loss=2.5785 pred=0.0287 kl=2.3464 recon=0.7147 spec=0.7317 | grad=5.51 | 23.7s
Epoch 39/200 | avg_loss=2.6180 pred=0.0287 kl=2.3480 recon=0.7280 spec=0.7405 | 26.6s
  New best model (loss=2.6180)
Epoch 40/200 | avg_loss=2.6203 pred=0.0283 kl=2.3584 recon=0.7266 spec=0.7436 | 26.6s
  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0040.pt
  [41/200] step 1100 | loss=2.7463 pred=0.0283 kl=2.3608 recon=0.7838 spec=0.7512 | grad=18.19 | 20.0s
Epoch 41/200 | avg_loss=2.6081 pred=0.0271 kl=2.3644 recon=0.7251 spec=0.7380 | 26.8s
  New best model (loss=2.6081)
Epoch 42/200 | avg_loss=2.5948 pred=0.0279 kl=2.3493 recon=0.7197 spec=0.7360 | 26.7s
  New best model (loss=2.5948)
  [43/200] step 1150 | loss=2.6223 pred=0.0281 kl=2.3435 recon=0.7132 spec=0.7629 | grad=10.47 | 16.1s
Epoch 43/200 | avg_loss=2.5626 pred=0.0275 kl=2.3467 recon=0.7075 spec=0.7311 | 26.9s
  New best model (loss=2.5626)
Epoch 44/200 | avg_loss=2.5482 pred=0.0262 kl=2.3584 recon=0.7020 spec=0.7296 | 26.8s
  New best model (loss=2.5482)
  [45/200] step 1200 | loss=2.5354 pred=0.0254 kl=2.3629 recon=0.6983 spec=0.7266 | grad=8.17 | 12.3s
Epoch 45/200 | avg_loss=2.5314 pred=0.0258 kl=2.3601 recon=0.6974 spec=0.7248 | 26.8s
  New best model (loss=2.5314)
Epoch 46/200 | avg_loss=2.5207 pred=0.0255 kl=2.3618 recon=0.6923 spec=0.7246 | 26.7s
  New best model (loss=2.5207)
  [47/200] step 1250 | loss=2.5220 pred=0.0252 kl=2.3601 recon=0.6899 spec=0.7290 | grad=4.47 | 8.3s
Epoch 47/200 | avg_loss=2.5180 pred=0.0253 kl=2.3646 recon=0.6911 spec=0.7246 | 26.6s
  New best model (loss=2.5180)
Epoch 48/200 | avg_loss=2.5174 pred=0.0252 kl=2.3674 recon=0.6896 spec=0.7262 | 26.5s
  New best model (loss=2.5174)
  [49/200] step 1300 | loss=2.5020 pred=0.0237 kl=2.3808 recon=0.6847 spec=0.7234 | grad=5.57 | 4.4s
Epoch 49/200 | avg_loss=2.5067 pred=0.0252 kl=2.3666 recon=0.6874 spec=0.7220 | 26.6s
  New best model (loss=2.5067)
  [50/200] step 1350 | loss=2.6388 pred=0.0287 kl=2.3445 recon=0.7289 spec=0.7525 | grad=14.39 | 26.7s
Epoch 50/200 | avg_loss=2.5203 pred=0.0259 kl=2.3627 recon=0.6931 spec=0.7231 | 26.8s
  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0050.pt
Epoch 51/200 | avg_loss=2.5382 pred=0.0275 kl=2.3595 recon=0.6992 spec=0.7258 | 26.7s
  [52/200] step 1400 | loss=2.5137 pred=0.0257 kl=2.3582 recon=0.6845 spec=0.7303 | grad=7.45 | 22.8s
Epoch 52/200 | avg_loss=2.5127 pred=0.0260 kl=2.3622 recon=0.6897 spec=0.7224 | 26.6s
Epoch 53/200 | avg_loss=2.4939 pred=0.0257 kl=2.3670 recon=0.6829 spec=0.7192 | 26.7s
  New best model (loss=2.4939)
  [54/200] step 1450 | loss=2.4921 pred=0.0246 kl=2.3704 recon=0.6866 spec=0.7137 | grad=5.12 | 18.9s
Epoch 54/200 | avg_loss=2.4788 pred=0.0255 kl=2.3646 recon=0.6781 spec=0.7156 | 26.8s
  New best model (loss=2.4788)
Epoch 55/200 | avg_loss=2.4712 pred=0.0239 kl=2.3831 recon=0.6744 spec=0.7165 | 26.7s
  New best model (loss=2.4712)
  [56/200] step 1500 | loss=2.4887 pred=0.0232 kl=2.3902 recon=0.6736 spec=0.7297 | grad=7.95 | 15.0s
Epoch 56/200 | avg_loss=2.4544 pred=0.0236 kl=2.3850 recon=0.6701 spec=0.7112 | 26.7s
  New best model (loss=2.4544)
Epoch 57/200 | avg_loss=2.4547 pred=0.0233 kl=2.3901 recon=0.6701 spec=0.7115 | 26.6s
  [58/200] step 1550 | loss=2.4111 pred=0.0234 kl=2.3922 recon=0.6599 spec=0.6960 | grad=4.94 | 11.0s
Epoch 58/200 | avg_loss=2.4372 pred=0.0232 kl=2.3893 recon=0.6639 spec=0.7081 | 26.7s
  New best model (loss=2.4372)
Epoch 59/200 | avg_loss=2.4653 pred=0.0232 kl=2.3853 recon=0.6725 spec=0.7156 | 27.4s
  [60/200] step 1600 | loss=2.4717 pred=0.0227 kl=2.3877 recon=0.6644 spec=0.7309 | grad=7.61 | 7.2s
Epoch 60/200 | avg_loss=2.4280 pred=0.0229 kl=2.3884 recon=0.6599 spec=0.7076 | 26.8s
  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0060.pt
  New best model (loss=2.4280)
Epoch 61/200 | avg_loss=2.4268 pred=0.0222 kl=2.3978 recon=0.6590 spec=0.7084 | 26.5s
  New best model (loss=2.4268)
  [62/200] step 1650 | loss=2.3963 pred=0.0226 kl=2.3918 recon=0.6497 spec=0.7002 | grad=3.71 | 3.4s
Epoch 62/200 | avg_loss=2.4323 pred=0.0228 kl=2.3914 recon=0.6612 spec=0.7088 | 26.7s
  [63/200] step 1700 | loss=2.4068 pred=0.0227 kl=2.3939 recon=0.6540 spec=0.7015 | grad=4.37 | 25.7s
Epoch 63/200 | avg_loss=2.4369 pred=0.0234 kl=2.3862 recon=0.6625 spec=0.7098 | 26.7s
Epoch 64/200 | avg_loss=2.3966 pred=0.0227 kl=2.3951 recon=0.6489 spec=0.7014 | 26.7s
  New best model (loss=2.3966)
  [65/200] step 1750 | loss=2.3868 pred=0.0220 kl=2.4025 recon=0.6484 spec=0.6960 | grad=4.48 | 22.0s
Epoch 65/200 | avg_loss=2.3947 pred=0.0225 kl=2.4031 recon=0.6473 spec=0.7024 | 26.8s
  New best model (loss=2.3947)
Epoch 66/200 | avg_loss=2.4162 pred=0.0219 kl=2.4036 recon=0.6549 spec=0.7071 | 26.6s
  [67/200] step 1800 | loss=2.3823 pred=0.0224 kl=2.4010 recon=0.6396 spec=0.7045 | grad=3.25 | 18.1s
Epoch 67/200 | avg_loss=2.3911 pred=0.0222 kl=2.3991 recon=0.6470 spec=0.7005 | 26.8s
  New best model (loss=2.3911)
Epoch 68/200 | avg_loss=2.3972 pred=0.0223 kl=2.3984 recon=0.6486 spec=0.7025 | 27.5s
  [69/200] step 1850 | loss=2.3773 pred=0.0217 kl=2.4010 recon=0.6350 spec=0.7077 | grad=5.08 | 14.0s
Epoch 69/200 | avg_loss=2.3685 pred=0.0217 kl=2.4027 recon=0.6373 spec=0.6988 | 26.6s
  New best model (loss=2.3685)
Epoch 70/200 | avg_loss=2.3709 pred=0.0212 kl=2.4103 recon=0.6396 spec=0.6976 | 26.6s
  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0070.pt
  [71/200] step 1900 | loss=2.3780 pred=0.0212 kl=2.4116 recon=0.6415 spec=0.6998 | grad=5.27 | 10.2s
Epoch 71/200 | avg_loss=2.3903 pred=0.0215 kl=2.4097 recon=0.6475 spec=0.6998 | 26.7s
Epoch 72/200 | avg_loss=2.3743 pred=0.0214 kl=2.4124 recon=0.6401 spec=0.6991 | 26.6s
  [73/200] step 1950 | loss=2.3665 pred=0.0216 kl=2.4074 recon=0.6433 spec=0.6895 | grad=6.87 | 6.3s
Epoch 73/200 | avg_loss=2.3476 pred=0.0208 kl=2.4145 recon=0.6311 spec=0.6936 | 26.7s
  New best model (loss=2.3476)
Epoch 74/200 | avg_loss=2.3751 pred=0.0206 kl=2.4173 recon=0.6415 spec=0.6983 | 26.7s
  [75/200] step 2000 | loss=2.3530 pred=0.0215 kl=2.4050 recon=0.6368 spec=0.6892 | grad=4.39 | 2.5s
Epoch 75/200 | avg_loss=2.3483 pred=0.0208 kl=2.4150 recon=0.6328 spec=0.6918 | 26.8s
  [76/200] step 2050 | loss=2.3388 pred=0.0200 kl=2.4332 recon=0.6292 spec=0.6907 | grad=3.30 | 24.7s
Epoch 76/200 | avg_loss=2.3433 pred=0.0203 kl=2.4257 recon=0.6302 spec=0.6923 | 26.7s
  New best model (loss=2.3433)
Epoch 77/200 | avg_loss=2.3291 pred=0.0201 kl=2.4292 recon=0.6261 spec=0.6883 | 27.2s
  New best model (loss=2.3291)
  [78/200] step 2100 | loss=2.3289 pred=0.0195 kl=2.4343 recon=0.6279 spec=0.6862 | grad=5.88 | 21.0s
Epoch 78/200 | avg_loss=2.3364 pred=0.0193 kl=2.4395 recon=0.6274 spec=0.6919 | 27.0s
Epoch 79/200 | avg_loss=2.3418 pred=0.0195 kl=2.4350 recon=0.6287 spec=0.6937 | 26.7s
  [80/200] step 2150 | loss=2.3271 pred=0.0202 kl=2.4211 recon=0.6260 spec=0.6871 | grad=5.95 | 17.1s
Epoch 80/200 | avg_loss=2.3267 pred=0.0199 kl=2.4251 recon=0.6241 spec=0.6896 | 26.8s
  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0080.pt
  New best model (loss=2.3267)
Epoch 81/200 | avg_loss=2.3267 pred=0.0196 kl=2.4266 recon=0.6243 spec=0.6895 | 26.6s
  New best model (loss=2.3267)
  [82/200] step 2200 | loss=2.3135 pred=0.0189 kl=2.4331 recon=0.6199 spec=0.6870 | grad=3.65 | 13.0s
Epoch 82/200 | avg_loss=2.3166 pred=0.0193 kl=2.4277 recon=0.6217 spec=0.6864 | 26.8s
  New best model (loss=2.3166)
Epoch 83/200 | avg_loss=2.3199 pred=0.0192 kl=2.4266 recon=0.6214 spec=0.6891 | 26.8s
  [84/200] step 2250 | loss=2.3533 pred=0.0191 kl=2.4279 recon=0.6260 spec=0.7052 | grad=4.00 | 9.2s
Epoch 84/200 | avg_loss=2.3257 pred=0.0191 kl=2.4266 recon=0.6228 spec=0.6912 | 26.6s
Epoch 85/200 | avg_loss=2.3389 pred=0.0193 kl=2.4254 recon=0.6291 spec=0.6913 | 26.8s
  [86/200] step 2300 | loss=2.3102 pred=0.0194 kl=2.4256 recon=0.6223 spec=0.6813 | grad=6.62 | 5.4s
Epoch 86/200 | avg_loss=2.3226 pred=0.0198 kl=2.4241 recon=0.6243 spec=0.6866 | 26.6s
Epoch 87/200 | avg_loss=2.3069 pred=0.0200 kl=2.4188 recon=0.6192 spec=0.6829 | 26.7s
  New best model (loss=2.3069)
  [88/200] step 2350 | loss=2.3216 pred=0.0198 kl=2.4191 recon=0.6240 spec=0.6863 | grad=7.57 | 1.5s
Epoch 88/200 | avg_loss=2.3111 pred=0.0203 kl=2.4208 recon=0.6184 spec=0.6865 | 26.8s
  [89/200] step 2400 | loss=2.3275 pred=0.0199 kl=2.4210 recon=0.6096 spec=0.7095 | grad=3.11 | 23.7s
Epoch 89/200 | avg_loss=2.2879 pred=0.0199 kl=2.4262 recon=0.6106 spec=0.6817 | 26.6s
  New best model (loss=2.2879)
Epoch 90/200 | avg_loss=2.2803 pred=0.0199 kl=2.4287 recon=0.6082 spec=0.6798 | 26.8s
  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3b/checkpoints/epoch_0090.pt
  New best model (loss=2.2803)
  [91/200] step 2450 | loss=2.2793 pred=0.0197 kl=2.4401 recon=0.6099 spec=0.6770 | grad=3.34 | 19.9s
Epoch 91/200 | avg_loss=2.2796 pred=0.0199 kl=2.4380 recon=0.6070 spec=0.6808 | 26.8s
  New best model (loss=2.2796)
Epoch 92/200 | avg_loss=2.2626 pred=0.0199 kl=2.4405 recon=0.6016 spec=0.6766 | 27.6s
  New best model (loss=2.2626)
  [93/200] step 2500 | loss=2.2485 pred=0.0191 kl=2.4473 recon=0.5994 spec=0.6708 | grad=2.98 | 15.7s
Epoch 93/200 | avg_loss=2.2596 pred=0.0194 kl=2.4440 recon=0.6001 spec=0.6770 | 26.6s
  New best model (loss=2.2596)
Epoch 94/200 | avg_loss=2.2549 pred=0.0192 kl=2.4421 recon=0.5986 spec=0.6761 | 26.6s
  New best model (loss=2.2549)
  [95/200] step 2550 | loss=2.2863 pred=0.0187 kl=2.4448 recon=0.5999 spec=0.6956 | grad=6.09 | 12.2s
Epoch 95/200 | avg_loss=2.2549 pred=0.0190 kl=2.4427 recon=0.5989 spec=0.6758 | 26.6s


     1→"""
     2→Inference pipeline for LeWM TTS.
     3→Text → JEPA predictor → mel reconstruction → Vocos vocoder → waveform
     4→"""
     5→
     6→import torch
     7→import torchaudio
     8→import numpy as np
     9→import argparse
    10→from pathlib import Path
    11→
    12→from model import LeWMTTS
    13→
    14→
    15→class LeWMTTSInference:
    16→    """Full inference pipeline."""
    17→
    18→    def __init__(self, checkpoint_path, device="cuda"):
    19→        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
    20→
    21→        # Load checkpoint
    22→        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
    23→        config = ckpt["config"]
    24→
    25→        # Build model
    26→        self.model = LeWMTTS(config).to(self.device)
    27→        self.model.load_state_dict(ckpt["model"])
    28→        self.model.eval()
    29→
    30→        # Vocos vocoder (directly decodes 100-mel)
    31→        from vocos import Vocos
    32→        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
    33→        self.vocos.eval()
    34→
    35→        self.config = config
    36→        self._start_emb_scale = 5.0
    37→        print(f"Model loaded from {checkpoint_path}")
    38→
    39→    def vocos_decode(self, mel):
    40→        """Decode mel [B, 100, T] → waveform via Vocos."""
    41→        with torch.no_grad():
    42→            waveform = self.vocos.decode(mel)  # [B, T_audio]
    43→        return waveform
    44→
    45→    def text_to_tokens(self, text):
    46→        tokens = list(text.encode("utf-8"))
    47→        return torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(self.device)
    48→
    49→    @torch.no_grad()
    50→    def synthesize(self, text, max_steps=300, temperature=0.0):
    51→        """Generate speech from text using AR prediction with KV-cache."""
    52→        text_tokens = self.text_to_tokens(text)
    53→        text_emb = self.model.text_encoder(text_tokens)
    54→
    55→        d = self.config["d_model"]
    56→        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
    57→        all_embeddings = [start_emb]
    58→
    59→        cache = self.model.init_ar_cache()
    60→        next_emb, cache = self.model.predict_next_cached(start_emb, text_emb, 0, cache)
    61→        if temperature > 0:
    62→            next_emb = next_emb + torch.randn_like(next_emb) * temperature
    63→        all_embeddings.append(next_emb)
    64→
    65→        prev_norms = [next_emb.norm().item()]
    66→
    67→        for step in range(1, max_steps):
    68→            next_emb, cache = self.model.predict_next_cached(next_emb, text_emb, step, cache)
    69→            if temperature > 0:
    70→                next_emb = next_emb + torch.randn_like(next_emb) * temperature
    71→            all_embeddings.append(next_emb)
    72→
    73→            norm = next_emb.norm().item()
    74→            prev_norms.append(norm)
    75→            if len(prev_norms) > 20:
    76→                recent = prev_norms[-20:]
    77→                if np.std(recent) < 0.01 * np.mean(recent) and step > 30:
    78→                    break
    79→
    80→        audio_embs = torch.cat(all_embeddings, dim=1)
    81→        mel = self.model.mel_decoder(audio_embs)
    82→        waveform = self.vocos_decode(mel.cpu())
    83→        return waveform.squeeze().numpy(), 24000
    84→
    85→    @torch.no_grad()
    86→    def reconstruct(self, mel_path):
    87→        """Teacher-forced reconstruction: encode real mel → decode back."""
    88→        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
    89→
    90→        z, mu, logvar = self.model.audio_encoder(mel)
    91→        mel_recon = self.model.mel_decoder(mu)
    92→
    93→        waveform = self.vocos_decode(mel_recon.cpu())
    94→        return waveform.squeeze().numpy(), 24000
    95→
    96→    @torch.no_grad()
    97→    def reconstruct_with_prediction(self, mel_path):
    98→        """Hybrid: encode real mel → predictor → decode predicted embeddings."""
    99→        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
   100→
   101→        z, mu, logvar = self.model.audio_encoder(mel)
   102→
   103→        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
   104→        text_emb = self.model.text_encoder(text_tokens)
   105→
   106→        predicted = self.model.predictor(mu[:, :-1], text_emb)
   107→        combined = torch.cat([mu[:, :1], predicted], dim=1)
   108→
   109→        mel_recon = self.model.mel_decoder(combined)
   110→        waveform = self.vocos_decode(mel_recon.cpu())
   111→        return waveform.squeeze().numpy(), 24000
   112→
   113→    def save_audio(self, waveform, sr, output_path):
   114→        # Normalize
   115→        mx = np.abs(waveform).max()
   116→        if mx > 0:
   117→            waveform = waveform / mx * 0.95
   118→        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
   119→        torchaudio.save(output_path, wav_tensor, sr)
   120→        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")
   121→
   122→
   123→def main():
   124→    parser = argparse.ArgumentParser()
   125→    parser.add_argument("--checkpoint", required=True)
   126→    parser.add_argument("--text", default=None, help="Text to synthesize")
   127→    parser.add_argument("--mel", default=None, help="Mel path for reconstruction test")
   128→    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "predict"],
   129→                        default="synthesize")
   130→    parser.add_argument("--output", default="output.wav")
   131→    parser.add_argument("--max_steps", type=int, default=300)
   132→    parser.add_argument("--temperature", type=float, default=0.0)
   133→    args = parser.parse_args()
   134→
   135→    tts = LeWMTTSInference(args.checkpoint)
   136→
   137→    if args.mode == "synthesize":
   138→        assert args.text, "Need --text for synthesis"
   139→        waveform, sr = tts.synthesize(args.text, args.max_steps, args.temperature)
   140→    elif args.mode == "reconstruct":
   141→        assert args.mel, "Need --mel for reconstruction"
   142→        waveform, sr = tts.reconstruct(args.mel)
   143→    elif args.mode == "predict":
   144→        assert args.mel, "Need --mel for prediction test"
   145→        waveform, sr = tts.reconstruct_with_prediction(args.mel)
   146→
   147→    tts.save_audio(waveform, sr, args.output)
   148→
   149→
   150→if __name__ == "__main__":
   151→    main()
   152→

"""
Inference pipeline for LeWM TTS.
Text → JEPA predictor → mel reconstruction → Vocos vocoder → waveform
"""

import torch
import torchaudio
import numpy as np
import argparse
from pathlib import Path

from model import LeWMTTS

class LeWMTTSInference:
    """Full inference pipeline."""

    def __init__(self, checkpoint_path, device="cuda"):
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")

        # Load checkpoint
        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
        config = ckpt["config"]

        # Build model
        self.model = LeWMTTS(config).to(self.device)
        self.model.load_state_dict(ckpt["model"])
        self.model.eval()

        # Vocos vocoder (directly decodes 100-mel)
        from vocos import Vocos
        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
        self.vocos.eval()

        self.config = config
        self._start_emb_scale = 5.0
        print(f"Model loaded from {checkpoint_path}")

    def vocos_decode(self, mel):
        """Decode mel [B, 100, T] → waveform via Vocos."""
        with torch.no_grad():
            waveform = self.vocos.decode(mel)  # [B, T_audio]
        return waveform

    def text_to_tokens(self, text):
        tokens = list(text.encode("utf-8"))
        return torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(self.device)

    @torch.no_grad()
    def synthesize(self, text, max_steps=300, temperature=0.0):
        """Generate speech from text using AR prediction with KV-cache."""
        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        d = self.config["d_model"]
        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
        all_embeddings = [start_emb]

        cache = self.model.init_ar_cache()
        next_emb, cache = self.model.predict_next_cached(start_emb, text_emb, 0, cache)
        if temperature > 0:
            next_emb = next_emb + torch.randn_like(next_emb) * temperature
        all_embeddings.append(next_emb)

        prev_norms = [next_emb.norm().item()]

        for step in range(1, max_steps):
            next_emb, cache = self.model.predict_next_cached(next_emb, text_emb, step, cache)
            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature
            all_embeddings.append(next_emb)

            norm = next_emb.norm().item()
            prev_norms.append(norm)
            if len(prev_norms) > 20:
                recent = prev_norms[-20:]
                if np.std(recent) < 0.01 * np.mean(recent) and step > 30:
                    break

        audio_embs = torch.cat(all_embeddings, dim=1)
        mel = self.model.mel_decoder(audio_embs)
        waveform = self.vocos_decode(mel.cpu())
        return waveform.squeeze().numpy(), 24000

    @torch.no_grad()
    def reconstruct(self, mel_path):
        """Teacher-forced reconstruction: encode real mel → decode back."""
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)

        z, mu, logvar = self.model.audio_encoder(mel)
        mel_recon = self.model.mel_decoder(mu)

        waveform = self.vocos_decode(mel_recon.cpu())
        return waveform.squeeze().numpy(), 24000

    @torch.no_grad()
    def reconstruct_with_prediction(self, mel_path):
        """Hybrid: encode real mel → predictor → decode predicted embeddings."""
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)

        z, mu, logvar = self.model.audio_encoder(mel)

        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
        text_emb = self.model.text_encoder(text_tokens)

        predicted = self.model.predictor(mu[:, :-1], text_emb)
        combined = torch.cat([mu[:, :1], predicted], dim=1)

        mel_recon = self.model.mel_decoder(combined)
        waveform = self.vocos_decode(mel_recon.cpu())
        return waveform.squeeze().numpy(), 24000

    def save_audio(self, waveform, sr, output_path):
        # Normalize
        mx = np.abs(waveform).max()
        if mx > 0:
            waveform = waveform / mx * 0.95
        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
        torchaudio.save(output_path, wav_tensor, sr)
        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--text", default=None, help="Text to synthesize")
    parser.add_argument("--mel", default=None, help="Mel path for reconstruction test")
    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "predict"],
                        default="synthesize")
    parser.add_argument("--output", default="output.wav")
    parser.add_argument("--max_steps", type=int, default=300)
    parser.add_argument("--temperature", type=float, default=0.0)
    args = parser.parse_args()

    tts = LeWMTTSInference(args.checkpoint)

    if args.mode == "synthesize":
        assert args.text, "Need --text for synthesis"
        waveform, sr = tts.synthesize(args.text, args.max_steps, args.temperature)
    elif args.mode == "reconstruct":
        assert args.mel, "Need --mel for reconstruction"
        waveform, sr = tts.reconstruct(args.mel)
    elif args.mode == "predict":
        assert args.mel, "Need --mel for prediction test"
        waveform, sr = tts.reconstruct_with_prediction(args.mel)

    tts.save_audio(waveform, sr, args.output)

if __name__ == "__main__":
    main()


     1→Device: cuda
     2→LeWM TTS model: 24.00M parameters
     3→Dataset: 6959 samples loaded
     4→Training: 6959 samples, 27 batches/epoch
     5→Epoch 1/200 | avg_loss=6.4157 pred=1.5698 kl=0.2692 recon=2.3557 spec=4.9796 | 27.5s
     6→  New best model (loss=6.4157)
     7→  [2/200] step 50 | loss=5.9234 pred=1.1925 kl=0.2006 recon=2.3602 spec=4.7394 | grad=1.48 | 22.9s
     8→Epoch 2/200 | avg_loss=6.1036 pred=1.3063 kl=0.2213 recon=2.3550 spec=4.8829 | 26.7s
     9→  New best model (loss=6.1036)
    10→Epoch 3/200 | avg_loss=5.6631 pred=0.9758 kl=0.1740 recon=2.3543 spec=4.6637 | 26.8s
    11→  New best model (loss=5.6631)
    12→  [4/200] step 100 | loss=4.8707 pred=0.7100 kl=0.2010 recon=2.3181 spec=3.6816 | grad=1.11 | 18.9s
    13→Epoch 4/200 | avg_loss=5.0403 pred=0.7430 kl=0.1915 recon=2.3460 spec=3.8993 | 26.7s
    14→  New best model (loss=5.0403)
    15→Epoch 5/200 | avg_loss=4.1751 pred=0.6059 kl=0.2271 recon=2.2503 spec=2.6327 | 26.9s
    16→  New best model (loss=4.1751)
    17→  [6/200] step 150 | loss=3.1632 pred=0.5672 kl=0.2311 recon=1.6025 spec=1.9807 | grad=1.91 | 15.1s
    18→Epoch 6/200 | avg_loss=3.2544 pred=0.5495 kl=0.2322 recon=1.6716 spec=2.0602 | 26.6s
    19→  New best model (loss=3.2544)
    20→Epoch 7/200 | avg_loss=2.5653 pred=0.4760 kl=0.2680 recon=1.3342 spec=1.5015 | 26.6s
    21→  New best model (loss=2.5653)
    22→  [8/200] step 200 | loss=2.2587 pred=0.4182 kl=0.3118 recon=1.2373 spec=1.1950 | grad=8.12 | 11.2s
    23→Epoch 8/200 | avg_loss=2.2319 pred=0.4148 kl=0.3167 recon=1.2378 spec=1.1467 | 26.7s
    24→  New best model (loss=2.2319)
    25→Epoch 9/200 | avg_loss=2.0277 pred=0.3794 kl=0.3523 recon=1.1345 spec=1.0126 | 26.8s
    26→  New best model (loss=2.0277)
    27→  [10/200] step 250 | loss=1.9283 pred=0.3482 kl=0.3789 recon=1.0729 spec=0.9971 | grad=5.03 | 7.3s
    28→Epoch 10/200 | avg_loss=1.8836 pred=0.3400 kl=0.3910 recon=1.0475 spec=0.9737 | 26.8s
    29→  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3/checkpoints/epoch_0010.pt
    30→  New best model (loss=1.8836)
    31→Epoch 11/200 | avg_loss=1.7388 pred=0.2963 kl=0.4481 recon=0.9699 spec=0.9217 | 26.8s
    32→  New best model (loss=1.7388)
    33→  [12/200] step 300 | loss=1.6745 pred=0.2646 kl=0.4951 recon=0.9533 spec=0.8856 | grad=1.69 | 3.5s
    34→Epoch 12/200 | avg_loss=1.6508 pred=0.2440 kl=0.5296 recon=0.9451 spec=0.8929 | 26.6s
    35→  New best model (loss=1.6508)
    36→  [13/200] step 350 | loss=1.6325 pred=0.1979 kl=0.6267 recon=0.9727 spec=0.8834 | grad=10.42 | 25.7s
    37→Epoch 13/200 | avg_loss=1.6044 pred=0.2082 kl=0.5997 recon=0.9364 spec=0.8822 | 26.6s
    38→  New best model (loss=1.6044)
    39→Epoch 14/200 | avg_loss=1.5552 pred=0.1817 kl=0.6592 recon=0.9173 spec=0.8679 | 26.8s
    40→  New best model (loss=1.5552)
    41→  [15/200] step 400 | loss=1.5585 pred=0.1451 kl=0.7535 recon=0.9383 spec=0.8945 | grad=9.06 | 21.9s
    42→Epoch 15/200 | avg_loss=1.5166 pred=0.1529 kl=0.7337 recon=0.9074 spec=0.8595 | 26.7s
    43→  New best model (loss=1.5166)
    44→Epoch 16/200 | avg_loss=1.4953 pred=0.1296 kl=0.8143 recon=0.9060 spec=0.8563 | 26.8s
    45→  New best model (loss=1.4953)
    46→  [17/200] step 450 | loss=1.4022 pred=0.1020 kl=0.9326 recon=0.8479 spec=0.8270 | grad=1.42 | 18.0s
    47→Epoch 17/200 | avg_loss=1.4262 pred=0.1062 kl=0.9142 recon=0.8650 spec=0.8345 | 26.6s
    48→  New best model (loss=1.4262)
    49→Epoch 18/200 | avg_loss=1.3860 pred=0.0832 kl=1.0482 recon=0.8448 spec=0.8242 | 26.7s
    50→  New best model (loss=1.3860)
    51→  [19/200] step 500 | loss=1.3669 pred=0.0708 kl=1.1489 recon=0.8329 spec=0.8203 | grad=3.57 | 13.9s
    52→Epoch 19/200 | avg_loss=1.3673 pred=0.0700 kl=1.1505 recon=0.8358 spec=0.8166 | 26.7s
    53→  New best model (loss=1.3673)
    54→Epoch 20/200 | avg_loss=1.3566 pred=0.0603 kl=1.2370 recon=0.8295 spec=0.8132 | 26.6s
    55→  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3/checkpoints/epoch_0020.pt
    56→  New best model (loss=1.3566)
    57→  [21/200] step 550 | loss=1.4187 pred=0.0541 kl=1.3104 recon=0.8739 spec=0.8481 | grad=6.37 | 10.1s
    58→Epoch 21/200 | avg_loss=1.3588 pred=0.0531 kl=1.3187 recon=0.8329 spec=0.8105 | 27.0s
    59→Epoch 22/200 | avg_loss=1.3616 pred=0.0474 kl=1.3792 recon=0.8339 spec=0.8125 | 27.1s
    60→  [23/200] step 600 | loss=1.4522 pred=0.0438 kl=1.4116 recon=0.8963 spec=0.8676 | grad=4.54 | 6.3s
    61→Epoch 23/200 | avg_loss=1.4340 pred=0.0469 kl=1.3991 recon=0.8903 spec=0.8364 | 26.7s
    62→Epoch 24/200 | avg_loss=1.3592 pred=0.0477 kl=1.3923 recon=0.8288 spec=0.8019 | 26.9s
    63→  [25/200] step 650 | loss=1.3464 pred=0.0445 kl=1.3972 recon=0.8127 spec=0.8104 | grad=4.69 | 2.4s
    64→Epoch 25/200 | avg_loss=1.3340 pred=0.0441 kl=1.3962 recon=0.8080 spec=0.7930 | 26.6s
    65→  New best model (loss=1.3340)
    66→  [26/200] step 700 | loss=1.3405 pred=0.0442 kl=1.3807 recon=0.8050 spec=0.8040 | grad=2.78 | 24.7s
    67→Epoch 26/200 | avg_loss=1.3337 pred=0.0434 kl=1.3884 recon=0.8070 spec=0.7898 | 26.7s
    68→  New best model (loss=1.3337)
    69→Epoch 27/200 | avg_loss=1.3329 pred=0.0444 kl=1.3660 recon=0.8040 spec=0.7882 | 26.9s
    70→  New best model (loss=1.3329)
    71→  [28/200] step 750 | loss=1.3280 pred=0.0462 kl=1.3407 recon=0.7963 spec=0.7850 | grad=4.03 | 20.8s
    72→Epoch 28/200 | avg_loss=1.3587 pred=0.0461 kl=1.3449 recon=0.8220 spec=0.7964 | 26.6s
    73→Epoch 29/200 | avg_loss=1.3300 pred=0.0471 kl=1.3270 recon=0.7972 spec=0.7825 | 26.7s
    74→  New best model (loss=1.3300)
    75→  [30/200] step 800 | loss=1.3473 pred=0.0481 kl=1.3105 recon=0.8069 spec=0.7906 | grad=4.27 | 17.0s
    76→Epoch 30/200 | avg_loss=1.3220 pred=0.0476 kl=1.3130 recon=0.7890 spec=0.7773 | 26.7s
    77→  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3/checkpoints/epoch_0030.pt
    78→  New best model (loss=1.3220)
    79→Epoch 31/200 | avg_loss=1.2998 pred=0.0500 kl=1.2973 recon=0.7671 spec=0.7677 | 26.7s
    80→  New best model (loss=1.2998)
    81→  [32/200] step 850 | loss=1.3111 pred=0.0488 kl=1.3004 recon=0.7676 spec=0.7849 | grad=3.71 | 13.1s
    82→Epoch 32/200 | avg_loss=1.2901 pred=0.0495 kl=1.2976 recon=0.7569 spec=0.7631 | 26.7s
    83→  New best model (loss=1.2901)
    84→Epoch 33/200 | avg_loss=1.3259 pred=0.0522 kl=1.2794 recon=0.7828 spec=0.7741 | 26.8s
    85→  [34/200] step 900 | loss=1.2806 pred=0.0530 kl=1.2680 recon=0.7440 spec=0.7561 | grad=3.79 | 9.4s
    86→Epoch 34/200 | avg_loss=1.2724 pred=0.0526 kl=1.2675 recon=0.7375 spec=0.7525 | 26.8s
    87→  New best model (loss=1.2724)
    88→Epoch 35/200 | avg_loss=1.2755 pred=0.0520 kl=1.2611 recon=0.7387 spec=0.7521 | 26.8s
    89→  [36/200] step 950 | loss=1.3263 pred=0.0543 kl=1.2489 recon=0.7722 spec=0.7800 | grad=3.81 | 5.3s
    90→Epoch 36/200 | avg_loss=1.2786 pred=0.0544 kl=1.2427 recon=0.7372 spec=0.7537 | 26.7s
    91→Epoch 37/200 | avg_loss=1.2692 pred=0.0546 kl=1.2311 recon=0.7279 spec=0.7488 | 26.7s
    92→  New best model (loss=1.2692)
    93→  [38/200] step 1000 | loss=1.2633 pred=0.0540 kl=1.2248 recon=0.7237 spec=0.7445 | grad=2.15 | 1.5s
    94→Epoch 38/200 | avg_loss=1.2539 pred=0.0547 kl=1.2155 recon=0.7141 spec=0.7423 | 26.6s
    95→  New best model (loss=1.2539)
    96→  [39/200] step 1050 | loss=1.3050 pred=0.0559 kl=1.2036 recon=0.7424 spec=0.7795 | grad=6.34 | 23.6s
    97→Epoch 39/200 | avg_loss=1.2668 pred=0.0565 kl=1.2033 recon=0.7202 spec=0.7488 | 26.6s
    98→Epoch 40/200 | avg_loss=1.3123 pred=0.0591 kl=1.1856 recon=0.7549 spec=0.7625 | 26.6s
    99→  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3/checkpoints/epoch_0040.pt
   100→  [41/200] step 1100 | loss=1.2381 pred=0.0562 kl=1.1798 recon=0.6980 spec=0.7318 | grad=2.48 | 19.9s
   101→Epoch 41/200 | avg_loss=1.2511 pred=0.0568 kl=1.1804 recon=0.7075 spec=0.7375 | 26.7s
   102→  New best model (loss=1.2511)
   103→Epoch 42/200 | avg_loss=1.2452 pred=0.0571 kl=1.1767 recon=0.7021 spec=0.7365 | 26.7s
   104→  New best model (loss=1.2452)
   105→  [43/200] step 1150 | loss=1.2513 pred=0.0569 kl=1.1774 recon=0.6994 spec=0.7546 | grad=3.59 | 16.0s
   106→Epoch 43/200 | avg_loss=1.2366 pred=0.0571 kl=1.1771 recon=0.6953 spec=0.7330 | 26.7s
   107→  New best model (loss=1.2366)
   108→Epoch 44/200 | avg_loss=1.2326 pred=0.0572 kl=1.1745 recon=0.6925 spec=0.7309 | 26.8s
   109→  New best model (loss=1.2326)
   110→  [45/200] step 1200 | loss=1.2360 pred=0.0566 kl=1.1756 recon=0.6933 spec=0.7370 | grad=2.57 | 12.8s
   111→Epoch 45/200 | avg_loss=1.2331 pred=0.0568 kl=1.1751 recon=0.6925 spec=0.7325 | 27.5s
   112→Epoch 46/200 | avg_loss=1.2724 pred=0.0585 kl=1.1694 recon=0.7236 spec=0.7469 | 26.8s
   113→  [47/200] step 1250 | loss=1.2325 pred=0.0575 kl=1.1726 recon=0.6910 spec=0.7336 | grad=3.15 | 8.4s
   114→Epoch 47/200 | avg_loss=1.2261 pred=0.0578 kl=1.1746 recon=0.6864 spec=0.7289 | 26.7s
   115→  New best model (loss=1.2261)
   116→Epoch 48/200 | avg_loss=1.2238 pred=0.0567 kl=1.1763 recon=0.6855 spec=0.7279 | 26.7s
   117→  New best model (loss=1.2238)
   118→  [49/200] step 1300 | loss=1.2081 pred=0.0576 kl=1.1726 recon=0.6847 spec=0.6970 | grad=1.74 | 4.4s
   119→Epoch 49/200 | avg_loss=1.2132 pred=0.0565 kl=1.1774 recon=0.6774 spec=0.7232 | 26.5s
   120→  New best model (loss=1.2132)
   121→

Device: cuda
LeWM TTS model: 24.00M parameters
Dataset: 6959 samples loaded
Training: 6959 samples, 27 batches/epoch
Epoch 1/200 | avg_loss=6.4157 pred=1.5698 kl=0.2692 recon=2.3557 spec=4.9796 | 27.5s
  New best model (loss=6.4157)
  [2/200] step 50 | loss=5.9234 pred=1.1925 kl=0.2006 recon=2.3602 spec=4.7394 | grad=1.48 | 22.9s
Epoch 2/200 | avg_loss=6.1036 pred=1.3063 kl=0.2213 recon=2.3550 spec=4.8829 | 26.7s
  New best model (loss=6.1036)
Epoch 3/200 | avg_loss=5.6631 pred=0.9758 kl=0.1740 recon=2.3543 spec=4.6637 | 26.8s
  New best model (loss=5.6631)
  [4/200] step 100 | loss=4.8707 pred=0.7100 kl=0.2010 recon=2.3181 spec=3.6816 | grad=1.11 | 18.9s
Epoch 4/200 | avg_loss=5.0403 pred=0.7430 kl=0.1915 recon=2.3460 spec=3.8993 | 26.7s
  New best model (loss=5.0403)
Epoch 5/200 | avg_loss=4.1751 pred=0.6059 kl=0.2271 recon=2.2503 spec=2.6327 | 26.9s
  New best model (loss=4.1751)
  [6/200] step 150 | loss=3.1632 pred=0.5672 kl=0.2311 recon=1.6025 spec=1.9807 | grad=1.91 | 15.1s
Epoch 6/200 | avg_loss=3.2544 pred=0.5495 kl=0.2322 recon=1.6716 spec=2.0602 | 26.6s
  New best model (loss=3.2544)
Epoch 7/200 | avg_loss=2.5653 pred=0.4760 kl=0.2680 recon=1.3342 spec=1.5015 | 26.6s
  New best model (loss=2.5653)
  [8/200] step 200 | loss=2.2587 pred=0.4182 kl=0.3118 recon=1.2373 spec=1.1950 | grad=8.12 | 11.2s
Epoch 8/200 | avg_loss=2.2319 pred=0.4148 kl=0.3167 recon=1.2378 spec=1.1467 | 26.7s
  New best model (loss=2.2319)
Epoch 9/200 | avg_loss=2.0277 pred=0.3794 kl=0.3523 recon=1.1345 spec=1.0126 | 26.8s
  New best model (loss=2.0277)
  [10/200] step 250 | loss=1.9283 pred=0.3482 kl=0.3789 recon=1.0729 spec=0.9971 | grad=5.03 | 7.3s
Epoch 10/200 | avg_loss=1.8836 pred=0.3400 kl=0.3910 recon=1.0475 spec=0.9737 | 26.8s
  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3/checkpoints/epoch_0010.pt
  New best model (loss=1.8836)
Epoch 11/200 | avg_loss=1.7388 pred=0.2963 kl=0.4481 recon=0.9699 spec=0.9217 | 26.8s
  New best model (loss=1.7388)
  [12/200] step 300 | loss=1.6745 pred=0.2646 kl=0.4951 recon=0.9533 spec=0.8856 | grad=1.69 | 3.5s
Epoch 12/200 | avg_loss=1.6508 pred=0.2440 kl=0.5296 recon=0.9451 spec=0.8929 | 26.6s
  New best model (loss=1.6508)
  [13/200] step 350 | loss=1.6325 pred=0.1979 kl=0.6267 recon=0.9727 spec=0.8834 | grad=10.42 | 25.7s
Epoch 13/200 | avg_loss=1.6044 pred=0.2082 kl=0.5997 recon=0.9364 spec=0.8822 | 26.6s
  New best model (loss=1.6044)
Epoch 14/200 | avg_loss=1.5552 pred=0.1817 kl=0.6592 recon=0.9173 spec=0.8679 | 26.8s
  New best model (loss=1.5552)
  [15/200] step 400 | loss=1.5585 pred=0.1451 kl=0.7535 recon=0.9383 spec=0.8945 | grad=9.06 | 21.9s
Epoch 15/200 | avg_loss=1.5166 pred=0.1529 kl=0.7337 recon=0.9074 spec=0.8595 | 26.7s
  New best model (loss=1.5166)
Epoch 16/200 | avg_loss=1.4953 pred=0.1296 kl=0.8143 recon=0.9060 spec=0.8563 | 26.8s
  New best model (loss=1.4953)
  [17/200] step 450 | loss=1.4022 pred=0.1020 kl=0.9326 recon=0.8479 spec=0.8270 | grad=1.42 | 18.0s
Epoch 17/200 | avg_loss=1.4262 pred=0.1062 kl=0.9142 recon=0.8650 spec=0.8345 | 26.6s
  New best model (loss=1.4262)
Epoch 18/200 | avg_loss=1.3860 pred=0.0832 kl=1.0482 recon=0.8448 spec=0.8242 | 26.7s
  New best model (loss=1.3860)
  [19/200] step 500 | loss=1.3669 pred=0.0708 kl=1.1489 recon=0.8329 spec=0.8203 | grad=3.57 | 13.9s
Epoch 19/200 | avg_loss=1.3673 pred=0.0700 kl=1.1505 recon=0.8358 spec=0.8166 | 26.7s
  New best model (loss=1.3673)
Epoch 20/200 | avg_loss=1.3566 pred=0.0603 kl=1.2370 recon=0.8295 spec=0.8132 | 26.6s
  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3/checkpoints/epoch_0020.pt
  New best model (loss=1.3566)
  [21/200] step 550 | loss=1.4187 pred=0.0541 kl=1.3104 recon=0.8739 spec=0.8481 | grad=6.37 | 10.1s
Epoch 21/200 | avg_loss=1.3588 pred=0.0531 kl=1.3187 recon=0.8329 spec=0.8105 | 27.0s
Epoch 22/200 | avg_loss=1.3616 pred=0.0474 kl=1.3792 recon=0.8339 spec=0.8125 | 27.1s
  [23/200] step 600 | loss=1.4522 pred=0.0438 kl=1.4116 recon=0.8963 spec=0.8676 | grad=4.54 | 6.3s
Epoch 23/200 | avg_loss=1.4340 pred=0.0469 kl=1.3991 recon=0.8903 spec=0.8364 | 26.7s
Epoch 24/200 | avg_loss=1.3592 pred=0.0477 kl=1.3923 recon=0.8288 spec=0.8019 | 26.9s
  [25/200] step 650 | loss=1.3464 pred=0.0445 kl=1.3972 recon=0.8127 spec=0.8104 | grad=4.69 | 2.4s
Epoch 25/200 | avg_loss=1.3340 pred=0.0441 kl=1.3962 recon=0.8080 spec=0.7930 | 26.6s
  New best model (loss=1.3340)
  [26/200] step 700 | loss=1.3405 pred=0.0442 kl=1.3807 recon=0.8050 spec=0.8040 | grad=2.78 | 24.7s
Epoch 26/200 | avg_loss=1.3337 pred=0.0434 kl=1.3884 recon=0.8070 spec=0.7898 | 26.7s
  New best model (loss=1.3337)
Epoch 27/200 | avg_loss=1.3329 pred=0.0444 kl=1.3660 recon=0.8040 spec=0.7882 | 26.9s
  New best model (loss=1.3329)
  [28/200] step 750 | loss=1.3280 pred=0.0462 kl=1.3407 recon=0.7963 spec=0.7850 | grad=4.03 | 20.8s
Epoch 28/200 | avg_loss=1.3587 pred=0.0461 kl=1.3449 recon=0.8220 spec=0.7964 | 26.6s
Epoch 29/200 | avg_loss=1.3300 pred=0.0471 kl=1.3270 recon=0.7972 spec=0.7825 | 26.7s
  New best model (loss=1.3300)
  [30/200] step 800 | loss=1.3473 pred=0.0481 kl=1.3105 recon=0.8069 spec=0.7906 | grad=4.27 | 17.0s
Epoch 30/200 | avg_loss=1.3220 pred=0.0476 kl=1.3130 recon=0.7890 spec=0.7773 | 26.7s
  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3/checkpoints/epoch_0030.pt
  New best model (loss=1.3220)
Epoch 31/200 | avg_loss=1.2998 pred=0.0500 kl=1.2973 recon=0.7671 spec=0.7677 | 26.7s
  New best model (loss=1.2998)
  [32/200] step 850 | loss=1.3111 pred=0.0488 kl=1.3004 recon=0.7676 spec=0.7849 | grad=3.71 | 13.1s
Epoch 32/200 | avg_loss=1.2901 pred=0.0495 kl=1.2976 recon=0.7569 spec=0.7631 | 26.7s
  New best model (loss=1.2901)
Epoch 33/200 | avg_loss=1.3259 pred=0.0522 kl=1.2794 recon=0.7828 spec=0.7741 | 26.8s
  [34/200] step 900 | loss=1.2806 pred=0.0530 kl=1.2680 recon=0.7440 spec=0.7561 | grad=3.79 | 9.4s
Epoch 34/200 | avg_loss=1.2724 pred=0.0526 kl=1.2675 recon=0.7375 spec=0.7525 | 26.8s
  New best model (loss=1.2724)
Epoch 35/200 | avg_loss=1.2755 pred=0.0520 kl=1.2611 recon=0.7387 spec=0.7521 | 26.8s
  [36/200] step 950 | loss=1.3263 pred=0.0543 kl=1.2489 recon=0.7722 spec=0.7800 | grad=3.81 | 5.3s
Epoch 36/200 | avg_loss=1.2786 pred=0.0544 kl=1.2427 recon=0.7372 spec=0.7537 | 26.7s
Epoch 37/200 | avg_loss=1.2692 pred=0.0546 kl=1.2311 recon=0.7279 spec=0.7488 | 26.7s
  New best model (loss=1.2692)
  [38/200] step 1000 | loss=1.2633 pred=0.0540 kl=1.2248 recon=0.7237 spec=0.7445 | grad=2.15 | 1.5s
Epoch 38/200 | avg_loss=1.2539 pred=0.0547 kl=1.2155 recon=0.7141 spec=0.7423 | 26.6s
  New best model (loss=1.2539)
  [39/200] step 1050 | loss=1.3050 pred=0.0559 kl=1.2036 recon=0.7424 spec=0.7795 | grad=6.34 | 23.6s
Epoch 39/200 | avg_loss=1.2668 pred=0.0565 kl=1.2033 recon=0.7202 spec=0.7488 | 26.6s
Epoch 40/200 | avg_loss=1.3123 pred=0.0591 kl=1.1856 recon=0.7549 spec=0.7625 | 26.6s
  Saved checkpoint: /home/ubuntu/lewm-tts/output_v3/checkpoints/epoch_0040.pt
  [41/200] step 1100 | loss=1.2381 pred=0.0562 kl=1.1798 recon=0.6980 spec=0.7318 | grad=2.48 | 19.9s
Epoch 41/200 | avg_loss=1.2511 pred=0.0568 kl=1.1804 recon=0.7075 spec=0.7375 | 26.7s
  New best model (loss=1.2511)
Epoch 42/200 | avg_loss=1.2452 pred=0.0571 kl=1.1767 recon=0.7021 spec=0.7365 | 26.7s
  New best model (loss=1.2452)
  [43/200] step 1150 | loss=1.2513 pred=0.0569 kl=1.1774 recon=0.6994 spec=0.7546 | grad=3.59 | 16.0s
Epoch 43/200 | avg_loss=1.2366 pred=0.0571 kl=1.1771 recon=0.6953 spec=0.7330 | 26.7s
  New best model (loss=1.2366)
Epoch 44/200 | avg_loss=1.2326 pred=0.0572 kl=1.1745 recon=0.6925 spec=0.7309 | 26.8s
  New best model (loss=1.2326)
  [45/200] step 1200 | loss=1.2360 pred=0.0566 kl=1.1756 recon=0.6933 spec=0.7370 | grad=2.57 | 12.8s
Epoch 45/200 | avg_loss=1.2331 pred=0.0568 kl=1.1751 recon=0.6925 spec=0.7325 | 27.5s
Epoch 46/200 | avg_loss=1.2724 pred=0.0585 kl=1.1694 recon=0.7236 spec=0.7469 | 26.8s
  [47/200] step 1250 | loss=1.2325 pred=0.0575 kl=1.1726 recon=0.6910 spec=0.7336 | grad=3.15 | 8.4s
Epoch 47/200 | avg_loss=1.2261 pred=0.0578 kl=1.1746 recon=0.6864 spec=0.7289 | 26.7s
  New best model (loss=1.2261)
Epoch 48/200 | avg_loss=1.2238 pred=0.0567 kl=1.1763 recon=0.6855 spec=0.7279 | 26.7s
  New best model (loss=1.2238)
  [49/200] step 1300 | loss=1.2081 pred=0.0576 kl=1.1726 recon=0.6847 spec=0.6970 | grad=1.74 | 4.4s
Epoch 49/200 | avg_loss=1.2132 pred=0.0565 kl=1.1774 recon=0.6774 spec=0.7232 | 26.5s
  New best model (loss=1.2132)