elements/session_states.py

# Utils
import uuid

# Streamlit
import streamlit as st

# Nix
from nix.models.TTS import NixTTSInference

# --------------------- SESSION STATE MANAGEMENT -------------------------

def init_session_state():
    # Model
    if "init_model" not in st.session_state:
        st.session_state.init_model = True
        st.session_state.random_str = uuid.uuid1().hex
        st.session_state.model_variant = "Stochastic"
        st.session_state.TTS = NixTTSInference("assets/nix-ljspeech-sdp-v0.1")

def update_model():
    if st.session_state.model_variant == "Deterministic":
        st.session_state.TTS = NixTTSInference("assets/nix-ljspeech-v0.1")
    elif st.session_state.model_variant == "Stochastic":
        st.session_state.TTS = NixTTSInference("assets/nix-ljspeech-sdp-v0.1")

def update_session_state(
    state_id,
    state_value,
):
    st.session_state[f"{state_id}"] = state_value