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 True:
        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',
            "placeholder_token": '<Alice>',
            "initializer_token": 'girl',
            "repeats": '100',
            "train_data_dir": 'resources/Alices',
            "output_dir": 'outputs/textual_inversion',
            "height": 512,
            "width": 512,
            "learning_rate": 5e-4,
            "max_train_steps": 1000,
            "save_steps": 200,
            "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 False:
        config['text_inversion'] = {
            "width": 512,
            "height": 512,
            "prompt": '<Alice> at the lake',
            "negative_prompt": '',
            "output_dir": 'outputs/text_inversion_txt2img',
        }
        gui_text_inversion = StableDiffusionUI_txt2img(
            **config['text_inversion']
        )
        start_end('txt2img inference', gui_train_text_inversion)