import os
import sys
import uuid

from webui_im2im import ControlnetRequestImg2Img
import numpy as np
import base64
import cv2
import os, sys
from gevent import pywsgi, monkey
from multiprocessing import Process, Queue
import glob

# from gpt4v_caption import caption_image

# monkey.patch_all()
# sys.setrecursionlimit(20000)

# 将当前工作目录切换到当前目录
project_dir = os.path.dirname(os.path.abspath(__file__))
os.chdir(project_dir)
sys.path.append(project_dir)


import global_variable as global_var
import json



def train_hair_lora():
    import os
    import json

    try:
        hair_train_dir = "/mnt/database2/jiangqian/0808/online_train_datas_2"
        hair_material_dir_list = os.listdir(hair_train_dir)

        for single_hair_material in hair_material_dir_list:
            # 获取请求参数
            task_id = str(uuid.uuid4())

            request_data = {}
            hair_material_dir = os.path.join(hair_train_dir, single_hair_material)
            print("--------------------hair_material_dir:", hair_material_dir)
            request_data['hair_material_dir'] = hair_material_dir

            img_dir = os.path.join(hair_material_dir, 'images')

            model_dir = os.path.join(hair_material_dir, 'model')
            if not os.path.exists(model_dir):
                os.makedirs(model_dir)

            #判断img_dir下面是否只有一个文件夹
            img_dir_list = os.listdir(img_dir)
            train_image_dir = os.path.join(img_dir, img_dir_list[0])

            #判断文件夹下面是否有图片
            request_data['train_image_dir'] = train_image_dir

            # 将请求数据放入队列
            train_thread(request_data)

            # 返回结果
            print("头发lora训练开始")
            print("\n\n\n\n")

    except Exception as e:
        print(e)

def train_thread(task_dict):
    try:
        hair_material_dir = task_dict['hair_material_dir']
        images_dir = os.path.join(hair_material_dir, 'images')
        model_dir = os.path.join(hair_material_dir, 'model')
        train_image_dir = task_dict['train_image_dir']
        tag = ""

        sample_dir = os.path.join(model_dir, 'sample')
        if not os.path.exists(sample_dir):
            os.makedirs(sample_dir)
        sample_txt = os.path.join(sample_dir, 'prompt.txt')
        with open(sample_txt, 'w') as f_s:
            f_s.write('titor hairstyle, easyphoto, faceless, no human, white background, simple background, '
                      + tag
                      + ' --n low quality, worst quality, bad anatomy, bad composition, poor, low effort --h 512 '
                        '--w 512 --s 30 --l 7')

        #训练头发lora
        cmd_train = (
            'docker run --rm --gpus all -v /home/student/Documents/workspace_cxt_tianjing_hair/miaoya/kohya_ss_home:/home/chinatszrn -v '
            '/mnt:/mnt -e PATH=/home/chinatszrn/.local/bin -w /home/chinatszrn/kohya_ss '
            '--net=host chinatszrn/ubuntu:kohya_ss accelerate launch --num_cpu_threads_per_process=2 "./train_network.py" --enable_bucket '
            '--min_bucket_reso=256 --max_bucket_reso=1800 --pretrained_model_name_or_path="/mnt/nas_hdd/米亚像馆/models/Stable-diffusion/majicmixRealistic_v7.safetensors" '
            f'--train_data_dir={images_dir} --resolution="768,768" '
            f'--output_dir={model_dir} '
            '--network_alpha="64" --save_model_as=safetensors --network_module=networks.lora --text_encoder_lr=5e-05 '
            '--unet_lr=0.0001 --network_dim=128 --output_name="hairstyle_lora" --lr_scheduler_num_cycles="20" '
            '--no_half_vae --learning_rate="0.0001" --lr_scheduler="cosine" --lr_warmup_steps="650" --train_batch_size="1" '
            '--max_train_steps="2000" --save_every_n_epochs="100" --mixed_precision="fp16" --save_precision="fp16" '
            '--caption_extension=".txt" --sample_sampler=ddim '
            f'--sample_prompts={sample_txt} --sample_every_n_epochs="1" '
            '--seed="1234" --cache_latents --optimizer_type="AdamW8bit" --max_data_loader_n_workers="0" --bucket_reso_steps=64 '
            '--xformers --bucket_no_upscale --noise_offset=0.0 --tokenizer_cache_dir="/home/chinatszrn/.cache/clip"')
        print("cmd_train:", cmd_train)
        os.system(cmd_train)

        lora_path = os.path.join(model_dir, 'hairstyle_lora.safetensors')

    except Exception as e:
        print(e)
        return


if __name__ == '__main__':
    global_var.webui_lora_dir = '/home/student/Documents/workspace_cxt_tianjing_hair/miaoya/webui_home/stable-diffusion-webui/models/Lora'

    train_hair_lora()