File size: 4,709 Bytes
3542be4
 
 
d472855
3542be4
 
 
 
c689a76
0c23f7f
27419c1
3542be4
953a099
402fe71
 
 
 
 
 
 
 
 
 
 
 
 
 
2bd18a2
953a099
 
c7df0ab
953a099
 
c7df0ab
953a099
6d6bac1
953a099
 
 
 
402fe71
 
 
 
 
 
 
 
 
0c23f7f
 
 
 
 
3515ae5
402fe71
 
 
 
 
 
 
d472855
402fe71
d472855
402fe71
 
 
 
 
 
 
 
a50b44f
 
402fe71
6dbbd62
 
06642da
4948a0e
6dbbd62
 
776de3e
 
6dbbd62
 
 
27419c1
 
 
402fe71
2b32e3d
 
 
 
 
 
 
 
 
 
5826348
4469e93
 
402fe71
72dca33
402fe71
72dca33
402fe71
72dca33
2b32e3d
72dca33
2b32e3d
f091221
 
 
6dbbd62
f091221
 
 
 
2b32e3d
2098e9c
2b32e3d
 
 
402fe71
 
 
 
 
 
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
122
123
124
125
126
127
128
129
130
131
132
133
import spaces
import gradio as gr
import torch
from diffusers import ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline, ControlNetModel, AutoencoderKL
from PIL import Image
import os
import time

from utils.utils import load_cn_model, load_cn_config, load_tagger_model, load_lora_model, resize_image_aspect_ratio, base_generation
from utils.prompt_utils import execute_prompt, remove_color, remove_duplicates
from utils.tagger import modelLoad, analysis



path = os.getcwd()
cn_dir = f"{path}/controlnet"
tagger_dir = f"{path}/tagger"
lora_dir = f"{path}/lora"
os.makedirs(cn_dir, exist_ok=True)
os.makedirs(tagger_dir, exist_ok=True)
os.makedirs(lora_dir, exist_ok=True)

load_cn_model(cn_dir)
load_cn_config(cn_dir)
load_tagger_model(tagger_dir)
load_lora_model(lora_dir)

def load_model(lora_dir, cn_dir):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16
    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
    controlnet = ControlNetModel.from_pretrained(cn_dir, torch_dtype=dtype, use_safetensors=True)
    pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
        "cagliostrolab/animagine-xl-3.1", controlnet=controlnet, vae=vae, torch_dtype=torch.float16
    )
    pipe.load_lora_weights(lora_dir, weight_name="style-lineart_02.safetensors")
    pipe = pipe.to(device)
    return pipe


@spaces.GPU
def predict(input_image_path, prompt, negative_prompt, controlnet_scale):
    pipe = load_model(lora_dir, cn_dir) 
    input_image_pil = Image.open(input_image_path)
    base_size = input_image_pil.size
    resize_image = resize_image_aspect_ratio(input_image_pil)
    white_base_pil = base_generation(resize_image.size, (255, 255, 255, 255)).convert("RGB")
    generator = torch.manual_seed(0)
    last_time = time.time()
    prompt = "masterpiece, best quality, monochrome, lineart, white background, " + prompt
    execute_tags = ["sketch", "transparent background"]
    prompt = execute_prompt(execute_tags, prompt)
    prompt = remove_duplicates(prompt)        
    prompt = remove_color(prompt)
    print(prompt)

    output_image = pipe(
        image=white_base_pil,
        control_image=resize_image,
        strength=1.0,
        prompt=prompt,
        negative_prompt = negative_prompt,
        controlnet_conditioning_scale=float(controlnet_scale),
        generator=generator,
        num_inference_steps=30,
        eta=1.0,
    ).images[0]
    print(f"Time taken: {time.time() - last_time}")
    output_image = output_image.resize(base_size, Image.LANCZOS)
    return output_image



class Img2Img:
    def __init__(self):
        self.demo = self.layout()
        self.post_filter = True
        self.tagger_model = None
        self.input_image_path = None

    def process_prompt_analysis(self, input_image_path):
        if self.tagger_model is None:
            self.tagger_model = modelLoad(tagger_dir)
        tags = analysis(input_image_path, tagger_dir, self.tagger_model)
        tags_list = tags      
        if self.post_filter:
            tags_list = remove_color(tags)
        return tags_list


    def layout(self):
        css = """
        #intro{
            max-width: 32rem;
            text-align: center;
            margin: 0 auto;
        }
        """
        with gr.Blocks(css=css) as demo:
            with gr.Row():
                with gr.Column():
                    self.input_image_path = gr.Image(label="input_image", type='filepath')
                    self.prompt = gr.Textbox(label="prompt", lines=3)
                    self.negative_prompt = gr.Textbox(label="negative_prompt", lines=3, value="lowres, error, extra digit, fewer digits, cropped, worst quality,low quality, normal quality, jpeg artifacts, blurry")

                    prompt_analysis_button = gr.Button("prompt_analysis")

                    self.controlnet_scale = gr.Slider(minimum=0.5, maximum=1.25, value=1.0, step=0.01, label="controlnet_scale")
                    
                    generate_button = gr.Button("generate")
                with gr.Column():
                    self.output_image = gr.Image(type="pil", label="output_image")


            prompt_analysis_button.click(
                        self.process_prompt_analysis,
                        inputs=[self.input_image_path],
                        outputs=self.prompt
            )


            generate_button.click(
                fn=predict,
                inputs=[self.input_image_path, self.prompt, self.negative_prompt, self.controlnet_scale],
                outputs=self.output_image
            )
        return demo



img2img = Img2Img()
img2img.demo.launch(share=True)