from diffusers import DiffusionPipeline, ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL
from diffusers.utils import load_image
from PIL import Image
import os, time
import torch
import numpy as np
import cv2
import gradio as gr
from torchvision import transforms 

lora_path = "/home/user/lora"
output_dir = "/home/user/output"

model = DiffusionPipeline.from_pretrained(
    "SimianLuo/LCM_Dreamshaper_v7",
    safety_checker = None,
)
model.to(torch_device="cpu", torch_dtype=torch.float32).to("mps")

low_threshold = 100
high_threshold = 200

safetensors_files_without_suffix = []

def process(prompt, negative_prompt, width, height, steps, scale, number, loras):
    if len(loras) > 0:
        for item in loras:
            model.load_lora_weights(lora_path, weight_name=f"{item}.safetensors")
        
    for _ in range(number):
        _seed = int.from_bytes(os.urandom(2), "big")
        generator = torch.manual_seed(_seed)
        image = model(
            prompt=prompt, 
            negative_prompt=negative_prompt,
            width=width, 
            height=height,
            num_inference_steps=steps,
            generator=generator,
            guidance_scale=scale
        ).images[0]

        timestamp = time.strftime("%Y%m%d-%H%M%S")
        output_path = os.path.join(output_dir, f"{timestamp}.png")
        image.save(output_path)

        # images.append(image)
    # return images

block = gr.Blocks().queue()

def update_loras_list():
    filenames = os.listdir(lora_path)
    safetensors_files = [filename for filename in filenames if filename.endswith(".safetensors")]
    safetensors_files_without_suffix = [filename[:-len(".safetensors")] for filename in safetensors_files]
    return gr.Dropdown.update(choices=safetensors_files_without_suffix)

with gr.Blocks(css="style.css") as block:
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="正向提示词", lines=3, value="Warhammer 40k, a new robot in Armor, god of empire detailed face octane painting, Devian art, concept art, The gate of the medieval castle")
            negative_prompt = gr.Textbox(label="反向提示词", value="low quality")
            loras = gr.Dropdown(safetensors_files_without_suffix, multiselect=True, label="Loras", info="增加Loras风格")
            with gr.Row():
                load_loras_btn = gr.Button(value="加载")
                clear_btn = gr.ClearButton(components=[loras], value="清空")
                    
        with gr.Column():
            with gr.Row():
                steps = gr.Slider(label="迭代步数", minimum=1, maximum=50, value=4, step=1)
                scale = gr.Slider(label="引导系数", minimum=1, maximum=50, value=4, step=0.5)
                number = gr.Slider(label="生成数量", minimum=1, maximum=10, value=1, step=1)

            with gr.Row():
                width = gr.Slider(label="宽度", minimum=384, maximum=1024, value=512, step=128)
                height = gr.Slider(label="高度", minimum=384, maximum=1024, value=512, step=128)

            with gr.Row():
                # result_gallery = gr.Gallery(label='输出', show_label=True, elem_id="gallery", columns=[2], height='auto')
                run_button = gr.Button(value="运行", elem_id="blue-button",)

    ips = [prompt, negative_prompt, width, height, steps, scale, number, loras]
    run_button.click(fn=process, inputs=ips)
    load_loras_btn.click(fn=update_loras_list, outputs=loras)


def main():
    block.launch(debug=True, share=False)

if __name__ == "__main__":
    main()