import gradio as gr
import sys
from utils import load_lora_model
from bg_alpha import adjust_transparency
from lineart import get_pipe
import os
from PIL import Image
import spaces

path = os.getcwd()
output_dir = f"{path}/output"
lora_dir = f"{path}/models/lora"

load_lora_model(lora_dir)

pipe = get_pipe(lora_dir)

@spaces.GPU()
def generate(prompt, negative_prompt):
    default_pos = "((white background)), lineart, <lora:sdxl_BWLine:1.0>, monochrome, "
    default_neg = ""
    prompt = default_pos + prompt 
    negative_prompt = default_neg + negative_prompt

    width, height = 1024, 1024
    color = (255, 255, 255)
    white_bg =  Image.new("RGB", (width, height), color)

    image = pipe(
                prompt=prompt,
                negative_prompt = negative_prompt,
                image=[white_bg],
                num_inference_steps=50,
                controlnet_conditioning_scale=[0.1]
            ).images[0]
    
    return image


class webui:
    def __init__(self):
        self.demo = gr.Blocks()

    def process(self, pos_prompt, neg_prompt):
        image = generate(pos_prompt, neg_prompt)
        image = adjust_transparency(image)
        return [image]

    def launch(self, share):
        with self.demo:
            with gr.Column():
                pos_prompt = gr.Textbox(value="1girl, cute, kawaii, medium breasts, medium hair, smile, mini skirt, best quality, very aesthetic,", max_lines=1000, label="Positive prompt")                    
                neg_prompt = gr.Textbox(value="bold line, multiple people,", max_lines=1000, label="Negative prompt")

                submit = gr.Button(value="Start", variant="primary")
                
                output_0 = gr.Image(label="Output", format="png")
                
            submit.click(
                self.process, 
                inputs=[pos_prompt, neg_prompt], #[input_image, pos_prompt, neg_prompt, alpha_th, thickness, reference_image], 
                outputs=[output_0]
            )

        self.demo.queue()
        self.demo.launch(share=share)


if __name__ == "__main__":
    ui = webui()
    if len(sys.argv) > 1:
        if sys.argv[1] == "share":
            ui.launch(share=True)
        else:
            ui.launch(share=False)
    else:
        ui.launch(share=False)