
    m\iW                        S r SSKrSSKrSSKrSSKrSSKrSSKrSSKJr  SSK	J
r
JrJrJr  SSKrSSKrSSKrSSKJr    SS\
\R(                  \R*                  4   S\S\S	\
\R(                  \R*                  4   4S
 jjrSS\
\R(                  \R*                  4   S\S	\
\R(                  \R*                  4   4S jjr " S S5      rS	\4S jrS	\4S jrS	\4S jr SS\
\R(                  \R*                  4   S\\   S	\4S jjr\" SS9r      S S\
\R(                  \R*                  4   S\
\\4   S\S\\   S\!S\\   S\\   S	\4S jjr"g)!z
Audio saving and transcoding utility module

Independent audio file operations outside of handler, supporting:
- Save audio tensor/numpy to files (default FLAC format, fast)
- Format conversion (FLAC/WAV/MP3)
- Batch processing
    N)Path)UnionOptionalListTuple)logger
audio_datafade_in_samplesfade_out_samplesreturnc                 0   US::  a  US::  a  U $ [        U [        R                  5      nU(       a   U R                  5       nUR                  S   nOU R                  5       nUR                  S   nUS:  a  [        X5      nU(       a?  [        R                  " SSXdR                  UR                  S9nUSSU24   U-  USSU24'   O8[        R                  " SSU[        R                  S9nUSSU24   U-  USSU24'   US:  a  [        X%5      nU(       aD  [        R                  " SSXR                  UR                  S9nUSXX-
  S24   U-  USXX-
  S24'   U$ [        R                  " SSU[        R                  S9nUSXX-
  S24   U-  USXX-
  S24'   U$ )	am  Apply linear fade in and/or fade out to audio data.

Args:
    audio_data: Audio data as torch.Tensor [channels, samples] or numpy.ndarray.
    fade_in_samples: Number of samples for fade in ramp (0 = no fade in).
    fade_out_samples: Number of samples for fade out ramp (0 = no fade out).

Returns:
    Audio data with fades applied, in the same format as input.
r   g        g      ?)dtypedevice.N)r   )
isinstancetorchTensorcloneshapecopyminlinspacer   r   npfloat32)	r	   r
   r   	is_tensoraudiototal_samples	actual_inramp
actual_outs	            2/mnt/workspace/ACE-Step-1.5/acestep/audio_utils.py
apply_fader"      s    ! 0A 5:u||4I  "B!B7	>>#sI[[QVQ]Q]^D%*3

?%;d%BE#z	z/";;sC"**ED%*3

?%;d%BE#z	z/"!)9
>>#sJkkRWR^R^_D6;CA[A\<\6]`d6dE#}1223
 L ;;sC2::FD6;CA[A\<\6]`d6dE#}1223L    	target_dbc                 v   [        U [        R                  5      (       a  U R                  5       nSnOU R	                  5       nSnU(       a+  [        R
                  " [        R                  " U5      5      nO*[        R
                  " [        R                  " U5      5      nUS:  a  U $ SUS-  -  nXT-  nX&-  nU$ )z
Apply peak normalization to audio data.

Args:
    audio_data: Audio data as torch.Tensor or numpy.ndarray
    target_db: Target peak level in dB (default: -1.0)
    
Returns:
    Normalized audio data in the same format as input
TFgư>
   g      4@)r   r   r   r   r   maxabsr   )r	   r$   r   r   peak
target_ampgains          r!   normalize_audior,   H   s     *ell++  "	!	 yy5)*vvbffUm$ d{ 	D()J D LELr#   c                      \ rS rSrSrSr1 SkrSrSS1rSS\	4S jjr
  S S
\R                  S\S\S\\	   S\\   SS	4S jjr     S!S\\R                  \R&                  4   S\\	\4   S\S\\	   S\S\\	   S\\   S\	4S jjr   S"S\\	\4   S\\	\4   S\	S\S\\	   S\\   S\	4S jjr      S#S\\\R                     \R                  4   S\\	\4   S\	S\S\\	   S\S\\	   S\\   S\\	   4S jjrSrg	)$
AudioSaverr   z*Audio saving and transcoding utility class128k>   r0   192k256k320k逻  iD  default_formatc                     UR                  5       U l        U R                  S;  a"  [        R                  " SU S35        SU l        gg)zv
Initialize audio saver

Args:
    default_format: Default save format ('flac', 'wav', 'mp3', 'wav32', 'opus', 'aac')
flacwavmp3wav32opusaacUnsupported format z, using 'flac'r8   N)lowerr5   r   warning)selfr5   s     r!   __init__AudioSaver.__init__z   sI     -224&TTNN00@OP"(D Ur#   Naudio_tensoroutput_pathinput_sample_ratemp3_bitratemp3_sample_rater   c                    [        U=(       d    U R                  5      R                  5       R                  5       nX`R                  ;  a  U R                  n [        U=(       d    U R                  5      nXpR                  ;  a  U R                  nUn[        U5      [        U5      :w  a3  [        R                  R                  U[        U5      [        U5      5      n[        R                  " SSS9 n	[        U	R                  5      n
SSS5         [        R                   " [        W
5      U[        U5      SSS9  SS	S
SSS[        U
5      SSS[        [        U5      5      SU[        U5      /n["        R$                  " USSSS9  [&        R(                  " SU SU SU S35          U
R7                  SS9  g! [         a    U R                  n GN\f = f! , (       d  f       N= f! [*         a  n[-        S5      UeSnAf["        R.                   a  n[-        S5      UeSnAf["        R0                   aJ  nUR2                  (       a  UR2                  R5                  SSS9O
[        U5      n[-        SU 35      UeSnAff = f! [         a    [&        R8                  " SU
 35         gf = f!  W
R7                  SS9  f ! [         a    [&        R8                  " SW
 35         f f = f= f) z=Save MP3 with explicit ffmpeg settings and 128k/48k defaults..wavF)suffixdeleteNT	soundfilechannels_firstbackendffmpegz-yz-hide_bannerz	-loglevelerrorz-iz-codec:a
libmp3lamez-arz-b:ax   )checkcapture_outputtimeout[AudioSaver] Saved audio to z (mp3, zHz, )zRffmpeg executable not found. Install ffmpeg or add it to PATH to export MP3 files.z.ffmpeg MP3 export timed out after 120 seconds.utf-8ignore)errorszffmpeg MP3 export failed: )
missing_okz2[AudioSaver] Failed to remove temporary WAV file: )strMP3_DEFAULT_BITRATEstripr?   MP3_ALLOWED_BITRATESintMP3_DEFAULT_SAMPLE_RATE	ExceptionMP3_ALLOWED_SAMPLE_RATES
torchaudio
functionalresampletempfileNamedTemporaryFiler   namesave
subprocessrunr   debugFileNotFoundErrorRuntimeErrorTimeoutExpiredCalledProcessErrorstderrdecodeunlinkr@   )rA   rD   rE   rF   rG   rH   bitratetarget_sample_ratetensor_to_savetemp_wavtemp_wav_pathcmdert   s                 r!   	_save_mp3AudioSaver._save_mp3   s    k=T%=%=>DDFLLN333..G	>!$_%T8T8T!U %B%BB!%!=!=% !S);%<<'22;;L#N_J`befxbyzN((uE /M F	eOOM"&'## $Wc-(Ls3123K C NN3d4MLL7}GL^K__cdkcllmnoe$$$5O  	>!%!=!=	> FE* ! 	|stz{{(( 	XOPVWW,, 	MBC((QXX__WX_>PSTUPVF!;F8DE1L	M  e!STaSbcdee$$$5 e!STaSbcdes   F8 2GBG& (J 8GG
G#&
J 0G<<J HJ 6AI;;J  J, #J)(J),K'.J>=K'>#K$!K'#K$$K'r	   sample_rateformatrO   c           
         U=(       d    U R                   R                  5       nUS;  a2  [        R                  " SU SU R                    35        U R                   n[	        U5      nUS:X  a  SOSU 3nUR
                  R                  5       S;  a  UR                  U5      nO[US:X  a0  UR
                  R                  5       S:X  a  UR                  S5      nO%US	:X  a  UR
                  R                  5       S
:X  a   [        U[        R                  5      (       a  U(       a%  [        R                  " U5      R                  5       n	O[        R                  " U5      R                  5       n	U	R                  5       S:X  a,  U	R                  S   U	R                  S   :  a  U	R                  n	OeUR!                  5       R                  5       n	U(       d@  U	R                  5       S:X  a,  U	R                  S   U	R                  S   :  a  U	R                  n	U	R#                  5       n	 US:X  a  U R%                  U	UUUUS9  ['        U5      $ US;   a"  [(        R*                  " ['        U5      U	USSS9  OUS;   a  US:X  ah   SSKn
U	R/                  SS5      R1                  5       nU
R3                  ['        U5      XSSS9  [        R4                  " SU SU S35        ['        U5      $ [(        R*                  " ['        U5      U	USSS9  O [(        R*                  " ['        U5      U	USS 9  [        R4                  " SU S!U S"U S35        ['        U5      $ ! [6         a&  n[        R8                  " SU S35        Sn SnANSnAff = f! [6         a  nUS:X  a  [        R8                  " S#U 35        e  SSKn
U	R/                  SS5      R1                  5       nUS:X  a  SnSnOUR;                  5       nSnU
R3                  ['        U5      XXS$9  [        R4                  " S%U S!U S"U S35        ['        U5      s SnA$ ! [6         a"  n[        R8                  " S&U S'U 35        e SnAff = fSnAff = f)(a7  
Save audio data to file

Args:
    audio_data: Audio data, torch.Tensor [channels, samples] or numpy.ndarray
    output_path: Output file path (extension can be omitted)
    sample_rate: Sample rate
    format: Audio format ('flac', 'wav', 'mp3', 'wav32', 'opus', 'aac'), defaults to default_format
    channels_first: If True, tensor format is [channels, samples], else [samples, channels]
    mp3_bitrate: Optional MP3 bitrate override (128k/192k/256k/320k)
    mp3_sample_rate: Optional MP3 sample rate override (44100/48000)

Returns:
    Actual saved file path
r7   r>   z, using r;   rJ   .)z.flacrJ   z.mp3z.opusz.aac.m4az.wav32r=   r      r      r:   )rG   rH   )r<   r=   TrQ   rN   )r8   r9   r;   NFLOATWAV)subtyper   rX   z	 (wav32, zHz)zFailed to save wav32: z, falling back to standard wavr9   rM   )rO   z (z, z1[AudioSaver] MP3 export failed without fallback: )r   r   z/[AudioSaver] Fallback soundfile Saved audio to z#[AudioSaver] Failed to save audio: z -> Fallback failed: )r5   r?   r   r@   r   rK   with_suffixr   r   ndarrayr   
from_numpyfloatdimr   Tcpu
contiguousr~   r^   rf   rl   rM   	transposenumpywritero   rd   rR   upper)rA   r	   rE   r   r   rO   rG   rH   extrD   sfaudio_npr}   	sf_formatr   inner_es                   r!   
save_audioAudioSaver.save_audio   sE   2 /D//668GGNN0ATAT@UVW((F ;' ')f6(|##%-__%11#6Kw;#5#5#;#;#=#I&226:[u_!3!3!9!9!;v!E j"**--$//
;AAC  %//
;AAC##%*|/A/A!/D|GYGYZ[G\/\$0NN\ &>>+113L!l&6&6&8A&=%%a(<+=+=a+@@#/>>L $..0N	  +$3   ;''?*$ #'$ 33 W$'. $0#9#9!Q#?#E#E#G [!18RYbgh'CK=PYZeYffi%jk";// $ #'' $ #'	 LL7}BvhbQ\P]]`ab{##- % ''=aS@^%_`!&'0  	PQRPSTU&'11!Q7==? W$ %I%G &I"G[)8dN{m[]^d]eeghsgttwxy;'' B1#EZ[bZcde'	so   %#N 	4N >A&M %A,N 
NM=8N =NN 
Q/ Q*0BP;5Q/;
Q'Q""Q''Q**Q/
input_pathoutput_formatremove_inputc           
      D   [        U5      n[        U5      nUR                  5       (       d  [        SU 35      e[        R                  " [        U5      5      u  pxU R                  UUUUSUUS9nU(       a)  UR                  5         [        R                  " SU 35        U$ )a  
Convert audio format

Args:
    input_path: Input audio file path
    output_path: Output audio file path
    output_format: Target format ('flac', 'wav', 'mp3', 'wav32', 'opus', 'aac')
    remove_input: Whether to delete input file
    mp3_bitrate: Optional MP3 bitrate override (128k/192k/256k/320k)
    mp3_sample_rate: Optional MP3 sample rate override (44100/48000)

Returns:
    Output file path
zInput file not found: Tr   r   rO   rG   rH   z![AudioSaver] Removed input file: )
r   existsrp   rf   loadr^   r   rv   r   ro   )	rA   r   rE   r   r   rG   rH   rD   r   s	            r!   convert_audioAudioSaver.convert_audioT  s    . *%
;'  ""#&<ZL$IJJ %/OOC
O$D! oo# #+ & 
 LL<ZLIJr#   audio_batch
output_dirfile_prefixc	                    [        U5      nUR                  SSS9  [        U[        R                  5      (       a?  UR                  5       S:X  a+  [        UR                  S   5       V	s/ s H  oU	   PM	     n
n	O[        U[        5      (       a  Un
OU/n
/ n[        U
5       H5  u  pX# SU	S 3-  nU R                  UUUUUUUS9nUR                  U5        M7     U$ s  sn	f )a  
Save audio batch

Args:
    audio_batch: Audio batch, List[tensor] or tensor [batch, channels, samples]
    output_dir: Output directory
    file_prefix: File prefix
    sample_rate: Sample rate
    format: Audio format
    channels_first: Tensor format flag
    mp3_bitrate: Optional MP3 bitrate override (128k/192k/256k/320k)
    mp3_sample_rate: Optional MP3 sample rate override (44100/48000)

Returns:
    List of saved file paths
T)parentsexist_ok   r   _04dr   )r   mkdirr   r   r   r   ranger   list	enumerater   append)rA   r   r   r   r   r   rO   rG   rH   i
audio_listsaved_pathsr   rE   
saved_paths                  r!   
save_batchAudioSaver.save_batch  s    6 *%
5 k5<<00[__5F!5K278I8I!8L2MN2MQa.2MJNJT**$J%J!*-HA$a#w'??K'-' / ) J z* . ) Os   )Cr5   )r8   )NNr4   NTNN)FNN)r   r4   NTNN)__name__
__module____qualname____firstlineno____doc__r_   ra   rc   re   r^   rB   r   r   r   rb   r   r~   r   r   r   boolr   r   r   r   __static_attributes__ r#   r!   r.   r.   r   s2   4 ;# %u~
)s 
)" &*)-8ell8e 8e 	8e
 c]8e "#8e 
8e| ! $#%))-R%,,

23R 39%R 	R
 R R c]R "#R 
Rr #%))-0#t)$0 39%0 	0
 0 c]0 "#0 
0l #  $#%))-54-u||;<5 #t)$5 	5
 5 5 5 c]5 "#5 
c5 5r#   r.   c                    [        U SS5      (       d  g[        U SS5      (       d  g[        U SS5      nUb  UR                  (       d  g[        R                  " 5       nSn[	        UR                  R                  5       5       GH6  nUR                  U   nUR                  S5      nU(       d  M-  / n[        R                  R                  U5      (       a  UR                  U5        O[        R                  R                  U5      (       a_  S HY  n[        R                  R                  Xh5      n	[        R                  R                  U	5      (       d  MH  UR                  U	5        M[     U HH  n	 [        U	S	5       n
 U
R                  S5      nU(       d  OUR                  U5        M,  SSS5        S
nMJ     GM9     U(       a  UR#                  5       $ S$ ! , (       d  f       N2= f! [          a     M  f = f)a  Compute an MD5 hash identifying the currently loaded LoRA adapter weights.

Iterates over the handler's LoRA service registry to find adapter weight
file paths, then hashes each file to produce a combined fingerprint.

Args:
    dit_handler: DiT handler instance with LoRA state attributes.

Returns:
    Hex digest string uniquely identifying the loaded LoRA weights,
    or empty string if no LoRA is active.
lora_loadedF use_lora_lora_serviceNpath)zadapter_model.safetensorszadapter_model.binzlokr_weights.safetensorsrbTi   )getattrregistryhashlibsha256sortedkeysgetosr   isfiler   isdirjoinopenreadupdateOSError	hexdigest)dit_handlerlora_servicehash_obj	found_anyadapter_namemeta	lora_path
candidatesfnamefpathfchunks               r!   get_lora_weights_hashr     s    ;u55;
E22;>L<#8#8~~HI|4499;<$$\2HHV$	 
77>>)$$i(WW]]9%%
 Y677>>%((%%e,  E	%&! !w$! .	  ' !	  ) =@ $-84"4 '&  s*   1G(=.G+
G(
G%!G((
G65G6c                 T   U c  g [        U [        5      (       a  [        R                  R	                  U 5      (       aH  [        U S5       n[        R                  " UR                  5       5      R                  5       sSSS5        $ [        R                  " U R                  S5      5      R                  5       $ [        U S5      (       aF  [        R                  " [        U R                  5      R                  S5      5      R                  5       $ [        R                  " [        U 5      R                  S5      5      R                  5       $ ! , (       d  f       N= f! [         a?    [        R                  " [        U 5      R                  S5      5      R                  5       s $ f = f)z
Get hash identifier for an audio file.

Args:
    audio_file: Path to audio file (str) or file-like object

Returns:
    Hash string or empty string
Nr   r   rZ   rk   )r   r^   r   r   r   r   r   r   r   r   encodehasattrrk   rd   )
audio_filer   s     r!   get_audio_file_hashr     s'    
Kj#&&ww~~j))*d+q">>!&&(3==? ,+>>*"3"3G"<=GGIIZ((>>#joo"6"="=g"FGQQSS~~c*o44W=>HHJJ ,+  K~~c*o44W=>HHJJKs>   AE 2E=	E 2E :AE ;E 
EE AF'&F'c                     [         R                  " U SSS9n[        R                  " UR	                  S5      5      nUR                  5       nUSS  SUSS  SUSS	  SUS	S
  SUS
S  3	nU$ )z
Generate deterministic UUID from generation parameters.
Same parameters will always generate the same UUID.

Args:
    params_dict: Dictionary of parameters

Returns:
    UUID string
TF)	sort_keysensure_asciirZ   r      -             )jsondumpsr   r   r   r   )params_dictparams_jsonr   hash_hexuuid_strs        r!   generate_uuid_from_paramsr     s     **[DuMK~~k009:H!!#H1Q-(1R.!18Br?2C1XbQS_DUUVW_`bceWfVghHOr#   seedc                 b   [        U [        R                  5      (       a  U R                  5       R	                  5       nOU n[
        R                  " UR                  5       5      R                  5       nUb9  U SU 3n[
        R                  " UR                  5       5      R                  5       $ U$ )z
Generate UUID from audio data (for caching/deduplication)

Args:
    audio_data: Audio data
    seed: Optional seed value

Returns:
    UUID string
r   )
r   r   r   r   r   r   r   tobytesr   r   )r	   r   r   	data_hashcombineds        r!   generate_uuid_from_audio_datar  '  s     *ell++>>#))+ x//12<<>I[$(~~hoo/0::<<r#   r8   r   rE   r   r   rO   rG   rH   c           	      8    [         R                  U UUUUUU5      $ )a  
Convenience function: save audio (using default configuration)

Args:
    audio_data: Audio data
    output_path: Output path
    sample_rate: Sample rate
    format: Format (default flac)
    channels_first: Tensor format flag
    mp3_bitrate: Optional MP3 bitrate override (128k/192k/256k/320k)
    mp3_sample_rate: Optional MP3 sample rate override (44100/48000)

Returns:
    Saved file path
)_default_saverr   )r	   rE   r   r   rO   rG   rH   s          r!   r   r   I  s,    0 $$ r#   )r   r   )g      )Nr   )#r   ior   r   rm   r   ri   pathlibr   typingr   r   r   r   r   r   r   rf   logurur   r   r   rb   r"   r   r,   r.   r^   r   r   r   r  r  r   r   r   r#   r!   <module>r	     s   
  	     / /    
 ,ellBJJ./,, , 5<<#$	,^&ellBJJ&> ? &E &]bchcocoqsq{q{c{]| &TI IX
95# 95xKs K4c * ellBJJ./
3- 	> 62  !%%) ellBJJ./ sDy!    SM	 
   #  c]  	 r#   