|
import gradio as gr
|
|
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)
|
|
|
|
class webui:
|
|
def __init__(self):
|
|
self.demo = gr.Blocks()
|
|
|
|
def undercoat(self, input_image, pos_prompt, neg_prompt, alpha_th):
|
|
org_line_image = input_image
|
|
image = pil2cv(input_image)
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA)
|
|
|
|
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.ANTIALIAS))
|
|
|
|
|
|
gen_image = generate(pipe, detectors, pos_prompt, neg_prompt)
|
|
output = process(gen_image.resize((image.shape[1], image.shape[0]), Image.ANTIALIAS) , org_line_image, alpha_th)
|
|
|
|
output = output.resize((image.shape[1], image.shape[0]) , Image.ANTIALIAS)
|
|
|
|
|
|
output = Image.alpha_composite(output, org_line_image)
|
|
name = randomname(10)
|
|
output.save(f"{output_dir}/output_{name}.png")
|
|
|
|
file_name = f"{output_dir}/output_{name}.png"
|
|
|
|
return output, file_name
|
|
|
|
|
|
|
|
def launch(self, share):
|
|
with self.demo:
|
|
with gr.Row():
|
|
with gr.Column():
|
|
input_image = gr.Image(type="pil", image_mode="RGBA")
|
|
|
|
pos_prompt = gr.Textbox(max_lines=1000, label="positive prompt")
|
|
neg_prompt = gr.Textbox(max_lines=1000, label="negative prompt")
|
|
|
|
alpha_th = gr.Slider(maximum = 255, value=100, label = "alpha threshold")
|
|
|
|
submit = gr.Button(value="Start")
|
|
with gr.Row():
|
|
with gr.Column():
|
|
with gr.Tab("output"):
|
|
output_0 = gr.Image()
|
|
|
|
output_file = gr.File()
|
|
submit.click(
|
|
self.undercoat,
|
|
inputs=[input_image, pos_prompt, neg_prompt, alpha_th],
|
|
outputs=[output_0, output_file]
|
|
)
|
|
|
|
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)
|
|
|
