from ui.config import config
from ui import (
        StableDiffusionUI_text_inversion,
        StableDiffusionUI_dreambooth,
        pipeline_superres,
        pipeline,
        StableDiffusionUI_convert,
    )
from ui.StableDiffusionUI_txt2img import StableDiffusionUI_txt2img
from ui.StableDiffusionUI_img2img import StableDiffusionUI_img2img
from ui.SuperResolutionUI import SuperResolutionUI
def start_end(name,func):
    print(f'------- test {name} start -------')
    func.on_run_button_click('test')
    print(f'------- test {name} finished -------')




if __name__ == '__main__':
    # 加载配置信息 具体配置请查看ui.config文件中

    # ------------------------------文生图------------------------------
    if False:
        config['txt2img']["prompt"] = 'extremely detailed CG unity 8k wallpaper,black long hair,cute face,1 adult girl,happy, green skirt dress, flower pattern in dress,solo,green gown,art of light novel,in field'
        config['txt2img']["negative_prompt"] = 'lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry'
        config['txt2img']["model_name"] = 'MoososCap/NOVEL-MODEL'
        config['txt2img']["output_dir"] = 'outputs/txt2img'
        gui_txt2img = StableDiffusionUI_txt2img(
            **config['txt2img']
        )
        start_end('txt2img', gui_txt2img)

    # ------------------------------图生图------------------------------
    if False:
        config['img2img']["prompt"] = 'red dress'
        config['img2img']["negative_prompt"] = 'lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry'
        config['img2img']["model_name"] =  'MoososCap/NOVEL-MODEL'
        config['img2img']["image_path"] = 'resources/cat2.jpg'
        config['img2img']["mask_path"] = 'resources/mask8.jpg'
        config['img2img']["output_dir"] = 'outputs/img2img'
        gui_img2img = StableDiffusionUI_img2img(
            **config['img2img']
        )
        start_end('img2img', gui_img2img)

    # ------------------------------文生图定制化训练--------------------------
    if False:
        config['train_text_inversion']["learnable_property"]= 'object'
        config['train_text_inversion']["placeholder_token"]= '<Alice>'
        config['train_text_inversion']["initializer_token"]= 'girl'
        config['train_text_inversion']["repeats"]= '100'
        config['train_text_inversion']["train_data_dir"]= 'resources/Alices'
        config['train_text_inversion']["output_dir"]= 'outputs/textual_inversion'
        config['train_text_inversion']["height"]= 512
        config['train_text_inversion']["width"]= 512
        config['train_text_inversion']["learning_rate"]= 5e-4
        config['train_text_inversion']["max_train_steps"]= 1000
        config['train_text_inversion']["save_steps"]= 200
        config['train_text_inversion']["model_name"]= 'MoososCap/NOVEL-MODEL'
        gui_train_text_inversion = StableDiffusionUI_text_inversion(
            **config['train_text_inversion']
        )
        start_end('txt2img train', gui_train_text_inversion)

    # 定制化训练文生图推理
    if True:
        config['text_inversion']["width"] = 512
        config['text_inversion']["height"] = 512
        config['text_inversion']["prompt"] = '<Alice> at the lake'
        config['text_inversion']["negative_prompt"] = ''
        config['text_inversion']["output_dir"] = 'outputs/text_inversion_txt2img'

        gui_text_inversion = StableDiffusionUI_txt2img(
            **config['text_inversion']
        )
        start_end('txt2img inference', gui_text_inversion)