soundscaps

app.py

@@ -21,6 +21,11 @@ import nltk
 from num2words import num2words
 from num2word_greek.numbers2words import convert_numbers
 from audionar import VitsModel, VitsTokenizer
+from audiocraft import AudioGen
+
+
+
+audiogen = AudioGen().eval().to('cpu')
 
 nltk.download('punkt', download_dir='./')
 nltk.download('punkt_tab', download_dir='./')
@@ -443,97 +448,118 @@ language_names = ['Ancient greek',
 
 
 def audionar_tts(text=None,
-                 lang='romanian'):
+                 lang='romanian',
+                 soundscape='',
+                 cache_lim=24):
 
     # https://huggingface.co/dkounadis/artificial-styletts2/blob/main/msinference.py
 
-    lang = lang.lower()
-
-    # https://huggingface.co/spaces/mms-meta/MMS
-
-    if 'hun' in lang:
-
-        lang_code = 'hun'
 
-    elif any([i in lang for i in ['ser', 'bosn', 'herzegov', 'montenegr', 'macedon']]):
+    lang_map = {
+            'ancient greek': 'grc',
+            'english': 'eng',
+            'deutsch': 'deu',
+            'french': 'fra',
+            'hungarian': 'hun',
+            'romanian': 'ron',
+            'serbian (approx.)': 'rmc-script_latin',
+        }
+    lang_code = lang_map.get(lang.lower(), lang.lower().split()[0].strip())
 
-        # romani carpathian (has also Vlax) - cooler voice
-        lang_code = 'rmc-script_latin'
+    global cached_lang_code, cached_net_g, cached_tokenizer
 
-    elif 'rom' in lang:
+    if 'cached_lang_code' not in globals() or cached_lang_code != lang_code:
+        cached_lang_code = lang_code
+        cached_net_g = VitsModel.from_pretrained(f'facebook/mms-tts-{lang_code}').eval()
+        cached_tokenizer = VitsTokenizer.from_pretrained(f'facebook/mms-tts-{lang_code}')
 
-        lang_code = 'ron'
+    net_g = cached_net_g
+    tokenizer = cached_tokenizer
 
-    elif 'ger' in lang or 'deu' in lang or 'allem' in lang:
+    total_audio = []
+    
+    final_audio = None
+    speech_audio = None
 
-        lang_code = 'deu'
 
-    elif 'french' in lang:
+    if text and text.strip():
+        
 
-        lang_code = 'fra'
+        text = only_greek_or_only_latin(text, lang=lang_code)
+        text = transliterate_number(text, lang=lang_code)
+        text = fix_vocals(text, lang=lang_code)
 
-    elif 'eng' in lang:
 
-        lang_code = 'eng'
+        sentences = textwrap.wrap(text, width=439)
 
-    elif 'ancient greek' in lang:
+        total_audio_parts = []
+        for sentence in sentences:
+            inputs = cached_tokenizer(sentence, return_tensors="pt")
+            with torch.no_grad():
+                audio_part = cached_net_g(
+                    input_ids=inputs.input_ids.to(device),
+                    attention_mask=inputs.attention_mask.to(device),
+                    lang_code=lang_code,
+                )[0, :]
+            total_audio_parts.append(audio_part)
 
-        lang_code = 'grc'
+        speech_audio = torch.cat(total_audio_parts).cpu().numpy()
 
-    else:
+    # AudioGen
+    if soundscape and soundscape.strip():
 
-        lang_code = lang.split()[0].strip()   # latin & future option
         
-    # LATIN / GRC / CYRILLIC    
+        speech_duration_secs = len(speech_audio) / 16000 if speech_audio is not None else 0
+        target_duration = max(speech_duration_secs + 0.74, 2.0)
 
-    text = only_greek_or_only_latin(text, lang=lang_code)  # assure gr-chars if lang=='grc' / latin if lang!='grc'
 
-    # NUMERALS (^ in math expression found & substituted here before arriving to fix_vocals)
+        background_audio = audiogen.generate(
+            soundscape,
+            duration=target_duration,
+            cache_lim=max(4, int(cache_lim))  # at least allow 10 A/R stEps
+         ).numpy()
 
-    text = transliterate_number(text, lang=lang_code)
+        if speech_audio is not None:
 
-    # PRONOUNC.
+            len_speech = len(speech_audio)
+            len_background = len(background_audio)
+            
+            if len_background > len_speech:
+                padding = np.zeros(len_background - len_speech,
+                                   dtype=np.float32)
+                speech_audio = np.concatenate([speech_audio, padding])
+            elif len_speech > len_background:
+                padding = np.zeros(len_speech - len_background,
+                                   dtype=np.float32)
+                background_audio = np.concatenate([background_audio, padding])
+
+
+            speech_audio_stereo = speech_audio[None, :]
+            background_audio_stereo = background_audio[None, :]
+
+
+            final_audio = np.concatenate([
+                0.49 * speech_audio_stereo + 0.51 * background_audio_stereo,
+                0.51 * background_audio_stereo + 0.49 * speech_audio_stereo
+            ], 0)
+        else:
+            final_audio = background_audio
+
+    # If no soundscape, use the speech audio as is.
+    elif speech_audio is not None:
+        final_audio = speech_audio
     
-    text = fix_vocals(text, lang=lang_code)
+    # If both inputs are empty, create a 2s silent audio file.
+    if final_audio is None:
+        final_audio = np.zeros(16000 * 2, dtype=np.float32)
 
-    # VITS
-
-    global cached_lang_code, cached_net_g, cached_tokenizer
-
-    if 'cached_lang_code' not in globals() or cached_lang_code != lang_code:
-        cached_lang_code = lang_code
-        cached_net_g = VitsModel.from_pretrained(f'facebook/mms-tts-{lang_code}').eval().to(device)
-        cached_tokenizer = VitsTokenizer.from_pretrained(f'facebook/mms-tts-{lang_code}')
-
-    net_g = cached_net_g
-    tokenizer = cached_tokenizer
-
-    total_audio = []
+    wavfile = '_vits_.wav'
+    audiofile.write(wavfile, final_audio, 16000)
     
-    if not isinstance(text, list):
-        text = textwrap.wrap(text, width=439)
-
-    for _t in text:
-        inputs = tokenizer(_t, return_tensors="pt")
-        with torch.no_grad():
-            x = net_g(input_ids=inputs.input_ids.to(device),
-                      attention_mask=inputs.attention_mask.to(device),
-                      lang_code=lang_code,
-                      )[0, :]
-            total_audio.append(x)
+    return wavfile
 
-        print(f'\n\n_______________________________ {_t} {x.shape=}')
 
-    x = torch.cat(total_audio).cpu().numpy()
-
-    tmp_file = f'_speech.wav'
-
-    audiofile.write(tmp_file, x, 16000)
-
-    return tmp_file
-
-
-# --
+# -- EXPRESSIO
 
 
 device = 0 if torch.cuda.is_available() else "cpu"
@@ -838,7 +864,334 @@ def plot_expression(arousal, dominance, valence):
     # plt.show()
 
 # TTS
-VOICES = [f'wav/{vox}' for vox in os.listdir('wav')]
+# VOICES = [f'wav/{vox}' for vox in os.listdir('wav')]
+# add unidecode (to parse non-roman characters for the StyleTTS2 
+# # for the VITS it should better skip the unknown letters - dont use unidecode())
+# at generation fill the state of "last tts"
+#    at record fill the state of "last record" and place in list of voice/langs for TTS
+VOICES = ['jv_ID_google-gmu_04982.wav', 
+        'it_IT_mls_1595.wav', 
+        'en_US_vctk_p303.wav', 
+        'en_US_vctk_p306.wav', 
+        'it_IT_mls_8842.wav', 
+        'en_US_cmu_arctic_ksp.wav', 
+        'jv_ID_google-gmu_05970.wav', 
+        'en_US_vctk_p318.wav', 
+        'ha_NE_openbible.wav', 
+        'ne_NP_ne-google_0883.wav', 
+        'en_US_vctk_p280.wav', 
+        'bn_multi_1010.wav', 
+        'en_US_vctk_p259.wav', 
+        'it_IT_mls_844.wav', 
+        'en_US_vctk_p269.wav', 
+        'en_US_vctk_p285.wav', 
+        'de_DE_m-ailabs_angela_merkel.wav', 
+        'en_US_vctk_p316.wav', 
+        'en_US_vctk_p362.wav', 
+        'jv_ID_google-gmu_06207.wav', 
+        'tn_ZA_google-nwu_9061.wav', 
+        'fr_FR_tom.wav', 
+        'en_US_vctk_p233.wav', 
+        'it_IT_mls_4975.wav', 
+        'en_US_vctk_p236.wav', 
+        'bn_multi_01232.wav', 
+        'bn_multi_5958.wav', 
+        'it_IT_mls_9185.wav', 
+        'en_US_vctk_p248.wav', 
+        'en_US_vctk_p287.wav', 
+        'it_IT_mls_9772.wav', 
+        'te_IN_cmu-indic_sk.wav', 
+        'tn_ZA_google-nwu_8333.wav', 
+        'en_US_vctk_p260.wav', 
+        'en_US_vctk_p247.wav', 
+        'en_US_vctk_p329.wav', 
+        'en_US_cmu_arctic_fem.wav', 
+        'en_US_cmu_arctic_rms.wav', 
+        'en_US_vctk_p308.wav', 
+        'jv_ID_google-gmu_08736.wav', 
+        'en_US_vctk_p245.wav', 
+        'fr_FR_m-ailabs_nadine_eckert_boulet.wav', 
+        'jv_ID_google-gmu_03314.wav', 
+        'en_US_vctk_p239.wav', 
+        'jv_ID_google-gmu_05540.wav', 
+        'it_IT_mls_7440.wav', 
+        'en_US_vctk_p310.wav', 
+        'en_US_vctk_p237.wav', 
+        'en_US_hifi-tts_92.wav', 
+        'en_US_cmu_arctic_aew.wav', 
+        'ne_NP_ne-google_2099.wav', 
+        'en_US_vctk_p226.wav', 
+        'af_ZA_google-nwu_1919.wav', 
+        'jv_ID_google-gmu_03727.wav', 
+        'en_US_vctk_p317.wav', 
+        'tn_ZA_google-nwu_0378.wav', 
+        'nl_pmk.wav', 
+        'en_US_vctk_p286.wav', 
+        'tn_ZA_google-nwu_3342.wav', 
+        # 'en_US_vctk_p343.wav', 
+        'de_DE_m-ailabs_ramona_deininger.wav', 
+        'jv_ID_google-gmu_03424.wav', 
+        'en_US_vctk_p341.wav', 
+        'jv_ID_google-gmu_03187.wav', 
+        'ne_NP_ne-google_3960.wav', 
+        'jv_ID_google-gmu_06080.wav', 
+        'ne_NP_ne-google_3997.wav', 
+        # 'en_US_vctk_p267.wav', 
+        'en_US_vctk_p240.wav', 
+        'ne_NP_ne-google_5687.wav', 
+        'ne_NP_ne-google_9407.wav', 
+        'jv_ID_google-gmu_05667.wav', 
+        'jv_ID_google-gmu_01519.wav', 
+        'ne_NP_ne-google_7957.wav', 
+        'it_IT_mls_4705.wav', 
+        'ne_NP_ne-google_6329.wav', 
+        'it_IT_mls_1725.wav', 
+        'tn_ZA_google-nwu_8914.wav', 
+        'en_US_ljspeech.wav', 
+        'tn_ZA_google-nwu_4850.wav', 
+        'en_US_vctk_p238.wav', 
+        'en_US_vctk_p302.wav', 
+        'jv_ID_google-gmu_08178.wav', 
+        'en_US_vctk_p313.wav', 
+        'af_ZA_google-nwu_2418.wav', 
+        'bn_multi_00737.wav', 
+        'en_US_vctk_p275.wav', # y
+        'af_ZA_google-nwu_0184.wav', 
+        'jv_ID_google-gmu_07638.wav', 
+        'ne_NP_ne-google_6587.wav', 
+        'ne_NP_ne-google_0258.wav', 
+        'en_US_vctk_p232.wav', 
+        'en_US_vctk_p336.wav', 
+        'jv_ID_google-gmu_09039.wav', 
+        'en_US_vctk_p312.wav', 
+        'af_ZA_google-nwu_8148.wav', 
+        'en_US_vctk_p326.wav', 
+        'en_US_vctk_p264.wav', 
+        'en_US_vctk_p295.wav', 
+        # 'en_US_vctk_p298.wav', 
+        'es_ES_m-ailabs_victor_villarraza.wav', 
+        'pl_PL_m-ailabs_nina_brown.wav', 
+        'tn_ZA_google-nwu_9365.wav', 
+        'en_US_vctk_p294.wav', 
+        'jv_ID_google-gmu_00658.wav', 
+        'jv_ID_google-gmu_08305.wav', 
+        'en_US_vctk_p330.wav', 
+        'gu_IN_cmu-indic_cmu_indic_guj_dp.wav', 
+        'jv_ID_google-gmu_05219.wav', 
+        'en_US_vctk_p284.wav', 
+        'de_DE_m-ailabs_eva_k.wav', 
+        # 'bn_multi_00779.wav', 
+        'en_UK_apope.wav', 
+        'en_US_vctk_p345.wav', 
+        'it_IT_mls_6744.wav', 
+        'en_US_vctk_p347.wav', 
+        'en_US_m-ailabs_mary_ann.wav', 
+        'en_US_m-ailabs_elliot_miller.wav', 
+        'en_US_vctk_p279.wav', 
+        'ru_RU_multi_nikolaev.wav', 
+        'bn_multi_4811.wav', 
+        'tn_ZA_google-nwu_7693.wav', 
+        'bn_multi_01701.wav', 
+        'en_US_vctk_p262.wav', 
+        # 'en_US_vctk_p266.wav', 
+        'en_US_vctk_p243.wav', 
+        'en_US_vctk_p297.wav', 
+        'en_US_vctk_p278.wav', 
+        'jv_ID_google-gmu_02059.wav', 
+        'en_US_vctk_p231.wav', 
+        'te_IN_cmu-indic_kpn.wav', 
+        'en_US_vctk_p250.wav', 
+        'it_IT_mls_4974.wav', 
+        'en_US_cmu_arctic_awbrms.wav', 
+        # 'en_US_vctk_p263.wav', 
+        'nl_femal.wav', 
+        'tn_ZA_google-nwu_6116.wav', 
+        'jv_ID_google-gmu_06383.wav', 
+        'en_US_vctk_p225.wav', 
+        'en_US_vctk_p228.wav', 
+        'it_IT_mls_277.wav', 
+        'tn_ZA_google-nwu_7866.wav', 
+        'en_US_vctk_p300.wav', 
+        'ne_NP_ne-google_0649.wav', 
+        'es_ES_carlfm.wav', 
+        'jv_ID_google-gmu_06510.wav', 
+        'de_DE_m-ailabs_rebecca_braunert_plunkett.wav', 
+        'en_US_vctk_p340.wav', 
+        'en_US_cmu_arctic_gka.wav', 
+        'ne_NP_ne-google_2027.wav', 
+        'jv_ID_google-gmu_09724.wav', 
+        'en_US_vctk_p361.wav', 
+        'ne_NP_ne-google_6834.wav', 
+        'jv_ID_google-gmu_02326.wav', 
+        'fr_FR_m-ailabs_zeckou.wav', 
+        'tn_ZA_google-nwu_1932.wav', 
+        # 'female-20-happy.wav', 
+        'tn_ZA_google-nwu_1483.wav', 
+        'de_DE_thorsten-emotion_amused.wav', 
+        'ru_RU_multi_minaev.wav', 
+        'sw_lanfrica.wav', 
+        'en_US_vctk_p271.wav', 
+        'tn_ZA_google-nwu_0441.wav', 
+        'it_IT_mls_6001.wav', 
+        'en_US_vctk_p305.wav', 
+        'it_IT_mls_8828.wav', 
+        'jv_ID_google-gmu_08002.wav', 
+        'it_IT_mls_2033.wav', 
+        'tn_ZA_google-nwu_3629.wav', 
+        'it_IT_mls_6348.wav', 
+        'en_US_cmu_arctic_axb.wav', 
+        'it_IT_mls_8181.wav', 
+        'en_US_vctk_p230.wav', 
+        'af_ZA_google-nwu_7214.wav', 
+        'nl_nathalie.wav', 
+        'it_IT_mls_8207.wav', 
+        'ko_KO_kss.wav', 
+        'af_ZA_google-nwu_6590.wav', 
+        'jv_ID_google-gmu_00264.wav', 
+        'tn_ZA_google-nwu_6234.wav', 
+        'jv_ID_google-gmu_05522.wav', 
+        'en_US_cmu_arctic_lnh.wav', 
+        'en_US_vctk_p272.wav', 
+        'en_US_cmu_arctic_slp.wav', 
+        'en_US_vctk_p299.wav', 
+        'en_US_hifi-tts_9017.wav', 
+        'it_IT_mls_4998.wav', 
+        'it_IT_mls_6299.wav', 
+        'en_US_cmu_arctic_rxr.wav', 
+        'female-46-neutral.wav', 
+        'jv_ID_google-gmu_01392.wav', 
+        'tn_ZA_google-nwu_8512.wav', 
+        'en_US_vctk_p244.wav', 
+        # 'bn_multi_3108.wav', 
+        # 'it_IT_mls_7405.wav', 
+        # 'bn_multi_3713.wav', 
+        # 'yo_openbible.wav', 
+        # 'jv_ID_google-gmu_01932.wav', 
+        'en_US_vctk_p270.wav', 
+        'tn_ZA_google-nwu_6459.wav', 
+        'bn_multi_4046.wav', 
+        'en_US_vctk_p288.wav', 
+        'en_US_vctk_p251.wav', 
+        'es_ES_m-ailabs_tux.wav', 
+        'tn_ZA_google-nwu_6206.wav', 
+        'bn_multi_9169.wav', 
+        # 'en_US_vctk_p293.wav', 
+        # 'en_US_vctk_p255.wav', 
+        'af_ZA_google-nwu_8963.wav',
+        # 'en_US_vctk_p265.wav', 
+        'gu_IN_cmu-indic_cmu_indic_guj_ad.wav', 
+        'jv_ID_google-gmu_07335.wav', 
+        'en_US_vctk_p323.wav', 
+        'en_US_vctk_p281.wav', 
+        'en_US_cmu_arctic_bdl.wav', 
+        'en_US_m-ailabs_judy_bieber.wav', 
+        'it_IT_mls_10446.wav', 
+        'en_US_vctk_p261.wav', 
+        'en_US_vctk_p292.wav', 
+        'te_IN_cmu-indic_ss.wav', 
+        'en_US_vctk_p311.wav', 
+        'it_IT_mls_12428.wav', 
+        'en_US_cmu_arctic_aup.wav', 
+        'jv_ID_google-gmu_04679.wav', 
+        'it_IT_mls_4971.wav', 
+        'en_US_cmu_arctic_ljm.wav', 
+        'fa_haaniye.wav', 
+        'en_US_vctk_p339.wav', 
+        'tn_ZA_google-nwu_7896.wav', 
+        'en_US_vctk_p253.wav', 
+        'it_IT_mls_5421.wav', 
+        # 'ne_NP_ne-google_0546.wav', 
+        'vi_VN_vais1000.wav', 
+        'en_US_vctk_p229.wav', 
+        'en_US_vctk_p254.wav', 
+        'en_US_vctk_p258.wav', 
+        'it_IT_mls_7936.wav', 
+        'en_US_vctk_p301.wav', 
+        'tn_ZA_google-nwu_0045.wav', 
+        'it_IT_mls_659.wav', 
+        'tn_ZA_google-nwu_7674.wav', 
+        'it_IT_mls_12804.wav', 
+        'el_GR_rapunzelina.wav', 
+        'en_US_hifi-tts_6097.wav', 
+        'en_US_vctk_p257.wav', 
+        'jv_ID_google-gmu_07875.wav', 
+        'it_IT_mls_1157.wav', 
+        'it_IT_mls_643.wav', 
+        'en_US_vctk_p304.wav', 
+        'ru_RU_multi_hajdurova.wav', 
+        'it_IT_mls_8461.wav', 
+        'bn_multi_3958.wav', 
+        'it_IT_mls_1989.wav', 
+        'en_US_vctk_p249.wav', 
+        # 'bn_multi_0834.wav', 
+        'en_US_vctk_p307.wav', 
+        'es_ES_m-ailabs_karen_savage.wav', 
+        'fr_FR_m-ailabs_bernard.wav', 
+        'en_US_vctk_p252.wav', 
+        'en_US_cmu_arctic_jmk.wav', 
+        'en_US_vctk_p333.wav', 
+        'tn_ZA_google-nwu_4506.wav', 
+        'ne_NP_ne-google_0283.wav', 
+        'de_DE_m-ailabs_karlsson.wav', 
+        'en_US_cmu_arctic_awb.wav', 
+        'en_US_vctk_p246.wav', 
+        'en_US_cmu_arctic_clb.wav', 
+        'en_US_vctk_p364.wav', 
+        'nl_flemishguy.wav', 
+        'en_US_vctk_p276.wav', # y
+        # 'en_US_vctk_p274.wav', 
+        'fr_FR_m-ailabs_gilles_g_le_blanc.wav', 
+        'it_IT_mls_7444.wav', 
+        'style_o22050.wav', 
+        'en_US_vctk_s5.wav', 
+        'en_US_vctk_p268.wav', 
+        'it_IT_mls_6807.wav', 
+        'it_IT_mls_2019.wav', 
+        'male-60-angry.wav', 
+        'af_ZA_google-nwu_8924.wav', 
+        'en_US_vctk_p374.wav', 
+        'en_US_vctk_p363.wav', 
+        'it_IT_mls_644.wav', 
+        'ne_NP_ne-google_3614.wav', 
+        'en_US_vctk_p241.wav', 
+        'ne_NP_ne-google_3154.wav', 
+        'en_US_vctk_p234.wav', 
+        'it_IT_mls_8384.wav', 
+        'fr_FR_m-ailabs_ezwa.wav', 
+        'it_IT_mls_5010.wav', 
+        'en_US_vctk_p351.wav', 
+        'en_US_cmu_arctic_eey.wav', 
+        'jv_ID_google-gmu_04285.wav', 
+        'jv_ID_google-gmu_06941.wav', 
+        'hu_HU_diana-majlinger.wav', 
+        'tn_ZA_google-nwu_2839.wav', 
+        'bn_multi_03042.wav', 
+        'tn_ZA_google-nwu_5628.wav', 
+        'it_IT_mls_4649.wav', 
+        'af_ZA_google-nwu_7130.wav', 
+        'en_US_cmu_arctic_slt.wav', 
+        'jv_ID_google-gmu_04175.wav', 
+        'gu_IN_cmu-indic_cmu_indic_guj_kt.wav', 
+        'jv_ID_google-gmu_00027.wav', 
+        'jv_ID_google-gmu_02884.wav', 
+        'en_US_vctk_p360.wav', 
+        'en_US_vctk_p334.wav', 
+        'male-27-sad.wav', 
+        'tn_ZA_google-nwu_1498.wav', 
+        'fi_FI_harri-tapani-ylilammi.wav', 
+        'bn_multi_rm.wav', 
+        'ne_NP_ne-google_2139.wav', 
+        'pl_PL_m-ailabs_piotr_nater.wav', 
+        'fr_FR_siwis.wav', 
+        'nl_bart-de-leeuw.wav', 
+        'jv_ID_google-gmu_04715.wav', 
+        'en_US_vctk_p283.wav', 
+        'en_US_vctk_p314.wav', 
+        'en_US_vctk_p335.wav', 
+        'jv_ID_google-gmu_07765.wav', 
+        'en_US_vctk_p273.wav'
+        ]
+
 _tts = StyleTTS2().to('cpu')
 
 def only_greek_or_only_latin(text, lang='grc'):
@@ -968,23 +1321,78 @@ def only_greek_or_only_latin(text, lang='grc'):
 
 
 def other_tts(text='Hallov worlds Far over the',
-              ref_s='wav/af_ZA_google-nwu_0184.wav'):
+              ref_s='wav/af_ZA_google-nwu_0184.wav',
+              soundscape='birds fomig',
+              cache_lim=64):
+    
+    total_audio = []
+    
+    final_audio = None
+    speech_audio = None
 
-    text = only_greek_or_only_latin(text, lang='eng')
 
-    x = _tts.inference(text, ref_s=ref_s)[0:1, 0, :]
+    if text and text.strip():
+        
+        text = only_greek_or_only_latin(text, lang='eng')
 
-    x = torch.cat([.99 * x,
-                   .94 * x], 0).cpu().numpy()  # Stereo
+        speech_audio = _tts.inference(text, ref_s=ref_s)[0, 0, :].numpy()  # 24 Khz
+        
+        if speech_audio.shape[0] > 10:
 
-    # x /= np.abs(x).max() + 1e-7  ~ Volume normalisation @api.py:tts_multi_sentence() OR demo.py
+            speech_audio = audresample.resample(signal=speech_audio.astype(np.float32),
+                                                original_rate=24000,
+                                                target_rate=16000)[0, :]   # 16 KHz
 
-    tmp_file = f'_speech.wav'  # N x clients (cleanup vs tmp file / client)
+    # AudioGen
+    if soundscape and soundscape.strip():
 
-    audiofile.write(tmp_file, x, 24000)
+        
+        speech_duration_secs = len(speech_audio) / 16000 if speech_audio is not None else 0
+        target_duration = max(speech_duration_secs + 0.74, 2.0)
 
-    return tmp_file
 
+        background_audio = audiogen.generate(
+            soundscape,
+            duration=target_duration,
+            cache_lim=max(4, int(cache_lim))  # at least allow 10 A/R stEps
+         ).numpy()
+
+        if speech_audio is not None:
+
+            len_speech = len(speech_audio)
+            len_background = len(background_audio)
+            
+            if len_background > len_speech:
+                padding = np.zeros(len_background - len_speech,
+                                   dtype=np.float32)
+                speech_audio = np.concatenate([speech_audio, padding])
+            elif len_speech > len_background:
+                padding = np.zeros(len_speech - len_background,
+                                   dtype=np.float32)
+                background_audio = np.concatenate([background_audio, padding])
+
+            # Convert to 2D arrays for stereo blending
+            speech_audio_stereo = speech_audio[None, :]
+            background_audio_stereo = background_audio[None, :]
+
+
+            final_audio = np.concatenate([
+                0.49 * speech_audio_stereo + 0.51 * background_audio_stereo,
+                0.51 * background_audio_stereo + 0.49 * speech_audio_stereo
+            ],0)
+        else:
+            final_audio = background_audio
+
+    elif speech_audio is not None:
+        final_audio = speech_audio
+
+    # If both inputs are empty, create a 2s silent audio file.
+    if final_audio is None:
+        final_audio = np.zeros(16000 * 2, dtype=np.float32)
+    print('\n=============F I N A L\n', final_audio.shape, final_audio.dtype, final_audio.min(), np.isnan(final_audio).sum())
+    wavfile = '_audionar_.wav'
+    audiofile.write(wavfile, final_audio, 16000)    
+    return wavfile
 
 def update_selected_voice(voice_filename):
     return 'wav/' + voice_filename + '.wav'
@@ -1035,11 +1443,19 @@ with gr.Blocks(theme='huggingface', css=css_buttons) as demo:
         # Main input and output components
         with gr.Row():
             text_input = gr.Textbox(
-                label="Enter text for TTS:",
+                label="TYpe text for TTS:",
                 placeholder="Type your message here...",
                 lines=4,
                 value="Farover the misty mountains cold too dungeons deep and caverns old.",
             )
+            soundscape_input = gr.Textbox(lines=1,
+                   value="frogs",
+                   label="AudioGen Txt"
+            ),
+            kv_input = gr.Number(
+                label="kv Period",
+                value=24,
+            )
             generate_button = gr.Button("Generate Audio", variant="primary")
 
         output_audio = gr.Audio(label="TTS Output")
@@ -1066,7 +1482,7 @@ with gr.Blocks(theme='huggingface', css=css_buttons) as demo:
 
         generate_button.click(
             fn=other_tts,
-            inputs=[text_input, selected_voice],
+            inputs=[text_input, selected_voice, soundscape_input, kv_input],
             outputs=output_audio
         )
 
@@ -1108,11 +1524,9 @@ with gr.Blocks(theme='huggingface', css=css_buttons) as demo:
                     value='Η γρηγορη καφετι αλεπου πειδαει πανω απο τον τεμπελη σκυλο.',
                     label="Type text for TTS"
                 )
-                lang_dropdown = gr.Dropdown(
-                    choices=language_names,
-                    label="TTS language",
-                    value="Ancient greek",
-                )
+                lang_dropdown = gr.Dropdown(choices=language_names, label="TTS language", value="Ancient greek")
+                soundscape_input = gr.Textbox(lines=1, value="dogs barg", label="AudioGen Txt")
+                kv_input = gr.Number(label="kv Period", value=70)
             
             # Create a button to trigger the TTS function
             tts_button = gr.Button("Generate Audio")
@@ -1123,7 +1537,7 @@ with gr.Blocks(theme='huggingface', css=css_buttons) as demo:
             # Link the button click event to the mms_tts function
             tts_button.click(
                 fn=audionar_tts,
-                inputs=[text_input, lang_dropdown],
+                inputs=[text_input, lang_dropdown, soundscape_input, kv_input],
                 outputs=audio_output
             )
 

audiocraft.py

@@ -0,0 +1,724 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from omegaconf import OmegaConf
+import numpy as np
+from huggingface_hub import hf_hub_download
+import os
+from torch.nn.utils import weight_norm
+from transformers import T5EncoderModel, T5Tokenizer  # type: ignore
+from einops import rearrange
+
+torch.backends.cuda.enable_mem_efficient_sdp(True)
+
+
+
+N_REPEAT = 2  # num (virtual batch_size) clones of audio sounds
+
+def _shift(x):
+    #print(x.shape, 'BATCH Independent SHIFT\n AudioGen')
+    for i, _slice in enumerate(x):
+        n = x.shape[2]
+        offset = np.random.randint(.24 * n, max(1, .74 * n))  # high should be above >= 0 TBD
+        print(offset)
+        x[i, :, :] = torch.roll(_slice, offset, dims=1)  # _slice 2D
+    return x
+
+class AudioGen(torch.nn.Module):
+
+    # https://huggingface.co/facebook/audiogen-medium
+
+    def __init__(self):
+
+        super().__init__()
+        _file_1 = hf_hub_download(
+            repo_id='facebook/audiogen-medium',
+            filename="compression_state_dict.bin",
+            cache_dir=os.environ.get('AUDIOCRAFT_CACHE_DIR', None),
+            library_name="audiocraft",
+            library_version= '1.3.0a1')  # Found at __init__.py #audiocraft.__version__)
+        pkg = torch.load(_file_1, map_location='cpu')# kwargs = OmegaConf.create(pkg['xp.cfg'])
+        self.compression_model = EncodecModel()
+        self.compression_model.load_state_dict(pkg['best_state'], strict=False)
+        self.compression_model.eval()  # ckpt has also unused encoder weights
+        self._chunk_len = 476
+        _file_2 = hf_hub_download(
+            repo_id='facebook/audiogen-medium',
+            filename="state_dict.bin",
+            cache_dir=os.environ.get('AUDIOCRAFT_CACHE_DIR', None),
+            library_name="audiocraft",
+            library_version= '1.3.0a1')  # Found at __init__.py #audiocraft.__version__)
+        pkg = torch.load(_file_2, map_location='cpu')
+        cfg = OmegaConf.create(pkg['xp.cfg'])  # CFG inside torch bin
+        _best = pkg['best_state']
+        _best['t5.output_proj.weight'] = _best.pop('condition_provider.conditioners.description.output_proj.weight')#.to(torch.float)
+        _best['t5.output_proj.bias'] = _best.pop('condition_provider.conditioners.description.output_proj.bias')#.to(torch.float)
+        self.lm = LMModel()
+        self.lm.load_state_dict(pkg['best_state'], strict=True)
+        self.lm.eval()
+
+
+    @torch.no_grad()
+    def generate(self,
+                 prompt='dogs mewo',
+                 duration=2.24,  # seconds of audio
+                 cache_lim=71,   # flush kv cache after cache_lim tok
+                 ):
+        torch.manual_seed(42)  # https://github.com/facebookresearch/audiocraft/issues/111#issuecomment-1614732858
+        self.lm.cache_lim = cache_lim
+        self.lm.n_draw = int(.8 * duration) + 1  # different beam every 0.47 seconds of audio
+        with torch.autocast(device_type='cpu', dtype=torch.bfloat16):
+            gen_tokens = self.lm.generate(
+                text_condition=[prompt] * N_REPEAT  + [''] * N_REPEAT,#['dogs', 'dogs...!', '', '']
+                max_tokens=int(.04 * duration / N_REPEAT * self.compression_model.frame_rate) + 12) # [bs, 4, 74*self.lm.n_draw]
+
+        # OOM if vocode all tokens
+        x = []
+
+
+        for i in range(7, gen_tokens.shape[2], self._chunk_len):  # min soundscape 2s assures 10 tokens
+
+            decoded_chunk = self.compression_model.decode(gen_tokens[:, :, i-7:i+self._chunk_len])
+
+            x.append(decoded_chunk)
+            
+        x = torch.cat(x, 2)  # [bs, 1, 114000]
+
+        x = _shift(x)  # clone() to have xN
+
+        return x.reshape(-1) #x / (x.abs().max() + 1e-7)
+
+
+class EncodecModel(nn.Module):
+
+    def __init__(self):
+
+        super().__init__()
+        self.decoder = SEANetDecoder()
+        self.quantizer = ResidualVectorQuantizer()
+        self.frame_rate = 50
+
+
+    def decode(self, codes):
+        # B,K,T -> B,C,T
+        emb = self.quantizer.decode(codes)
+        return self.decoder(emb)
+
+
+class StreamableLSTM(nn.Module):
+
+    def __init__(self,
+                 dimension,
+                 num_layers=2,
+                 skip=True):
+        
+        super().__init__()
+        self.skip = skip
+        self.lstm = nn.LSTM(dimension, dimension, num_layers)
+
+    def forward(self, x):
+        x = x.permute(2, 0, 1)
+        y, _ = self.lstm(x)
+        if self.skip:
+            y = y + x
+        y = y.permute(1, 2, 0)
+        return y
+
+
+
+class SEANetResnetBlock(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 kernel_sizes = [3, 1], 
+                 pad_mode = 'reflect', 
+                 compress = 2):
+        
+        super().__init__()
+
+        hidden = dim // compress
+        block = []
+        for i, kernel_size in enumerate(kernel_sizes):
+            in_chs = dim if i == 0 else hidden
+            out_chs = dim if i == len(kernel_sizes) - 1 else hidden
+            block += [nn.ELU(),
+                      StreamableConv1d(in_chs,
+                                       out_chs,
+                                       kernel_size=kernel_size,
+                                       pad_mode=pad_mode)]
+        self.block = nn.Sequential(*block)
+
+    def forward(self, x):
+        return x + self.block(x)
+
+
+
+
+
+class SEANetDecoder(nn.Module):
+#  channels=1 dimension=128 n_filters=64 n_residual_layers=1 ratios=[8, 5, 4, 2] 
+# activation='ELU' activation_params={'alpha': 1.0}, final_activation=None
+# final_activation_params=None norm='weight_norm' 
+# norm_params={} kernel_size=7 last_kernel_size=7 residual_kernel_size=3 dilation_base=2 
+# causal=False pad_mode='constant' 
+# true_skip=True compress=2 lstm=2 disable_norm_outer_blocks=0 trim_right_ratio=1.0 
+
+    def __init__(self,
+                 channels = 1,
+                 dimension = 128,
+                 n_filters = 64,
+                 n_residual_layers = 1,
+                 ratios = [8, 5, 4, 2], 
+                 kernel_size = 7,
+                 last_kernel_size = 7, 
+                 residual_kernel_size = 3, 
+                 pad_mode = 'constant', 
+                 compress = 2,
+                 lstm = 2):
+        
+        super().__init__()
+
+
+        mult = int(2 ** len(ratios))
+        model = [
+            StreamableConv1d(dimension, mult * n_filters,
+                             kernel_size,
+                             pad_mode=pad_mode)
+        ]
+
+        if lstm:
+            print('\n\n\n\nLSTM IN SEANET\n\n\n\n')
+            model += [StreamableLSTM(mult * n_filters,
+                                     num_layers=lstm)]
+
+        # Upsample to raw audio scale
+        for i, ratio in enumerate(ratios):
+            
+
+            model += [
+                nn.ELU(),
+                StreamableConvTranspose1d(mult * n_filters,
+                                          mult * n_filters // 2,
+                                          kernel_size=ratio * 2,
+                                          stride=ratio),
+            ]
+            # Add residual layers
+            for j in range(n_residual_layers):
+
+                model += [
+                    SEANetResnetBlock(mult * n_filters // 2, 
+                                      kernel_sizes=[residual_kernel_size, 1],
+                                      pad_mode=pad_mode, 
+                                      compress=compress)]
+
+            mult //= 2
+        
+        # Add final layers
+        model += [
+            nn.ELU(),
+            StreamableConv1d(n_filters,
+                             channels,
+                             last_kernel_size,
+                             pad_mode=pad_mode)]
+        self.model=nn.Sequential(*model)
+
+    def forward(self, z):
+        return self.model(z)
+
+
+
+
+def unpad1d(x, paddings):
+    padding_left, padding_right = paddings
+    end = x.shape[-1] - padding_right
+    return x[..., padding_left: end]
+
+
+class NormConv1d(nn.Module):
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        
+        self.conv = weight_norm(nn.Conv1d(*args, **kwargs))  # norm = weight_norm
+        
+    def forward(self, x):
+        return self.conv(x)
+
+
+
+
+
+class NormConvTranspose1d(nn.Module):
+    
+    def __init__(self, *args, causal: bool = False, norm: str = 'none',
+                 norm_kwargs = {}, **kwargs):
+        super().__init__()
+        
+        self.convtr = weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+        
+    def forward(self, x):
+        return self.convtr(x)
+        
+
+
+
+
+
+class StreamableConv1d(nn.Module):
+
+    def __init__(self,
+                 in_channels, 
+                 out_channels,
+                 kernel_size, 
+                 stride=1, 
+                 groups=1, 
+                 bias=True,
+                 pad_mode='reflect'):
+        super().__init__()
+        if (stride != 1) or (groups != 1):
+            raise ValueError
+        self.conv = NormConv1d(in_channels,
+                               out_channels,
+                               kernel_size,
+                               stride,
+                               groups=groups,
+                               bias=bias)
+        self.pad_mode = pad_mode
+
+    def forward(self, x):
+        kernel_size = self.conv.conv.kernel_size[0]
+        kernel_size = (kernel_size - 1) * self.conv.conv.dilation[0] + 1
+        padding_total = kernel_size - self.conv.conv.stride[0]
+        padding_right = padding_total // 2
+        padding_left = padding_total - padding_right
+
+        # x = pad1d(x, (padding_left, padding_right), mode=self.pad_mode)
+        x = F.pad(x, (padding_left, padding_right), self.pad_mode)
+        return self.conv(x)
+
+
+class StreamableConvTranspose1d(nn.Module):
+
+    def __init__(self, in_channels: int, out_channels: int,
+                 kernel_size: int, stride: int = 1, causal: bool = False,
+                 norm: str = 'none', trim_right_ratio: float = 1.,
+                 norm_kwargs = {}):
+        super().__init__()
+        self.convtr = NormConvTranspose1d(in_channels,
+                                          out_channels,
+                                          kernel_size,
+                                          stride)
+        
+
+    def forward(self, x):
+
+        padding_total = self.convtr.convtr.kernel_size[0] - self.convtr.convtr.stride[0]
+
+        y = self.convtr(x)
+
+        # Asymmetric padding required for odd strides
+        # print('\n   \n\n\nn\n\n\nnANTICAUSAL T\n\n\n')
+        padding_right = padding_total // 2
+        padding_left = padding_total - padding_right
+
+        y = unpad1d(y, (padding_left, padding_right))
+        return y
+
+
+# VQ
+
+class EuclideanCodebook(nn.Module):
+    def __init__(self,
+                 dim,
+                 codebook_size):
+        super().__init__()
+        self.register_buffer("embed", torch.zeros(codebook_size, dim))
+
+
+
+
+class VectorQuantization(nn.Module):
+    
+    def __init__(self,
+                 dim,
+                 codebook_size):
+        
+        super().__init__()
+        self._codebook = EuclideanCodebook(dim=dim,
+                                           codebook_size=codebook_size)
+
+    def decode(self, _ind):
+        return F.embedding(_ind, self._codebook.embed)
+
+
+class ResidualVectorQuantization(nn.Module):
+
+    def __init__(self, *, num_quantizers, **kwargs):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [VectorQuantization(**kwargs) for _ in range(num_quantizers)]
+        )
+
+    def decode(self, _ind):
+        x = 0.0
+        for i, _code in enumerate(_ind):
+            x = x + self.layers[i].decode(_code)
+        return x.transpose(1, 2)
+
+
+class ResidualVectorQuantizer(nn.Module):
+    
+# dimension=128 n_q=4 q_dropout=False bins=2048 decay=0.99 kmeans_init=True 
+# kmeans_iters=50 threshold_ema_dead_code=2
+# orthogonal_reg_weight=0.0 orthogonal_reg_active_codes_only=False 
+# orthogonal_reg_max_codes=None
+
+    def __init__(
+        self,
+        dimension = 128,
+        n_q = 4,
+        bins = 2048
+        ):
+        
+        super().__init__()
+        self.vq = ResidualVectorQuantization(dim=dimension,
+                                             codebook_size=bins,
+                                             num_quantizers=n_q)
+
+    def decode(self, codes):
+        # codes is [B, K, T], with T frames, K nb of codebooks, vq.decode expects [K, B, T].
+        return self.vq.decode(codes.transpose(0, 1))    
+
+
+class T5(nn.Module):
+
+    def __init__(self):
+
+        super().__init__()
+        self.output_proj = nn.Linear(1024,  # t5-large
+                                     1536)  # lm hidden
+        self.t5_tokenizer = T5Tokenizer.from_pretrained('t5-large', legacy=True)
+        t5 = T5EncoderModel.from_pretrained('t5-large').train(mode=False)
+
+        # this makes sure that the t5 is not part
+        # of the saved checkpoint
+        self.__dict__['t5'] = t5.to('cpu')
+
+    def forward(self, prompt):
+        with torch.set_grad_enabled(False): #, torch.autocast(device_type='cpu', dtype=torch.float32):
+
+            bs = len(prompt) // 2
+            d = self.t5_tokenizer(prompt,
+                                    return_tensors='pt',
+                                    padding=True).to(self.output_proj.bias.device)
+            d['attention_mask'][bs:, :] = 0  # null condition t5 attn_mask should be zero
+
+            x = self.t5(input_ids=d['input_ids'],
+                            attention_mask=d['attention_mask']).last_hidden_state  # no kv
+        # Float 16
+        # > self.output_proj() is outside of autocast of t5 - however inside the autocast of lm thus computed in torch.float16
+        x = self.output_proj(x)  # nn.Linear() - produces different result if there is no duplicate txt condition here
+        x[bs:, :, :] = 0  # venv/../site-packages/audiocraft/modules/conditioners.py -> tokenize()
+        return x
+
+
+class LMModel(nn.Module):
+
+    def __init__(self,
+                 n_q = 4,
+                 card = 2048,
+                 dim = 1536
+                 ):
+        super().__init__()
+        self.cache_lim = -1
+        self.t5 = T5()
+        self.card = card # 2048
+        self.n_draw = 1  # draw > 1 tokens of different CFG scale
+                         # batch size > 1 is slower from n_draw as calls transformer on larger batch
+        self.emb = nn.ModuleList([nn.Embedding(self.card + 1, dim) for _ in range(n_q)])  # EMBEDDING HAS 2049
+        self.transformer = StreamingTransformer()
+        self.out_norm = nn.LayerNorm(dim, eps=1e-5)
+        self.linears = nn.ModuleList([nn.Linear(dim, self.card, bias=False) for _ in range(n_q)])  # LINEAR DOESNT HAVE 2049
+
+    def forward(self,
+                sequence,
+                condition_tensors=None,
+                cache_position=None
+                ):
+
+        bs, n_q, time_frames = sequence.shape # [bs, 4, time]
+
+        input_ = sum([self.emb[k](sequence[:, k]) for k in range(n_q)])
+
+        out = self.transformer(torch.cat([input_, input_], 0),  # duplicate null condition (bs x 2) for ClassifierFreeGuidance
+                               cross_attention_src=condition_tensors,
+                               cache_position=cache_position)
+
+        out = self.out_norm(out)
+
+        logits = torch.stack([self.linears[k](out) for k in range(n_q)], dim=1) # [2*bs, 4, 1,      2048]
+        logits = 3 * logits[:bs, :, :, :] - self._scale * logits[bs:, :, :, :]                 # [  bs, 4, n_draw, 2048]
+
+        #bs, n_q, n_draw, vocab = logits.shape
+        tokens = torch.multinomial(torch.softmax(logits.view(bs * self.n_draw * n_q, 2048), dim=1),
+                                   num_samples=1)
+        return tokens.view(bs, n_q, self.n_draw).transpose(1, 2)
+
+    @torch.no_grad()
+    def generate(self,
+                 max_tokens=None,
+                 text_condition=None
+                 ):
+        x = self.t5(text_condition)
+        bs = x.shape[0] // 2  # has null conditions - bs*2*N_REPEAT applys in builders.py
+        self._scale = .3 * torch.rand(1, 1, self.n_draw, 1, device=x.device) + 1.94
+        cache_position = 0
+
+        out_codes = torch.full((bs,
+                                self.n_draw,
+                                4,
+                                4 + 3 + max_tokens),  # 4 + max_tokens + 4-1 to have sufficient to index the 1st antidiagonal of 4x4 + 4 xtra tokens
+                               self.card,
+                               dtype=torch.long,
+                               device=x.device) # [bs, n_draw, 4, dur]
+
+        # A/R
+        for offset in range(0, max_tokens + 4 - 1):  # max_tokens + n_q - 1
+
+            # extract diagonal via indexing out_codes[ [0, 1, 2, 3], [0, 1, 2, 3] ]
+            next_token = self.forward(out_codes[:, 0, [0, 1, 2, 3], torch.tensor([3, 2, 1, 0]) + offset][:, :, None],  # index diagonal & exapnd to [bs, n_q, dur=1]
+                                      #gen_sequence[:, 0, :, offset-1:offset],  # DIAGINDEXING for setting prediction of lm into gen_sequence THE GENSEQUENCE has to be un-delayed in the end [Because it has to be de-delayed for the vocoder then is actually only the lm input that requires to see the delay thus we could just feed by diaggather] so it matches gen_codes -1 a[[0, 1, 2, 3], torch.tensor([0, 1, 2, 3]) + 5]  the gen_sequence is indexed by vertical column and fed to lm however the prediction of lm is place diagonally with delay to the gen_sequence
+                                      condition_tensors=x,  # utilisation of the attention mask of txt condition ?
+                                      cache_position=cache_position)  # [bs, n_draw, 4]
+
+            # Fill of next_token should be also placed on antidiagonal [not column]
+
+            #   Do Not Overwrite 2048 of TRIU/TRIL = START/END => Do Not Fill them by Predicted Tokens
+            # 0-th antidiagonal should be full of card = [2048, 2048, 2048, 2048]
+            #
+            #   [2048, 2048, 2048, 2048,    0,    1,    2,    3,    4,    5,    6, 2048, 2048, 2048],
+            #   [2048, 2048, 2048, 2048, 2048,    0,    1,    2,    3,    4,    5,    6, 2048, 2048],
+            #   [2048, 2048, 2048, 2048, 2048, 2048,    0,    1,    2,    3,    4,    5,    6, 2048],
+            #   [2048, 2048, 2048, 2048, 2048, 2048, 2048,    0,    1,    2,    3,    4,    5,    6]]
+            # NO OVerWriting
+            if offset == 0:
+
+                next_token[:, :, 1:4] = 2048  # self.card - bottom 3 entries of the antidiagonal should remain 2048
+
+            elif offset == 1:
+
+                next_token[:, :, 2:4] = 2048  # bottom 2 entries of the antidiagonal should remain 2048
+
+            elif offset == 2:
+
+                next_token[:, :, 3:4] = 2048
+
+            elif offset == max_tokens:
+
+                next_token[:, :, 0:1] = 2048  # top 1 entry of the antidiagonal should stay to 2048
+
+            elif offset == (max_tokens + 1):
+
+                next_token[:, :, 0:2] = 2048
+
+            elif offset == (max_tokens + 2):
+
+                next_token[:, :, 0:3] = 2048
+
+            else:  # offset 3,4,5,6,7...... max_tokens-1   # FILL Complete n_q = 4 ANTIDIAGONAL ENTRIES
+
+                pass #print('No delete anti-diag')
+
+            out_codes[:, :, [0, 1, 2, 3], torch.tensor([3, 2, 1, 0]) + offset + 1] = next_token
+            # Sink Attn
+            if (offset > 0) and (offset % self.cache_lim) == 0:
+                n_preserve = 4
+                self.transformer._flush(n_preserve=n_preserve)
+                cache_position = n_preserve
+            else:
+                cache_position += 1
+
+        # [bs, n_draw, 4, time+xtra] -> [bs, 4, n_draw, time] ->  [bs, 4, time * n_draw]
+        out_codes = out_codes[:, :, :, 4:max_tokens+4].transpose(1, 2).reshape(bs, 4, self.n_draw * max_tokens)
+
+        # flush for next API call
+        self.transformer._flush()
+
+        return out_codes  # SKIP THE 4 fill 2048
+
+
+
+
+def create_sin_embedding(positions,
+                         dim,
+                         max_period=10000
+                         ):
+    # assert dim % 2 == 0
+    half_dim = dim // 2
+    positions = positions.to(torch.float)
+    adim = torch.arange(half_dim, device=positions.device,
+                        dtype=torch.float).view(1, 1, -1)
+    max_period_tensor = torch.full([],
+                                   max_period,
+                                   device=positions.device,
+                                   dtype=torch.float)  # avoid sync point
+    phase = positions / (max_period_tensor ** (adim / (half_dim - 1)))
+    # OFFICIAL is torch.float32 HOWEVER self_attn.in_prod_weight = torch.float16
+    return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1)
+
+
+class StreamingMultiheadAttention(nn.Module):
+
+    def __init__(self,
+                 embed_dim,
+                 num_heads,
+                 cross_attention=False,
+                 ):
+
+        super().__init__()
+
+        self.cross_attention = cross_attention
+        # if not self.cross_attention then it has kvcachingn
+        self.k_history = None
+        # cleanup history through LM inside GENERATION - Each 0,..,47 mha has different kv history
+        self.v_history = None
+        self.num_heads = num_heads
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.register_buffer('in_proj_weight', torch.ones((3 * embed_dim, embed_dim),
+                                                          dtype=torch.float))
+
+    def forward(self,
+                query,
+                key=None,
+                value=None):
+        layout = "b h t d"
+        if self.cross_attention:
+
+            # Different queries, keys, values > split in_proj_weight
+
+            dim = self.in_proj_weight.shape[0] // 3
+
+            q = nn.functional.linear(query, self.in_proj_weight[:dim])
+            k = nn.functional.linear(key,   self.in_proj_weight[dim: 2 * dim])
+            v = nn.functional.linear(value, self.in_proj_weight[2 * dim:])
+
+            q, k, v = [
+                rearrange(x, f"b t (h d) -> {layout}", h=self.num_heads) for x in [q, k, v]]
+
+        else:
+
+            # Here <else> = self_attention for audio with itself (above is cross attention txt)
+
+            # HISTORY - DIFFERENT FOR EACH TRANSF LAYER
+
+            # here we have different floating values from official
+            projected = nn.functional.linear(query, self.in_proj_weight, None)
+            # print(query.sum(), projected.sum() , self.in_proj_weight.sum(), 'Lc')   # verified official AudioGen values
+            bound_layout = "b h p t d"
+            packed = rearrange(
+                projected, f"b t (p h d) -> {bound_layout}", p=3, h=self.num_heads)
+            q, k, v = packed.unbind(dim=2)
+            if self.k_history is not None:
+                # IF ctrl^c during live_demo the assigning of each of kv is non-atomic k!=v
+                # thus it will try to continue with incompatible k/v dims!
+                self.k_history = torch.cat([self.k_history, k], 2)
+                self.v_history = torch.cat([self.v_history, v], 2)
+            else:
+                self.k_history = k
+                self.v_history = v
+
+            # Assign Completed k / v to k / v
+
+            k = self.k_history
+            v = self.v_history
+
+            # -> kv CACHE ONLY APPLIES if not self.cross_attention
+
+        x = torch.nn.functional.scaled_dot_product_attention(
+            q, k, v, attn_mask=None, is_causal=False, dropout_p=0.0)
+
+        x = rearrange(x, f"{layout} -> b t (h d)", h=self.num_heads)
+        x = self.out_proj(x)
+        return x
+
+
+class StreamingTransformerLayer(nn.Module):
+
+    def __init__(self,
+                 d_model,
+                 num_heads,
+                 dim_feedforward):
+        
+        super().__init__()
+
+        self.self_attn = StreamingMultiheadAttention(embed_dim=d_model,
+                                                     num_heads=num_heads)
+        self.linear1 = nn.Linear(d_model, dim_feedforward, bias=False)
+        self.linear2 = nn.Linear(dim_feedforward, d_model, bias=False)
+        self.cross_attention = StreamingMultiheadAttention(embed_dim=d_model,
+                                                           num_heads=num_heads,
+                                                           cross_attention=True)
+        self.norm_cross = nn.LayerNorm(d_model, eps=1e-5)
+        self.norm1 = nn.LayerNorm(d_model, eps=1e-5)
+        self.norm2 = nn.LayerNorm(d_model, eps=1e-5)
+
+    def forward(self,
+                x,
+                cross_attention_src=None):
+        x = x + self.self_attn(self.norm1(x))
+        x = x + self.cross_attention(query=self.norm_cross(x),
+                                     key=cross_attention_src,
+                                     value=cross_attention_src)  # txtcondition
+        x = x + self.linear2(F.gelu(self.linear1(self.norm2(x))))
+        return x
+
+
+class StreamingTransformer(nn.Module):
+
+    def __init__(self,
+                 d_model=1536,
+                 num_heads=24,
+                 num_layers=48,
+                 dim_feedforward=6144):
+        super().__init__()
+
+        self.layers = nn.ModuleList(
+                [
+                    StreamingTransformerLayer(d_model=d_model,
+                                              num_heads=num_heads,
+                                              dim_feedforward=dim_feedforward) for _ in range(num_layers)
+                    ]
+                )
+
+    def forward(self,
+                x,
+                cache_position=None,
+                cross_attention_src=None):
+
+        x = x + create_sin_embedding(
+                torch.zeros(x.shape[0], 1, 1, device=x.device) + cache_position, 1536)
+
+        for lay in self.layers:
+            x = lay(x,
+                    cross_attention_src=cross_attention_src)
+        return x
+
+    def _flush(self,
+               n_preserve=None):
+
+        for lay in self.layers:
+            if n_preserve is not None:
+                # cache position is difficult to choose to also preserve kv from end
+                lay.self_attn.k_history = lay.self_attn.k_history[:, :, :n_preserve, :]
+                lay.self_attn.v_history = lay.self_attn.v_history[:, :, :n_preserve, :]
+            else:
+                lay.self_attn.k_history = None
+                lay.self_attn.v_history = None
+
+
+if __name__ == '__main__':
+
+    import audiofile
+    model = AudioGen().to('cpu')
+    x = model.generate(prompt='swims in lake frogs', duration=6.4).cpu().numpy()
+    audiofile.write('_sound_.wav', x, 16000)

requirements.txt

@@ -1,8 +1,6 @@
-torch
+omegaconf
 nltk
-pydantic==2.10.6
 librosa
-transformers
 phonemizer
 audiofile
 matplotlib
@@ -11,4 +9,8 @@ num2words
 numpy<2.0.0
 gradio==5.27.0
 Numbers2Words-Greek
-
+einops
+torch
+pydantic==2.10.6
+transformers==4.49.0
+sentencepiece