File size: 14,109 Bytes
87d91a7
4c31ad4
2419492
4c31ad4
 
a066e08
2419492
067da29
 
2419492
 
 
4c31ad4
 
 
a7f51c5
 
3651eaa
 
 
 
 
 
2419492
be9073e
331eca2
0af1d8d
 
 
 
 
a045576
 
 
 
 
 
4c31ad4
 
 
 
 
 
 
 
 
 
 
 
331eca2
4c31ad4
510810d
4c31ad4
3651eaa
4c31ad4
 
 
0699667
4c31ad4
 
 
 
c058625
4c31ad4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2ca2323
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a4f5d8c
4c31ad4
 
 
 
 
 
 
 
 
 
 
 
238d2f7
 
4c31ad4
 
 
 
a7f51c5
067da29
 
 
331eca2
2ca2323
 
7f39ca4
 
 
 
 
 
 
 
 
331eca2
 
067da29
66af8e6
067da29
 
331eca2
 
 
4c31ad4
 
 
 
 
 
 
 
 
 
 
 
 
 
331eca2
a231571
7f39ca4
 
 
 
 
00f4438
 
 
331eca2
7f39ca4
 
 
 
 
 
 
00f4438
 
 
331eca2
c058625
 
7f39ca4
83f75b0
4c31ad4
c058625
8f2831f
 
 
c1e58aa
1069a6c
8f2831f
067da29
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4c31ad4
 
 
 
 
 
8f2831f
067da29
8f2831f
 
067da29
 
 
 
3073f7f
067da29
8f7b910
067da29
 
 
 
 
 
 
 
 
331eca2
35f5939
331eca2
8f2831f
 
038916c
a231571
4c31ad4
a045576
4c31ad4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a045576
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c058625
4c31ad4
c058625
a045576
 
 
 
 
 
 
4c31ad4
 
 
 
 
 
 
 
 
 
 
 
c058625
 
4c31ad4
 
c058625
3651eaa
238d2f7
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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
import gradio as gr
from huggingface_hub import login, HfFileSystem, HfApi, ModelCard
import os
import spaces
import random
import torch

is_shared_ui = True if "fffiloni/sdxl-control-loras" in os.environ['SPACE_ID'] else False

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

fs = HfFileSystem(token=hf_token)
api = HfApi()

device="cuda" if torch.cuda.is_available() else "cpu"

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
)

def check_use_custom_or_no(value):
    if value is True:
        return gr.update(visible=True)
    else:
        return gr.update(visible=False)

def get_files(file_paths):
    last_files = {}  # Dictionary to store the last file for each path

    for file_path in file_paths:
        # Split the file path into directory and file components
        directory, file_name = file_path.rsplit('/', 1)
    
        # Update the last file for the current path
        last_files[directory] = file_name
    
    # Extract the last files from the dictionary
    result = list(last_files.values())

    return result

def load_model(model_name):

    if model_name == "":
        gr.Warning("If you want to use a private model, you need to duplicate this space on your personal account.")
        raise gr.Error("You forgot to define Model ID.")

    # Get instance_prompt a.k.a trigger word
    card = ModelCard.load(model_name)
    repo_data = card.data.to_dict()
    instance_prompt = repo_data.get("instance_prompt")

    if instance_prompt is not None:
        print(f"Trigger word: {instance_prompt}")
    else:
        instance_prompt = "no trigger word needed"
        print(f"Trigger word: no trigger word needed")

    # List all ".safetensors" files in repo
    sfts_available_files = fs.glob(f"{model_name}/*safetensors")
    sfts_available_files = get_files(sfts_available_files)

    if sfts_available_files == []:
        sfts_available_files = ["NO SAFETENSORS FILE"]

    print(f"Safetensors available: {sfts_available_files}")

    return model_name, "Model Ready", gr.update(choices=sfts_available_files, value=sfts_available_files[0], visible=True), gr.update(value=instance_prompt, visible=True)

def custom_model_changed(model_name, previous_model):
    if model_name == "" and previous_model == "" :
        status_message = ""      
    elif model_name != previous_model:
        status_message = "model changed, please reload before any new run"
    else:
        status_message = "model ready"
    return status_message

def resize_image(input_path, output_path, target_height):
    # Open the input image
    img = Image.open(input_path)

    # Calculate the aspect ratio of the original image
    original_width, original_height = img.size
    original_aspect_ratio = original_width / original_height

    # Calculate the new width while maintaining the aspect ratio and the target height
    new_width = int(target_height * original_aspect_ratio)

    # Resize the image while maintaining the aspect ratio and fixing the height
    img = img.resize((new_width, target_height), Image.LANCZOS)

    # Save the resized image
    img.save(output_path)

    return output_path

def infer(use_custom_model, model_name, weight_name, custom_lora_weight, image_in, prompt, negative_prompt, preprocessor, controlnet_conditioning_scale, guidance_scale, inf_steps, seed, progress=gr.Progress(track_tqdm=True)):

    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(device)
    
    prompt = prompt
    negative_prompt = negative_prompt

    if seed < 0 :
        seed = random.randint(0, 423538377342)
    
    generator = torch.Generator(device=device).manual_seed(seed)

    if image_in == None:
        raise gr.Error("You forgot to upload a source image.")
    
    image_in = resize_image(image_in, "resized_input.jpg", 1024)
    
    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:
        
        if model_name == "":
            raise gr.Error("you forgot to set a custom model name.")
        
        custom_model = model_name

        # This is where you load your trained weights
        if weight_name == "NO SAFETENSORS FILE": 
            pipe.load_lora_weights(
                custom_model,     
                low_cpu_mem_usage = True,
                use_auth_token = True
            )
    
        else:
            pipe.load_lora_weights(
                custom_model,
                weight_name = weight_name,        
                low_cpu_mem_usage = True,
                use_auth_token = True
            )
    
        lora_scale=custom_lora_weight

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

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

    return f"result.png", seed

css="""
#col-container{
    margin: 0 auto;
    max-width: 720px;
    text-align: left;
}
div#warning-duplicate {
    background-color: #ebf5ff;
    padding: 0 10px 5px;
    margin: 20px 0;
}
div#warning-duplicate > .gr-prose > h2, div#warning-duplicate > .gr-prose > p {
    color: #0f4592!important;
}
div#warning-duplicate strong {
    color: #0f4592;
}
p.actions {
    display: flex;
    align-items: center;
    margin: 20px 0;
}
div#warning-duplicate .actions a {
    display: inline-block;
    margin-right: 10px;
}
button#load_model_btn{
    height: 46px;
}
#status_info{
    font-size: 0.9em;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        if is_shared_ui:
            top_description = gr.HTML(f'''
                <div class="gr-prose">
                    <h2><svg xmlns="http://www.w3.org/2000/svg" width="18px" height="18px" style="margin-right: 0px;display: inline-block;"fill="none"><path fill="#fff" d="M7 13.2a6.3 6.3 0 0 0 4.4-10.7A6.3 6.3 0 0 0 .6 6.9 6.3 6.3 0 0 0 7 13.2Z"/><path fill="#fff" fill-rule="evenodd" d="M7 0a6.9 6.9 0 0 1 4.8 11.8A6.9 6.9 0 0 1 0 7 6.9 6.9 0 0 1 7 0Zm0 0v.7V0ZM0 7h.6H0Zm7 6.8v-.6.6ZM13.7 7h-.6.6ZM9.1 1.7c-.7-.3-1.4-.4-2.2-.4a5.6 5.6 0 0 0-4 1.6 5.6 5.6 0 0 0-1.6 4 5.6 5.6 0 0 0 1.6 4 5.6 5.6 0 0 0 4 1.7 5.6 5.6 0 0 0 4-1.7 5.6 5.6 0 0 0 1.7-4 5.6 5.6 0 0 0-1.7-4c-.5-.5-1.1-.9-1.8-1.2Z" clip-rule="evenodd"/><path fill="#000" fill-rule="evenodd" d="M7 2.9a.8.8 0 1 1 0 1.5A.8.8 0 0 1 7 3ZM5.8 5.7c0-.4.3-.6.6-.6h.7c.3 0 .6.2.6.6v3.7h.5a.6.6 0 0 1 0 1.3H6a.6.6 0 0 1 0-1.3h.4v-3a.6.6 0 0 1-.6-.7Z" clip-rule="evenodd"/></svg>
                    Note: you might want to use a <strong>private</strong> custom LoRa model</h2>
                    <p class="main-message">
                        To do so, <strong>duplicate the Space</strong> and run it on your own profile using <strong>your own access token</strong> and eventually a GPU (T4-small or A10G-small) for faster inference without waiting in the queue.<br />
                    </p>
                    <p class="actions">
                        <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}?duplicate=true">
                            <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-lg-dark.svg" alt="Duplicate this Space" />
                        </a>
                        to start using private models and skip the queue
                    </p>
                </div>
            ''', elem_id="warning-duplicate")
        gr.HTML("""
<h2 style="text-align: center;">SD-XL Control LoRas</h2>
<p style="text-align: center;">Use StableDiffusion XL with <a href="https://huggingface.co/collections/diffusers/sdxl-controlnets-64f9c35846f3f06f5abe351f">Diffusers' SDXL ControlNets</a></p>

        """)

        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.")
        
        with gr.Box(visible=False) as custom_model_box:
            with gr.Row():
                with gr.Column():
                    if not is_shared_ui:
                        your_username = api.whoami()["name"]
                        my_models = api.list_models(author=your_username, filter=["diffusers", "stable-diffusion-xl", 'lora'])
                        model_names = [item.modelId for item in my_models]
    
                    if not is_shared_ui:
                        custom_model = gr.Dropdown(
                            label = "Your custom model ID",
                            info="You can pick one of your private models",
                            choices = model_names,
                            allow_custom_value = True
                            #placeholder = "username/model_id"
                        )
                    else:
                        custom_model = gr.Textbox(
                            label="Your custom model ID", 
                            placeholder="your_username/your_trained_model_name", 
                            info="Make sure your model is set to PUBLIC"
                        )
                    
                    weight_name = gr.Dropdown(
                        label="Safetensors file", 
                        #value="pytorch_lora_weights.safetensors", 
                        info="specify which one if model has several .safetensors files",
                        allow_custom_value=True,
                        visible = False
                    )
                with gr.Column():
                    with gr.Group():
                        load_model_btn = gr.Button("Load my model", elem_id="load_model_btn")
                        previous_model = gr.Textbox(
                            visible = False
                        )
                        model_status = gr.Textbox(
                            label = "model status",
                            show_label = False,
                            elem_id = "status_info"
                        )
                    trigger_word = gr.Textbox(label="Trigger word", interactive=False, visible=False)
        
        image_in = gr.Image(source="upload", type="filepath")
        
        with gr.Row():
            
            with gr.Column():
                with gr.Group():
                    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")
                with gr.Group():
                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=7.5)
                    inf_steps = gr.Slider(label="Inference Steps", minimum="25", maximum="50", step=1, value=25)
                    custom_lora_weight = gr.Slider(label="Custom model weights", minimum=0.1, maximum=0.9, step=0.1, value=0.9)

            with gr.Column():
                with gr.Group():
                    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)
                with gr.Group():
                    seed = gr.Slider(
                        label="Seed",
                        info = "-1 denotes a random seed",
                        minimum=-1,
                        maximum=423538377342,
                        step=1,
                        value=-1
                    )
                    last_used_seed = gr.Number(
                        label = "Last used seed",
                        info = "the seed used in the last generation",
                    )
        
        
        
        submit_btn = gr.Button("Submit")
        
        result = gr.Image(label="Result")
    
    use_custom_model.change(
        fn = check_custom_or_no,
        inputs =[use_custom_model],
        outputs = [custom_model_box],
        queue = False
    )
    custom_model.blur(
        fn=custom_model_changed,
        inputs = [custom_model, previous_model],
        outputs = [model_status],
        queue = False
    )
    load_model_btn.click(
        fn = load_model,
        inputs=[custom_model],
        outputs = [previous_model, model_status, weight_name, trigger_word],
        queue = False
    )
    submit_btn.click(
        fn = infer,
        inputs = [use_custom_model, custom_model, weight_name, custom_lora_weight, image_in, prompt, negative_prompt, preprocessor, controlnet_conditioning_scale, guidance_scale, inf_steps, seed],
        outputs = [result, last_used_seed]
    )

demo.queue(max_size=12).launch()