kadirnar commited on
Commit
82fe494
1 Parent(s): 239ea0c

Upload 8 files

Browse files
app.py CHANGED
@@ -1,8 +1,20 @@
 
1
  from utils.image2image import stable_diffusion_img2img
2
  from utils.text2image import stable_diffusion_text2img
3
  from utils.inpaint import stable_diffusion_inpaint
 
 
 
 
 
 
 
 
 
 
4
  import gradio as gr
5
 
 
6
  stable_model_list = [
7
  "runwayml/stable-diffusion-v1-5",
8
  "stabilityai/stable-diffusion-2",
@@ -27,7 +39,7 @@ stable_negative_prompt_list = [
27
  ]
28
  app = gr.Blocks()
29
  with app:
30
- gr.Markdown("# **<h2 align='center'>Stable Diffusion + ControlNet WebUI<h2>**")
31
  gr.Markdown(
32
  """
33
  <h5 style='text-align: center'>
@@ -178,6 +190,288 @@ with app:
178
 
179
  inpaint_predict = gr.Button(value='Generator')
180
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
181
  with gr.Tab('Generator'):
182
  with gr.Column():
183
  output_image = gr.Image(label='Image')
@@ -222,4 +516,91 @@ with app:
222
  outputs = [output_image],
223
  )
224
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
225
  app.launch()
 
1
+
2
  from utils.image2image import stable_diffusion_img2img
3
  from utils.text2image import stable_diffusion_text2img
4
  from utils.inpaint import stable_diffusion_inpaint
5
+
6
+ from controlnet.controlnet_canny import stable_diffusion_controlnet_img2img
7
+ from controlnet.controlnet_depth import stable_diffusion_controlnet_img2img
8
+ from controlnet.controlnet_hed import stable_diffusion_controlnet_img2img
9
+ from controlnet.controlnet_mlsd import stable_diffusion_controlnet_img2img
10
+ from controlnet.controlnet_pose import stable_diffusion_controlnet_img2img
11
+ from controlnet.controlnet_scribble import stable_diffusion_controlnet_img2img
12
+ from controlnet.controlnet_seg import stable_diffusion_controlnet_img2img
13
+
14
+
15
  import gradio as gr
16
 
17
+
18
  stable_model_list = [
19
  "runwayml/stable-diffusion-v1-5",
20
  "stabilityai/stable-diffusion-2",
 
39
  ]
40
  app = gr.Blocks()
41
  with app:
42
+ gr.Markdown("# **<h2 align='center'>Stable Diffusion WebUI<h2>**")
43
  gr.Markdown(
44
  """
45
  <h5 style='text-align: center'>
 
190
 
191
  inpaint_predict = gr.Button(value='Generator')
192
 
193
+
194
+ with gr.Tab('ControlNet'):
195
+ with gr.Tab('Canny'):
196
+ controlnet_image_file = gr.Image(label='Image')
197
+
198
+ controlnet_model_id = gr.Dropdown(
199
+ choices=stable_inpiant_model_list,
200
+ value=stable_inpiant_model_list[0],
201
+ label='Stable Model Id'
202
+ )
203
+
204
+ controlnet_prompt = gr.Textbox(
205
+ lines=1,
206
+ value=stable_prompt_list[0],
207
+ label='Prompt'
208
+ )
209
+
210
+ controlnet_negative_prompt = gr.Textbox(
211
+ lines=1,
212
+ value=stable_negative_prompt_list[0],
213
+ label='Negative Prompt'
214
+ )
215
+
216
+ with gr.Accordion("Advanced Options", open=False):
217
+ controlnet_guidance_scale = gr.Slider(
218
+ minimum=0.1,
219
+ maximum=15,
220
+ step=0.1,
221
+ value=7.5,
222
+ label='Guidance Scale'
223
+ )
224
+
225
+ controlnet_num_inference_step = gr.Slider(
226
+ minimum=1,
227
+ maximum=100,
228
+ step=1,
229
+ value=50,
230
+ label='Num Inference Step'
231
+ )
232
+
233
+ controlnet_canny_predict = gr.Button(value='Generator')
234
+
235
+ with gr.Tab('Hed'):
236
+ controlnet_image_file = gr.Image(label='Image')
237
+
238
+ controlnet_model_id = gr.Dropdown(
239
+ choices=stable_inpiant_model_list,
240
+ value=stable_inpiant_model_list[0],
241
+ label='Stable Model Id'
242
+ )
243
+
244
+ controlnet_prompt = gr.Textbox(
245
+ lines=1,
246
+ value=stable_prompt_list[0],
247
+ label='Prompt'
248
+ )
249
+
250
+ controlnet_negative_prompt = gr.Textbox(
251
+ lines=1,
252
+ value=stable_negative_prompt_list[0],
253
+ label='Negative Prompt'
254
+ )
255
+
256
+ with gr.Accordion("Advanced Options", open=False):
257
+ controlnet_guidance_scale = gr.Slider(
258
+ minimum=0.1,
259
+ maximum=15,
260
+ step=0.1,
261
+ value=7.5,
262
+ label='Guidance Scale'
263
+ )
264
+
265
+ controlnet_num_inference_step = gr.Slider(
266
+ minimum=1,
267
+ maximum=100,
268
+ step=1,
269
+ value=50,
270
+ label='Num Inference Step'
271
+ )
272
+
273
+ controlnet_hed_predict = gr.Button(value='Generator')
274
+
275
+ with gr.Tab('MLSD line'):
276
+ controlnet_image_file = gr.Image(label='Image')
277
+
278
+ controlnet_model_id = gr.Dropdown(
279
+ choices=stable_inpiant_model_list,
280
+ value=stable_inpiant_model_list[0],
281
+ label='Stable Model Id'
282
+ )
283
+
284
+ controlnet_prompt = gr.Textbox(
285
+ lines=1,
286
+ value=stable_prompt_list[0],
287
+ label='Prompt'
288
+ )
289
+
290
+ controlnet_negative_prompt = gr.Textbox(
291
+ lines=1,
292
+ value=stable_negative_prompt_list[0],
293
+ label='Negative Prompt'
294
+ )
295
+
296
+ with gr.Accordion("Advanced Options", open=False):
297
+ controlnet_guidance_scale = gr.Slider(
298
+ minimum=0.1,
299
+ maximum=15,
300
+ step=0.1,
301
+ value=7.5,
302
+ label='Guidance Scale'
303
+ )
304
+
305
+ controlnet_num_inference_step = gr.Slider(
306
+ minimum=1,
307
+ maximum=100,
308
+ step=1,
309
+ value=50,
310
+ label='Num Inference Step'
311
+ )
312
+
313
+ controlnet_mlsd_predict = gr.Button(value='Generator')
314
+
315
+ with gr.Tab('Segmentation'):
316
+ controlnet_image_file = gr.Image(label='Image')
317
+
318
+ controlnet_model_id = gr.Dropdown(
319
+ choices=stable_inpiant_model_list,
320
+ value=stable_inpiant_model_list[0],
321
+ label='Stable Model Id'
322
+ )
323
+
324
+ controlnet_prompt = gr.Textbox(
325
+ lines=1,
326
+ value=stable_prompt_list[0],
327
+ label='Prompt'
328
+ )
329
+
330
+ controlnet_negative_prompt = gr.Textbox(
331
+ lines=1,
332
+ value=stable_negative_prompt_list[0],
333
+ label='Negative Prompt'
334
+ )
335
+
336
+ with gr.Accordion("Advanced Options", open=False):
337
+ controlnet_guidance_scale = gr.Slider(
338
+ minimum=0.1,
339
+ maximum=15,
340
+ step=0.1,
341
+ value=7.5,
342
+ label='Guidance Scale'
343
+ )
344
+
345
+ controlnet_num_inference_step = gr.Slider(
346
+ minimum=1,
347
+ maximum=100,
348
+ step=1,
349
+ value=50,
350
+ label='Num Inference Step'
351
+ )
352
+
353
+ controlnet_seg_predict = gr.Button(value='Generator')
354
+
355
+ with gr.Tab('Depth'):
356
+ controlnet_image_file = gr.Image(label='Image')
357
+
358
+ controlnet_model_id = gr.Dropdown(
359
+ choices=stable_inpiant_model_list,
360
+ value=stable_inpiant_model_list[0],
361
+ label='Stable Model Id'
362
+ )
363
+
364
+ controlnet_prompt = gr.Textbox(
365
+ lines=1,
366
+ value=stable_prompt_list[0],
367
+ label='Prompt'
368
+ )
369
+
370
+ controlnet_negative_prompt = gr.Textbox(
371
+ lines=1,
372
+ value=stable_negative_prompt_list[0],
373
+ label='Negative Prompt'
374
+ )
375
+
376
+ with gr.Accordion("Advanced Options", open=False):
377
+ controlnet_guidance_scale = gr.Slider(
378
+ minimum=0.1,
379
+ maximum=15,
380
+ step=0.1,
381
+ value=7.5,
382
+ label='Guidance Scale'
383
+ )
384
+
385
+ controlnet_num_inference_step = gr.Slider(
386
+ minimum=1,
387
+ maximum=100,
388
+ step=1,
389
+ value=50,
390
+ label='Num Inference Step'
391
+ )
392
+
393
+ controlnet_depth_predict = gr.Button(value='Generator')
394
+
395
+ with gr.Tab('Scribble'):
396
+ controlnet_image_file = gr.Image(label='Image')
397
+
398
+ controlnet_model_id = gr.Dropdown(
399
+ choices=stable_inpiant_model_list,
400
+ value=stable_inpiant_model_list[0],
401
+ label='Stable Model Id'
402
+ )
403
+
404
+ controlnet_prompt = gr.Textbox(
405
+ lines=1,
406
+ value=stable_prompt_list[0],
407
+ label='Prompt'
408
+ )
409
+
410
+ controlnet_negative_prompt = gr.Textbox(
411
+ lines=1,
412
+ value=stable_negative_prompt_list[0],
413
+ label='Negative Prompt'
414
+ )
415
+
416
+ with gr.Accordion("Advanced Options", open=False):
417
+ controlnet_guidance_scale = gr.Slider(
418
+ minimum=0.1,
419
+ maximum=15,
420
+ step=0.1,
421
+ value=7.5,
422
+ label='Guidance Scale'
423
+ )
424
+
425
+ controlnet_num_inference_step = gr.Slider(
426
+ minimum=1,
427
+ maximum=100,
428
+ step=1,
429
+ value=50,
430
+ label='Num Inference Step'
431
+ )
432
+
433
+ controlnet_scribble_predict = gr.Button(value='Generator')
434
+
435
+ with gr.Tab('Pose'):
436
+ controlnet_image_file = gr.Image(label='Image')
437
+
438
+ controlnet_model_id = gr.Dropdown(
439
+ choices=stable_inpiant_model_list,
440
+ value=stable_inpiant_model_list[0],
441
+ label='Stable Model Id'
442
+ )
443
+
444
+ controlnet_prompt = gr.Textbox(
445
+ lines=1,
446
+ value=stable_prompt_list[0],
447
+ label='Prompt'
448
+ )
449
+
450
+ controlnet_negative_prompt = gr.Textbox(
451
+ lines=1,
452
+ value=stable_negative_prompt_list[0],
453
+ label='Negative Prompt'
454
+ )
455
+
456
+ with gr.Accordion("Advanced Options", open=False):
457
+ controlnet_guidance_scale = gr.Slider(
458
+ minimum=0.1,
459
+ maximum=15,
460
+ step=0.1,
461
+ value=7.5,
462
+ label='Guidance Scale'
463
+ )
464
+
465
+ controlnet_num_inference_step = gr.Slider(
466
+ minimum=1,
467
+ maximum=100,
468
+ step=1,
469
+ value=50,
470
+ label='Num Inference Step'
471
+ )
472
+
473
+ controlnet_pose_predict = gr.Button(value='Generator')
474
+
475
  with gr.Tab('Generator'):
476
  with gr.Column():
477
  output_image = gr.Image(label='Image')
 
516
  outputs = [output_image],
517
  )
518
 
519
+
520
+ controlnet_canny_predict.click(
521
+ fn = stable_diffusion_controlnet_img2img,
522
+ inputs = [
523
+ inpaint_image_file,
524
+ inpaint_model_id,
525
+ inpaint_prompt,
526
+ inpaint_negative_prompt,
527
+ inpaint_guidance_scale,
528
+ inpaint_num_inference_step,
529
+ ],
530
+ outputs = [output_image],
531
+ )
532
+
533
+ controlnet_hed_predict.click(
534
+ fn = stable_diffusion_controlnet_img2img,
535
+ inputs = [
536
+ inpaint_image_file,
537
+ inpaint_model_id,
538
+ inpaint_prompt,
539
+ inpaint_negative_prompt,
540
+ inpaint_guidance_scale,
541
+ inpaint_num_inference_step,
542
+ ],
543
+ outputs = [output_image],
544
+ )
545
+ controlnet_mlsd_predict.click(
546
+ fn = stable_diffusion_controlnet_img2img,
547
+ inputs = [
548
+ inpaint_image_file,
549
+ inpaint_model_id,
550
+ inpaint_prompt,
551
+ inpaint_negative_prompt,
552
+ inpaint_guidance_scale,
553
+ inpaint_num_inference_step,
554
+ ],
555
+ outputs = [output_image],
556
+ )
557
+ controlnet_seg_predict.click(
558
+ fn = stable_diffusion_controlnet_img2img,
559
+ inputs = [
560
+ inpaint_image_file,
561
+ inpaint_model_id,
562
+ inpaint_prompt,
563
+ inpaint_negative_prompt,
564
+ inpaint_guidance_scale,
565
+ inpaint_num_inference_step,
566
+ ],
567
+ outputs = [output_image],
568
+ )
569
+ controlnet_depth_predict.click(
570
+ fn = stable_diffusion_controlnet_img2img,
571
+ inputs = [
572
+ inpaint_image_file,
573
+ inpaint_model_id,
574
+ inpaint_prompt,
575
+ inpaint_negative_prompt,
576
+ inpaint_guidance_scale,
577
+ inpaint_num_inference_step,
578
+ ],
579
+ outputs = [output_image],
580
+ )
581
+ controlnet_scribble_predict.click(
582
+ fn = stable_diffusion_controlnet_img2img,
583
+ inputs = [
584
+ inpaint_image_file,
585
+ inpaint_model_id,
586
+ inpaint_prompt,
587
+ inpaint_negative_prompt,
588
+ inpaint_guidance_scale,
589
+ inpaint_num_inference_step,
590
+ ],
591
+ outputs = [output_image],
592
+ )
593
+ controlnet_pose_predict.click(
594
+ fn = stable_diffusion_controlnet_img2img,
595
+ inputs = [
596
+ inpaint_image_file,
597
+ inpaint_model_id,
598
+ inpaint_prompt,
599
+ inpaint_negative_prompt,
600
+ inpaint_guidance_scale,
601
+ inpaint_num_inference_step,
602
+ ],
603
+ outputs = [output_image],
604
+ )
605
+
606
  app.launch()
controlnet/controlnet_canny.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusers import ( StableDiffusionControlNetPipeline,
2
+ ControlNetModel, UniPCMultistepScheduler)
3
+
4
+ from PIL import Image
5
+ import numpy as np
6
+ import torch
7
+ import cv2
8
+
9
+
10
+ def controlnet_canny(
11
+ image_path:str,
12
+ low_th:int,
13
+ high_th:int,
14
+ ):
15
+ image = Image.open(image_path)
16
+ image = np.array(image)
17
+
18
+ image = cv2.Canny(image, low_th, high_th)
19
+ image = image[:, :, None]
20
+ image = np.concatenate([image, image, image], axis=2)
21
+ image = Image.fromarray(image)
22
+
23
+ controlnet = ControlNetModel.from_pretrained(
24
+ "lllyasviel/sd-controlnet-canny",
25
+ torch_dtype=torch.float16
26
+ )
27
+ return controlnet, image
28
+
29
+
30
+ def stable_diffusion_controlnet_img2img(
31
+ stable_model_path:str,
32
+ image_path:str,
33
+ prompt:str,
34
+ negative_prompt:str,
35
+ num_samples:int,
36
+ guidance_scale:int,
37
+ num_inference_step:int,
38
+ low_th:int,
39
+ high_th:int
40
+ ):
41
+
42
+ controlnet, image = controlnet_canny(
43
+ image_path=image_path,
44
+ low_th=low_th,
45
+ high_th=high_th
46
+ )
47
+ pipe = StableDiffusionControlNetPipeline.from_pretrained(
48
+ pretrained_model_name_or_path=stable_model_path,
49
+ controlnet=controlnet,
50
+ safety_checker=None,
51
+ torch_dtype=torch.float16,
52
+ )
53
+ pipe.to("cuda")
54
+ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
55
+ pipe.enable_xformers_memory_efficient_attention()
56
+
57
+ output = pipe(
58
+ prompt = prompt,
59
+ image = image,
60
+ negative_prompt = negative_prompt,
61
+ num_images_per_prompt = num_samples,
62
+ num_inference_steps = num_inference_step,
63
+ guidance_scale = guidance_scale,
64
+ ).images
65
+
66
+ return output
controlnet/controlnet_depth.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusers import ( StableDiffusionControlNetPipeline,
2
+ ControlNetModel, UniPCMultistepScheduler,
3
+ DDIMScheduler)
4
+
5
+ from transformers import pipeline
6
+ from PIL import Image
7
+ import numpy as np
8
+ import torch
9
+
10
+
11
+ def controlnet_depth(image_path:str):
12
+ depth_estimator = pipeline('depth-estimation')
13
+
14
+ image = Image.open(image_path)
15
+ image = depth_estimator(image)['depth']
16
+ image = np.array(image)
17
+ image = image[:, :, None]
18
+ image = np.concatenate([image, image, image], axis=2)
19
+ image = Image.fromarray(image)
20
+
21
+ controlnet = ControlNetModel.from_pretrained(
22
+ "fusing/stable-diffusion-v1-5-controlnet-depth", torch_dtype=torch.float16
23
+ )
24
+
25
+ return controlnet, image
26
+
27
+ def stable_diffusion_controlnet_img2img(
28
+ stable_model_path:str,
29
+ image_path:str,
30
+ prompt:str,
31
+ negative_prompt:str,
32
+ num_samples:int,
33
+ guidance_scale:int,
34
+ num_inference_step:int,
35
+ ):
36
+
37
+ controlnet, image = controlnet_depth(image_path=image_path)
38
+
39
+ pipe = StableDiffusionControlNetPipeline.from_pretrained(
40
+ pretrained_model_name_or_path=stable_model_path,
41
+ controlnet=controlnet,
42
+ safety_checker=None,
43
+ torch_dtype=torch.float16
44
+ )
45
+
46
+ pipe.to("cuda")
47
+ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
48
+ pipe.enable_xformers_memory_efficient_attention()
49
+
50
+ output = pipe(
51
+ prompt = prompt,
52
+ image = image,
53
+ negative_prompt = negative_prompt,
54
+ num_images_per_prompt = num_samples,
55
+ num_inference_steps = num_inference_step,
56
+ guidance_scale = guidance_scale,
57
+ ).images
58
+
59
+ return output
controlnet/controlnet_hed.py ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusers import ( StableDiffusionControlNetPipeline,
2
+ ControlNetModel, UniPCMultistepScheduler,
3
+ DDIMScheduler)
4
+
5
+ from controlnet_aux import HEDdetector
6
+ from PIL import Image
7
+ import numpy as np
8
+ import torch
9
+ import cv2
10
+
11
+
12
+ def controlnet_hed(image_path:str):
13
+ hed = HEDdetector.from_pretrained('lllyasviel/ControlNet')
14
+
15
+ image = Image.open(image_path)
16
+ image = hed(image)
17
+
18
+ controlnet = ControlNetModel.from_pretrained(
19
+ "fusing/stable-diffusion-v1-5-controlnet-hed",
20
+ torch_dtype=torch.float16
21
+ )
22
+ return controlnet, image
23
+
24
+
25
+ def stable_diffusion_controlnet_img2img(
26
+ stable_model_path:str,
27
+ image_path:str,
28
+ prompt:str,
29
+ negative_prompt:str,
30
+ num_samples:int,
31
+ guidance_scale:int,
32
+ num_inference_step:int,
33
+ ):
34
+
35
+ controlnet, image = controlnet_hed(image_path=image_path)
36
+
37
+ pipe = StableDiffusionControlNetPipeline.from_pretrained(
38
+ pretrained_model_name_or_path=stable_model_path,
39
+ controlnet=controlnet,
40
+ safety_checker=None,
41
+ torch_dtype=torch.float16
42
+ )
43
+
44
+ pipe.to("cuda")
45
+ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
46
+ pipe.enable_xformers_memory_efficient_attention()
47
+
48
+ output = pipe(
49
+ prompt = prompt,
50
+ image = image,
51
+ negative_prompt = negative_prompt,
52
+ num_images_per_prompt = num_samples,
53
+ num_inference_steps = num_inference_step,
54
+ guidance_scale = guidance_scale,
55
+ ).images
56
+
57
+ return output
controlnet/controlnet_mlsd.py ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusers import ( StableDiffusionControlNetPipeline,
2
+ ControlNetModel, UniPCMultistepScheduler,
3
+ DDIMScheduler)
4
+
5
+ from controlnet_aux import MLSDdetector
6
+ from PIL import Image
7
+ import numpy as np
8
+ import torch
9
+ import cv2
10
+
11
+
12
+ def controlnet_mlsd(image_path:str):
13
+ mlsd = MLSDdetector.from_pretrained('lllyasviel/ControlNet')
14
+
15
+ image = Image.open(image_path)
16
+ image = mlsd(image)
17
+
18
+ controlnet = ControlNetModel.from_pretrained(
19
+ "fusing/stable-diffusion-v1-5-controlnet-mlsd",
20
+ torch_dtype=torch.float16
21
+ )
22
+
23
+ return controlnet, image
24
+
25
+ def stable_diffusion_controlnet_img2img(
26
+ stable_model_path:str,
27
+ image_path:str,
28
+ prompt:str,
29
+ negative_prompt:str,
30
+ num_samples:int,
31
+ guidance_scale:int,
32
+ num_inference_step:int,
33
+ ):
34
+
35
+ controlnet, image = controlnet_mlsd(image_path=image_path)
36
+
37
+ pipe = StableDiffusionControlNetPipeline.from_pretrained(
38
+ pretrained_model_name_or_path=stable_model_path,
39
+ controlnet=controlnet,
40
+ safety_checker=None,
41
+ torch_dtype=torch.float16
42
+ )
43
+
44
+ pipe.to("cuda")
45
+ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
46
+ pipe.enable_xformers_memory_efficient_attention()
47
+
48
+ output = pipe(
49
+ prompt = prompt,
50
+ image = image,
51
+ negative_prompt = negative_prompt,
52
+ num_images_per_prompt = num_samples,
53
+ num_inference_steps = num_inference_step,
54
+ guidance_scale = guidance_scale,
55
+ ).images
56
+
57
+ return output
controlnet/controlnet_pose.py ADDED
@@ -0,0 +1,55 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusers import ( StableDiffusionControlNetPipeline,
2
+ ControlNetModel, UniPCMultistepScheduler)
3
+
4
+ from controlnet_aux import OpenposeDetector
5
+
6
+ from PIL import Image
7
+ import torch
8
+
9
+
10
+ def controlnet_pose(image_path:str):
11
+ openpose = OpenposeDetector.from_pretrained('lllyasviel/ControlNet')
12
+
13
+ image = Image.open(image_path)
14
+ image = openpose(image)
15
+
16
+ controlnet = ControlNetModel.from_pretrained(
17
+ "fusing/stable-diffusion-v1-5-controlnet-openpose",
18
+ torch_dtype=torch.float16
19
+ )
20
+
21
+ return controlnet, image
22
+
23
+ def stable_diffusion_controlnet_img2img(
24
+ stable_model_path:str,
25
+ image_path:str,
26
+ prompt:str,
27
+ negative_prompt:str,
28
+ num_samples:int,
29
+ guidance_scale:int,
30
+ num_inference_step:int,
31
+ ):
32
+
33
+ controlnet, image = controlnet_pose(image_path=image_path)
34
+
35
+ pipe = StableDiffusionControlNetPipeline.from_pretrained(
36
+ pretrained_model_name_or_path=stable_model_path,
37
+ controlnet=controlnet,
38
+ safety_checker=None,
39
+ torch_dtype=torch.float16
40
+ )
41
+
42
+ pipe.to("cuda")
43
+ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
44
+ pipe.enable_xformers_memory_efficient_attention()
45
+
46
+ output = pipe(
47
+ prompt = prompt,
48
+ image = image,
49
+ negative_prompt = negative_prompt,
50
+ num_images_per_prompt = num_samples,
51
+ num_inference_steps = num_inference_step,
52
+ guidance_scale = guidance_scale,
53
+ ).images
54
+
55
+ return output
controlnet/controlnet_scribble.py ADDED
@@ -0,0 +1,55 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusers import ( StableDiffusionControlNetPipeline,
2
+ ControlNetModel, UniPCMultistepScheduler,
3
+ DDIMScheduler)
4
+
5
+ from controlnet_aux import HEDdetector
6
+
7
+ from PIL import Image
8
+ import torch
9
+
10
+
11
+ def controlnet_scribble(image_path:str):
12
+ hed = HEDdetector.from_pretrained('lllyasviel/ControlNet')
13
+
14
+ image = Image.open(image_path)
15
+ image = hed(image, scribble=True)
16
+
17
+ controlnet = ControlNetModel.from_pretrained(
18
+ "fusing/stable-diffusion-v1-5-controlnet-scribble", torch_dtype=torch.float16
19
+ )
20
+
21
+ return controlnet, image
22
+
23
+ def stable_diffusion_controlnet_img2img(
24
+ stable_model_path:str,
25
+ image_path:str,
26
+ prompt:str,
27
+ negative_prompt:str,
28
+ num_samples:int,
29
+ guidance_scale:int,
30
+ num_inference_step:int,
31
+ ):
32
+
33
+ controlnet, image = controlnet_scribble(image_path=image_path)
34
+
35
+ pipe = StableDiffusionControlNetPipeline.from_pretrained(
36
+ pretrained_model_name_or_path=stable_model_path,
37
+ controlnet=controlnet,
38
+ safety_checker=None,
39
+ torch_dtype=torch.float16
40
+ )
41
+
42
+ pipe.to("cuda")
43
+ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
44
+ pipe.enable_xformers_memory_efficient_attention()
45
+
46
+ output = pipe(
47
+ prompt = prompt,
48
+ image = image,
49
+ negative_prompt = negative_prompt,
50
+ num_images_per_prompt = num_samples,
51
+ num_inference_steps = num_inference_step,
52
+ guidance_scale = guidance_scale,
53
+ ).images
54
+
55
+ return output
controlnet/controlnet_seg.py ADDED
@@ -0,0 +1,113 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoImageProcessor, UperNetForSemanticSegmentation
2
+ import torch
3
+ from diffusers import (StableDiffusionControlNetPipeline,
4
+ ControlNetModel, UniPCMultistepScheduler)
5
+
6
+
7
+ from PIL import Image
8
+ import numpy as np
9
+ import torch
10
+
11
+
12
+ def ade_palette():
13
+ """ADE20K palette that maps each class to RGB values."""
14
+ return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
15
+ [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
16
+ [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
17
+ [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
18
+ [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
19
+ [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
20
+ [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220],
21
+ [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224],
22
+ [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255],
23
+ [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7],
24
+ [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],
25
+ [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255],
26
+ [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0],
27
+ [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255],
28
+ [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255],
29
+ [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],
30
+ [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0],
31
+ [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0],
32
+ [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255],
33
+ [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255],
34
+ [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],
35
+ [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255],
36
+ [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255],
37
+ [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255],
38
+ [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0],
39
+ [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],
40
+ [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255],
41
+ [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112],
42
+ [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160],
43
+ [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163],
44
+ [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],
45
+ [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0],
46
+ [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255],
47
+ [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204],
48
+ [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255],
49
+ [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],
50
+ [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194],
51
+ [102, 255, 0], [92, 0, 255]]
52
+
53
+
54
+ def controlnet_mlsd(image_path:str):
55
+ image_processor = AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-small")
56
+ image_segmentor = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-small")
57
+
58
+ image = Image.open(image_path).convert('RGB')
59
+ pixel_values = image_processor(image, return_tensors="pt").pixel_values
60
+
61
+ with torch.no_grad():
62
+ outputs = image_segmentor(pixel_values)
63
+
64
+ seg = image_processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
65
+
66
+ color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
67
+ palette = np.array(ade_palette())
68
+
69
+ for label, color in enumerate(palette):
70
+ color_seg[seg == label, :] = color
71
+
72
+ color_seg = color_seg.astype(np.uint8)
73
+ image = Image.fromarray(color_seg)
74
+ controlnet = ControlNetModel.from_pretrained(
75
+ "fusing/stable-diffusion-v1-5-controlnet-seg", torch_dtype=torch.float16
76
+ )
77
+
78
+ return controlnet, image
79
+
80
+
81
+ def stable_diffusion_controlnet_img2img(
82
+ stable_model_path:str,
83
+ image_path:str,
84
+ prompt:str,
85
+ negative_prompt:str,
86
+ num_samples:int,
87
+ guidance_scale:int,
88
+ num_inference_step:int,
89
+ ):
90
+
91
+ controlnet, image = controlnet_mlsd(image_path=image_path)
92
+
93
+ pipe = StableDiffusionControlNetPipeline.from_pretrained(
94
+ pretrained_model_name_or_path=stable_model_path,
95
+ controlnet=controlnet,
96
+ safety_checker=None,
97
+ torch_dtype=torch.float16
98
+ )
99
+
100
+ pipe.to("cuda")
101
+ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
102
+ pipe.enable_xformers_memory_efficient_attention()
103
+
104
+ output = pipe(
105
+ prompt = prompt,
106
+ image = image,
107
+ negative_prompt = negative_prompt,
108
+ num_images_per_prompt = num_samples,
109
+ num_inference_steps = num_inference_step,
110
+ guidance_scale = guidance_scale,
111
+ ).images
112
+
113
+ return output