Update inference.py

inference.py

@@ -3,11 +3,11 @@ from torch import Tensor
 import torch.nn as nn
 import torch.nn.functional as F
 import logging
-import json
 from typing import Optional, Dict, Type, Union, List, Tuple
 from pathlib import Path
 from dataclasses import dataclass
 import os
+import io
 
 from mars5.model import CodecLM, ResidualTransformer
 from vocos import Vocos
@@ -16,7 +16,7 @@ from mars5.diffuser import MultinomialDiffusion, DSH, perform_simple_inference
 from mars5.minbpe.regex import RegexTokenizer, GPT4_SPLIT_PATTERN
 from mars5.minbpe.codebook import CodebookTokenizer
 from mars5.ar_generate import ar_generate
-from mars5.utils import nuke_weight_norm
+from mars5.utils import nuke_weight_norm, construct_padding_mask
 from mars5.trim import trim
 from huggingface_hub import ModelHubMixin, hf_hub_download
 from safetensors import safe_open