Hugging Face's logo

Spaces:
OzzyGT
/
diffusers-recolor
Running on Zero

App Files Community
diffusers-recolor
File size: 3,966 Bytes 
58f8532
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
444a999
58f8532
 
 
 
 
 
 
 
 
 
 
 
 
444a999
58f8532
 
 
 
 
 
 
 
 
 
 
 
 
d1d527c
58f8532
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
444a999
58f8532
 
 
 
 
 
 
 
 
 
 
 
 
 
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
134
135
136
137
138
139
140
141
142
143
144
145
 
import gradio as gr
import spaces
import torch
from diffusers import AutoencoderKL, ControlNetModel, TCDScheduler
from gradio_imageslider import ImageSlider
from image_gen_aux import LineArtPreprocessor
from PIL import Image, ImageEnhance

from controlnet_union import ControlNetModel_Union
from pipeline_sdxl_recolor import StableDiffusionXLRecolorPipeline

lineart_preprocessor = LineArtPreprocessor.from_pretrained("OzzyGT/lineart").to("cuda")

controlnet = [
    ControlNetModel.from_pretrained(
        "OzzyGT/ControlNet-recolorXL", torch_dtype=torch.float16, variant="fp16"
    ),
    ControlNetModel_Union.from_pretrained(
        "OzzyGT/controlnet-union-promax-sdxl-1.0",
        torch_dtype=torch.float16,
        variant="fp16",
    ),
]

vae = AutoencoderKL.from_pretrained(
    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
).to("cuda")

pipe = StableDiffusionXLRecolorPipeline.from_pretrained(
    "recoilme/ColorfulXL-Lightning",
    torch_dtype=torch.float16,
    vae=vae,
    controlnet=controlnet,
    variant="fp16",
).to("cuda")

pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)

pipe.load_ip_adapter(
    "h94/IP-Adapter",
    subfolder="sdxl_models",
    weight_name="ip-adapter_sdxl_vit-h.safetensors",
    image_encoder_folder="models/image_encoder",
)

scale = {
    "up": {"block_0": [1.0, 0.0, 1.0]},
}
pipe.set_ip_adapter_scale(scale)

prompt = "high quality color photo, sharp, detailed, 4k, colorized, remastered"
negative_prompt = "blurry, low resolution, bad quality, pixelated, black and white, b&w, grayscale, monochrome, sepia"

(
    prompt_embeds,
    negative_prompt_embeds,
    pooled_prompt_embeds,
    negative_pooled_prompt_embeds,
) = pipe.encode_prompt(prompt, negative_prompt, "cuda", True)


@spaces.GPU(duration=16)
def recolor_image(image):
    source_image = image["background"]

    lineart_image = lineart_preprocessor(source_image, resolution_scale=0.7)[0]

    for image in pipe(
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        pooled_prompt_embeds=pooled_prompt_embeds,
        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
        image=[source_image, lineart_image],
        ip_adapter_image=source_image,
        num_inference_steps=8,
        guidance_scale=2.0,
        controlnet_conditioning_scale=[1.0, 0.5],
        control_guidance_end=[1.0, 0.9],
    ):
        yield source_image, image

    image = image.convert("RGBA")
    source_image = source_image.convert("RGBA")

    enhancer = ImageEnhance.Color(image)
    image = enhancer.enhance(4.0)

    alpha = image.split()[3]
    alpha = alpha.point(lambda p: p * 0.20)
    image.putalpha(alpha)

    merged_image = Image.alpha_composite(source_image, image)

    yield source_image, merged_image


def clear_result():
    return gr.update(value=None)


css = """
.gradio-container {
    width: 1024px !important;
}
"""


title = """<h1 align="center">Diffusers Image Recolor</h1>
<div align="center">Upload a grayscale image to colorize it.</div>
<div align="center">This space is a PoC made for the guide <a href='https://huggingface.co/blog/OzzyGT/diffusers-recolor'>Recoloring photos with diffusers</a>.</div>
"""

with gr.Blocks(css=css) as demo:
    gr.HTML(title)

    run_button = gr.Button("Generate")

    with gr.Row():
        input_image = gr.ImageEditor(
            type="pil",
            label="Input Image",
            crop_size=(1024, 1024),
            canvas_size=(1024, 1024),
            layers=False,
            eraser=False,
            brush=False,
            sources=["upload"],
            image_mode="RGB",
        )

        result = ImageSlider(interactive=False, label="Generated Image", type="pil")

    run_button.click(
        fn=clear_result,
        inputs=None,
        outputs=result,
    ).then(
        fn=recolor_image,
        inputs=[input_image],
        outputs=result,
    )


demo.launch(share=False)