diff --git a/src/extensions/core/uploadAudio.ts b/src/extensions/core/uploadAudio.ts
index 147f9958d..d43c793ad 100644
--- a/src/extensions/core/uploadAudio.ts
+++ b/src/extensions/core/uploadAudio.ts
@@ -88,11 +88,11 @@ app.registerExtension({
   name: 'Comfy.AudioWidget',
   async beforeRegisterNodeDef(nodeType, nodeData) {
     if (
-      // @ts-expect-error ComfyNode
-      ['LoadAudio', 'SaveAudio', 'PreviewAudio'].includes(nodeType.comfyClass)
+      ['LoadAudio', 'SaveAudio', 'PreviewAudio'].includes(
+        nodeType.prototype.comfyClass
+      )
     ) {
-      // @ts-expect-error InputSpec is not typed correctly
-      nodeData.input.required.audioUI = ['AUDIO_UI']
+      nodeData.input.required.audioUI = ['AUDIO_UI', {}]
     }
   },
   getCustomWidgets() {
@@ -150,8 +150,7 @@ app.registerExtension({
   name: 'Comfy.UploadAudio',
   async beforeRegisterNodeDef(nodeType, nodeData: ComfyNodeDef) {
     if (nodeData?.input?.required?.audio?.[1]?.audio_upload === true) {
-      // @ts-expect-error InputSpec is not typed correctly
-      nodeData.input.required.upload = ['AUDIOUPLOAD']
+      nodeData.input.required.upload = ['AUDIOUPLOAD', {}]
     }
   },
   getCustomWidgets() {
diff --git a/src/services/litegraphService.ts b/src/services/litegraphService.ts
index 822319336..8abba2c63 100644
--- a/src/services/litegraphService.ts
+++ b/src/services/litegraphService.ts
@@ -173,6 +173,7 @@ export const useLitegraphService = () => {
     // Note: Some extensions expects node.comfyClass to be set in
     // `beforeRegisterNodeDef`.
     node.prototype.comfyClass = nodeDefV1.name
+    node.comfyClass = nodeDefV1.name
     await extensionService.invokeExtensionsAsync(
       'beforeRegisterNodeDef',
       node,
@@ -180,7 +181,6 @@ export const useLitegraphService = () => {
     )
 
     const nodeDef = new ComfyNodeDefImpl(nodeDefV1)
-    node.comfyClass = nodeDef.name
     node.nodeData = nodeDef
     LiteGraph.registerNodeType(nodeId, node)
     // Note: Do not following assignments before `LiteGraph.registerNodeType`