File size: 4,035 Bytes
87d91a7
2419492
 
 
 
 
 
3651eaa
 
 
 
 
 
2419492
62b2a7f
0af1d8d
 
 
 
 
 
 
 
 
 
 
 
 
 
c1071da
510810d
089e6cb
3651eaa
7f39ca4
 
 
0699667
7f39ca4
b571647
0699667
83f75b0
672cfcb
c058625
7f39ca4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c058625
 
7f39ca4
83f75b0
7f39ca4
c058625
8f2831f
 
 
8d07e5f
1069a6c
8f2831f
 
 
 
 
 
 
 
 
b571647
7d3dd84
c058625
a613ef1
 
 
81e01f7
 
a613ef1
 
 
0699667
 
c058625
 
3651eaa
c058625
 
7f39ca4
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, 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 = "extra digit, fewer digits, cropped, worst quality, low quality, glitch, deformed, mutated, ugly, disfigured"

    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="Model to use", placeholder="username/my_model")
        image_in = gr.Image(source="upload", type="filepath")
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(label="Prompt")
                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)  
                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, preprocessor, controlnet_conditioning_scale, guidance_scale, seed],
        outputs = [result]
    )

demo.queue().launch()