noise-suppress-filter.c 28 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058
  1. #include <stdint.h>
  2. #include <inttypes.h>
  3. #include <util/circlebuf.h>
  4. #include <obs-module.h>
  5. #ifdef LIBSPEEXDSP_ENABLED
  6. #include <speex/speex_preprocess.h>
  7. #endif
  8. #ifdef LIBRNNOISE_ENABLED
  9. #ifdef _MSC_VER
  10. #define ssize_t intptr_t
  11. #endif
  12. #include <rnnoise.h>
  13. #include <media-io/audio-resampler.h>
  14. #endif
  15. bool nvafx_loaded = false;
  16. #ifdef LIBNVAFX_ENABLED
  17. #include "nvafx-load.h"
  18. #include <pthread.h>
  19. #endif
  20. /* -------------------------------------------------------- */
  21. #define do_log(level, format, ...) \
  22. blog(level, "[noise suppress: '%s'] " format, \
  23. obs_source_get_name(ng->context), ##__VA_ARGS__)
  24. #define warn(format, ...) do_log(LOG_WARNING, format, ##__VA_ARGS__)
  25. #define info(format, ...) do_log(LOG_INFO, format, ##__VA_ARGS__)
  26. #ifdef _DEBUG
  27. #define debug(format, ...) do_log(LOG_DEBUG, format, ##__VA_ARGS__)
  28. #else
  29. #define debug(format, ...)
  30. #endif
  31. /* -------------------------------------------------------- */
  32. #define S_SUPPRESS_LEVEL "suppress_level"
  33. #define S_NVAFX_INTENSITY "intensity"
  34. #define S_METHOD "method"
  35. #define S_METHOD_SPEEX "speex"
  36. #define S_METHOD_RNN "rnnoise"
  37. #define S_METHOD_NVAFX "nvafx"
  38. #define MT_ obs_module_text
  39. #define TEXT_SUPPRESS_LEVEL MT_("NoiseSuppress.SuppressLevel")
  40. #define TEXT_NVAFX_INTENSITY MT_("NoiseSuppress.Intensity")
  41. #define TEXT_METHOD MT_("NoiseSuppress.Method")
  42. #define TEXT_METHOD_SPEEX MT_("NoiseSuppress.Method.Speex")
  43. #define TEXT_METHOD_RNN MT_("NoiseSuppress.Method.RNNoise")
  44. #define TEXT_METHOD_NVAFX MT_("NoiseSuppress.Method.nvafx")
  45. #define MAX_PREPROC_CHANNELS 8
  46. /* RNNoise constants, these can't be changed */
  47. #define RNNOISE_SAMPLE_RATE 48000
  48. #define RNNOISE_FRAME_SIZE 480
  49. /* nvafx constants, these can't be changed */
  50. #define NVAFX_SAMPLE_RATE 48000
  51. #define NVAFX_FRAME_SIZE \
  52. 480 /* the sdk does not explicitly set this as a constant though it relies on it*/
  53. /* If the following constant changes, RNNoise breaks */
  54. #define BUFFER_SIZE_MSEC 10
  55. /* -------------------------------------------------------- */
  56. struct noise_suppress_data {
  57. obs_source_t *context;
  58. int suppress_level;
  59. uint64_t last_timestamp;
  60. uint64_t latency;
  61. size_t frames;
  62. size_t channels;
  63. struct circlebuf info_buffer;
  64. struct circlebuf input_buffers[MAX_PREPROC_CHANNELS];
  65. struct circlebuf output_buffers[MAX_PREPROC_CHANNELS];
  66. bool use_rnnoise;
  67. bool use_nvafx;
  68. bool nvafx_enabled;
  69. #ifdef LIBSPEEXDSP_ENABLED
  70. /* Speex preprocessor state */
  71. SpeexPreprocessState *spx_states[MAX_PREPROC_CHANNELS];
  72. #endif
  73. #ifdef LIBRNNOISE_ENABLED
  74. /* RNNoise state */
  75. DenoiseState *rnn_states[MAX_PREPROC_CHANNELS];
  76. /* Resampler */
  77. audio_resampler_t *rnn_resampler;
  78. audio_resampler_t *rnn_resampler_back;
  79. #endif
  80. #ifdef LIBNVAFX_ENABLED
  81. /* NVAFX handle, one per audio channel */
  82. NvAFX_Handle handle[MAX_PREPROC_CHANNELS];
  83. uint32_t sample_rate;
  84. float intensity_ratio;
  85. unsigned int num_samples_per_frame, num_channels;
  86. char *model;
  87. bool nvafx_initialized;
  88. /* Resampler */
  89. audio_resampler_t *nvafx_resampler;
  90. audio_resampler_t *nvafx_resampler_back;
  91. /* Initialization */
  92. bool nvafx_loading;
  93. pthread_t nvafx_thread;
  94. pthread_mutex_t nvafx_mutex;
  95. #endif
  96. /* PCM buffers */
  97. float *copy_buffers[MAX_PREPROC_CHANNELS];
  98. #ifdef LIBSPEEXDSP_ENABLED
  99. spx_int16_t *spx_segment_buffers[MAX_PREPROC_CHANNELS];
  100. #endif
  101. #ifdef LIBRNNOISE_ENABLED
  102. float *rnn_segment_buffers[MAX_PREPROC_CHANNELS];
  103. #endif
  104. #ifdef LIBNVAFX_ENABLED
  105. float *nvafx_segment_buffers[MAX_PREPROC_CHANNELS];
  106. #endif
  107. /* output data */
  108. struct obs_audio_data output_audio;
  109. DARRAY(float) output_data;
  110. };
  111. #ifdef LIBNVAFX_ENABLED
  112. /* global mutex for nvafx load functions since they aren't thread-safe */
  113. bool nvafx_initializer_mutex_initialized;
  114. pthread_mutex_t nvafx_initializer_mutex;
  115. #endif
  116. /* -------------------------------------------------------- */
  117. #define SUP_MIN -60
  118. #define SUP_MAX 0
  119. static const float c_32_to_16 = (float)INT16_MAX;
  120. static const float c_16_to_32 = ((float)INT16_MAX + 1.0f);
  121. /* -------------------------------------------------------- */
  122. static const char *noise_suppress_name(void *unused)
  123. {
  124. UNUSED_PARAMETER(unused);
  125. return obs_module_text("NoiseSuppress");
  126. }
  127. static void noise_suppress_destroy(void *data)
  128. {
  129. struct noise_suppress_data *ng = data;
  130. #ifdef LIBNVAFX_ENABLED
  131. if (ng->nvafx_enabled)
  132. pthread_mutex_lock(&ng->nvafx_mutex);
  133. #endif
  134. for (size_t i = 0; i < ng->channels; i++) {
  135. #ifdef LIBSPEEXDSP_ENABLED
  136. speex_preprocess_state_destroy(ng->spx_states[i]);
  137. #endif
  138. #ifdef LIBRNNOISE_ENABLED
  139. rnnoise_destroy(ng->rnn_states[i]);
  140. #endif
  141. #ifdef LIBNVAFX_ENABLED
  142. if (ng->handle[0]) {
  143. if (NvAFX_DestroyEffect(ng->handle[i]) !=
  144. NVAFX_STATUS_SUCCESS) {
  145. do_log(LOG_ERROR, "NvAFX_Release() failed");
  146. }
  147. }
  148. #endif
  149. circlebuf_free(&ng->input_buffers[i]);
  150. circlebuf_free(&ng->output_buffers[i]);
  151. }
  152. #ifdef LIBSPEEXDSP_ENABLED
  153. bfree(ng->spx_segment_buffers[0]);
  154. #endif
  155. #ifdef LIBRNNOISE_ENABLED
  156. bfree(ng->rnn_segment_buffers[0]);
  157. if (ng->rnn_resampler) {
  158. audio_resampler_destroy(ng->rnn_resampler);
  159. audio_resampler_destroy(ng->rnn_resampler_back);
  160. }
  161. #endif
  162. #ifdef LIBNVAFX_ENABLED
  163. bfree(ng->nvafx_segment_buffers[0]);
  164. if (ng->nvafx_resampler) {
  165. audio_resampler_destroy(ng->nvafx_resampler);
  166. audio_resampler_destroy(ng->nvafx_resampler_back);
  167. }
  168. bfree(ng->model);
  169. if (ng->nvafx_enabled) {
  170. if (ng->use_nvafx)
  171. pthread_join(ng->nvafx_thread, NULL);
  172. pthread_mutex_unlock(&ng->nvafx_mutex);
  173. pthread_mutex_destroy(&ng->nvafx_mutex);
  174. }
  175. #endif
  176. bfree(ng->copy_buffers[0]);
  177. circlebuf_free(&ng->info_buffer);
  178. da_free(ng->output_data);
  179. bfree(ng);
  180. }
  181. static void *nvafx_initialize(void *data)
  182. {
  183. #ifdef LIBNVAFX_ENABLED
  184. struct noise_suppress_data *ng = data;
  185. int err;
  186. if (!ng->use_nvafx || !nvafx_loaded)
  187. return NULL;
  188. pthread_mutex_lock(&ng->nvafx_mutex);
  189. pthread_mutex_lock(&nvafx_initializer_mutex);
  190. if (!ng->handle[0]) {
  191. ng->sample_rate = NVAFX_SAMPLE_RATE;
  192. for (size_t i = 0; i < ng->channels; i++) {
  193. err = NvAFX_CreateEffect(NVAFX_EFFECT_DENOISER,
  194. &ng->handle[i]);
  195. if (err != NVAFX_STATUS_SUCCESS) {
  196. do_log(LOG_ERROR,
  197. "NvAFX_CreateEffect() failed, error %i",
  198. err);
  199. goto failure;
  200. }
  201. err = NvAFX_SetU32(ng->handle[i],
  202. NVAFX_PARAM_DENOISER_SAMPLE_RATE,
  203. ng->sample_rate);
  204. if (err != NVAFX_STATUS_SUCCESS) {
  205. do_log(LOG_ERROR,
  206. "NvAFX_SetU32(Sample Rate: %f) failed, error %i",
  207. ng->sample_rate, err);
  208. goto failure;
  209. }
  210. // initial setting of intensity to 1.0f
  211. err = NvAFX_SetFloat(
  212. ng->handle[i],
  213. NVAFX_PARAM_DENOISER_INTENSITY_RATIO,
  214. ng->intensity_ratio);
  215. if (err != NVAFX_STATUS_SUCCESS) {
  216. do_log(LOG_ERROR,
  217. "NvAFX_SetFloat(Intensity Ratio: %f) failed, error %i",
  218. 1.0f, err);
  219. goto failure;
  220. }
  221. err = NvAFX_SetString(ng->handle[i],
  222. NVAFX_PARAM_DENOISER_MODEL_PATH,
  223. ng->model);
  224. if (err != NVAFX_STATUS_SUCCESS) {
  225. do_log(LOG_ERROR,
  226. "NvAFX_SetString() failed, error %i",
  227. err);
  228. goto failure;
  229. }
  230. err = NvAFX_Load(ng->handle[i]);
  231. if (err != NVAFX_STATUS_SUCCESS) {
  232. do_log(LOG_ERROR,
  233. "NvAFX_Load() failed with error %i",
  234. err);
  235. goto failure;
  236. }
  237. }
  238. if (ng->use_nvafx) {
  239. err = NvAFX_GetU32(ng->handle[0],
  240. NVAFX_PARAM_DENOISER_NUM_CHANNELS,
  241. &ng->num_channels);
  242. if (err != NVAFX_STATUS_SUCCESS) {
  243. do_log(LOG_ERROR,
  244. "NvAFX_GetU32() failed to get the number of channels, error %i",
  245. err);
  246. goto failure;
  247. }
  248. if (ng->num_channels != 1) {
  249. do_log(LOG_ERROR,
  250. "The number of channels is not 1 in the sdk any more ==> update code");
  251. goto failure;
  252. }
  253. NvAFX_Status err = NvAFX_GetU32(
  254. ng->handle[0],
  255. NVAFX_PARAM_DENOISER_NUM_SAMPLES_PER_FRAME,
  256. &ng->num_samples_per_frame);
  257. if (err != NVAFX_STATUS_SUCCESS) {
  258. do_log(LOG_ERROR,
  259. "NvAFX_GetU32() failed to get the number of samples per frame, error %i",
  260. err);
  261. goto failure;
  262. }
  263. if (ng->num_samples_per_frame != NVAFX_FRAME_SIZE) {
  264. do_log(LOG_ERROR,
  265. "The number of samples per frame has changed from 480 (= 10 ms) ==> update code");
  266. goto failure;
  267. }
  268. }
  269. }
  270. ng->nvafx_initialized = true;
  271. pthread_mutex_unlock(&nvafx_initializer_mutex);
  272. pthread_mutex_unlock(&ng->nvafx_mutex);
  273. return NULL;
  274. failure:
  275. ng->use_nvafx = false;
  276. pthread_mutex_unlock(&nvafx_initializer_mutex);
  277. pthread_mutex_unlock(&ng->nvafx_mutex);
  278. return NULL;
  279. #else
  280. UNUSED_PARAMETER(data);
  281. return NULL;
  282. #endif
  283. }
  284. static inline void alloc_channel(struct noise_suppress_data *ng,
  285. uint32_t sample_rate, size_t channel,
  286. size_t frames)
  287. {
  288. #ifdef LIBSPEEXDSP_ENABLED
  289. ng->spx_states[channel] =
  290. speex_preprocess_state_init((int)frames, sample_rate);
  291. #endif
  292. #ifdef LIBRNNOISE_ENABLED
  293. ng->rnn_states[channel] = rnnoise_create(NULL);
  294. #endif
  295. circlebuf_reserve(&ng->input_buffers[channel], frames * sizeof(float));
  296. circlebuf_reserve(&ng->output_buffers[channel], frames * sizeof(float));
  297. }
  298. static inline enum speaker_layout convert_speaker_layout(uint8_t channels)
  299. {
  300. switch (channels) {
  301. case 0:
  302. return SPEAKERS_UNKNOWN;
  303. case 1:
  304. return SPEAKERS_MONO;
  305. case 2:
  306. return SPEAKERS_STEREO;
  307. case 3:
  308. return SPEAKERS_2POINT1;
  309. case 4:
  310. return SPEAKERS_4POINT0;
  311. case 5:
  312. return SPEAKERS_4POINT1;
  313. case 6:
  314. return SPEAKERS_5POINT1;
  315. case 8:
  316. return SPEAKERS_7POINT1;
  317. default:
  318. return SPEAKERS_UNKNOWN;
  319. }
  320. }
  321. static void noise_suppress_update(void *data, obs_data_t *s)
  322. {
  323. struct noise_suppress_data *ng = data;
  324. uint32_t sample_rate = audio_output_get_sample_rate(obs_get_audio());
  325. size_t channels = audio_output_get_channels(obs_get_audio());
  326. size_t frames = (size_t)sample_rate / (1000 / BUFFER_SIZE_MSEC);
  327. const char *method = obs_data_get_string(s, S_METHOD);
  328. ng->suppress_level = (int)obs_data_get_int(s, S_SUPPRESS_LEVEL);
  329. ng->latency = 1000000000LL / (1000 / BUFFER_SIZE_MSEC);
  330. ng->use_rnnoise = strcmp(method, S_METHOD_RNN) == 0;
  331. ng->use_nvafx = ng->nvafx_enabled &&
  332. strcmp(method, S_METHOD_NVAFX) == 0;
  333. /* Process 10 millisecond segments to keep latency low */
  334. /* Also RNNoise only supports buffers of this exact size. */
  335. /* At 48kHz, NVAFX processes 480 samples which corresponds to 10 ms.*/
  336. ng->frames = frames;
  337. ng->channels = channels;
  338. #ifdef LIBNVAFX_ENABLED
  339. ng->intensity_ratio = (float)obs_data_get_double(s, S_NVAFX_INTENSITY);
  340. if (ng->use_nvafx) {
  341. pthread_mutex_lock(&ng->nvafx_mutex);
  342. if (ng->nvafx_initialized) {
  343. int err;
  344. for (size_t i = 0; i < ng->channels; i++) {
  345. err = NvAFX_SetFloat(
  346. ng->handle[i],
  347. NVAFX_PARAM_DENOISER_INTENSITY_RATIO,
  348. ng->intensity_ratio);
  349. if (err != NVAFX_STATUS_SUCCESS) {
  350. do_log(LOG_ERROR,
  351. "NvAFX_SetFloat(Intensity Ratio: %f) failed, error %i",
  352. ng->intensity_ratio, err);
  353. ng->use_nvafx = false;
  354. }
  355. }
  356. }
  357. pthread_mutex_unlock(&ng->nvafx_mutex);
  358. }
  359. #endif
  360. /* Ignore if already allocated */
  361. #if defined(LIBSPEEXDSP_ENABLED)
  362. if (!ng->use_rnnoise && !ng->use_nvafx && ng->spx_states[0])
  363. return;
  364. #endif
  365. #ifdef LIBNVAFX_ENABLED
  366. if (ng->use_nvafx && (ng->nvafx_initialized || ng->nvafx_loading))
  367. return;
  368. #endif
  369. #ifdef LIBRNNOISE_ENABLED
  370. if (ng->use_rnnoise && ng->rnn_states[0])
  371. return;
  372. #endif
  373. /* One speex/rnnoise state for each channel (limit 2) */
  374. ng->copy_buffers[0] = bmalloc(frames * channels * sizeof(float));
  375. #ifdef LIBSPEEXDSP_ENABLED
  376. ng->spx_segment_buffers[0] =
  377. bmalloc(frames * channels * sizeof(spx_int16_t));
  378. #endif
  379. #ifdef LIBRNNOISE_ENABLED
  380. ng->rnn_segment_buffers[0] =
  381. bmalloc(RNNOISE_FRAME_SIZE * channels * sizeof(float));
  382. #endif
  383. #ifdef LIBNVAFX_ENABLED
  384. ng->nvafx_segment_buffers[0] =
  385. bmalloc(NVAFX_FRAME_SIZE * channels * sizeof(float));
  386. #endif
  387. for (size_t c = 1; c < channels; ++c) {
  388. ng->copy_buffers[c] = ng->copy_buffers[c - 1] + frames;
  389. #ifdef LIBSPEEXDSP_ENABLED
  390. ng->spx_segment_buffers[c] =
  391. ng->spx_segment_buffers[c - 1] + frames;
  392. #endif
  393. #ifdef LIBRNNOISE_ENABLED
  394. ng->rnn_segment_buffers[c] =
  395. ng->rnn_segment_buffers[c - 1] + RNNOISE_FRAME_SIZE;
  396. #endif
  397. #ifdef LIBNVAFX_ENABLED
  398. ng->nvafx_segment_buffers[c] =
  399. ng->nvafx_segment_buffers[c - 1] + NVAFX_FRAME_SIZE;
  400. #endif
  401. }
  402. #ifdef LIBNVAFX_ENABLED
  403. if (!ng->nvafx_initialized && ng->use_nvafx && !ng->nvafx_loading) {
  404. ng->nvafx_loading = true;
  405. pthread_create(&ng->nvafx_thread, NULL, nvafx_initialize, ng);
  406. }
  407. #endif
  408. for (size_t i = 0; i < channels; i++)
  409. alloc_channel(ng, sample_rate, i, frames);
  410. #ifdef LIBRNNOISE_ENABLED
  411. if (sample_rate == RNNOISE_SAMPLE_RATE) {
  412. ng->rnn_resampler = NULL;
  413. ng->rnn_resampler_back = NULL;
  414. } else {
  415. struct resample_info src, dst;
  416. src.samples_per_sec = sample_rate;
  417. src.format = AUDIO_FORMAT_FLOAT_PLANAR;
  418. src.speakers = convert_speaker_layout((uint8_t)channels);
  419. dst.samples_per_sec = RNNOISE_SAMPLE_RATE;
  420. dst.format = AUDIO_FORMAT_FLOAT_PLANAR;
  421. dst.speakers = convert_speaker_layout((uint8_t)channels);
  422. ng->rnn_resampler = audio_resampler_create(&dst, &src);
  423. ng->rnn_resampler_back = audio_resampler_create(&src, &dst);
  424. }
  425. #endif
  426. #ifdef LIBNVAFX_ENABLED
  427. if (sample_rate == NVAFX_SAMPLE_RATE) {
  428. ng->nvafx_resampler = NULL;
  429. ng->nvafx_resampler_back = NULL;
  430. } else {
  431. struct resample_info src, dst;
  432. src.samples_per_sec = sample_rate;
  433. src.format = AUDIO_FORMAT_FLOAT_PLANAR;
  434. src.speakers = convert_speaker_layout((uint8_t)channels);
  435. dst.samples_per_sec = NVAFX_SAMPLE_RATE;
  436. dst.format = AUDIO_FORMAT_FLOAT_PLANAR;
  437. dst.speakers = convert_speaker_layout((uint8_t)channels);
  438. ng->nvafx_resampler = audio_resampler_create(&dst, &src);
  439. ng->nvafx_resampler_back = audio_resampler_create(&src, &dst);
  440. }
  441. #endif
  442. }
  443. #ifdef _MSC_VER
  444. #pragma warning(push)
  445. #pragma warning(disable : 4706)
  446. #endif
  447. bool load_nvafx(void)
  448. {
  449. #ifdef LIBNVAFX_ENABLED
  450. if (!load_lib()) {
  451. blog(LOG_INFO,
  452. "[noise suppress]: NVIDIA RTX denoiser disabled, redistributable not found");
  453. return false;
  454. }
  455. nvafx_initializer_mutex_initialized =
  456. pthread_mutex_init(&nvafx_initializer_mutex, NULL) == 0;
  457. #define LOAD_SYM_FROM_LIB(sym, lib, dll) \
  458. if (!(sym = (sym##_t)GetProcAddress(lib, #sym))) { \
  459. DWORD err = GetLastError(); \
  460. printf("[noise suppress]: Couldn't load " #sym " from " dll \
  461. ": %lu (0x%lx)", \
  462. err, err); \
  463. goto unload_everything; \
  464. }
  465. #define LOAD_SYM(sym) LOAD_SYM_FROM_LIB(sym, nv_audiofx, "NVAudioEffects.dll")
  466. LOAD_SYM(NvAFX_GetEffectList);
  467. LOAD_SYM(NvAFX_CreateEffect);
  468. LOAD_SYM(NvAFX_DestroyEffect);
  469. LOAD_SYM(NvAFX_SetU32);
  470. LOAD_SYM(NvAFX_SetString);
  471. LOAD_SYM(NvAFX_SetFloat);
  472. LOAD_SYM(NvAFX_GetU32);
  473. LOAD_SYM(NvAFX_GetString);
  474. LOAD_SYM(NvAFX_GetFloat);
  475. LOAD_SYM(NvAFX_Load);
  476. LOAD_SYM(NvAFX_Run);
  477. #undef LOAD_SYM
  478. int err;
  479. NvAFX_Handle h = NULL;
  480. err = NvAFX_CreateEffect(NVAFX_EFFECT_DENOISER, &h);
  481. if (err != NVAFX_STATUS_SUCCESS) {
  482. if (err == NVAFX_STATUS_GPU_UNSUPPORTED) {
  483. blog(LOG_INFO,
  484. "[noise suppress]: NVIDIA RTX denoiser disabled: unsupported GPU");
  485. } else {
  486. blog(LOG_ERROR,
  487. "[noise suppress]: NVIDIA RTX denoiser disabled: error %i",
  488. err);
  489. }
  490. goto unload_everything;
  491. }
  492. err = NvAFX_DestroyEffect(h);
  493. if (err != NVAFX_STATUS_SUCCESS) {
  494. blog(LOG_ERROR, "NvAFX_DestroyEffect() failed, error %i", err);
  495. goto unload_everything;
  496. }
  497. nvafx_loaded = true;
  498. blog(LOG_INFO, "[noise suppress]: NVIDIA RTX denoiser enabled");
  499. return true;
  500. unload_everything:
  501. release_lib();
  502. #endif
  503. return false;
  504. }
  505. #ifdef _MSC_VER
  506. #pragma warning(pop)
  507. #endif
  508. void unload_nvafx(void)
  509. {
  510. #ifdef LIBNVAFX_ENABLED
  511. release_lib();
  512. if (nvafx_initializer_mutex_initialized) {
  513. pthread_mutex_destroy(&nvafx_initializer_mutex);
  514. nvafx_initializer_mutex_initialized = false;
  515. }
  516. #endif
  517. }
  518. static void *noise_suppress_create(obs_data_t *settings, obs_source_t *filter)
  519. {
  520. struct noise_suppress_data *ng =
  521. bzalloc(sizeof(struct noise_suppress_data));
  522. ng->context = filter;
  523. #ifdef LIBNVAFX_ENABLED
  524. char sdk_path[MAX_PATH];
  525. if (!nvafx_get_sdk_path(sdk_path, sizeof(sdk_path))) {
  526. ng->nvafx_enabled = false;
  527. do_log(LOG_ERROR, "NVAFX redist is not installed.");
  528. } else {
  529. const char *file = "\\models\\denoiser_48k.trtpkg";
  530. size_t size = strlen(sdk_path) + strlen(file) + 1;
  531. char *buffer = (char *)bmalloc(size);
  532. strcpy(buffer, sdk_path);
  533. strcat(buffer, file);
  534. ng->model = buffer;
  535. ng->nvafx_enabled = true;
  536. ng->nvafx_initialized = false;
  537. ng->nvafx_loading = false;
  538. pthread_mutex_init(&ng->nvafx_mutex, NULL);
  539. info("NVAFX SDK redist path was found here %s", sdk_path);
  540. }
  541. #endif
  542. noise_suppress_update(ng, settings);
  543. return ng;
  544. }
  545. static inline void process_speexdsp(struct noise_suppress_data *ng)
  546. {
  547. #ifdef LIBSPEEXDSP_ENABLED
  548. /* Set args */
  549. for (size_t i = 0; i < ng->channels; i++)
  550. speex_preprocess_ctl(ng->spx_states[i],
  551. SPEEX_PREPROCESS_SET_NOISE_SUPPRESS,
  552. &ng->suppress_level);
  553. /* Convert to 16bit */
  554. for (size_t i = 0; i < ng->channels; i++)
  555. for (size_t j = 0; j < ng->frames; j++) {
  556. float s = ng->copy_buffers[i][j];
  557. if (s > 1.0f)
  558. s = 1.0f;
  559. else if (s < -1.0f)
  560. s = -1.0f;
  561. ng->spx_segment_buffers[i][j] =
  562. (spx_int16_t)(s * c_32_to_16);
  563. }
  564. /* Execute */
  565. for (size_t i = 0; i < ng->channels; i++)
  566. speex_preprocess_run(ng->spx_states[i],
  567. ng->spx_segment_buffers[i]);
  568. /* Convert back to 32bit */
  569. for (size_t i = 0; i < ng->channels; i++)
  570. for (size_t j = 0; j < ng->frames; j++)
  571. ng->copy_buffers[i][j] =
  572. (float)ng->spx_segment_buffers[i][j] /
  573. c_16_to_32;
  574. #endif
  575. }
  576. static inline void process_rnnoise(struct noise_suppress_data *ng)
  577. {
  578. #ifdef LIBRNNOISE_ENABLED
  579. /* Adjust signal level to what RNNoise expects, resample if necessary */
  580. if (ng->rnn_resampler) {
  581. float *output[MAX_PREPROC_CHANNELS];
  582. uint32_t out_frames;
  583. uint64_t ts_offset;
  584. audio_resampler_resample(ng->rnn_resampler, (uint8_t **)output,
  585. &out_frames, &ts_offset,
  586. (const uint8_t **)ng->copy_buffers,
  587. (uint32_t)ng->frames);
  588. for (size_t i = 0; i < ng->channels; i++) {
  589. for (ssize_t j = 0, k = (ssize_t)out_frames -
  590. RNNOISE_FRAME_SIZE;
  591. j < RNNOISE_FRAME_SIZE; ++j, ++k) {
  592. if (k >= 0) {
  593. ng->rnn_segment_buffers[i][j] =
  594. output[i][k] * 32768.0f;
  595. } else {
  596. ng->rnn_segment_buffers[i][j] = 0;
  597. }
  598. }
  599. }
  600. } else {
  601. for (size_t i = 0; i < ng->channels; i++) {
  602. for (size_t j = 0; j < RNNOISE_FRAME_SIZE; ++j) {
  603. ng->rnn_segment_buffers[i][j] =
  604. ng->copy_buffers[i][j] * 32768.0f;
  605. }
  606. }
  607. }
  608. /* Execute */
  609. for (size_t i = 0; i < ng->channels; i++) {
  610. rnnoise_process_frame(ng->rnn_states[i],
  611. ng->rnn_segment_buffers[i],
  612. ng->rnn_segment_buffers[i]);
  613. }
  614. /* Revert signal level adjustment, resample back if necessary */
  615. if (ng->rnn_resampler) {
  616. float *output[MAX_PREPROC_CHANNELS];
  617. uint32_t out_frames;
  618. uint64_t ts_offset;
  619. audio_resampler_resample(
  620. ng->rnn_resampler_back, (uint8_t **)output, &out_frames,
  621. &ts_offset, (const uint8_t **)ng->rnn_segment_buffers,
  622. RNNOISE_FRAME_SIZE);
  623. for (size_t i = 0; i < ng->channels; i++) {
  624. for (ssize_t j = 0,
  625. k = (ssize_t)out_frames - ng->frames;
  626. j < (ssize_t)ng->frames; ++j, ++k) {
  627. if (k >= 0) {
  628. ng->copy_buffers[i][j] =
  629. output[i][k] / 32768.0f;
  630. } else {
  631. ng->copy_buffers[i][j] = 0;
  632. }
  633. }
  634. }
  635. } else {
  636. for (size_t i = 0; i < ng->channels; i++) {
  637. for (size_t j = 0; j < RNNOISE_FRAME_SIZE; ++j) {
  638. ng->copy_buffers[i][j] =
  639. ng->rnn_segment_buffers[i][j] /
  640. 32768.0f;
  641. }
  642. }
  643. }
  644. #else
  645. UNUSED_PARAMETER(ng);
  646. #endif
  647. }
  648. static inline void process_nvafx(struct noise_suppress_data *ng)
  649. {
  650. #ifdef LIBNVAFX_ENABLED
  651. if (nvafx_loaded && ng->use_nvafx && ng->nvafx_initialized) {
  652. /* Resample if necessary */
  653. if (ng->nvafx_resampler) {
  654. float *output[MAX_PREPROC_CHANNELS];
  655. uint32_t out_frames;
  656. uint64_t ts_offset;
  657. audio_resampler_resample(
  658. ng->nvafx_resampler, (uint8_t **)output,
  659. &out_frames, &ts_offset,
  660. (const uint8_t **)ng->copy_buffers,
  661. (uint32_t)ng->frames);
  662. for (size_t i = 0; i < ng->channels; i++) {
  663. for (ssize_t j = 0, k = (ssize_t)out_frames -
  664. NVAFX_FRAME_SIZE;
  665. j < NVAFX_FRAME_SIZE; ++j, ++k) {
  666. if (k >= 0) {
  667. ng->nvafx_segment_buffers[i][j] =
  668. output[i][k];
  669. } else {
  670. ng->nvafx_segment_buffers[i][j] =
  671. 0;
  672. }
  673. }
  674. }
  675. } else {
  676. for (size_t i = 0; i < ng->channels; i++) {
  677. for (size_t j = 0; j < NVAFX_FRAME_SIZE; ++j) {
  678. ng->nvafx_segment_buffers[i][j] =
  679. ng->copy_buffers[i][j];
  680. }
  681. }
  682. }
  683. /* Execute */
  684. for (size_t i = 0; i < ng->channels; i++) {
  685. NvAFX_Status err = NvAFX_Run(
  686. ng->handle[i], &ng->nvafx_segment_buffers[i],
  687. &ng->nvafx_segment_buffers[i],
  688. ng->num_samples_per_frame, ng->num_channels);
  689. if (err != NVAFX_STATUS_SUCCESS)
  690. do_log(LOG_ERROR,
  691. "NvAFX_Run() failed, error %i", err);
  692. }
  693. /* Revert signal level adjustment, resample back if necessary */
  694. if (ng->nvafx_resampler) {
  695. float *output[MAX_PREPROC_CHANNELS];
  696. uint32_t out_frames;
  697. uint64_t ts_offset;
  698. audio_resampler_resample(
  699. ng->nvafx_resampler_back, (uint8_t **)output,
  700. &out_frames, &ts_offset,
  701. (const uint8_t **)ng->nvafx_segment_buffers,
  702. NVAFX_FRAME_SIZE);
  703. for (size_t i = 0; i < ng->channels; i++) {
  704. for (ssize_t j = 0, k = (ssize_t)out_frames -
  705. ng->frames;
  706. j < (ssize_t)ng->frames; ++j, ++k) {
  707. if (k >= 0) {
  708. ng->copy_buffers[i][j] =
  709. output[i][k];
  710. } else {
  711. ng->copy_buffers[i][j] = 0;
  712. }
  713. }
  714. }
  715. } else {
  716. for (size_t i = 0; i < ng->channels; i++) {
  717. for (size_t j = 0; j < NVAFX_FRAME_SIZE; ++j) {
  718. ng->copy_buffers[i][j] =
  719. ng->nvafx_segment_buffers[i][j];
  720. }
  721. }
  722. }
  723. }
  724. #else
  725. UNUSED_PARAMETER(ng);
  726. #endif
  727. }
  728. static inline void process(struct noise_suppress_data *ng)
  729. {
  730. /* Pop from input circlebuf */
  731. for (size_t i = 0; i < ng->channels; i++)
  732. circlebuf_pop_front(&ng->input_buffers[i], ng->copy_buffers[i],
  733. ng->frames * sizeof(float));
  734. if (ng->use_rnnoise) {
  735. process_rnnoise(ng);
  736. } else if (ng->use_nvafx) {
  737. if (nvafx_loaded) {
  738. process_nvafx(ng);
  739. }
  740. } else {
  741. process_speexdsp(ng);
  742. }
  743. /* Push to output circlebuf */
  744. for (size_t i = 0; i < ng->channels; i++)
  745. circlebuf_push_back(&ng->output_buffers[i], ng->copy_buffers[i],
  746. ng->frames * sizeof(float));
  747. }
  748. struct ng_audio_info {
  749. uint32_t frames;
  750. uint64_t timestamp;
  751. };
  752. static inline void clear_circlebuf(struct circlebuf *buf)
  753. {
  754. circlebuf_pop_front(buf, NULL, buf->size);
  755. }
  756. static void reset_data(struct noise_suppress_data *ng)
  757. {
  758. for (size_t i = 0; i < ng->channels; i++) {
  759. clear_circlebuf(&ng->input_buffers[i]);
  760. clear_circlebuf(&ng->output_buffers[i]);
  761. }
  762. clear_circlebuf(&ng->info_buffer);
  763. }
  764. static struct obs_audio_data *
  765. noise_suppress_filter_audio(void *data, struct obs_audio_data *audio)
  766. {
  767. struct noise_suppress_data *ng = data;
  768. struct ng_audio_info info;
  769. size_t segment_size = ng->frames * sizeof(float);
  770. size_t out_size;
  771. #ifdef LIBSPEEXDSP_ENABLED
  772. if (!ng->spx_states[0])
  773. return audio;
  774. #endif
  775. #ifdef LIBRNNOISE_ENABLED
  776. if (!ng->rnn_states[0])
  777. return audio;
  778. #endif
  779. /* -----------------------------------------------
  780. * if timestamp has dramatically changed, consider it a new stream of
  781. * audio data. clear all circular buffers to prevent old audio data
  782. * from being processed as part of the new data. */
  783. if (ng->last_timestamp) {
  784. int64_t diff = llabs((int64_t)ng->last_timestamp -
  785. (int64_t)audio->timestamp);
  786. if (diff > 1000000000LL)
  787. reset_data(ng);
  788. }
  789. ng->last_timestamp = audio->timestamp;
  790. /* -----------------------------------------------
  791. * push audio packet info (timestamp/frame count) to info circlebuf */
  792. info.frames = audio->frames;
  793. info.timestamp = audio->timestamp;
  794. circlebuf_push_back(&ng->info_buffer, &info, sizeof(info));
  795. /* -----------------------------------------------
  796. * push back current audio data to input circlebuf */
  797. for (size_t i = 0; i < ng->channels; i++)
  798. circlebuf_push_back(&ng->input_buffers[i], audio->data[i],
  799. audio->frames * sizeof(float));
  800. /* -----------------------------------------------
  801. * pop/process each 10ms segments, push back to output circlebuf */
  802. while (ng->input_buffers[0].size >= segment_size)
  803. process(ng);
  804. /* -----------------------------------------------
  805. * peek front of info circlebuf, check to see if we have enough to
  806. * pop the expected packet size, if not, return null */
  807. memset(&info, 0, sizeof(info));
  808. circlebuf_peek_front(&ng->info_buffer, &info, sizeof(info));
  809. out_size = info.frames * sizeof(float);
  810. if (ng->output_buffers[0].size < out_size)
  811. return NULL;
  812. /* -----------------------------------------------
  813. * if there's enough audio data buffered in the output circlebuf,
  814. * pop and return a packet */
  815. circlebuf_pop_front(&ng->info_buffer, NULL, sizeof(info));
  816. da_resize(ng->output_data, out_size * ng->channels);
  817. for (size_t i = 0; i < ng->channels; i++) {
  818. ng->output_audio.data[i] =
  819. (uint8_t *)&ng->output_data.array[i * out_size];
  820. circlebuf_pop_front(&ng->output_buffers[i],
  821. ng->output_audio.data[i], out_size);
  822. }
  823. ng->output_audio.frames = info.frames;
  824. ng->output_audio.timestamp = info.timestamp - ng->latency;
  825. return &ng->output_audio;
  826. }
  827. static bool noise_suppress_method_modified(obs_properties_t *props,
  828. obs_property_t *property,
  829. obs_data_t *settings)
  830. {
  831. obs_property_t *p_suppress_level =
  832. obs_properties_get(props, S_SUPPRESS_LEVEL);
  833. obs_property_t *p_navfx_intensity =
  834. obs_properties_get(props, S_NVAFX_INTENSITY);
  835. const char *method = obs_data_get_string(settings, S_METHOD);
  836. bool enable_level = strcmp(method, S_METHOD_SPEEX) == 0;
  837. bool enable_intensity = strcmp(method, S_METHOD_NVAFX) == 0;
  838. obs_property_set_visible(p_suppress_level, enable_level);
  839. obs_property_set_visible(p_navfx_intensity, enable_intensity);
  840. UNUSED_PARAMETER(property);
  841. return true;
  842. }
  843. static void noise_suppress_defaults_v1(obs_data_t *s)
  844. {
  845. obs_data_set_default_int(s, S_SUPPRESS_LEVEL, -30);
  846. #if defined(LIBRNNOISE_ENABLED) && !defined(LIBSPEEXDSP_ENABLED)
  847. obs_data_set_default_string(s, S_METHOD, S_METHOD_RNN);
  848. #else
  849. obs_data_set_default_string(s, S_METHOD, S_METHOD_SPEEX);
  850. #endif
  851. #if defined(LIBNVAFX_ENABLED)
  852. obs_data_set_default_double(s, S_NVAFX_INTENSITY, 1.0);
  853. #endif
  854. }
  855. static void noise_suppress_defaults_v2(obs_data_t *s)
  856. {
  857. obs_data_set_default_int(s, S_SUPPRESS_LEVEL, -30);
  858. #if defined(LIBRNNOISE_ENABLED)
  859. obs_data_set_default_string(s, S_METHOD, S_METHOD_RNN);
  860. #else
  861. obs_data_set_default_string(s, S_METHOD, S_METHOD_SPEEX);
  862. #endif
  863. #if defined(LIBNVAFX_ENABLED)
  864. obs_data_set_default_double(s, S_NVAFX_INTENSITY, 1.0);
  865. #endif
  866. }
  867. static obs_properties_t *noise_suppress_properties(void *data)
  868. {
  869. obs_properties_t *ppts = obs_properties_create();
  870. #ifdef LIBNVAFX_ENABLED
  871. struct noise_suppress_data *ng = (struct noise_suppress_data *)data;
  872. #else
  873. UNUSED_PARAMETER(data);
  874. #endif
  875. #if defined(LIBRNNOISE_ENABLED) && defined(LIBSPEEXDSP_ENABLED)
  876. obs_property_t *method = obs_properties_add_list(
  877. ppts, S_METHOD, TEXT_METHOD, OBS_COMBO_TYPE_LIST,
  878. OBS_COMBO_FORMAT_STRING);
  879. obs_property_list_add_string(method, TEXT_METHOD_SPEEX, S_METHOD_SPEEX);
  880. obs_property_list_add_string(method, TEXT_METHOD_RNN, S_METHOD_RNN);
  881. #ifdef LIBNVAFX_ENABLED
  882. if (ng->nvafx_enabled)
  883. obs_property_list_add_string(method, TEXT_METHOD_NVAFX,
  884. S_METHOD_NVAFX);
  885. #endif
  886. obs_property_set_modified_callback(method,
  887. noise_suppress_method_modified);
  888. #endif
  889. #ifdef LIBSPEEXDSP_ENABLED
  890. obs_property_t *speex_slider = obs_properties_add_int_slider(
  891. ppts, S_SUPPRESS_LEVEL, TEXT_SUPPRESS_LEVEL, SUP_MIN, SUP_MAX,
  892. 1);
  893. obs_property_int_set_suffix(speex_slider, " dB");
  894. #endif
  895. #ifdef LIBNVAFX_ENABLED
  896. obs_properties_add_float_slider(ppts, S_NVAFX_INTENSITY,
  897. TEXT_NVAFX_INTENSITY, 0.0f, 1.0f,
  898. 0.01f);
  899. if (!nvafx_loaded) {
  900. obs_property_list_item_disable(method, 2, true);
  901. }
  902. #endif
  903. return ppts;
  904. }
  905. struct obs_source_info noise_suppress_filter = {
  906. .id = "noise_suppress_filter",
  907. .type = OBS_SOURCE_TYPE_FILTER,
  908. .output_flags = OBS_SOURCE_AUDIO | OBS_SOURCE_CAP_OBSOLETE,
  909. .get_name = noise_suppress_name,
  910. .create = noise_suppress_create,
  911. .destroy = noise_suppress_destroy,
  912. .update = noise_suppress_update,
  913. .filter_audio = noise_suppress_filter_audio,
  914. .get_defaults = noise_suppress_defaults_v1,
  915. .get_properties = noise_suppress_properties,
  916. };
  917. struct obs_source_info noise_suppress_filter_v2 = {
  918. .id = "noise_suppress_filter",
  919. .version = 2,
  920. .type = OBS_SOURCE_TYPE_FILTER,
  921. .output_flags = OBS_SOURCE_AUDIO,
  922. .get_name = noise_suppress_name,
  923. .create = noise_suppress_create,
  924. .destroy = noise_suppress_destroy,
  925. .update = noise_suppress_update,
  926. .filter_audio = noise_suppress_filter_audio,
  927. .get_defaults = noise_suppress_defaults_v2,
  928. .get_properties = noise_suppress_properties,
  929. };