§
    ’Þjö  ã                   ót   — d dl mZ ddlmZ d dlZd dlmZ d dlmZ	 ddl
mZ dZd	„ Zd
„ Z G d„ d¦  «        ZdS )é    )Ú
merge_helpé   )ÚHEARTMULA_HELPN)Úmps_device_or)Úfiles_locator©ÚHEARTMULA_LYRIC_PROMPTÚ3Bc            	      ó¢   — i dd“dd“dd“dd“dd	“d
d“dd“dd“dd“dd“ddg“dddddœ“dd“dt           “dd“dddddd d!d"œ“d#d“d$d%d&d'd(ddd)œ¥S )*NÚgroupÚmusicÚ
audio_onlyTÚimage_outputsFÚsliding_windowÚguidance_max_phasesr   Úno_negative_promptÚinference_stepsÚtemperatureÚimage_prompt_types_allowedÚ Úsupports_early_stopÚprofiles_dirÚheartmula_oss_3bÚ
alt_promptzKeywords / Tagsúpiano,happy,weddingé   )ÚlabelÚplaceholderÚlinesÚ
lm_decoderÚ!text_prompt_enhancer_instructionsÚprompt_enhancer_button_labelzCompose LyricsÚduration_sliderzMaximum Song Duration (seconds)zMaximum Song Durationé   éð   gš™™™™™¹?éx   )r   ÚnameÚminÚmaxÚ	incrementÚdefaultÚtop_k_sliderç      ø?é2   éÀÔ ç      ô?é
   )Úheartmula_cfg_scaleÚheartmula_topkÚheartmula_max_audio_length_msÚheartmula_codec_guidance_scaleÚheartmula_codec_stepsÚheartmula_codec_versionÚcompiler   © ó    úDC:\Users\johnn\OneDrive\Documents\AI\models\TTS\heartmula_handler.pyÚ_get_heartmula_model_defr<      s,  € ð$Ø�ð$à�dð$ð 	˜ð$ð 	˜%ð	$ð
 	˜qð$ð 	˜dð$ð 	˜5ð$ð 	�tð$ð 	% bð$ð 	˜tð$ð 	Ð+Ð,ð$ð 	Ø&Ø0Øð
ð 
ð$ð" 	�dð#$ð$ 	,Õ-Cð%$ð& 	'Ð(8ð'$ð( 	Ø6Ø+ØØØØð
ð 
ð)$ð8 	˜ð9$ð:  #ØØ)/Ø*.Ø!#Ø#%ØðG$ð $ð $ð $r:   c                 óv   — | pi                       dd¦  «        }|rd|› �nd}d}ddd|› d�d	|› d
�g}|dg|gdœgS )Nr7   r   Ú_zDeepBeepMeep/TTSzgen_config.jsonztokenizer.jsonÚcodec_configz.jsonÚHeartMula_codecz.safetensorsÚ	HeartMula)ÚrepoIdÚsourceFolderListÚfileList)Úget)Ú	model_defÚcodec_versionÚcodec_suffixÚrepo_idÚ	gen_filess        r;   Ú_get_heartmula_download_defrK   7   s…   € Ø�_ "×)Ò)Ð*CÀRÑHÔH€MØ*7Ð?Ð&�}Ð&Ð&Ð&¸R€LØ €GàØØ*�|Ð*Ð*Ð*Ø4˜,Ð4Ð4Ð4ð	€Ið Ø!, Ø"˜ð	
ð 	
ðð r:   c                   ó  — e Zd Zed„ ¦   «         Zed„ ¦   «         Zed„ ¦   «         Zed„ ¦   «         Zed„ ¦   «         Zed„ ¦   «         Z	edd„¦   «         Z
e	 	 	 	 	 	 	 	 	 	 dd„¦   «         Zed„ ¦   «         Zed„ ¦   «         Zed„ ¦   «         ZdS )Úfamily_handlerc                  ó   — dgS )Nr   r9   r9   r:   r;   Úquery_supported_typesz$family_handler.query_supported_typesK   s   € à"Ð#Ð#r:   c                  ó
   — i i fS ©Nr9   r9   r:   r;   Úquery_family_mapsz family_handler.query_family_mapsO   s   € à�2ˆvˆr:   c                  ó   — dS )NÚttsr9   r9   r:   r;   Úquery_model_familyz!family_handler.query_model_familyS   s   € àˆur:   c                  ó   — dddœS )N)i“  ÚMusic)i˜  ÚTTS)r   rT   r9   r9   r:   r;   Úquery_family_infosz!family_handler.query_family_infosW   s   € à(°Ð?Ð?Ð?r:   c                 ó   — dS )NÚ
heart_mular9   )Úbase_model_types    r;   Úget_lora_dirzfamily_handler.get_lora_dir[   s   € àˆ|r:   c                 óP   — t          i t          ¦   «         ¥ddi¥t          ¦  «        S )NÚdevice_explicitT)r   r<   r   )r\   rF   s     r;   Úquery_model_defzfamily_handler.query_model_def_   s+   € åÐQÕ5Ñ7Ô7ÐQÐ9JÈDÐQÐQÕSaÑbÔbÐbr:   Nc                 ó$   — t          |pi ¦  «        S rQ   )rK   )ÚcomputeListr\   rF   s      r;   Úquery_model_filesz family_handler.query_model_filesc   s   € å*¨9¨?¸Ñ;Ô;Ð;r:   Fr   c                 ól  — ddl m} t          j        ¦   «         }d }t	          | t
          t          f¦  «        rt          | ¦  «        dk    r| d         }n| }d }|rt          j        |d¬¦  «        }|€|} ||t          t          j        d¦  «        ¦  «        t          |||                     dd¦  «        |                     d	d
¦  «        |                     dd¦  «        |                     dd¦  «        |                     dd¦  «        |                     dd¦  «        |j        |j        ¬¦  «        }d|j        j        d         j        _        d|j        j        j        _        |j        |j        j        d         |j        dœ}|dgdgdœ|j        rddgnddœ}t-          |¦  «        dv rddi|d<   ||fS )Nr   )ÚHeartMuLaPipeliner   F)Úerror_if_noneÚcpur2   r-   r3   r.   r4   r/   r6   r1   r5   r0   r7   r   )Ú	ckpt_rootÚdeviceÚversionÚ	VAE_dtypeÚheartmula_weights_pathÚ	cfg_scaleÚtopkÚmax_audio_length_msÚcodec_stepsÚcodec_guidance_scalerG   Úcuda_graphsÚflash_attention)ÚtransformerÚtransformer2Úcodecru   rt   )rt   ru   )ÚpipeÚcoTenantsMapÚ
cudaGraphs)r   r   é   é   éÈ   Úbudgets)ÚHeartMula.pipelinere   ÚflÚget_download_locationÚ
isinstanceÚlistÚtupleÚlenÚlocate_filer   Útorchri   ÚHEARTMULA_VERSIONrE   ÚgraphsÚflashÚmulaÚdecoderÚlayersÚ_compile_meÚbackbonerv   Úint)Úmodel_filenameÚ
model_typer\   rF   ÚquantizeTransformerÚtext_encoder_quantizationÚdtyperk   Úmixed_precision_transformerÚsave_quantizedÚsubmodel_no_listÚtext_encoder_filenameÚprofileÚlm_featuresÚkwargsre   rh   Úweights_candidaterl   Úpipelinerw   s                        r;   Ú
load_modelzfamily_handler.load_modelg   s  € ð$ 	:Ð9Ð9Ð9Ð9Ð9åÔ,Ñ.Ô.ˆ	Ø ÐÝ�n¥t­U mÑ4Ô4ð 	/Ý�>Ñ"Ô" QÒ&Ð&Ø$2°1Ô$5Ð!øà .ÐØ!%ÐØð 	;Ý%'¤^Ø!°ð&ñ &ô &Ð"ð &Ð-Ø):Ð&Ø$Ð$ØÝ ¥¤¨eÑ!4Ô!4Ñ5Ô5Ý%ØØ#9Ø—m’mÐ$9¸3Ñ?Ô?Ø—’Ð/°Ñ4Ô4Ø )§¢Ð.MÈvÑ VÔ VØ!ŸšÐ&=¸rÑBÔBØ!*§¢Ð/OÐQUÑ!VÔ!VØ#Ÿ-š-Ð(AÀ2ÑFÔFØ#Ô*Ø'Ô-ð
ñ 
ô 
ˆð  7<ˆŒÔ˜aÔ Ô'Ô3Ø49ˆŒÔÔ%Ô1à#œ=Ø$œMÔ1°!Ô4Ø”^ð
ð 
ˆð à .Ð/Ø!. ðð ð
 >IÔ=OÐZ˜=¨.Ð9Ð9ÐUZð
ð 
ˆõ ˆw‰<Œ<˜<Ð'Ð'Ø-¨sÐ3ˆD�‰Oà˜ˆ~Ðr:   c                 óô   — d|vrd|d<   ddi}|                      ¦   «         D ]\  }}|                     ||¦  «         Œ|dk     r4|                     dd¦  «        |d<   |                     dd	¦  «        |d
<   d S d S )Nr   r   Úaudio_prompt_typeg…ëQ¸…@r2   r-   Úguidance_scaler3   r.   Útop_k)ÚitemsÚ
setdefaultrE   )r\   Úsettings_versionrF   Úui_defaultsÚdefaultsÚkeyÚvalues          r;   Úfix_settingszfamily_handler.fix_settings¯   sª   € à˜{Ð*Ð*Ø(*ˆK˜Ñ%ð   ð
ˆð #Ÿ.š.Ñ*Ô*ð 	/ð 	/‰JˆC�Ø×"Ò" 3¨Ñ.Ô.Ð.Ð.à˜dÒ"Ð"Ø,5¯MªMÐ:OÐQTÑ,UÔ,UˆKÐ(Ñ)Ø#,§=¢=Ð1AÀ2Ñ#FÔ#FˆK˜Ñ Ð Ð ð #Ð"r:   c                 óì   — |                      di ¦  «        }|                     ddd|                      dd¦  «        dddd|                      d	d
¦  «        |                      dd¦  «        ddœ¦  «         d S )Nr#   r   r   r   r+   r&   r   g      ð?r2   r-   r3   r.   ÚFG)r    r   Úrepeat_generationÚduration_secondsÚvideo_lengthÚnum_inference_stepsÚnegative_promptr   r¡   r¢   Úmulti_prompts_gen_type)rE   Úupdate)r\   rF   r¦   Úduration_defs       r;   Úupdate_default_settingsz&family_handler.update_default_settings¾   s‘   € à —}’}Ð%6¸Ñ;Ô;ˆØ×Òà%'Ø3Ø%&Ø$0×$4Ò$4°YÀÑ$DÔ$DØ !Ø'(Ø#%Ø"Ø"+§-¢-Ð0EÀsÑ"KÔ"KØ"ŸšÐ'7¸Ñ<Ô<Ø*.ðð ñ	
ô 	
ð 	
ð 	
ð 	
r:   c                 óö   — |                      dd¦  «        }|�2t          t          |¦  «                             ¦   «         ¦  «        dk    rdS |                      d¦  «        €|                      d¦  «        �dS d S )Nr   r   r   z.Keywords prompt cannot be empty for HeartMuLa.Úaudio_guideÚaudio_guide2z/HeartMuLa does not support reference audio yet.)rE   r„   ÚstrÚstrip)r\   rF   ÚinputsÚ
one_promptr   s        r;   Úvalidate_generative_promptz)family_handler.validate_generative_promptÑ   ss   € à—Z’Z ¨bÑ1Ô1ˆ
ØÐ¥¥S¨¡_¤_×%:Ò%:Ñ%<Ô%<Ñ!=Ô!=ÀÒ!BÐ!BØCÐCØ�:Š:�mÑ$Ô$Ð0°F·J²J¸~Ñ4NÔ4NÐ4ZØDÐDØˆtr:   rQ   )
FNNNFFNNr   N)Ú__name__Ú
__module__Ú__qualname__ÚstaticmethodrO   rR   rU   rY   r]   r`   rc   rž   rª   rµ   r½   r9   r:   r;   rM   rM   J   ss  € € € € € Øð$ð $ñ „\ð$ð ðð ñ „\ðð ðð ñ „\ðð ð@ð @ñ „\ð@ð ðð ñ „\ðð ðcð cñ „\ðcð ð<ð <ð <ñ „\ð<ð ð "Ø"&ØØØ$)ØØØ"ØØðEð Eð Eñ „\ðEðN ðGð Gñ „\ðGð ð
ð 
ñ „\ð
ð$ ðð ñ „\ðð ð r:   rM   )Úmodels.help_sharedr   Úhelpr   r†   Ú
shared.mpsr   Úshared.utilsr   r   Úprompt_enhancersr	   r‡   r<   rK   rM   r9   r:   r;   ú<module>rÇ      sÉ   ðà )Ð )Ð )Ð )Ð )Ð )Ø  Ð  Ð  Ð  Ð  Ð  à €€€à $Ð $Ð $Ð $Ð $Ð $Ø ,Ð ,Ð ,Ð ,Ð ,Ð ,à 4Ð 4Ð 4Ð 4Ð 4Ð 4àÐ ð%ð %ð %ðPð ð ð&Nð Nð Nð Nð Nñ Nô Nð Nð Nð Nr:   