main.py

from argparse import ArgumentParser
from types import SimpleNamespace

import streamlit as st

from app.SessionState import get as get_app_state
from app.components import Janus
from app.login import login_widget
from app.model import instantiate_generator, checkpoint_widget
from app.utils import Session, get_user_history
from app.globals import TEXT_GENERATION_ATTRIBUTES 

if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--device', default=None)
    args = parser.parse_args()
    
    # Create a session state
    application_state = get_app_state(
        login_state=SimpleNamespace(username=None, successful=False, registering=False),
        current_session=Session(mode="Exploratory",
                                id=1,
                                name="",
                                description="",
                                generations=[],
                                favorites=set(),
                                attributes=TEXT_GENERATION_ATTRIBUTES),
        model_settings={},
        checkpoint_info={},
        session_history=[],
        app_state={},
    )

    # The main title
    st.title(":rocket: Janus: Write with Mercury")

    # Login the user
    login_state = login_widget(login_state=application_state.login_state)

    # Load back the user's information
    user_history = get_user_history(login_state.username)
    application_state.current_session.id = user_history.session_id
    application_state.session_history = user_history.session_history

    # Create a generator using a model checkpoint
    checkpoint_info = checkpoint_widget()
    application_state.checkpoint_info = checkpoint_info.__dict__
    generator = instantiate_generator(**checkpoint_info.__dict__, device=args.device)

    # Create the application
    janus = Janus(
        generator=generator,
        current_session=application_state.current_session,
        model_settings=application_state.model_settings,
        session_history=application_state.session_history,
        username=login_state.username,
    )
    janus.start()