diff --git a/src/voxcpm/core.py b/src/voxcpm/core.py
index f5e82b44..c48b7e18 100644
--- a/src/voxcpm/core.py
+++ b/src/voxcpm/core.py
@@ -91,7 +91,7 @@ def from_pretrained(
         hf_model_id: str = "openbmb/VoxCPM2",
         load_denoiser: bool = True,
         zipenhancer_model_id: str = "iic/speech_zipenhancer_ans_multiloss_16k_base",
-        cache_dir: str = None,
+        cache_dir: Optional[str] = None,
         local_files_only: bool = False,
         optimize: bool = True,
         lora_config: Optional[LoRAConfig] = None,
@@ -160,9 +160,9 @@ def generate_streaming(self, *args, **kwargs) -> Generator[np.ndarray, None, Non
     def _generate(
         self,
         text: str,
-        prompt_wav_path: str = None,
-        prompt_text: str = None,
-        reference_wav_path: str = None,
+        prompt_wav_path: Optional[str] = None,
+        prompt_text: Optional[str] = None,
+        reference_wav_path: Optional[str] = None,
         cfg_value: float = 2.0,
         inference_timesteps: int = 10,
         min_len: int = 2,
diff --git a/src/voxcpm/modules/locdit/local_dit.py b/src/voxcpm/modules/locdit/local_dit.py
index 521c7c77..aaba8f66 100644
--- a/src/voxcpm/modules/locdit/local_dit.py
+++ b/src/voxcpm/modules/locdit/local_dit.py
@@ -1,4 +1,5 @@
 import torch
+from typing import Optional
 from ..minicpm4 import MiniCPMModel, MiniCPM4Config
 import torch.nn as nn
 import math
@@ -27,7 +28,7 @@ def __init__(
         self,
         in_channels: int,
         time_embed_dim: int,
-        out_dim: int = None,
+        out_dim: Optional[int] = None,
     ):
         super().__init__()
 
diff --git a/src/voxcpm/modules/minicpm4/config.py b/src/voxcpm/modules/minicpm4/config.py
index 332dd11c..821dc390 100644
--- a/src/voxcpm/modules/minicpm4/config.py
+++ b/src/voxcpm/modules/minicpm4/config.py
@@ -1,5 +1,5 @@
 from pydantic import BaseModel
-from typing import List
+from typing import List, Optional
 
 
 class RopeScalingConfig(BaseModel):
@@ -26,5 +26,5 @@ class MiniCPM4Config(BaseModel):
     dim_model_base: int
     scale_depth: float
     rope_theta: float
-    kv_channels: int = None
+    kv_channels: Optional[int] = None
     no_rope: bool = False