File size: 4,231 Bytes
87d91a7
2419492
 
 
 
 
 
3651eaa
 
 
 
 
 
2419492
62b2a7f
0af1d8d
 
 
 
 
 
 
 
 
 
 
 
 
 
c1071da
510810d
089e6cb
3651eaa
093e5a8
7f39ca4
 
0699667
7f39ca4
b571647
0699667
83f75b0
093e5a8
c058625
7f39ca4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c058625
 
7f39ca4
83f75b0
7f39ca4
c058625
8f2831f
 
 
8d07e5f
1069a6c
8f2831f
 
 
 
 
 
 
 
 
b571647
bf07a86
c058625
a613ef1
 
 
093e5a8
81e01f7
 
093e5a8
a613ef1
bf07a86
0699667
c058625
 
3651eaa
c058625
 
093e5a8
c058625
 
3651eaa
c058625
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
import gradio as gr
from huggingface_hub import login
import os

hf_token = os.environ.get("HF_TOKEN")
login(token=hf_token)

from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL
from diffusers.utils import load_image
from PIL import Image
import torch
import numpy as np
import cv2

#vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
controlnet = ControlNetModel.from_pretrained(
    "diffusers/controlnet-canny-sdxl-1.0",
    torch_dtype=torch.float16
)

pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    controlnet=controlnet,
    #vae=vae,
    torch_dtype=torch.float16, 
    variant="fp16",
    use_safetensors=True
)
pipe.to("cuda")
generator = torch.Generator(device="cuda")

#pipe.enable_model_cpu_offload()

def infer(use_custom_model, model_name, image_in, prompt, negative_prompt, preprocessor, controlnet_conditioning_scale, guidance_scale, seed):
    if use_custom_model:
        custom_model = model_name

        # This is where you load your trained weights
        pipe.load_lora_weights(custom_model, use_auth_token=True)
    
    prompt = prompt
    negative_prompt = negative_prompt

    if preprocessor == "canny":

        image = load_image(image_in)

        image = np.array(image)
        image = cv2.Canny(image, 100, 200)
        image = image[:, :, None]
        image = np.concatenate([image, image, image], axis=2)
        image = Image.fromarray(image)

    if use_custom_model:
        lora_scale= 0.9

        images = pipe(
            prompt, 
            negative_prompt=negative_prompt, 
            image=image, 
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guidance_scale = guidance_scale,
            num_inference_steps=50,
            generator=generator.manual_seed(seed),
            cross_attention_kwargs={"scale": lora_scale}
        ).images
    else:
        images = pipe(
            prompt, 
            negative_prompt=negative_prompt, 
            image=image, 
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guidance_scale = guidance_scale,
            num_inference_steps=50,
            generator=generator.manual_seed(seed),
        ).images

    images[0].save(f"result.png")

    return f"result.png"

css="""
#col-container{
    margin: 0 auto;
    max-width: 680px;
    text-align: left;
}
"""
with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("""
# SD-XL Control LoRas
Use StableDiffusion XL with ControlNet pretrained LoRas 

        """)
        use_custom_model = gr.Checkbox(label="Use a public custom model ?(optional)", value=False, info="To use a private model, you'll prefer to duplicate the space with your own access token.")
        model_name = gr.Textbox(label="Custom Model to use", placeholder="username/my_custom_public_model")
        image_in = gr.Image(source="upload", type="filepath")
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(label="Prompt")
                negative_prompt = gr.Textbox(label="Negative prompt", value="extra digit, fewer digits, cropped, worst quality, low quality, glitch, deformed, mutated, ugly, disfigured")
                guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=7.5)
            with gr.Column():
                preprocessor = gr.Dropdown(label="Preprocessor", choices=["canny"], value="canny", interactive=False, info="For the moment, only canny is available")  
                controlnet_conditioning_scale = gr.Slider(label="Controlnet conditioning Scale", minimum=0.1, maximum=0.9, step=0.01, value=0.5, type="float")
                seed = gr.Slider(label="seed", minimum=0, maximum=500000, step=1, value=42)

        submit_btn = gr.Button("Submit")
        result = gr.Image(label="Result")

    submit_btn.click(
        fn = infer,
        inputs = [use_custom_model, model_name, image_in, prompt, negative_prompt, preprocessor, controlnet_conditioning_scale, guidance_scale, seed],
        outputs = [result]
    )

demo.queue().launch()