
    fGiNQ              
          U d Z ddlZddlmZmZmZmZ ddlZddlZddl	m
Z
 ddlZddlZddlZddlmZ 	 ddlZ eed          s
d Zee_        n# e$ r dZY nw xY w	 dd	lmZ dZn# e$ rZdZeZY dZ[ndZ[ww xY wd
ej        dej        defdZdedededeej        ef         fdZd)dej        dededej        fdZ da!ee         e"d<    ej#                    Z$ e
e%          &                                j'        d         dz  Z(	 e()                    dd          5 Z* ej+        e*          Z,ddd           n# 1 swxY w Y   n# e$ r i Z,Y nw xY w e-e,.                    d          pd          /                                Z0e0re0a1e-e"d<   n ed          a1e-e"d<    ee,.                    dd                     Z2d*d"e-defd#Z3d*dej        ded"e-dej        fd$Z4d+d"ee-         ddfd%Z5 G d& d'          Z6d'd(gZ7dS ),u  轻量级内存说话人注册表，用于处理 PCM 切片。

- 在内存中维护说话人 embedding 缓存。
- 给定 PCM 切片，如果相似度超过阈值则返回已有说话人 ID；
    否则分配新的 ID 并缓存。
- 设计为被 `gzzm.app` 导入以支持下游的说话人分离逻辑。

注意：当前使用一个简单的 embedding（对帧求均值和标准差）作为占位实现。
如有更好的说话人模型，请替换 `_extract_embedding` 以提取真实 embedding。
    N)ListOptionalTupleUnion)Path   )choose_devicelist_audio_backendsc                      g S N r       src/gzzm/speaker_id.py_dummy_backendsr      s    Ir   )SpeakerRecognitionabreturnc                     t           j                            |           t           j                            |          z  dz   }t          t          j        | |          |z            S )N:0yE>)nplinalgnormfloatdot)r   r   denoms      r   _cosine_simr   /   sJ    Y^^A!2!22d:E1%&&&r   datasample_ratechannelsc                     t          |           dz  dk    rt          d          |dk    rt          d          |dk    rt          d          t          j        | t          j                  }|dk    rd|j        |z  dk    rt          d          |                    d	|                              d
                              t          j                  }|                    t          j	                  dz  }||fS )N   r   z'PCM bytes length must be even for int16zInvalid sample_ratezInvalid channelsdtyper   z"PCM size not divisible by channelsaxisg      @)
len
ValueErrorr   
frombufferint16sizereshapemeanastypefloat32)r   r   r    pcmwavs        r   _pcm_bytes_to_mono_float32r3   4   s    
4yy1}BCCCa.///1}}+,,,
-BH
-
-
-C!||8h!##ABBBkk"h'',,!,44;;BHEE
**RZ
 
 7
*Cr   >  r2   src_srdst_src           	         | j         dk    r t          j        dt          j                  S ||k    r!|                     t          j        d          S | j        d         t          |          z  }t          dt          t          ||z                                }t          j
        d|| j        d         d          }t          j
        d||d          }t          j        |||                               t          j                  S )	Nr   )r   r#   F)copyr           )numendpoint)r,   r   zerosr0   r/   shaper   maxintroundlinspaceinterp)r2   r5   r6   durout_lenx_oldx_news          r   _resample_linearrG   E   s    
x1}}xBJ////zz"*5z111
)A,v
&C!SsV|,,--..GKScilUCCCEKSg>>>E9UE3''..rz:::r   _MODELgzzm_config.jsonrutf-8encodingmodel_device_speaker _MODEL_DEVICE speaker_initial_baseline_seconds      $@cpudevicec                     t           t          d          t          t          5  t          | t
          k    rt          j        dd| i          a| at          cd d d            S # 1 swxY w Y   d S )Nz.speechbrain is required for speaker embeddingsz!speechbrain/spkrec-ecapa-voxcelebrT   )sourcerun_opts)r   ImportError_IMPORT_ERROR_MODEL_LOCKrH   rP   from_hparamsrT   s    r   
_get_modelr]   f   s    !JKKQ^^ 
 	 	 N- 7 7'4:"F+  F #M	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	s   3A$$A(+A(c                    | j         dk    r t          j        dt          j                  S t	          | |d          }t          |          }t          j        |                              d          }t          j	                    5  |
                    |                    |                    }ddd           n# 1 swxY w Y   |                                                                                                                    t          j                  }t          j                            |          dz   }||z  S )uG   通过 SpeechBrain 提取 ECAPA embedding（spkrec-ecapa-voxceleb）。r      r#   r4   r\   Nr   )r,   r   r<   r0   rG   r]   torch
from_numpy	unsqueezeno_gradencode_batchtosqueezerS   numpyr/   r   r   )	r2   r   rT   wav16modeltensorembemb_npr   s	            r   _extract_embeddingrm   x   s:   
x1}}x2:....S+u55Ef%%%Ee$$..q11F	 4 4  6!2!2334 4 4 4 4 4 4 4 4 4 4 4 4 4 4[[]]  &&((//
;;F9>>&!!D(DD=s   )B<<C C c                 4    t          | pt                     dS )zLPreload speaker recognition model at startup to avoid first-request latency.r\   N)r]   rP   r\   s    r   preload_speaker_modelro      s    f-......r   c                   :   e Zd Z	 	 d$eddddddd	ed
ededee         dededededdfdZ		 	 	 d%de
dededee         dedeeeef         z  fdZdej        defdZdedefdZdee         dee         fd Z	 	 d&de
dededee         def
d!Zd'd"Zdefd#ZdS )(SpeakerRegistry      ?333333?NgMbP?i  g333333?T)model_device disable_soft_reuse_for_persistedenergy_thresholdmin_duration_ms	agg_alphaupdate_on_matchsimilarity_thresholdsoft_marginrt   ru   rv   rw   rx   ry   r   c          	         	 g | _         g | _        d| _        t          j                    | _        || _        t          d|          | _        |	 t          t                                                    j        d         dz  }	|	                    dd          5 }
t          j        |
          }ddd           n# 1 swxY w Y   n# t           $ r i }Y nw xY wt#          |                    d|                    d	d
                              }t#          |          | _        t)          |          | _        t-          |          | _        t)          |          | _        t#          |          | _        t5          |          | _        g | _        	 t)          t:                              dd                    dz  | _        n# t           $ r
 d| _        Y nw xY wt5          t:                              dt:                              dd                              }t          |          | _        	 | j                             dd           n# t           $ r Y nw xY wtC                      | _"        dS )u  内存中说话人注册表构造函数。

        新增参数说明（仅列出新增参数的含义与推荐档位）：

        参数（新增）:
            energy_threshold: 当片段的 RMS 能量低于该阈值时视为静音并跳过。值越大越严格（更多片段被视为静音）。
            min_duration_ms: 只有时长 >= 该值（毫秒）的片段才会被用于提取 embedding。增大可减少短切片带来的噪声。
            agg_alpha: 命中后使用移动平均聚合缓存 embedding 时新 embedding 的权重，范围 0.0-1.0。值越小越保守（更偏历史）。
            update_on_match: 命中后是否更新缓存 embedding（启用可逐步聚合说话人特征）。

        推荐档位（示例，按实际数据/场景调整）:
            - 高速度（低延迟/更少模型调用）:
                energy_threshold=5e-3, min_duration_ms=800, agg_alpha=0.6, update_on_match=False
            - 均衡（默认，实时场景常用）:
                energy_threshold=1e-3, min_duration_ms=400, agg_alpha=0.6, update_on_match=True
            - 高精确度（更敏感，适合离线或批处理）:
                energy_threshold=5e-4, min_duration_ms=250, agg_alpha=0.3, update_on_match=True

        使用建议:
            - 实时低延迟服务优先选择“高速度”或“均衡”。
            - 做离线评估或要求较高识别质量时使用“高精确度”。
            - 在引入这些值后，请在目标环境下用小样本集合验证并微调阈值。
        r   r9   NrI   rJ   rK   rL   speaker_disable_soft_reuseru   FrQ   rR        @@     @speaker_store_dirspeaker_storeT)parentsexist_ok)#_embs_ids_next_id	threadingLock_lock_thrr>   _softr   __file__resolver   openjsonload	Exceptionboolget_disable_soft_for_persistedr   _energy_thrr?   _min_dur_ms
_agg_alpha_update_on_matchstr_model_device
_durations_GZZM_CONFIG_initial_baseline_ms
_store_dirmkdirset_persisted_ids)selfrz   r{   rt   ru   rv   rw   rx   ry   cfg_path_f_cfgstores                r   __init__zSpeakerRegistry.__init__   s   L	
 (*
+-	^%%
(	k**
 ,3>>1133;A>ASS]]3]99 )R9R==D) ) ) ) ) ) ) ) ) ) ) ) ) ) )   /35txx@bdi7j7jkk0 0, ,00P+Q+Q( !!122//  	** $_ 5 5 !.. (*	0(-l.>.>?acg.h.h(i(ilr(rD%% 	0 	0 	0(/D%%%	0 L$$%8,:J:J?\k:l:lmmnnu++	O!!$!>>>> 	 	 	D	 "ees[   AC B:.C :B>>C B>C CC0G   GG-I
 

IIr   F	pcm_bytesr   r    rT   
return_simc           	         t          |||          \  }}|p| j        }|j        dk    rd}d}	ntt          t	          j        t	          j        t	          j        |          d                              }t          |j        d                   t          |          z  dz  }	|| j	        k     s|	t          | j
                  k     r|rdndS t          |||          }
t          j                            |
          dk     r|rdndS | j        5  d	}d
}d	}t          t!          | j        | j                            D ]$\  }\  }}t'          |
|          }||k    r|}|}|}%||| j        k    r| j        rs|p|t-          | j                  k     r| j        |         nd}|dk    r:|	dk    r4t          |	          |	|z   z  }t1          dt3          |d                    }n#t1          dt3          d| j                            }||
z  d|z
  | j        |         z  z   }t          j                            |          dz   }||z                      t          j                  | j        |<   	 |t-          | j                  k     r| j        |         nd|	z   | j        |<   nM# t:          $ r@ |t-          | j                  k     r|	| j        |<   n| j                            |	           Y nw xY w|r||fn|cd	d	d	           S | j        dz  }d}| j         r|| j!        v rd}||r|| j        |z
  k    r| j        rs|p|t-          | j                  k     r| j        |         nd}|dk    r:|	dk    r4t          |	          |	|z   z  }t1          dt3          |d                    }n#t1          dt3          d| j                            }||
z  d|z
  | j        |         z  z   }t          j                            |          dz   }||z                      t          j                  | j        |<   	 |t-          | j                  k     r| j        |         nd|	z   | j        |<   nM# t:          $ r@ |t-          | j                  k     r|	| j        |<   n| j                            |	           Y nw xY w|r||fn|cd	d	d	           S |||t-          | j                  k     r| j        |         nd}t1          d| j        | j        dz  z
            }|	| j"        k    rK||	k     rD||k    r=| j        r!|dk    rt          |	          |	|z   z  nd}t1          dt3          |d                    }||
z  d|z
  | j        |         z  z   }t          j                            |          dz   }||z                      t          j                  | j        |<   	 |t-          | j                  k     r| j        |         nd|	z   | j        |<   nM# t:          $ r@ |t-          | j                  k     r|	| j        |<   n| j                            |	           Y nw xY w|r||fn|cd	d	d	           S | j#        }| xj#        dz  c_#        | j                            |           | j                            |
           	 | j                            |	           n # t:          $ r | j        |	gz   | _        Y nw xY w|r||fn|cd	d	d	           S # 1 swxY w Y   d	S )u  如果相似度超过阈值则返回已有说话人 ID，否则新增并返回新 ID。

        参数：
            pcm_bytes: 原始 PCM int16 字节流。
            sample_rate: PCM 的采样率。
            channels: 通道数。
            device: 用于说话人模型的 torch 设备字符串（"cpu" | "cuda" | "cuda:0" ...）。
            return_sim: 如果为 True，同时返回最佳相似度分数。
        r   r9   r&   r~   )r%   r9   r%   r\   gư>Ng      g?gffffff?g      ?r   g      ?TFg?r"   r   )$r3   r   r,   r   r   sqrtr.   squarer=   r   r   rm   r   r   r   	enumeratezipr   r   r   r   r   r(   r   r>   minr   r/   r0   r   appendr   r   r   r   r   )r   r   r   r    rT   r   r2   srrmsdur_msrk   best_idbest_simbest_idxidx	cached_id
cached_embsim
stored_durr   new_embr   effective_softcan_soft_reusemin_enhance_simnew_ids                             r   register_or_matchz!SpeakerRegistry.register_or_match   s   " -YXNNR -4- 8q==CFF	#Q ? ? ?@@AAC39Q<((5994v=F$"""t7G1H1H(H(H *2992 b8889>>#%% *2992Z [	@ [	@15GH&*H09#di:T:T0U0U # #,,i!#z22>>"H'G"H "x49'<'<( ;X-A>FT_I]I]>]>]!:!:cfJ!A~~&1**!&MMVj-@Ac!Tll33Sdo%>%>?? 3wC!Gtz(7K+KLG9>>'22T9D,3dN+B+B2:+N+NDJx(;RZ]`aeap]q]qRqRqT_X5N5Nwz  E  5E11$ ; ; ;#c$/&:&:::8>DOH55 O226:::; /9E**gC[	@ [	@ [	@ [	@ [	@ [	@ [	@ [	@H "Z#-N!N/ 'W@S5S5S!&"~"(tySaGa:b:b( ;X-A>FT_I]I]>]>]!:!:cfJ!A~~&1**!&MMVj-@Ac!Tll33Sdo%>%>?? 3wC!Gtz(7K+KLG9>>'22T9D,3dN+B+B2:+N+NDJx(;RZ]`aeap]q]qRqRqT_X5N5Nwz  E  5E11$ ; ; ;#c$/&:&:::8>DOH55 O226:::	;
 /9E**gw[	@ [	@ [	@ [	@ [	@ [	@ [	@ [	@~ "x';:BSEYEY:Y:YT_X66_b
"%c49
Q+G"H"HT666:;N;NS[_nSnSn, ?EORS^^E&MMVj-@AAY\c!Tll33#$s7a4:h;O/O"P!y~~g66=07$/F/Frz/R/R
8,?V^adeietauauVuVu9R9R{~  CI  9IDOH55( ? ? ?'#do*>*>>><B 9 9 $ 6 6v > > >	?
 3=IGX..'a[	@ [	@ [	@ [	@ [	@ [	@ [	@ [	@f ]FMMQMMIV$$$Jc"""=&&v.... = = ="&/VH"<=)3?FH%%w[	@ [	@ [	@ [	@ [	@ [	@ [	@ [	@ [	@ [	@ [	@ [	@ [	@ [	@ [	@ [	@ [	@ [	@s   E[?*2J[?AK'$[?&K''
[?>D-[?,2Q[?AR)&[?(R))
[? D[?2X[?AY
[?Y
[?$A[?0[
[?[(%[?'[((
[??\\rk   c                     t          t          j                              }| j        | dz  }	 t	          j        ||           n# t          $ r  w xY w|S )u>   将 embedding 持久化到磁盘并返回其 UUID 字符串。.npy)r   uuiduuid4r   r   saver   )r   rk   uidpaths       r   persist_embeddingz!SpeakerRegistry.persist_embeddings  sf    $*,,C-	GD# 	 	 		 
s   A Ar   c                    |sdS | j         | dz  }|                                sdS 	 t          j        |          }n# t          $ r Y dS w xY w| j        5  || j        v r	 ddd           dS | j                            |           |                    t          j	                  }t          j
                            |          dz   }|dk    r||z  }| j                            |           | j                            |           	 | j                            | j                   n*# t          $ r | j                            d           Y nw xY wddd           n# 1 swxY w Y   dS )uU   通过 UUID 将持久化的 embedding 加载到内存缓存；成功则返回 True。Fr   NTr   r   r   )r   existsr   r   r   r   r   r   r/   r0   r   r   r   r   addr   r   )r   r   r   rk   r   s        r   load_persistedzSpeakerRegistry.load_persisted~  s    	5C-{{}} 	5	'$--CC 	 	 	55	Z 	0 	0di	0 	0 	0 	0 	0 	0 	0 	0 IS!!!**RZ((C9>>#&&-DqyyDjJc"""##C(((0&&t'@AAAA 0 0 0&&w/////0	0 	0 	0 	0 	0 	0 	0 	0 	0 	0 	0 	0 	0 	0 	0" tsL   > 
AAE.BE	D)(E)$EEEEE#&E#uidsc                     g }|pg D ]B}	 |                      |          }n# t          $ r d}Y nw xY w|r|                    |           C|S )uC   尝试加载多个持久化 id；返回成功加载的 id 列表。F)r   r   r   )r   r   loadedr   oks        r   load_persisted_bulkz#SpeakerRegistry.load_persisted_bulk  sx    :2 	# 	#C((--    #c"""s    //c                     t          |||          \  }}|p| j        }t          |||          }|                     |          }	 |                     |           n# t
          $ r Y nw xY w|S )ue   从原始 PCM 字节提取 embedding，持久化到磁盘，加载到内存缓存，并返回 UUID。r\   )r3   r   rm   r   r   r   )	r   r   r   r    rT   r2   r   rk   r   s	            r   persist_from_pcmz SpeakerRegistry.persist_from_pcm  s     -YXNNR-4- b888$$S))	$$$$ 	 	 	D	
s   A 
A)(A)c                    | j         5  | j                                         | j                                         d| _        	 | j                                         n# t          $ r Y nw xY wd d d            d S # 1 swxY w Y   d S )Nr   )r   r   clearr   r   r   r   r   s    r   r   zSpeakerRegistry.clear  s    Z 	 	JIOODM%%''''   	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	s5   :A:AA:
A*'A:)A**A::A>A>c                     | j         5  t          | j                  t          | j                  | j        dcd d d            S # 1 swxY w Y   d S )N)countids	threshold)r   r(   r   listr   r   s    r   statszSpeakerRegistry.stats  s    Z 	 	TYDI!Y 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	s   /AAA)rr   rs   )r   NF)r   N)r   N)__name__
__module____qualname__rP   r   r   r   r   r?   r   bytestupler   r   ndarrayr   r   r   r   r   r   dictr   r   r   r   rq   rq      s>        '+ _$
 * <@"&" $_$ _$ _$#_$ _$
 _$ +34._$  _$ _$ _$ _$ 
_$ _$ _$ _$J  $ B@ B@B@ B@ 	B@
 B@ B@ 
uS%Z 	 B@ B@ B@ B@H	RZ 	C 	 	 	 	# $    <
S	 
d3i 
 
 
 
   $   	
  
   &   t      r   rq   ro   )r4   )rS   r   )8__doc__r   typingr   r   r   r   r   ospathlibr   r   rg   r   r`   device_utilsr	   
torchaudiohasattrr   r
   r   speechbrain.pretrainedr   rY   excr   r   r   r   r?   r3   rG   rH   __annotations__r   rZ   r   r   r   _config_pathr   r   r   r   r   r   strip_cfg_devrP   _INITIAL_BASELINE_SECONDSr]   rm   ro   rq   __all__r   r   r   <module>r      sF  	 	 	     / / / / / / / / / / / /  				             ' ' ' ' ' '	7:455 9	 	 	 *9
&   JJJ999999
 MM	    MMMMMM'2: '"* ' ' ' ' '
U   PUVXV`beVePf    "	; 	;"* 	;c 	;3 	;2: 	; 	; 	; 	; (,#$ + + +in tH~~%%''/25GG			3		1	1 %R ty}}% % % % % % % % % % % % % % %   LLL 3| 677=2>>DDFF +!M3!!!!&r**M3***!E,"2"23UW["\"\]]  s '9    $ BJ S # RTR\     / /(3- /4 / / / /
~ ~ ~ ~ ~ ~ ~ ~B
 5
6s_   A AAA' 'A:,A55A:	E	  D=1E	 =EE	 EE	 	EE