Hugging Face's logo

Spaces:
mattyamonaca
/
starline
Runtime error

App Files Community
starline
File size: 4,979 Bytes 
e3dd038
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
134c8c2
e3dd038
 
 
 
 
 
 
 
134c8c2
e3dd038
 
 
134c8c2
e3dd038
 
 
 
 
 
134c8c2
 
e3dd038
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
134c8c2
e3dd038
 
 
 
 
134c8c2
 
 
e3dd038
 
 
 
 
 
 
 
 
 
64a6d23
e3dd038
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
 
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
import zipfile

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)


def zip_png_files(folder_path):
    # Zipファイルの名前を設定（フォルダ名と同じにします）
    zip_path = os.path.join(folder_path, 'output.zip')
    
    # zipfileオブジェクトを作成し、書き込みモードで開く
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        # フォルダ内のすべてのファイルをループ処理
        for foldername, subfolders, filenames in os.walk(folder_path):
            for filename in filenames:
                # PNGファイルのみを対象にする
                if filename.endswith('.png'):
                    # ファイルのフルパスを取得
                    file_path = os.path.join(foldername, filename)
                    # zipファイルに追加
                    zipf.write(file_path, arcname=os.path.relpath(file_path, folder_path))


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

    def undercoat(self, input_image, pos_prompt, neg_prompt, alpha_th, thickness, reference_flg, reference_img):
        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(reference_flg)
        detectors = get_cn_detector(input_image.resize((1024, 1024), Image.ANTIALIAS))
        

        gen_image = generate(pipe, detectors, pos_prompt, neg_prompt, reference_flg, reference_img)
        color_img, unfinished = process(gen_image.resize((image.shape[1], image.shape[0]), Image.ANTIALIAS) , org_line_image, alpha_th, thickness)
        #color_img = color_img.resize((image.shape[1], image.shape[0]) , Image.ANTIALIAS)


        output_img = Image.alpha_composite(color_img, org_line_image)
        name = randomname(10)
        if not os.path.exists(f"{output_dir}"):
            os.makedirs(f"{output_dir}")
        os.makedirs(f"{output_dir}/{name}")
        output_img.save(f"{output_dir}/{name}/output_image.png")
        org_line_image.save(f"{output_dir}/{name}/line_image.png")
        color_img.save(f"{output_dir}/{name}/color_image.png")
        unfinished.save(f"{output_dir}/{name}/unfinished_image.png")

        outputs = [output_img, org_line_image, color_img, unfinished]
        zip_png_files(f"{output_dir}/{name}")
        filename = f"{output_dir}/{name}/output.zip"

        return outputs, filename



    def launch(self, share):
        with self.demo:
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(type="pil", image_mode="RGBA", label="lineart")
                    pos_prompt = gr.Textbox(value="1girl, blue hair, pink shirts, bestquality, 4K", max_lines=1000, label="positive prompt")                    
                    neg_prompt = gr.Textbox(value=" (worst quality, low quality:1.2), (lowres:1.2), (bad anatomy:1.2), (greyscale, monochrome:1.4)", max_lines=1000, label="negative prompt")

                    alpha_th = gr.Slider(maximum = 255, value=100, label = "alpha threshold")
                    thickness = gr.Number(value=5, label="Thickness of correction area (Odd numbers need to be entered)")
                
                    reference_image = gr.Image(type="pil", image_mode="RGB", label="reference_image")
                    reference_flg = gr.Checkbox(value=True, label="reference_flg")
                    #gr.Slider(maximum = 21, value=3, step=2, label = "Thickness of correction area")

                    submit = gr.Button(value="Start")
                with gr.Row():
                    with gr.Column():
                        with gr.Tab("output"):
                            output_0 = gr.Gallery(format="png")
                        output_file = gr.File()
            submit.click(
                self.undercoat, 
                inputs=[input_image, pos_prompt, neg_prompt, alpha_th, thickness, reference_flg, reference_image], 
                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)