Update app.py

app.py

@@ -23,6 +23,8 @@ class TTS:
             last_chpt1 = './log/1038_eunsik_01/Glow_TTS_00289602.pt'
         elif model_variant == 'KSS':
             last_chpt1 = './log/KSS/Glow_TTS_00280641.pt'
+        elif model_variant == '태연':
+            last_chpt1 = './log/Taeyeon/Glow_TTS_400000.pt'
         check_point = torch.load(last_chpt1, map_location = device)
         self.flowgenerator.load_state_dict(check_point['generator'])
         self.flowgenerator.decoder.skip()
@@ -31,6 +33,8 @@ class TTS:
             last_chpt2 = './log/1038_eunsik_01/HiFI_GAN_00257000.pt'
         elif model_variant == 'KSS':
             last_chpt2 = './log/KSS/HiFi_GAN_00135000.pt'
+        elif model_variant == '태연':
+            last_chpt2 = './log/Taeyeon/HiFi_GAN_337000.pt'
         check_point = torch.load(last_chpt2, map_location = device)
         self.voicegenerator.load_state_dict(check_point['gen_model'])
         self.voicegenerator.eval()
@@ -54,14 +58,16 @@ def init_session_state():
     # Model
     if "init_model" not in st.session_state:
         st.session_state.init_model = True
-        st.session_state.model_variant = "감기걸린 은식"
-        st.session_state.TTS = TTS("감기걸린 은식")
+        st.session_state.model_variant = "태연"
+        st.session_state.TTS = TTS("태연")
 
 def update_model():
     if st.session_state.model_variant == "KSS":
         st.session_state.TTS = TTS("KSS")
     elif st.session_state.model_variant == "감기걸린 은식":
         st.session_state.TTS = TTS("감기걸린 은식")
+    elif st.seesion_state.model_varaiant == '태연':
+        st.session_state.TTS = TTS("태연")
 
 def update_session_state(state_id, state_value):
     st.session_state[f"{state_id}"] = state_value
@@ -89,19 +95,18 @@ col1, col2 = st.columns(2)
 with col1:
     input_text = st.text_input(
         "한글로만 입력해주세요",
-        value = "밥은 먹고 다녀",
+        value = "밥은 먹고 다녀?",
     )
 with col2:
-    model_variant = st.selectbox("목소리 선택해주세요", options = ["KSS", "감기걸린 은식"], index = 1)
+    model_variant = st.selectbox("목소리 선택해주세요", options = ["KSS", "감기걸린 은식", "태연"], index = 1)
     
 button_change = st.button("Change Vocie") 
 if button_change == True:
     if model_variant != st.session_state.model_variant:
-    # Update variant choice
-    update_session_state("model_variant", model_variant)
-    # Re-load model
-    update_model()
-    st.snow()
+        with st.spinner('Wait for it...'):
+            update_session_state("model_variant", model_variant)
+            update_model()
+        st.success('Done!', icon="✅")
     
 noise_scale = st.slider('noise를 추가합니다.', 0., 2., value = 0.33, step = 0.01)
 length_scale = st.slider('속도를 조절합니다.', 0., 2., value = 1., step = 0.01)