import sys
from starline import process

from utils import load_cn_model, load_cn_config, randomname
from convertor import pil2cv, cv2pil

from sd_model import get_cn_pipeline, generate, get_cn_detector
import cv2
import os
import numpy as np
from PIL import Image


path = os.getcwd()
output_dir = f"{path}/output"
input_dir = f"{path}/input"
cn_lineart_dir = f"{path}/controlnet/lineart"

load_cn_model(cn_lineart_dir)
load_cn_config(cn_lineart_dir)

pos_prompt="brown hair, brown eyes, white T-shirt"
neg_prompt=""
input_image=Image.open("/home/alfredplpl/Downloads/IMG_0617.PNG")

org_line_image=input_image
image = pil2cv(input_image)
image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA)
alpha_th=100

index = np.where(image[:, :, 3] == 0)
image[index] = [255, 255, 255, 255]
input_image = cv2pil(image)

pipe = get_cn_pipeline()
detectors = get_cn_detector(input_image.resize((1024, 1024), Image.Resampling.BICUBIC))


gen_image = generate(pipe, detectors, pos_prompt, neg_prompt)
output = process(gen_image.resize((image.shape[1], image.shape[0]), Image.Resampling.BICUBIC) , org_line_image, alpha_th)

output = output.resize((image.shape[1], image.shape[0]) , Image.Resampling.BICUBIC)


output = Image.alpha_composite(output, org_line_image)
name = randomname(10)
output.save(f"{output_dir}/output_{name}.png")
#output = pil2cv(output)
file_name = f"{output_dir}/output_{name}.png"