Improve efficiency and fix concurrent issue of audio dumping

UnityPy/export/AudioClipConverter.py

@@ -3,7 +3,8 @@
 import ctypes
 import os
 import platform
-from typing import TYPE_CHECKING, Dict, Union
+from threading import Lock
+from typing import TYPE_CHECKING, Dict
 
 from UnityPy.streams import EndianBinaryWriter
 
@@ -27,15 +28,14 @@
 
 
 def get_fmod_path(
-    system: Union["Windows", "Linux", "Darwin"], arch: ["x64", "x86", "arm", "arm64"]
+    system: str,  # "Windows", "Linux", "Darwin"
+    arch: str,  # "x64", "x86", "arm", "arm64"
 ) -> str:
     if system == "Darwin":
         # universal dylib
         return "lib/FMOD/Darwin/libfmod.dylib"
-
     if system == "Windows":
         return f"lib/FMOD/Windows/{arch}/fmod.dll"
-
     if system == "Linux":
         if arch == "x64":
             arch = "x86_64"
@@ -111,40 +111,56 @@ def extract_audioclip_samples(
     return dump_samples(audio, audio_data, convert_pcm_float)
 
 
+SYSTEM_INSTANCES = {}  # (channels, flags) -> (pyfmodex_system_instance, lock)
+SYSTEM_GLOBAL_LOCK = Lock()
+
+def get_pyfmodex_system_instance(channels: int, flags: int):
+    global pyfmodex, SYSTEM_INSTANCES, SYSTEM_GLOBAL_LOCK
+    with SYSTEM_GLOBAL_LOCK:
+        instance_key = (channels, flags)
+        if instance_key in SYSTEM_INSTANCES:
+            return SYSTEM_INSTANCES[instance_key]
+
+        system = pyfmodex.System()
+        system.init(channels, flags, None)
+        lock = Lock()
+        SYSTEM_INSTANCES[instance_key] = (system, lock)
+        return system, lock
+
+
 def dump_samples(
     clip: AudioClip, audio_data: bytes, convert_pcm_float: bool = True
 ) -> Dict[str, bytes]:
+    global pyfmodex
     if pyfmodex is None:
         import_pyfmodex()
     if not pyfmodex:
         return {}
-    # init system
-    system = pyfmodex.System()
-    system.init(clip.m_Channels, pyfmodex.flags.INIT_FLAGS.NORMAL, None)
-
-    sound = system.create_sound(
-        bytes(audio_data),
-        pyfmodex.flags.MODE.OPENMEMORY,
-        exinfo=pyfmodex.structure_declarations.CREATESOUNDEXINFO(
-            length=len(audio_data),
-            numchannels=clip.m_Channels,
-            defaultfrequency=clip.m_Frequency,
-        ),
-    )
 
-    # iterate over subsounds
-    samples = {}
-    for i in range(sound.num_subsounds):
-        if i > 0:
-            filename = "%s-%i.wav" % (clip.m_Name, i)
-        else:
-            filename = "%s.wav" % clip.m_Name
-        subsound = sound.get_subsound(i)
-        samples[filename] = subsound_to_wav(subsound, convert_pcm_float)
-        subsound.release()
+    system, lock = get_pyfmodex_system_instance(clip.m_Channels, pyfmodex.flags.INIT_FLAGS.NORMAL)
+    with lock:
+        sound = system.create_sound(
+            bytes(audio_data),
+            pyfmodex.flags.MODE.OPENMEMORY,
+            exinfo=pyfmodex.structure_declarations.CREATESOUNDEXINFO(
+                length=len(audio_data),
+                numchannels=clip.m_Channels,
+                defaultfrequency=clip.m_Frequency,
+            ),
+        )
+
+        # iterate over subsounds
+        samples = {}
+        for i in range(sound.num_subsounds):
+            if i > 0:
+                filename = "%s-%i.wav" % (clip.m_Name, i)
+            else:
+                filename = "%s.wav" % clip.m_Name
+            subsound = sound.get_subsound(i)
+            samples[filename] = subsound_to_wav(subsound, convert_pcm_float)
+            subsound.release()
 
-    sound.release()
-    system.release()
+        sound.release()
     return samples
 
 
@@ -162,10 +178,12 @@ def subsound_to_wav(subsound, convert_pcm_float: bool = True) -> bytes:
         pyfmodex.enums.SOUND_FORMAT.PCM24,
         pyfmodex.enums.SOUND_FORMAT.PCM32,
     ]:
+        # format is PCM integer
         audio_format = 1
         wav_data_length = sound_data_length
         convert_pcm_float = False
     elif sound_format == pyfmodex.enums.SOUND_FORMAT.PCMFLOAT:
+        # format is IEEE 754 float
         if convert_pcm_float:
             audio_format = 1
             bits = 16
@@ -205,11 +223,11 @@ def subsound_to_wav(subsound, convert_pcm_float: bool = True) -> bytes:
         if convert_pcm_float:
             if np is not None:
                 ptr_data = np.frombuffer(ptr_data, dtype=np.float32)
-                ptr_data = (ptr_data * 2**15).astype(np.int16).tobytes()
+                ptr_data = (ptr_data * (1 << 15)).astype(np.int16).tobytes()
             else:
                 ptr_data = struct.unpack("<%df" % (len(ptr_data) // 4), ptr_data)
                 ptr_data = struct.pack(
-                    "<%dh" % len(ptr_data), *[int(f * 2**15) for f in ptr_data]
+                    "<%dh" % len(ptr_data), *[int(f * (1 << 15)) for f in ptr_data]
                 )
 
         w.write(ptr_data)