import os
# check main envs


def check_install(verbose = True):
    try:
        import paddle
    except:
        print("please install paddle==2.4.2 with version dtk-23.04  before running")
        exit()
    try:
        print('checking install.......')
        import safetensors
        from ppdiffusers.utils import image_grid
        from paddlenlp.transformers.clip.feature_extraction import CLIPFeatureExtractor
        from paddlenlp.transformers import FeatureExtractionMixin
        import ipywidgets
        import PIL
        import tqdm
        print('检测完成，库完整')
    except (ModuleNotFoundError, ImportError, AttributeError):
        if verbose: print('检测到库不完整, 正在安装库')
        os.system("pip install -U pip  -i https://mirror.baidu.com/pypi/simple")
        os.system("pip install -U OmegaConf --user")
        os.system("pip install ppdiffusers==0.9.0 --user")
        os.system("pip install paddlenlp==2.4.9 --user")
        os.system("pip install -U safetensors --user")
        os.system("pip install ipython==8.14.0")
        os.system("pip install ipywidgets==8.0.7")
        os.system("pip install pillow==9.5.0")
        os.system("pip install tqdm==4.65.0")

def start_end(name,func):
    print(f'------- test {name} start -------')
    func.on_run_button_click('test')
    print(f'------- test {name} finished -------')

if __name__=='__main__':
    check_install()

    from ui import gui_train_text_inversion,gui_txt2img,gui_img2img

    start_end('txt2img',gui_txt2img)

    start_end('img2img',gui_img2img)

    start_end('txt2img train',gui_train_text_inversion)