kadirnar commited on
Commit
a22dacf
1 Parent(s): b5fb0c0

Upload 25 files

Browse files
Files changed (25) hide show
  1. diffusion_webui/__init__.py +1 -0
  2. diffusion_webui/diffusion_models/__init__.py +0 -0
  3. diffusion_webui/diffusion_models/controlnet/__init__.py +0 -0
  4. diffusion_webui/diffusion_models/controlnet/controlnet_canny.py +183 -0
  5. diffusion_webui/diffusion_models/controlnet/controlnet_depth.py +187 -0
  6. diffusion_webui/diffusion_models/controlnet/controlnet_hed.py +181 -0
  7. diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/__init__.py +0 -0
  8. diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/controlnet_inpaint_app.py +203 -0
  9. diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/pipeline_stable_diffusion_controlnet_inpaint.py +607 -0
  10. diffusion_webui/diffusion_models/controlnet/controlnet_mlsd.py +173 -0
  11. diffusion_webui/diffusion_models/controlnet/controlnet_pose.py +189 -0
  12. diffusion_webui/diffusion_models/controlnet/controlnet_scribble.py +188 -0
  13. diffusion_webui/diffusion_models/controlnet/controlnet_seg.py +353 -0
  14. diffusion_webui/diffusion_models/stable_diffusion/__init__.py +0 -0
  15. diffusion_webui/diffusion_models/stable_diffusion/__pycache__/__init__.cpython-38.pyc +0 -0
  16. diffusion_webui/diffusion_models/stable_diffusion/__pycache__/img2img_app.cpython-38.pyc +0 -0
  17. diffusion_webui/diffusion_models/stable_diffusion/__pycache__/inpaint_app.cpython-38.pyc +0 -0
  18. diffusion_webui/diffusion_models/stable_diffusion/__pycache__/text2img_app.cpython-38.pyc +0 -0
  19. diffusion_webui/diffusion_models/stable_diffusion/img2img_app.py +153 -0
  20. diffusion_webui/diffusion_models/stable_diffusion/inpaint_app.py +148 -0
  21. diffusion_webui/diffusion_models/stable_diffusion/text2img_app.py +170 -0
  22. diffusion_webui/helpers.py +33 -0
  23. diffusion_webui/utils/__init__.py +0 -0
  24. diffusion_webui/utils/model_list.py +33 -0
  25. diffusion_webui/utils/scheduler_list.py +47 -0
diffusion_webui/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ __version__ = "1.4.1"
diffusion_webui/diffusion_models/__init__.py ADDED
File without changes
diffusion_webui/diffusion_models/controlnet/__init__.py ADDED
File without changes
diffusion_webui/diffusion_models/controlnet/controlnet_canny.py ADDED
@@ -0,0 +1,183 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import gradio as gr
3
+ import numpy as np
4
+ import torch
5
+ from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
6
+ from PIL import Image
7
+
8
+ from diffusion_webui.utils.model_list import (
9
+ controlnet_canny_model_list,
10
+ stable_model_list,
11
+ )
12
+ from diffusion_webui.utils.scheduler_list import (
13
+ SCHEDULER_LIST,
14
+ get_scheduler_list,
15
+ )
16
+
17
+
18
+ class StableDiffusionControlNetCannyGenerator:
19
+ def __init__(self):
20
+ self.pipe = None
21
+
22
+ def load_model(self, stable_model_path, controlnet_model_path, scheduler):
23
+ if self.pipe is None:
24
+ controlnet = ControlNetModel.from_pretrained(
25
+ controlnet_model_path, torch_dtype=torch.float16
26
+ )
27
+ self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
28
+ pretrained_model_name_or_path=stable_model_path,
29
+ controlnet=controlnet,
30
+ safety_checker=None,
31
+ torch_dtype=torch.float16,
32
+ )
33
+
34
+ self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
35
+ self.pipe.to("cuda")
36
+ self.pipe.enable_xformers_memory_efficient_attention()
37
+
38
+ return self.pipe
39
+
40
+ def controlnet_canny(
41
+ self,
42
+ image_path: str,
43
+ ):
44
+ image = Image.open(image_path)
45
+ image = np.array(image)
46
+
47
+ image = cv2.Canny(image, 100, 200)
48
+ image = image[:, :, None]
49
+ image = np.concatenate([image, image, image], axis=2)
50
+ image = Image.fromarray(image)
51
+
52
+ return image
53
+
54
+ def generate_image(
55
+ self,
56
+ image_path: str,
57
+ stable_model_path: str,
58
+ controlnet_model_path: str,
59
+ prompt: str,
60
+ negative_prompt: str,
61
+ num_images_per_prompt: int,
62
+ guidance_scale: int,
63
+ num_inference_step: int,
64
+ scheduler: str,
65
+ seed_generator: int,
66
+ ):
67
+ pipe = self.load_model(
68
+ stable_model_path=stable_model_path,
69
+ controlnet_model_path=controlnet_model_path,
70
+ scheduler=scheduler,
71
+ )
72
+
73
+ image = self.controlnet_canny(image_path=image_path)
74
+
75
+ if seed_generator == 0:
76
+ random_seed = torch.randint(0, 1000000, (1,))
77
+ generator = torch.manual_seed(random_seed)
78
+ else:
79
+ generator = torch.manual_seed(seed_generator)
80
+
81
+ output = pipe(
82
+ prompt=prompt,
83
+ image=image,
84
+ negative_prompt=negative_prompt,
85
+ num_images_per_prompt=num_images_per_prompt,
86
+ num_inference_steps=num_inference_step,
87
+ guidance_scale=guidance_scale,
88
+ generator=generator,
89
+ ).images
90
+
91
+ return output
92
+
93
+ def app():
94
+ with gr.Blocks():
95
+ with gr.Row():
96
+ with gr.Column():
97
+ controlnet_canny_image_file = gr.Image(
98
+ type="filepath", label="Image"
99
+ )
100
+
101
+ controlnet_canny_prompt = gr.Textbox(
102
+ lines=1,
103
+ placeholder="Prompt",
104
+ show_label=False,
105
+ )
106
+
107
+ controlnet_canny_negative_prompt = gr.Textbox(
108
+ lines=1,
109
+ placeholder="Negative Prompt",
110
+ show_label=False,
111
+ )
112
+ with gr.Row():
113
+ with gr.Column():
114
+ controlnet_canny_stable_model_id = gr.Dropdown(
115
+ choices=stable_model_list,
116
+ value=stable_model_list[0],
117
+ label="Stable Model Id",
118
+ )
119
+
120
+ controlnet_canny_guidance_scale = gr.Slider(
121
+ minimum=0.1,
122
+ maximum=15,
123
+ step=0.1,
124
+ value=7.5,
125
+ label="Guidance Scale",
126
+ )
127
+ controlnet_canny_num_inference_step = gr.Slider(
128
+ minimum=1,
129
+ maximum=100,
130
+ step=1,
131
+ value=50,
132
+ label="Num Inference Step",
133
+ )
134
+ controlnet_canny_num_images_per_prompt = gr.Slider(
135
+ minimum=1,
136
+ maximum=10,
137
+ step=1,
138
+ value=1,
139
+ label="Number Of Images",
140
+ )
141
+ with gr.Row():
142
+ with gr.Column():
143
+ controlnet_canny_model_id = gr.Dropdown(
144
+ choices=controlnet_canny_model_list,
145
+ value=controlnet_canny_model_list[0],
146
+ label="ControlNet Model Id",
147
+ )
148
+
149
+ controlnet_canny_scheduler = gr.Dropdown(
150
+ choices=SCHEDULER_LIST,
151
+ value=SCHEDULER_LIST[0],
152
+ label="Scheduler",
153
+ )
154
+
155
+ controlnet_canny_seed_generator = gr.Number(
156
+ value=0,
157
+ label="Seed Generator",
158
+ )
159
+ controlnet_canny_predict = gr.Button(value="Generator")
160
+
161
+ with gr.Column():
162
+ output_image = gr.Gallery(
163
+ label="Generated images",
164
+ show_label=False,
165
+ elem_id="gallery",
166
+ ).style(grid=(1, 2))
167
+
168
+ controlnet_canny_predict.click(
169
+ fn=StableDiffusionControlNetCannyGenerator().generate_image,
170
+ inputs=[
171
+ controlnet_canny_image_file,
172
+ controlnet_canny_stable_model_id,
173
+ controlnet_canny_model_id,
174
+ controlnet_canny_prompt,
175
+ controlnet_canny_negative_prompt,
176
+ controlnet_canny_num_images_per_prompt,
177
+ controlnet_canny_guidance_scale,
178
+ controlnet_canny_num_inference_step,
179
+ controlnet_canny_scheduler,
180
+ controlnet_canny_seed_generator,
181
+ ],
182
+ outputs=[output_image],
183
+ )
diffusion_webui/diffusion_models/controlnet/controlnet_depth.py ADDED
@@ -0,0 +1,187 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import numpy as np
3
+ import torch
4
+ from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
5
+ from PIL import Image
6
+ from transformers import pipeline
7
+
8
+ from diffusion_webui.utils.model_list import (
9
+ controlnet_depth_model_list,
10
+ stable_model_list,
11
+ )
12
+ from diffusion_webui.utils.scheduler_list import (
13
+ SCHEDULER_LIST,
14
+ get_scheduler_list,
15
+ )
16
+
17
+
18
+ class StableDiffusionControlNetDepthGenerator:
19
+ def __init__(self):
20
+ self.pipe = None
21
+
22
+ def load_model(self, stable_model_path, controlnet_model_path, scheduler):
23
+ if self.pipe is None:
24
+ controlnet = ControlNetModel.from_pretrained(
25
+ controlnet_model_path, torch_dtype=torch.float16
26
+ )
27
+
28
+ self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
29
+ pretrained_model_name_or_path=stable_model_path,
30
+ controlnet=controlnet,
31
+ safety_checker=None,
32
+ torch_dtype=torch.float16,
33
+ )
34
+
35
+ self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
36
+ self.pipe.to("cuda")
37
+ self.pipe.enable_xformers_memory_efficient_attention()
38
+
39
+ return self.pipe
40
+
41
+ def controlnet_depth(self, image_path: str):
42
+ depth_estimator = pipeline("depth-estimation")
43
+ image = Image.open(image_path)
44
+ image = depth_estimator(image)["depth"]
45
+ image = np.array(image)
46
+ image = image[:, :, None]
47
+ image = np.concatenate([image, image, image], axis=2)
48
+ image = Image.fromarray(image)
49
+
50
+ return image
51
+
52
+ def generate_image(
53
+ self,
54
+ image_path: str,
55
+ stable_model_path: str,
56
+ depth_model_path: str,
57
+ prompt: str,
58
+ negative_prompt: str,
59
+ num_images_per_prompt: int,
60
+ guidance_scale: int,
61
+ num_inference_step: int,
62
+ scheduler: str,
63
+ seed_generator: int,
64
+ ):
65
+ image = self.controlnet_depth(image_path)
66
+
67
+ pipe = self.load_model(
68
+ stable_model_path=stable_model_path,
69
+ controlnet_model_path=depth_model_path,
70
+ scheduler=scheduler,
71
+ )
72
+
73
+ if seed_generator == 0:
74
+ random_seed = torch.randint(0, 1000000, (1,))
75
+ generator = torch.manual_seed(random_seed)
76
+ else:
77
+ generator = torch.manual_seed(seed_generator)
78
+
79
+ output = pipe(
80
+ prompt=prompt,
81
+ image=image,
82
+ negative_prompt=negative_prompt,
83
+ num_images_per_prompt=num_images_per_prompt,
84
+ num_inference_steps=num_inference_step,
85
+ guidance_scale=guidance_scale,
86
+ generator=generator,
87
+ ).images
88
+
89
+ return output
90
+
91
+ def app():
92
+ with gr.Blocks():
93
+ with gr.Row():
94
+ with gr.Column():
95
+ controlnet_depth_image_file = gr.Image(
96
+ type="filepath", label="Image"
97
+ )
98
+
99
+ controlnet_depth_prompt = gr.Textbox(
100
+ lines=1,
101
+ show_label=False,
102
+ placeholder="Prompt",
103
+ )
104
+
105
+ controlnet_depth_negative_prompt = gr.Textbox(
106
+ lines=1,
107
+ show_label=False,
108
+ placeholder="Negative Prompt",
109
+ )
110
+
111
+ with gr.Row():
112
+ with gr.Column():
113
+ controlnet_depth_stable_model_id = gr.Dropdown(
114
+ choices=stable_model_list,
115
+ value=stable_model_list[0],
116
+ label="Stable Model Id",
117
+ )
118
+ controlnet_depth_guidance_scale = gr.Slider(
119
+ minimum=0.1,
120
+ maximum=15,
121
+ step=0.1,
122
+ value=7.5,
123
+ label="Guidance Scale",
124
+ )
125
+
126
+ controlnet_depth_num_inference_step = gr.Slider(
127
+ minimum=1,
128
+ maximum=100,
129
+ step=1,
130
+ value=50,
131
+ label="Num Inference Step",
132
+ )
133
+
134
+ controlnet_depth_num_images_per_prompt = gr.Slider(
135
+ minimum=1,
136
+ maximum=10,
137
+ step=1,
138
+ value=1,
139
+ label="Number Of Images",
140
+ )
141
+ with gr.Row():
142
+ with gr.Column():
143
+ controlnet_depth_model_id = gr.Dropdown(
144
+ choices=controlnet_depth_model_list,
145
+ value=controlnet_depth_model_list[0],
146
+ label="ControlNet Model Id",
147
+ )
148
+
149
+ controlnet_depth_scheduler = gr.Dropdown(
150
+ choices=SCHEDULER_LIST,
151
+ value=SCHEDULER_LIST[0],
152
+ label="Scheduler",
153
+ )
154
+
155
+ controlnet_depth_seed_generator = gr.Number(
156
+ minimum=0,
157
+ maximum=1000000,
158
+ step=1,
159
+ value=0,
160
+ label="Seed Generator",
161
+ )
162
+
163
+ controlnet_depth_predict = gr.Button(value="Generator")
164
+
165
+ with gr.Column():
166
+ output_image = gr.Gallery(
167
+ label="Generated images",
168
+ show_label=False,
169
+ elem_id="gallery",
170
+ ).style(grid=(1, 2))
171
+
172
+ controlnet_depth_predict.click(
173
+ fn=StableDiffusionControlNetDepthGenerator().generate_image,
174
+ inputs=[
175
+ controlnet_depth_image_file,
176
+ controlnet_depth_stable_model_id,
177
+ controlnet_depth_model_id,
178
+ controlnet_depth_prompt,
179
+ controlnet_depth_negative_prompt,
180
+ controlnet_depth_num_images_per_prompt,
181
+ controlnet_depth_guidance_scale,
182
+ controlnet_depth_num_inference_step,
183
+ controlnet_depth_scheduler,
184
+ controlnet_depth_seed_generator,
185
+ ],
186
+ outputs=output_image,
187
+ )
diffusion_webui/diffusion_models/controlnet/controlnet_hed.py ADDED
@@ -0,0 +1,181 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from controlnet_aux import HEDdetector
4
+ from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
5
+ from PIL import Image
6
+
7
+ from diffusion_webui.utils.model_list import (
8
+ controlnet_hed_model_list,
9
+ stable_model_list,
10
+ )
11
+ from diffusion_webui.utils.scheduler_list import (
12
+ SCHEDULER_LIST,
13
+ get_scheduler_list,
14
+ )
15
+
16
+
17
+ class StableDiffusionControlNetHEDGenerator:
18
+ def __init__(self):
19
+ self.pipe = None
20
+
21
+ def load_model(self, stable_model_path, controlnet_model_path, scheduler):
22
+ if self.pipe is None:
23
+ controlnet = ControlNetModel.from_pretrained(
24
+ controlnet_model_path, torch_dtype=torch.float16
25
+ )
26
+
27
+ self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
28
+ pretrained_model_name_or_path=stable_model_path,
29
+ controlnet=controlnet,
30
+ safety_checker=None,
31
+ torch_dtype=torch.float16,
32
+ )
33
+
34
+ self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
35
+ self.pipe.to("cuda")
36
+ self.pipe.enable_xformers_memory_efficient_attention()
37
+
38
+ return self.pipe
39
+
40
+ def controlnet_hed(self, image_path: str):
41
+ hed = HEDdetector.from_pretrained("lllyasviel/ControlNet")
42
+ image = Image.open(image_path)
43
+ image = hed(image)
44
+
45
+ return image
46
+
47
+ def generate_image(
48
+ self,
49
+ image_path: str,
50
+ stable_model_path: str,
51
+ controlnet_hed_model_path: str,
52
+ prompt: str,
53
+ negative_prompt: str,
54
+ num_images_per_prompt: int,
55
+ guidance_scale: int,
56
+ num_inference_step: int,
57
+ sheduler: str,
58
+ seed_generator: int,
59
+ ):
60
+
61
+ image = self.controlnet_hed(image_path=image_path)
62
+
63
+ pipe = self.load_model(
64
+ stable_model_path=stable_model_path,
65
+ controlnet_model_path=controlnet_hed_model_path,
66
+ scheduler=sheduler,
67
+ )
68
+
69
+ if seed_generator == 0:
70
+ random_seed = torch.randint(0, 1000000, (1,))
71
+ generator = torch.manual_seed(random_seed)
72
+ else:
73
+ generator = torch.manual_seed(seed_generator)
74
+
75
+ output = pipe(
76
+ prompt=prompt,
77
+ image=image,
78
+ negative_prompt=negative_prompt,
79
+ num_images_per_prompt=num_images_per_prompt,
80
+ num_inference_steps=num_inference_step,
81
+ guidance_scale=guidance_scale,
82
+ generator=generator,
83
+ ).images
84
+
85
+ return output
86
+
87
+ def app():
88
+ with gr.Blocks():
89
+ with gr.Row():
90
+ with gr.Column():
91
+ controlnet_hed_image_file = gr.Image(
92
+ type="filepath", label="Image"
93
+ )
94
+ controlnet_hed_prompt = gr.Textbox(
95
+ lines=1,
96
+ show_label=False,
97
+ placeholder="Prompt",
98
+ )
99
+
100
+ controlnet_hed_negative_prompt = gr.Textbox(
101
+ lines=1,
102
+ show_label=False,
103
+ placeholder="Negative Prompt",
104
+ )
105
+
106
+ with gr.Row():
107
+ with gr.Column():
108
+ controlnet_hed_stable_model_id = gr.Dropdown(
109
+ choices=stable_model_list,
110
+ value=stable_model_list[0],
111
+ label="Stable Model Id",
112
+ )
113
+ controlnet_hed_guidance_scale = gr.Slider(
114
+ minimum=0.1,
115
+ maximum=15,
116
+ step=0.1,
117
+ value=7.5,
118
+ label="Guidance Scale",
119
+ )
120
+ controlnet_hed_num_inference_step = gr.Slider(
121
+ minimum=1,
122
+ maximum=100,
123
+ step=1,
124
+ value=50,
125
+ label="Num Inference Step",
126
+ )
127
+
128
+ controlnet_hed_num_images_per_prompt = gr.Slider(
129
+ minimum=1,
130
+ maximum=10,
131
+ step=1,
132
+ value=1,
133
+ label="Number Of Images",
134
+ )
135
+
136
+ with gr.Row():
137
+ with gr.Column():
138
+ controlnet_hed_model_id = gr.Dropdown(
139
+ choices=controlnet_hed_model_list,
140
+ value=controlnet_hed_model_list[0],
141
+ label="ControlNet Model Id",
142
+ )
143
+ controlnet_hed_scheduler = gr.Dropdown(
144
+ choices=SCHEDULER_LIST,
145
+ value=SCHEDULER_LIST[0],
146
+ label="Scheduler",
147
+ )
148
+
149
+ controlnet_hed_seed_generator = gr.Number(
150
+ minimum=0,
151
+ maximum=1000000,
152
+ step=1,
153
+ value=0,
154
+ label="Seed Generator",
155
+ )
156
+
157
+ controlnet_hed_predict = gr.Button(value="Generator")
158
+
159
+ with gr.Column():
160
+ output_image = gr.Gallery(
161
+ label="Generated images",
162
+ show_label=False,
163
+ elem_id="gallery",
164
+ ).style(grid=(1, 2))
165
+
166
+ controlnet_hed_predict.click(
167
+ fn=StableDiffusionControlNetHEDGenerator().generate_image,
168
+ inputs=[
169
+ controlnet_hed_image_file,
170
+ controlnet_hed_stable_model_id,
171
+ controlnet_hed_model_id,
172
+ controlnet_hed_prompt,
173
+ controlnet_hed_negative_prompt,
174
+ controlnet_hed_num_images_per_prompt,
175
+ controlnet_hed_guidance_scale,
176
+ controlnet_hed_num_inference_step,
177
+ controlnet_hed_scheduler,
178
+ controlnet_hed_seed_generator,
179
+ ],
180
+ outputs=[output_image],
181
+ )
diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/__init__.py ADDED
File without changes
diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/controlnet_inpaint_app.py ADDED
@@ -0,0 +1,203 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import gradio as gr
3
+ import numpy as np
4
+ import torch
5
+ from diffusers import (
6
+ ControlNetModel,
7
+ StableDiffusionControlNetPipeline,
8
+ )
9
+ from PIL import Image
10
+
11
+ from diffusion_webui.utils.model_list import (
12
+ controlnet_canny_model_list,
13
+ stable_model_list,
14
+ )
15
+ from diffusion_webui.utils.scheduler_list import (
16
+ SCHEDULER_LIST,
17
+ get_scheduler_list,
18
+ )
19
+
20
+
21
+ class StableDiffusionControlInpaintNetCannyGenerator:
22
+ def __init__(self):
23
+ self.pipe = None
24
+
25
+ def controlnet_canny_inpaint(
26
+ self,
27
+ image_path: str,
28
+ ):
29
+ image = Image.open(image_path)
30
+ image = np.array(image)
31
+
32
+ image = cv2.Canny(image, 100, 200)
33
+ image = image[:, :, None]
34
+ image = np.concatenate([image, image, image], axis=2)
35
+ image = Image.fromarray(image)
36
+
37
+ return image
38
+
39
+ def load_model(self, stable_model_path, controlnet_model_path, scheduler):
40
+ if self.pipe is None:
41
+ controlnet = ControlNetModel.from_pretrained(
42
+ controlnet_model_path, torch_dtype=torch.float16
43
+ )
44
+ self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
45
+ pretrained_model_name_or_path=stable_model_path,
46
+ controlnet=controlnet,
47
+ safety_checker=None,
48
+ torch_dtype=torch.float16,
49
+ )
50
+
51
+ self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
52
+ self.pipe.to("cuda")
53
+ self.pipe.enable_xformers_memory_efficient_attention()
54
+
55
+ return self.pipe
56
+
57
+ def generate_image(
58
+ self,
59
+ image_path: str,
60
+ stable_model_path: str,
61
+ controlnet_model_path: str,
62
+ prompt: str,
63
+ negative_prompt: str,
64
+ num_images_per_prompt: int,
65
+ guidance_scale: int,
66
+ num_inference_step: int,
67
+ scheduler: str,
68
+ seed_generator: int,
69
+ ):
70
+
71
+ image = self.controlnet_canny_inpaint(
72
+ image_path=image_path, controlnet_model_path=controlnet_model_path
73
+ )
74
+
75
+ pipe = self.load_model(
76
+ stable_model_path=stable_model_path,
77
+ controlnet_model_path=controlnet_model_path,
78
+ scheduler=scheduler,
79
+ )
80
+
81
+ if seed_generator == 0:
82
+ random_seed = torch.randint(0, 1000000, (1,))
83
+ generator = torch.manual_seed(random_seed)
84
+ else:
85
+ generator = torch.manual_seed(seed_generator)
86
+
87
+ output = pipe(
88
+ prompt=prompt,
89
+ image=image,
90
+ negative_prompt=negative_prompt,
91
+ num_images_per_prompt=num_images_per_prompt,
92
+ num_inference_steps=num_inference_step,
93
+ guidance_scale=guidance_scale,
94
+ generator=generator,
95
+ ).images
96
+
97
+ return output
98
+
99
+ def app():
100
+ with gr.Blocks():
101
+ with gr.Row():
102
+ with gr.Column():
103
+ controlnet_canny_inpaint_image_file = gr.Image(
104
+ type="filepath", label="Image"
105
+ )
106
+
107
+ controlnet_canny_inpaint_prompt = gr.Textbox(
108
+ lines=1, placeholder="Prompt", show_label=False
109
+ )
110
+
111
+ controlnet_canny_inpaint_negative_prompt = gr.Textbox(
112
+ lines=1,
113
+ show_label=False,
114
+ placeholder="Negative Prompt",
115
+ )
116
+ with gr.Row():
117
+ with gr.Column():
118
+ controlnet_canny_inpaint_stable_model_id = (
119
+ gr.Dropdown(
120
+ choices=stable_model_list,
121
+ value=stable_model_list[0],
122
+ label="Stable Model Id",
123
+ )
124
+ )
125
+
126
+ controlnet_canny_inpaint_guidance_scale = gr.Slider(
127
+ minimum=0.1,
128
+ maximum=15,
129
+ step=0.1,
130
+ value=7.5,
131
+ label="Guidance Scale",
132
+ )
133
+
134
+ controlnet_canny_inpaint_num_inference_step = (
135
+ gr.Slider(
136
+ minimum=1,
137
+ maximum=100,
138
+ step=1,
139
+ value=50,
140
+ label="Num Inference Step",
141
+ )
142
+ )
143
+ controlnet_canny_inpaint_num_images_per_prompt = (
144
+ gr.Slider(
145
+ minimum=1,
146
+ maximum=10,
147
+ step=1,
148
+ value=1,
149
+ label="Number Of Images",
150
+ )
151
+ )
152
+ with gr.Row():
153
+ with gr.Column():
154
+ controlnet_canny_inpaint_model_id = gr.Dropdown(
155
+ choices=controlnet_canny_model_list,
156
+ value=controlnet_canny_model_list[0],
157
+ label="Controlnet Model Id",
158
+ )
159
+ controlnet_canny_inpaint_scheduler = (
160
+ gr.Dropdown(
161
+ choices=SCHEDULER_LIST,
162
+ value=SCHEDULER_LIST[0],
163
+ label="Scheduler",
164
+ )
165
+ )
166
+
167
+ controlnet_canny_inpaint_seed_generator = (
168
+ gr.Slider(
169
+ minimum=0,
170
+ maximum=1000000,
171
+ step=1,
172
+ value=0,
173
+ label="Seed Generator",
174
+ )
175
+ )
176
+
177
+ controlnet_canny_inpaint_predict = gr.Button(
178
+ value="Generator"
179
+ )
180
+
181
+ with gr.Column():
182
+ output_image = gr.Gallery(
183
+ label="Generated images",
184
+ show_label=False,
185
+ elem_id="gallery",
186
+ ).style(grid=(1, 2))
187
+
188
+ controlnet_canny_inpaint_predict.click(
189
+ fn=StableDiffusionControlInpaintNetCannyGenerator().generate_image,
190
+ inputs=[
191
+ controlnet_canny_inpaint_image_file,
192
+ controlnet_canny_inpaint_stable_model_id,
193
+ controlnet_canny_inpaint_model_id,
194
+ controlnet_canny_inpaint_prompt,
195
+ controlnet_canny_inpaint_negative_prompt,
196
+ controlnet_canny_inpaint_num_images_per_prompt,
197
+ controlnet_canny_inpaint_guidance_scale,
198
+ controlnet_canny_inpaint_num_inference_step,
199
+ controlnet_canny_inpaint_scheduler,
200
+ controlnet_canny_inpaint_seed_generator,
201
+ ],
202
+ outputs=[output_image],
203
+ )
diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/pipeline_stable_diffusion_controlnet_inpaint.py ADDED
@@ -0,0 +1,607 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2023 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import numpy as np
16
+ import PIL.Image
17
+ import torch
18
+ from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import *
19
+
20
+ EXAMPLE_DOC_STRING = """
21
+ Examples:
22
+ ```py
23
+ >>> # !pip install opencv-python transformers accelerate
24
+ >>> from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, UniPCMultistepScheduler
25
+ >>> from diffusers.utils import load_image
26
+ >>> import numpy as np
27
+ >>> import torch
28
+
29
+ >>> import cv2
30
+ >>> from PIL import Image
31
+ >>> # download an image
32
+ >>> image = load_image(
33
+ ... "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
34
+ ... )
35
+ >>> image = np.array(image)
36
+ >>> mask_image = load_image(
37
+ ... "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
38
+ ... )
39
+ >>> mask_image = np.array(mask_image)
40
+ >>> # get canny image
41
+ >>> canny_image = cv2.Canny(image, 100, 200)
42
+ >>> canny_image = canny_image[:, :, None]
43
+ >>> canny_image = np.concatenate([canny_image, canny_image, canny_image], axis=2)
44
+ >>> canny_image = Image.fromarray(canny_image)
45
+
46
+ >>> # load control net and stable diffusion v1-5
47
+ >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
48
+ >>> pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
49
+ ... "runwayml/stable-diffusion-inpainting", controlnet=controlnet, torch_dtype=torch.float16
50
+ ... )
51
+
52
+ >>> # speed up diffusion process with faster scheduler and memory optimization
53
+ >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
54
+ >>> # remove following line if xformers is not installed
55
+ >>> pipe.enable_xformers_memory_efficient_attention()
56
+
57
+ >>> pipe.enable_model_cpu_offload()
58
+
59
+ >>> # generate image
60
+ >>> generator = torch.manual_seed(0)
61
+ >>> image = pipe(
62
+ ... "futuristic-looking doggo",
63
+ ... num_inference_steps=20,
64
+ ... generator=generator,
65
+ ... image=image,
66
+ ... control_image=canny_image,
67
+ ... mask_image=mask_image
68
+ ... ).images[0]
69
+ ```
70
+ """
71
+
72
+
73
+ def prepare_mask_and_masked_image(image, mask):
74
+ """
75
+ Prepares a pair (image, mask) to be consumed by the Stable Diffusion pipeline. This means that those inputs will be
76
+ converted to ``torch.Tensor`` with shapes ``batch x channels x height x width`` where ``channels`` is ``3`` for the
77
+ ``image`` and ``1`` for the ``mask``.
78
+ The ``image`` will be converted to ``torch.float32`` and normalized to be in ``[-1, 1]``. The ``mask`` will be
79
+ binarized (``mask > 0.5``) and cast to ``torch.float32`` too.
80
+ Args:
81
+ image (Union[np.array, PIL.Image, torch.Tensor]): The image to inpaint.
82
+ It can be a ``PIL.Image``, or a ``height x width x 3`` ``np.array`` or a ``channels x height x width``
83
+ ``torch.Tensor`` or a ``batch x channels x height x width`` ``torch.Tensor``.
84
+ mask (_type_): The mask to apply to the image, i.e. regions to inpaint.
85
+ It can be a ``PIL.Image``, or a ``height x width`` ``np.array`` or a ``1 x height x width``
86
+ ``torch.Tensor`` or a ``batch x 1 x height x width`` ``torch.Tensor``.
87
+ Raises:
88
+ ValueError: ``torch.Tensor`` images should be in the ``[-1, 1]`` range. ValueError: ``torch.Tensor`` mask
89
+ should be in the ``[0, 1]`` range. ValueError: ``mask`` and ``image`` should have the same spatial dimensions.
90
+ TypeError: ``mask`` is a ``torch.Tensor`` but ``image`` is not
91
+ (ot the other way around).
92
+ Returns:
93
+ tuple[torch.Tensor]: The pair (mask, masked_image) as ``torch.Tensor`` with 4
94
+ dimensions: ``batch x channels x height x width``.
95
+ """
96
+ if isinstance(image, torch.Tensor):
97
+ if not isinstance(mask, torch.Tensor):
98
+ raise TypeError(
99
+ f"`image` is a torch.Tensor but `mask` (type: {type(mask)} is not"
100
+ )
101
+
102
+ # Batch single image
103
+ if image.ndim == 3:
104
+ assert (
105
+ image.shape[0] == 3
106
+ ), "Image outside a batch should be of shape (3, H, W)"
107
+ image = image.unsqueeze(0)
108
+
109
+ # Batch and add channel dim for single mask
110
+ if mask.ndim == 2:
111
+ mask = mask.unsqueeze(0).unsqueeze(0)
112
+
113
+ # Batch single mask or add channel dim
114
+ if mask.ndim == 3:
115
+ # Single batched mask, no channel dim or single mask not batched but channel dim
116
+ if mask.shape[0] == 1:
117
+ mask = mask.unsqueeze(0)
118
+
119
+ # Batched masks no channel dim
120
+ else:
121
+ mask = mask.unsqueeze(1)
122
+
123
+ assert (
124
+ image.ndim == 4 and mask.ndim == 4
125
+ ), "Image and Mask must have 4 dimensions"
126
+ assert (
127
+ image.shape[-2:] == mask.shape[-2:]
128
+ ), "Image and Mask must have the same spatial dimensions"
129
+ assert (
130
+ image.shape[0] == mask.shape[0]
131
+ ), "Image and Mask must have the same batch size"
132
+
133
+ # Check image is in [-1, 1]
134
+ if image.min() < -1 or image.max() > 1:
135
+ raise ValueError("Image should be in [-1, 1] range")
136
+
137
+ # Check mask is in [0, 1]
138
+ if mask.min() < 0 or mask.max() > 1:
139
+ raise ValueError("Mask should be in [0, 1] range")
140
+
141
+ # Binarize mask
142
+ mask[mask < 0.5] = 0
143
+ mask[mask >= 0.5] = 1
144
+
145
+ # Image as float32
146
+ image = image.to(dtype=torch.float32)
147
+ elif isinstance(mask, torch.Tensor):
148
+ raise TypeError(
149
+ f"`mask` is a torch.Tensor but `image` (type: {type(image)} is not"
150
+ )
151
+ else:
152
+ # preprocess image
153
+ if isinstance(image, (PIL.Image.Image, np.ndarray)):
154
+ image = [image]
155
+
156
+ if isinstance(image, list) and isinstance(image[0], PIL.Image.Image):
157
+ image = [np.array(i.convert("RGB"))[None, :] for i in image]
158
+ image = np.concatenate(image, axis=0)
159
+ elif isinstance(image, list) and isinstance(image[0], np.ndarray):
160
+ image = np.concatenate([i[None, :] for i in image], axis=0)
161
+
162
+ image = image.transpose(0, 3, 1, 2)
163
+ image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0
164
+
165
+ # preprocess mask
166
+ if isinstance(mask, (PIL.Image.Image, np.ndarray)):
167
+ mask = [mask]
168
+
169
+ if isinstance(mask, list) and isinstance(mask[0], PIL.Image.Image):
170
+ mask = np.concatenate(
171
+ [np.array(m.convert("L"))[None, None, :] for m in mask], axis=0
172
+ )
173
+ mask = mask.astype(np.float32) / 255.0
174
+ elif isinstance(mask, list) and isinstance(mask[0], np.ndarray):
175
+ mask = np.concatenate([m[None, None, :] for m in mask], axis=0)
176
+
177
+ mask[mask < 0.5] = 0
178
+ mask[mask >= 0.5] = 1
179
+ mask = torch.from_numpy(mask)
180
+
181
+ masked_image = image * (mask < 0.5)
182
+
183
+ return mask, masked_image
184
+
185
+
186
+ class StableDiffusionControlNetInpaintPipeline(
187
+ StableDiffusionControlNetPipeline
188
+ ):
189
+ r"""
190
+ Pipeline for text-guided image inpainting using Stable Diffusion with ControlNet guidance.
191
+
192
+ This model inherits from [`StableDiffusionControlNetPipeline`]. Check the superclass documentation for the generic methods the
193
+ library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
194
+
195
+ Args:
196
+ vae ([`AutoencoderKL`]):
197
+ Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
198
+ text_encoder ([`CLIPTextModel`]):
199
+ Frozen text-encoder. Stable Diffusion uses the text portion of
200
+ [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
201
+ the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
202
+ tokenizer (`CLIPTokenizer`):
203
+ Tokenizer of class
204
+ [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
205
+ unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
206
+ controlnet ([`ControlNetModel`]):
207
+ Provides additional conditioning to the unet during the denoising process
208
+ scheduler ([`SchedulerMixin`]):
209
+ A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
210
+ [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
211
+ safety_checker ([`StableDiffusionSafetyChecker`]):
212
+ Classification module that estimates whether generated images could be considered offensive or harmful.
213
+ Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.
214
+ feature_extractor ([`CLIPFeatureExtractor`]):
215
+ Model that extracts features from generated images to be used as inputs for the `safety_checker`.
216
+ """
217
+
218
+ def prepare_mask_latents(
219
+ self,
220
+ mask,
221
+ masked_image,
222
+ batch_size,
223
+ height,
224
+ width,
225
+ dtype,
226
+ device,
227
+ generator,
228
+ do_classifier_free_guidance,
229
+ ):
230
+ # resize the mask to latents shape as we concatenate the mask to the latents
231
+ # we do that before converting to dtype to avoid breaking in case we're using cpu_offload
232
+ # and half precision
233
+ mask = torch.nn.functional.interpolate(
234
+ mask,
235
+ size=(
236
+ height // self.vae_scale_factor,
237
+ width // self.vae_scale_factor,
238
+ ),
239
+ )
240
+ mask = mask.to(device=device, dtype=dtype)
241
+
242
+ masked_image = masked_image.to(device=device, dtype=dtype)
243
+
244
+ # encode the mask image into latents space so we can concatenate it to the latents
245
+ if isinstance(generator, list):
246
+ masked_image_latents = [
247
+ self.vae.encode(masked_image[i : i + 1]).latent_dist.sample(
248
+ generator=generator[i]
249
+ )
250
+ for i in range(batch_size)
251
+ ]
252
+ masked_image_latents = torch.cat(masked_image_latents, dim=0)
253
+ else:
254
+ masked_image_latents = self.vae.encode(
255
+ masked_image
256
+ ).latent_dist.sample(generator=generator)
257
+ masked_image_latents = (
258
+ self.vae.config.scaling_factor * masked_image_latents
259
+ )
260
+
261
+ # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method
262
+ if mask.shape[0] < batch_size:
263
+ if not batch_size % mask.shape[0] == 0:
264
+ raise ValueError(
265
+ "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"
266
+ f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number"
267
+ " of masks that you pass is divisible by the total requested batch size."
268
+ )
269
+ mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1)
270
+ if masked_image_latents.shape[0] < batch_size:
271
+ if not batch_size % masked_image_latents.shape[0] == 0:
272
+ raise ValueError(
273
+ "The passed images and the required batch size don't match. Images are supposed to be duplicated"
274
+ f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed."
275
+ " Make sure the number of images that you pass is divisible by the total requested batch size."
276
+ )
277
+ masked_image_latents = masked_image_latents.repeat(
278
+ batch_size // masked_image_latents.shape[0], 1, 1, 1
279
+ )
280
+
281
+ mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask
282
+ masked_image_latents = (
283
+ torch.cat([masked_image_latents] * 2)
284
+ if do_classifier_free_guidance
285
+ else masked_image_latents
286
+ )
287
+
288
+ # aligning device to prevent device errors when concating it with the latent model input
289
+ masked_image_latents = masked_image_latents.to(
290
+ device=device, dtype=dtype
291
+ )
292
+ return mask, masked_image_latents
293
+
294
+ @torch.no_grad()
295
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
296
+ def __call__(
297
+ self,
298
+ prompt: Union[str, List[str]] = None,
299
+ image: Union[torch.FloatTensor, PIL.Image.Image] = None,
300
+ control_image: Union[
301
+ torch.FloatTensor,
302
+ PIL.Image.Image,
303
+ List[torch.FloatTensor],
304
+ List[PIL.Image.Image],
305
+ ] = None,
306
+ mask_image: Union[torch.FloatTensor, PIL.Image.Image] = None,
307
+ height: Optional[int] = None,
308
+ width: Optional[int] = None,
309
+ num_inference_steps: int = 50,
310
+ guidance_scale: float = 7.5,
311
+ negative_prompt: Optional[Union[str, List[str]]] = None,
312
+ num_images_per_prompt: Optional[int] = 1,
313
+ eta: float = 0.0,
314
+ generator: Optional[
315
+ Union[torch.Generator, List[torch.Generator]]
316
+ ] = None,
317
+ latents: Optional[torch.FloatTensor] = None,
318
+ prompt_embeds: Optional[torch.FloatTensor] = None,
319
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None,
320
+ output_type: Optional[str] = "pil",
321
+ return_dict: bool = True,
322
+ callback: Optional[
323
+ Callable[[int, int, torch.FloatTensor], None]
324
+ ] = None,
325
+ callback_steps: int = 1,
326
+ cross_attention_kwargs: Optional[Dict[str, Any]] = None,
327
+ controlnet_conditioning_scale: float = 1.0,
328
+ ):
329
+ r"""
330
+ Function invoked when calling the pipeline for generation.
331
+ Args:
332
+ prompt (`str` or `List[str]`, *optional*):
333
+ The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
334
+ instead.
335
+ image (`PIL.Image.Image`):
336
+ `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will
337
+ be masked out with `mask_image` and repainted according to `prompt`.
338
+ control_image (`torch.FloatTensor`, `PIL.Image.Image`, `List[torch.FloatTensor]` or `List[PIL.Image.Image]`):
339
+ The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If
340
+ the type is specified as `Torch.FloatTensor`, it is passed to ControlNet as is. PIL.Image.Image` can
341
+ also be accepted as an image. The control image is automatically resized to fit the output image.
342
+ mask_image (`PIL.Image.Image`):
343
+ `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be
344
+ repainted, while black pixels will be preserved. If `mask_image` is a PIL image, it will be converted
345
+ to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L)
346
+ instead of 3, so the expected shape would be `(B, H, W, 1)`.
347
+ height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
348
+ The height in pixels of the generated image.
349
+ width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
350
+ The width in pixels of the generated image.
351
+ num_inference_steps (`int`, *optional*, defaults to 50):
352
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
353
+ expense of slower inference.
354
+ guidance_scale (`float`, *optional*, defaults to 7.5):
355
+ Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
356
+ `guidance_scale` is defined as `w` of equation 2. of [Imagen
357
+ Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
358
+ 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
359
+ usually at the expense of lower image quality.
360
+ negative_prompt (`str` or `List[str]`, *optional*):
361
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
362
+ `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
363
+ Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
364
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
365
+ The number of images to generate per prompt.
366
+ eta (`float`, *optional*, defaults to 0.0):
367
+ Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
368
+ [`schedulers.DDIMScheduler`], will be ignored for others.
369
+ generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
370
+ One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
371
+ to make generation deterministic.
372
+ latents (`torch.FloatTensor`, *optional*):
373
+ Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
374
+ generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
375
+ tensor will ge generated by sampling using the supplied random `generator`.
376
+ prompt_embeds (`torch.FloatTensor`, *optional*):
377
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
378
+ provided, text embeddings will be generated from `prompt` input argument.
379
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
380
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
381
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
382
+ argument.
383
+ output_type (`str`, *optional*, defaults to `"pil"`):
384
+ The output format of the generate image. Choose between
385
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
386
+ return_dict (`bool`, *optional*, defaults to `True`):
387
+ Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
388
+ plain tuple.
389
+ callback (`Callable`, *optional*):
390
+ A function that will be called every `callback_steps` steps during inference. The function will be
391
+ called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
392
+ callback_steps (`int`, *optional*, defaults to 1):
393
+ The frequency at which the `callback` function will be called. If not specified, the callback will be
394
+ called at every step.
395
+ cross_attention_kwargs (`dict`, *optional*):
396
+ A kwargs dictionary that if specified is passed along to the `AttnProcessor` as defined under
397
+ `self.processor` in
398
+ [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
399
+ controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):
400
+ The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added
401
+ to the residual in the original unet.
402
+ Examples:
403
+ Returns:
404
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
405
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
406
+ When returning a tuple, the first element is a list with the generated images, and the second element is a
407
+ list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
408
+ (nsfw) content, according to the `safety_checker`.
409
+ """
410
+ # 0. Default height and width to unet
411
+ height, width = self._default_height_width(height, width, control_image)
412
+
413
+ # 1. Check inputs. Raise error if not correct
414
+ self.check_inputs(
415
+ prompt,
416
+ control_image,
417
+ height,
418
+ width,
419
+ callback_steps,
420
+ negative_prompt,
421
+ prompt_embeds,
422
+ negative_prompt_embeds,
423
+ )
424
+
425
+ # 2. Define call parameters
426
+ if prompt is not None and isinstance(prompt, str):
427
+ batch_size = 1
428
+ elif prompt is not None and isinstance(prompt, list):
429
+ batch_size = len(prompt)
430
+ else:
431
+ batch_size = prompt_embeds.shape[0]
432
+
433
+ device = self._execution_device
434
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
435
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
436
+ # corresponds to doing no classifier free guidance.
437
+ do_classifier_free_guidance = guidance_scale > 1.0
438
+
439
+ # 3. Encode input prompt
440
+ prompt_embeds = self._encode_prompt(
441
+ prompt,
442
+ device,
443
+ num_images_per_prompt,
444
+ do_classifier_free_guidance,
445
+ negative_prompt,
446
+ prompt_embeds=prompt_embeds,
447
+ negative_prompt_embeds=negative_prompt_embeds,
448
+ )
449
+
450
+ # 4. Prepare image
451
+ control_image = self.prepare_image(
452
+ control_image,
453
+ width,
454
+ height,
455
+ batch_size * num_images_per_prompt,
456
+ num_images_per_prompt,
457
+ device,
458
+ self.controlnet.dtype,
459
+ )
460
+
461
+ if do_classifier_free_guidance:
462
+ control_image = torch.cat([control_image] * 2)
463
+
464
+ # 5. Prepare timesteps
465
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
466
+ timesteps = self.scheduler.timesteps
467
+
468
+ # 6. Prepare latent variables
469
+ num_channels_latents = self.controlnet.in_channels
470
+ latents = self.prepare_latents(
471
+ batch_size * num_images_per_prompt,
472
+ num_channels_latents,
473
+ height,
474
+ width,
475
+ prompt_embeds.dtype,
476
+ device,
477
+ generator,
478
+ latents,
479
+ )
480
+
481
+ # EXTRA: prepare mask latents
482
+ mask, masked_image = prepare_mask_and_masked_image(image, mask_image)
483
+ mask, masked_image_latents = self.prepare_mask_latents(
484
+ mask,
485
+ masked_image,
486
+ batch_size * num_images_per_prompt,
487
+ height,
488
+ width,
489
+ prompt_embeds.dtype,
490
+ device,
491
+ generator,
492
+ do_classifier_free_guidance,
493
+ )
494
+
495
+ # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
496
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
497
+
498
+ # 8. Denoising loop
499
+ num_warmup_steps = (
500
+ len(timesteps) - num_inference_steps * self.scheduler.order
501
+ )
502
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
503
+ for i, t in enumerate(timesteps):
504
+ # expand the latents if we are doing classifier free guidance
505
+ latent_model_input = (
506
+ torch.cat([latents] * 2)
507
+ if do_classifier_free_guidance
508
+ else latents
509
+ )
510
+ latent_model_input = self.scheduler.scale_model_input(
511
+ latent_model_input, t
512
+ )
513
+
514
+ down_block_res_samples, mid_block_res_sample = self.controlnet(
515
+ latent_model_input,
516
+ t,
517
+ encoder_hidden_states=prompt_embeds,
518
+ controlnet_cond=control_image,
519
+ return_dict=False,
520
+ )
521
+
522
+ down_block_res_samples = [
523
+ down_block_res_sample * controlnet_conditioning_scale
524
+ for down_block_res_sample in down_block_res_samples
525
+ ]
526
+ mid_block_res_sample *= controlnet_conditioning_scale
527
+
528
+ # predict the noise residual
529
+ latent_model_input = torch.cat(
530
+ [latent_model_input, mask, masked_image_latents], dim=1
531
+ )
532
+ noise_pred = self.unet(
533
+ latent_model_input,
534
+ t,
535
+ encoder_hidden_states=prompt_embeds,
536
+ cross_attention_kwargs=cross_attention_kwargs,
537
+ down_block_additional_residuals=down_block_res_samples,
538
+ mid_block_additional_residual=mid_block_res_sample,
539
+ ).sample
540
+
541
+ # perform guidance
542
+ if do_classifier_free_guidance:
543
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
544
+ noise_pred = noise_pred_uncond + guidance_scale * (
545
+ noise_pred_text - noise_pred_uncond
546
+ )
547
+
548
+ # compute the previous noisy sample x_t -> x_t-1
549
+ latents = self.scheduler.step(
550
+ noise_pred, t, latents, **extra_step_kwargs
551
+ ).prev_sample
552
+
553
+ # call the callback, if provided
554
+ if i == len(timesteps) - 1 or (
555
+ (i + 1) > num_warmup_steps
556
+ and (i + 1) % self.scheduler.order == 0
557
+ ):
558
+ progress_bar.update()
559
+ if callback is not None and i % callback_steps == 0:
560
+ callback(i, t, latents)
561
+
562
+ # If we do sequential model offloading, let's offload unet and controlnet
563
+ # manually for max memory savings
564
+ if (
565
+ hasattr(self, "final_offload_hook")
566
+ and self.final_offload_hook is not None
567
+ ):
568
+ self.unet.to("cpu")
569
+ self.controlnet.to("cpu")
570
+ torch.cuda.empty_cache()
571
+
572
+ if output_type == "latent":
573
+ image = latents
574
+ has_nsfw_concept = None
575
+ elif output_type == "pil":
576
+ # 8. Post-processing
577
+ image = self.decode_latents(latents)
578
+
579
+ # 9. Run safety checker
580
+ image, has_nsfw_concept = self.run_safety_checker(
581
+ image, device, prompt_embeds.dtype
582
+ )
583
+
584
+ # 10. Convert to PIL
585
+ image = self.numpy_to_pil(image)
586
+ else:
587
+ # 8. Post-processing
588
+ image = self.decode_latents(latents)
589
+
590
+ # 9. Run safety checker
591
+ image, has_nsfw_concept = self.run_safety_checker(
592
+ image, device, prompt_embeds.dtype
593
+ )
594
+
595
+ # Offload last model to CPU
596
+ if (
597
+ hasattr(self, "final_offload_hook")
598
+ and self.final_offload_hook is not None
599
+ ):
600
+ self.final_offload_hook.offload()
601
+
602
+ if not return_dict:
603
+ return (image, has_nsfw_concept)
604
+
605
+ return StableDiffusionPipelineOutput(
606
+ images=image, nsfw_content_detected=has_nsfw_concept
607
+ )
diffusion_webui/diffusion_models/controlnet/controlnet_mlsd.py ADDED
@@ -0,0 +1,173 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from controlnet_aux import MLSDdetector
4
+ from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
5
+ from PIL import Image
6
+
7
+ from diffusion_webui.utils.model_list import stable_model_list
8
+ from diffusion_webui.utils.scheduler_list import (
9
+ SCHEDULER_LIST,
10
+ get_scheduler_list,
11
+ )
12
+
13
+
14
+ class StableDiffusionControlNetMLSDGenerator:
15
+ def __init__(self):
16
+ self.pipe = None
17
+
18
+ def load_model(self, stable_model_path, controlnet_model_path, scheduler):
19
+ if self.pipe is None:
20
+ controlnet = ControlNetModel.from_pretrained(
21
+ controlnet_model_path, torch_dtype=torch.float16
22
+ )
23
+
24
+ self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
25
+ pretrained_model_name_or_path=stable_model_path,
26
+ controlnet=controlnet,
27
+ safety_checker=None,
28
+ torch_dtype=torch.float16,
29
+ )
30
+
31
+ self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
32
+ self.pipe.to("cuda")
33
+ self.pipe.enable_xformers_memory_efficient_attention()
34
+
35
+ return self.pipe
36
+
37
+ def controlnet_mlsd(self, image_path: str):
38
+ mlsd = MLSDdetector.from_pretrained("lllyasviel/ControlNet")
39
+
40
+ image = Image.open(image_path)
41
+ image = mlsd(image)
42
+
43
+ return image
44
+
45
+ def generate_image(
46
+ self,
47
+ image_path: str,
48
+ model_path: str,
49
+ prompt: str,
50
+ negative_prompt: str,
51
+ num_images_per_prompt: int,
52
+ guidance_scale: int,
53
+ num_inference_step: int,
54
+ scheduler: str,
55
+ seed_generator: int,
56
+ ):
57
+ image = self.controlnet_mlsd(image_path=image_path)
58
+
59
+ pipe = self.load_model(
60
+ stable_model_path=model_path,
61
+ controlnet_model_path="lllyasviel/sd-controlnet-mlsd",
62
+ scheduler=scheduler,
63
+ )
64
+
65
+ if seed_generator == 0:
66
+ random_seed = torch.randint(0, 1000000, (1,))
67
+ generator = torch.manual_seed(random_seed)
68
+ else:
69
+ generator = torch.manual_seed(seed_generator)
70
+
71
+ output = pipe(
72
+ prompt=prompt,
73
+ image=image,
74
+ negative_prompt=negative_prompt,
75
+ num_images_per_prompt=num_images_per_prompt,
76
+ num_inference_steps=num_inference_step,
77
+ guidance_scale=guidance_scale,
78
+ generator=generator,
79
+ ).images
80
+
81
+ return output
82
+
83
+ def app():
84
+ with gr.Blocks():
85
+ with gr.Row():
86
+ with gr.Column():
87
+ controlnet_mlsd_image_file = gr.Image(
88
+ type="filepath", label="Image"
89
+ )
90
+
91
+ controlnet_mlsd_prompt = gr.Textbox(
92
+ lines=1,
93
+ show_label=False,
94
+ placeholder="Prompt",
95
+ )
96
+
97
+ controlnet_mlsd_negative_prompt = gr.Textbox(
98
+ lines=1,
99
+ show_label=False,
100
+ placeholder="Negative Prompt",
101
+ )
102
+
103
+ with gr.Row():
104
+ with gr.Column():
105
+ controlnet_mlsd_model_id = gr.Dropdown(
106
+ choices=stable_model_list,
107
+ value=stable_model_list[0],
108
+ label="Stable Model Id",
109
+ )
110
+ controlnet_mlsd_guidance_scale = gr.Slider(
111
+ minimum=0.1,
112
+ maximum=15,
113
+ step=0.1,
114
+ value=7.5,
115
+ label="Guidance Scale",
116
+ )
117
+ controlnet_mlsd_num_inference_step = gr.Slider(
118
+ minimum=1,
119
+ maximum=100,
120
+ step=1,
121
+ value=50,
122
+ label="Num Inference Step",
123
+ )
124
+
125
+ with gr.Row():
126
+ with gr.Column():
127
+ controlnet_mlsd_scheduler = gr.Dropdown(
128
+ choices=SCHEDULER_LIST,
129
+ value=SCHEDULER_LIST[0],
130
+ label="Scheduler",
131
+ )
132
+
133
+ controlnet_mlsd_seed_generator = gr.Slider(
134
+ minimum=0,
135
+ maximum=1000000,
136
+ step=1,
137
+ value=0,
138
+ label="Seed Generator",
139
+ )
140
+ controlnet_mlsd_num_images_per_prompt = (
141
+ gr.Slider(
142
+ minimum=1,
143
+ maximum=10,
144
+ step=1,
145
+ value=1,
146
+ label="Number Of Images",
147
+ )
148
+ )
149
+
150
+ controlnet_mlsd_predict = gr.Button(value="Generator")
151
+
152
+ with gr.Column():
153
+ output_image = gr.Gallery(
154
+ label="Generated images",
155
+ show_label=False,
156
+ elem_id="gallery",
157
+ ).style(grid=(1, 2))
158
+
159
+ controlnet_mlsd_predict.click(
160
+ fn=StableDiffusionControlNetMLSDGenerator().generate_image,
161
+ inputs=[
162
+ controlnet_mlsd_image_file,
163
+ controlnet_mlsd_model_id,
164
+ controlnet_mlsd_prompt,
165
+ controlnet_mlsd_negative_prompt,
166
+ controlnet_mlsd_num_images_per_prompt,
167
+ controlnet_mlsd_guidance_scale,
168
+ controlnet_mlsd_num_inference_step,
169
+ controlnet_mlsd_scheduler,
170
+ controlnet_mlsd_seed_generator,
171
+ ],
172
+ outputs=output_image,
173
+ )
diffusion_webui/diffusion_models/controlnet/controlnet_pose.py ADDED
@@ -0,0 +1,189 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from controlnet_aux import OpenposeDetector
4
+ from diffusers import (
5
+ ControlNetModel,
6
+ StableDiffusionControlNetPipeline,
7
+ UniPCMultistepScheduler,
8
+ )
9
+ from PIL import Image
10
+
11
+ from diffusion_webui.utils.model_list import (
12
+ controlnet_pose_model_list,
13
+ stable_model_list,
14
+ )
15
+ from diffusion_webui.utils.scheduler_list import (
16
+ SCHEDULER_LIST,
17
+ get_scheduler_list,
18
+ )
19
+
20
+
21
+ class StableDiffusionControlNetPoseGenerator:
22
+ def __init__(self):
23
+ self.pipe = None
24
+
25
+ def load_model(self, stable_model_path, controlnet_model_path, scheduler):
26
+ if self.pipe is None:
27
+ controlnet = ControlNetModel.from_pretrained(
28
+ controlnet_model_path, torch_dtype=torch.float16
29
+ )
30
+
31
+ self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
32
+ pretrained_model_name_or_path=stable_model_path,
33
+ controlnet=controlnet,
34
+ safety_checker=None,
35
+ torch_dtype=torch.float16,
36
+ )
37
+
38
+ self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
39
+ self.pipe.to("cuda")
40
+ self.pipe.enable_xformers_memory_efficient_attention()
41
+
42
+ return self.pipe
43
+
44
+ def controlnet_pose(self, image_path: str):
45
+ openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
46
+
47
+ image = Image.open(image_path)
48
+ image = openpose(image)
49
+
50
+ return image
51
+
52
+ def generate_image(
53
+ self,
54
+ image_path: str,
55
+ stable_model_path: str,
56
+ controlnet_pose_model_path: str,
57
+ prompt: str,
58
+ negative_prompt: str,
59
+ num_images_per_prompt: int,
60
+ guidance_scale: int,
61
+ num_inference_step: int,
62
+ scheduler: str,
63
+ seed_generator: int,
64
+ ):
65
+
66
+ image = self.controlnet_pose(image_path=image_path)
67
+
68
+ pipe = self.load_model(
69
+ stable_model_path=stable_model_path,
70
+ controlnet_model_path=controlnet_pose_model_path,
71
+ scheduler=scheduler,
72
+ )
73
+
74
+ if seed_generator == 0:
75
+ random_seed = torch.randint(0, 1000000, (1,))
76
+ generator = torch.manual_seed(random_seed)
77
+ else:
78
+ generator = torch.manual_seed(seed_generator)
79
+
80
+ output = pipe(
81
+ prompt=prompt,
82
+ image=image,
83
+ negative_prompt=negative_prompt,
84
+ num_images_per_prompt=num_images_per_prompt,
85
+ num_inference_steps=num_inference_step,
86
+ guidance_scale=guidance_scale,
87
+ generator=generator,
88
+ ).images
89
+
90
+ return output
91
+
92
+ def app():
93
+ with gr.Blocks():
94
+ with gr.Row():
95
+ with gr.Column():
96
+ controlnet_pose_image_file = gr.Image(
97
+ type="filepath", label="Image"
98
+ )
99
+
100
+ controlnet_pose_prompt = gr.Textbox(
101
+ lines=1,
102
+ show_label=False,
103
+ placeholder="Prompt",
104
+ )
105
+
106
+ controlnet_pose_negative_prompt = gr.Textbox(
107
+ lines=1,
108
+ show_label=False,
109
+ placeholder="Negative Prompt",
110
+ )
111
+
112
+ with gr.Row():
113
+ with gr.Column():
114
+ controlnet_pose_stable_model_id = gr.Dropdown(
115
+ choices=stable_model_list,
116
+ value=stable_model_list[0],
117
+ label="Stable Model Id",
118
+ )
119
+ controlnet_pose_guidance_scale = gr.Slider(
120
+ minimum=0.1,
121
+ maximum=15,
122
+ step=0.1,
123
+ value=7.5,
124
+ label="Guidance Scale",
125
+ )
126
+
127
+ controlnet_pose_num_inference_step = gr.Slider(
128
+ minimum=1,
129
+ maximum=100,
130
+ step=1,
131
+ value=50,
132
+ label="Num Inference Step",
133
+ )
134
+
135
+ controlnet_pose_num_images_per_prompt = gr.Slider(
136
+ minimum=1,
137
+ maximum=10,
138
+ step=1,
139
+ value=1,
140
+ label="Number Of Images",
141
+ )
142
+
143
+ with gr.Row():
144
+ with gr.Column():
145
+ controlnet_pose_model_id = gr.Dropdown(
146
+ choices=controlnet_pose_model_list,
147
+ value=controlnet_pose_model_list[0],
148
+ label="ControlNet Model Id",
149
+ )
150
+
151
+ controlnet_pose_scheduler = gr.Dropdown(
152
+ choices=SCHEDULER_LIST,
153
+ value=SCHEDULER_LIST[0],
154
+ label="Scheduler",
155
+ )
156
+
157
+ controlnet_pose_seed_generator = gr.Number(
158
+ minimum=0,
159
+ maximum=1000000,
160
+ step=1,
161
+ value=0,
162
+ label="Seed Generator",
163
+ )
164
+
165
+ controlnet_pose_predict = gr.Button(value="Generator")
166
+
167
+ with gr.Column():
168
+ output_image = gr.Gallery(
169
+ label="Generated images",
170
+ show_label=False,
171
+ elem_id="gallery",
172
+ ).style(grid=(1, 2))
173
+
174
+ controlnet_pose_predict.click(
175
+ fn=StableDiffusionControlNetPoseGenerator().generate_image,
176
+ inputs=[
177
+ controlnet_pose_image_file,
178
+ controlnet_pose_stable_model_id,
179
+ controlnet_pose_model_id,
180
+ controlnet_pose_prompt,
181
+ controlnet_pose_negative_prompt,
182
+ controlnet_pose_num_images_per_prompt,
183
+ controlnet_pose_guidance_scale,
184
+ controlnet_pose_num_inference_step,
185
+ controlnet_pose_scheduler,
186
+ controlnet_pose_seed_generator,
187
+ ],
188
+ outputs=output_image,
189
+ )
diffusion_webui/diffusion_models/controlnet/controlnet_scribble.py ADDED
@@ -0,0 +1,188 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from controlnet_aux import HEDdetector
4
+ from diffusers import (
5
+ ControlNetModel,
6
+ StableDiffusionControlNetPipeline,
7
+ UniPCMultistepScheduler,
8
+ )
9
+ from PIL import Image
10
+
11
+ from diffusion_webui.utils.model_list import (
12
+ controlnet_scribble_model_list,
13
+ stable_model_list,
14
+ )
15
+ from diffusion_webui.utils.scheduler_list import (
16
+ SCHEDULER_LIST,
17
+ get_scheduler_list,
18
+ )
19
+
20
+
21
+ class StableDiffusionControlNetScribbleGenerator:
22
+ def __init__(self):
23
+ self.pipe = None
24
+
25
+ def load_model(self, stable_model_path, controlnet_model_path, scheduler):
26
+ if self.pipe is None:
27
+ controlnet = ControlNetModel.from_pretrained(
28
+ controlnet_model_path, torch_dtype=torch.float16
29
+ )
30
+
31
+ self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
32
+ pretrained_model_name_or_path=stable_model_path,
33
+ controlnet=controlnet,
34
+ safety_checker=None,
35
+ torch_dtype=torch.float16,
36
+ )
37
+
38
+ self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
39
+ self.pipe.to("cuda")
40
+ self.pipe.enable_xformers_memory_efficient_attention()
41
+
42
+ return self.pipe
43
+
44
+ def controlnet_scribble(self, image_path: str):
45
+ hed = HEDdetector.from_pretrained("lllyasviel/ControlNet")
46
+
47
+ image = Image.open(image_path)
48
+ image = hed(image, scribble=True)
49
+
50
+ return image
51
+
52
+ def generate_image(
53
+ self,
54
+ image_path: str,
55
+ stable_model_path: str,
56
+ controlnet_hed_model_path: str,
57
+ prompt: str,
58
+ negative_prompt: str,
59
+ num_images_per_prompt: int,
60
+ guidance_scale: int,
61
+ num_inference_step: int,
62
+ scheduler: str,
63
+ seed_generator: int,
64
+ ):
65
+
66
+ image = self.controlnet_scribble(image_path=image_path)
67
+
68
+ pipe = self.load_model(
69
+ stable_model_path=stable_model_path,
70
+ controlnet_model_path=controlnet_hed_model_path,
71
+ scheduler=scheduler,
72
+ )
73
+ if seed_generator == 0:
74
+ random_seed = torch.randint(0, 1000000, (1,))
75
+ generator = torch.manual_seed(random_seed)
76
+ else:
77
+ generator = torch.manual_seed(seed_generator)
78
+
79
+ output = pipe(
80
+ prompt=prompt,
81
+ image=image,
82
+ negative_prompt=negative_prompt,
83
+ num_images_per_prompt=num_images_per_prompt,
84
+ num_inference_steps=num_inference_step,
85
+ guidance_scale=guidance_scale,
86
+ generator=generator,
87
+ ).images
88
+
89
+ return output
90
+
91
+ def app():
92
+ with gr.Blocks():
93
+ with gr.Row():
94
+ with gr.Column():
95
+ controlnet_scribble_image_file = gr.Image(
96
+ type="filepath", label="Image"
97
+ )
98
+ controlnet_scribble_prompt = gr.Textbox(
99
+ lines=1,
100
+ show_label=False,
101
+ placeholder="Prompt",
102
+ )
103
+
104
+ controlnet_scribble_negative_prompt = gr.Textbox(
105
+ lines=1,
106
+ show_label=False,
107
+ placeholder="Negative Prompt",
108
+ )
109
+
110
+ with gr.Row():
111
+ with gr.Column():
112
+ controlnet_scribble_stable_model_id = gr.Dropdown(
113
+ choices=stable_model_list,
114
+ value=stable_model_list[0],
115
+ label="Stable Model Id",
116
+ )
117
+ controlnet_scribble_guidance_scale = gr.Slider(
118
+ minimum=0.1,
119
+ maximum=15,
120
+ step=0.1,
121
+ value=7.5,
122
+ label="Guidance Scale",
123
+ )
124
+
125
+ controlnet_scribble_num_inference_step = gr.Slider(
126
+ minimum=1,
127
+ maximum=100,
128
+ step=1,
129
+ value=50,
130
+ label="Num Inference Step",
131
+ )
132
+
133
+ controlnet_scribble_num_images_per_prompt = (
134
+ gr.Slider(
135
+ minimum=1,
136
+ maximum=10,
137
+ step=1,
138
+ value=1,
139
+ label="Number Of Images",
140
+ )
141
+ )
142
+ with gr.Row():
143
+ with gr.Column():
144
+ controlnet_scribble_model_id = gr.Dropdown(
145
+ choices=controlnet_scribble_model_list,
146
+ value=controlnet_scribble_model_list[0],
147
+ label="ControlNet Model Id",
148
+ )
149
+
150
+ controlnet_scribble_scheduler = gr.Dropdown(
151
+ choices=SCHEDULER_LIST,
152
+ value=SCHEDULER_LIST[0],
153
+ label="Scheduler",
154
+ )
155
+
156
+ controlnet_scribble_seed_generator = gr.Number(
157
+ minimum=0,
158
+ maximum=1000000,
159
+ step=1,
160
+ value=0,
161
+ label="Seed Generator",
162
+ )
163
+
164
+ controlnet_scribble_predict = gr.Button(value="Generator")
165
+
166
+ with gr.Column():
167
+ output_image = gr.Gallery(
168
+ label="Generated images",
169
+ show_label=False,
170
+ elem_id="gallery",
171
+ ).style(grid=(1, 2))
172
+
173
+ controlnet_scribble_predict.click(
174
+ fn=StableDiffusionControlNetScribbleGenerator().generate_image,
175
+ inputs=[
176
+ controlnet_scribble_image_file,
177
+ controlnet_scribble_stable_model_id,
178
+ controlnet_scribble_model_id,
179
+ controlnet_scribble_prompt,
180
+ controlnet_scribble_negative_prompt,
181
+ controlnet_scribble_num_images_per_prompt,
182
+ controlnet_scribble_guidance_scale,
183
+ controlnet_scribble_num_inference_step,
184
+ controlnet_scribble_scheduler,
185
+ controlnet_scribble_seed_generator,
186
+ ],
187
+ outputs=output_image,
188
+ )
diffusion_webui/diffusion_models/controlnet/controlnet_seg.py ADDED
@@ -0,0 +1,353 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import numpy as np
3
+ import torch
4
+ from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
5
+ from PIL import Image
6
+ from transformers import AutoImageProcessor, UperNetForSemanticSegmentation
7
+
8
+ from diffusion_webui.utils.model_list import stable_model_list
9
+ from diffusion_webui.utils.scheduler_list import (
10
+ SCHEDULER_LIST,
11
+ get_scheduler_list,
12
+ )
13
+
14
+
15
+ def ade_palette():
16
+ """ADE20K palette that maps each class to RGB values."""
17
+ return [
18
+ [120, 120, 120],
19
+ [180, 120, 120],
20
+ [6, 230, 230],
21
+ [80, 50, 50],
22
+ [4, 200, 3],
23
+ [120, 120, 80],
24
+ [140, 140, 140],
25
+ [204, 5, 255],
26
+ [230, 230, 230],
27
+ [4, 250, 7],
28
+ [224, 5, 255],
29
+ [235, 255, 7],
30
+ [150, 5, 61],
31
+ [120, 120, 70],
32
+ [8, 255, 51],
33
+ [255, 6, 82],
34
+ [143, 255, 140],
35
+ [204, 255, 4],
36
+ [255, 51, 7],
37
+ [204, 70, 3],
38
+ [0, 102, 200],
39
+ [61, 230, 250],
40
+ [255, 6, 51],
41
+ [11, 102, 255],
42
+ [255, 7, 71],
43
+ [255, 9, 224],
44
+ [9, 7, 230],
45
+ [220, 220, 220],
46
+ [255, 9, 92],
47
+ [112, 9, 255],
48
+ [8, 255, 214],
49
+ [7, 255, 224],
50
+ [255, 184, 6],
51
+ [10, 255, 71],
52
+ [255, 41, 10],
53
+ [7, 255, 255],
54
+ [224, 255, 8],
55
+ [102, 8, 255],
56
+ [255, 61, 6],
57
+ [255, 194, 7],
58
+ [255, 122, 8],
59
+ [0, 255, 20],
60
+ [255, 8, 41],
61
+ [255, 5, 153],
62
+ [6, 51, 255],
63
+ [235, 12, 255],
64
+ [160, 150, 20],
65
+ [0, 163, 255],
66
+ [140, 140, 140],
67
+ [250, 10, 15],
68
+ [20, 255, 0],
69
+ [31, 255, 0],
70
+ [255, 31, 0],
71
+ [255, 224, 0],
72
+ [153, 255, 0],
73
+ [0, 0, 255],
74
+ [255, 71, 0],
75
+ [0, 235, 255],
76
+ [0, 173, 255],
77
+ [31, 0, 255],
78
+ [11, 200, 200],
79
+ [255, 82, 0],
80
+ [0, 255, 245],
81
+ [0, 61, 255],
82
+ [0, 255, 112],
83
+ [0, 255, 133],
84
+ [255, 0, 0],
85
+ [255, 163, 0],
86
+ [255, 102, 0],
87
+ [194, 255, 0],
88
+ [0, 143, 255],
89
+ [51, 255, 0],
90
+ [0, 82, 255],
91
+ [0, 255, 41],
92
+ [0, 255, 173],
93
+ [10, 0, 255],
94
+ [173, 255, 0],
95
+ [0, 255, 153],
96
+ [255, 92, 0],
97
+ [255, 0, 255],
98
+ [255, 0, 245],
99
+ [255, 0, 102],
100
+ [255, 173, 0],
101
+ [255, 0, 20],
102
+ [255, 184, 184],
103
+ [0, 31, 255],
104
+ [0, 255, 61],
105
+ [0, 71, 255],
106
+ [255, 0, 204],
107
+ [0, 255, 194],
108
+ [0, 255, 82],
109
+ [0, 10, 255],
110
+ [0, 112, 255],
111
+ [51, 0, 255],
112
+ [0, 194, 255],
113
+ [0, 122, 255],
114
+ [0, 255, 163],
115
+ [255, 153, 0],
116
+ [0, 255, 10],
117
+ [255, 112, 0],
118
+ [143, 255, 0],
119
+ [82, 0, 255],
120
+ [163, 255, 0],
121
+ [255, 235, 0],
122
+ [8, 184, 170],
123
+ [133, 0, 255],
124
+ [0, 255, 92],
125
+ [184, 0, 255],
126
+ [255, 0, 31],
127
+ [0, 184, 255],
128
+ [0, 214, 255],
129
+ [255, 0, 112],
130
+ [92, 255, 0],
131
+ [0, 224, 255],
132
+ [112, 224, 255],
133
+ [70, 184, 160],
134
+ [163, 0, 255],
135
+ [153, 0, 255],
136
+ [71, 255, 0],
137
+ [255, 0, 163],
138
+ [255, 204, 0],
139
+ [255, 0, 143],
140
+ [0, 255, 235],
141
+ [133, 255, 0],
142
+ [255, 0, 235],
143
+ [245, 0, 255],
144
+ [255, 0, 122],
145
+ [255, 245, 0],
146
+ [10, 190, 212],
147
+ [214, 255, 0],
148
+ [0, 204, 255],
149
+ [20, 0, 255],
150
+ [255, 255, 0],
151
+ [0, 153, 255],
152
+ [0, 41, 255],
153
+ [0, 255, 204],
154
+ [41, 0, 255],
155
+ [41, 255, 0],
156
+ [173, 0, 255],
157
+ [0, 245, 255],
158
+ [71, 0, 255],
159
+ [122, 0, 255],
160
+ [0, 255, 184],
161
+ [0, 92, 255],
162
+ [184, 255, 0],
163
+ [0, 133, 255],
164
+ [255, 214, 0],
165
+ [25, 194, 194],
166
+ [102, 255, 0],
167
+ [92, 0, 255],
168
+ ]
169
+
170
+
171
+ class StableDiffusionControlNetSegGenerator:
172
+ def __init__(self):
173
+ self.pipe = None
174
+
175
+ def load_model(
176
+ self,
177
+ stable_model_path,
178
+ scheduler,
179
+ ):
180
+
181
+ if self.pipe is None:
182
+ controlnet = ControlNetModel.from_pretrained(
183
+ "lllyasviel/sd-controlnet-seg", torch_dtype=torch.float16
184
+ )
185
+ self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
186
+ pretrained_model_name_or_path=stable_model_path,
187
+ controlnet=controlnet,
188
+ safety_checker=None,
189
+ torch_dtype=torch.float16,
190
+ )
191
+
192
+ self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
193
+ self.pipe.to("cuda")
194
+ self.pipe.enable_xformers_memory_efficient_attention()
195
+
196
+ return self.pipe
197
+
198
+ def controlnet_seg(self, image_path: str):
199
+ image_processor = AutoImageProcessor.from_pretrained(
200
+ "openmmlab/upernet-convnext-small"
201
+ )
202
+ image_segmentor = UperNetForSemanticSegmentation.from_pretrained(
203
+ "openmmlab/upernet-convnext-small"
204
+ )
205
+
206
+ image = Image.open(image_path).convert("RGB")
207
+ pixel_values = image_processor(image, return_tensors="pt").pixel_values
208
+
209
+ with torch.no_grad():
210
+ outputs = image_segmentor(pixel_values)
211
+
212
+ seg = image_processor.post_process_semantic_segmentation(
213
+ outputs, target_sizes=[image.size[::-1]]
214
+ )[0]
215
+
216
+ color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
217
+ palette = np.array(ade_palette())
218
+
219
+ for label, color in enumerate(palette):
220
+ color_seg[seg == label, :] = color
221
+
222
+ color_seg = color_seg.astype(np.uint8)
223
+ image = Image.fromarray(color_seg)
224
+
225
+ return image
226
+
227
+ def generate_image(
228
+ self,
229
+ image_path: str,
230
+ model_path: str,
231
+ prompt: str,
232
+ negative_prompt: str,
233
+ num_images_per_prompt: int,
234
+ guidance_scale: int,
235
+ num_inference_step: int,
236
+ scheduler: str,
237
+ seed_generator: int,
238
+ ):
239
+
240
+ image = self.controlnet_seg(image_path=image_path)
241
+ pipe = self.load_model(
242
+ stable_model_path=model_path,
243
+ scheduler=scheduler,
244
+ )
245
+ if seed_generator == 0:
246
+ random_seed = torch.randint(0, 1000000, (1,))
247
+ generator = torch.manual_seed(random_seed)
248
+ else:
249
+ generator = torch.manual_seed(seed_generator)
250
+
251
+ output = pipe(
252
+ prompt=prompt,
253
+ image=image,
254
+ negative_prompt=negative_prompt,
255
+ num_images_per_prompt=num_images_per_prompt,
256
+ num_inference_steps=num_inference_step,
257
+ guidance_scale=guidance_scale,
258
+ generator=generator,
259
+ ).images
260
+
261
+ return output
262
+
263
+ def app():
264
+ with gr.Blocks():
265
+ with gr.Row():
266
+ with gr.Column():
267
+ controlnet_seg_image_file = gr.Image(
268
+ type="filepath", label="Image"
269
+ )
270
+
271
+ controlnet_seg_prompt = gr.Textbox(
272
+ lines=1,
273
+ show_label=False,
274
+ placeholder="Prompt",
275
+ )
276
+
277
+ controlnet_seg_negative_prompt = gr.Textbox(
278
+ lines=1,
279
+ show_label=False,
280
+ placeholder="Negative Prompt",
281
+ )
282
+
283
+ with gr.Row():
284
+ with gr.Column():
285
+ controlnet_seg_model_id = gr.Dropdown(
286
+ choices=stable_model_list,
287
+ value=stable_model_list[0],
288
+ label="Stable Model Id",
289
+ )
290
+ controlnet_seg_guidance_scale = gr.Slider(
291
+ minimum=0.1,
292
+ maximum=15,
293
+ step=0.1,
294
+ value=7.5,
295
+ label="Guidance Scale",
296
+ )
297
+
298
+ controlnet_seg_num_inference_step = gr.Slider(
299
+ minimum=1,
300
+ maximum=100,
301
+ step=1,
302
+ value=50,
303
+ label="Num Inference Step",
304
+ )
305
+
306
+ with gr.Row():
307
+ with gr.Column():
308
+ controlnet_seg_scheduler = gr.Dropdown(
309
+ choices=SCHEDULER_LIST,
310
+ value=SCHEDULER_LIST[0],
311
+ label="Scheduler",
312
+ )
313
+ controlnet_seg_num_images_per_prompt = (
314
+ gr.Slider(
315
+ minimum=1,
316
+ maximum=10,
317
+ step=1,
318
+ value=1,
319
+ label="Number Of Images",
320
+ )
321
+ )
322
+ controlnet_seg_seed_generator = gr.Slider(
323
+ minimum=0,
324
+ maximum=1000000,
325
+ step=1,
326
+ value=0,
327
+ label="Seed Generator",
328
+ )
329
+
330
+ controlnet_seg_predict = gr.Button(value="Generator")
331
+
332
+ with gr.Column():
333
+ output_image = gr.Gallery(
334
+ label="Generated images",
335
+ show_label=False,
336
+ elem_id="gallery",
337
+ ).style(grid=(1, 2))
338
+
339
+ controlnet_seg_predict.click(
340
+ fn=StableDiffusionControlNetSegGenerator().generate_image,
341
+ inputs=[
342
+ controlnet_seg_image_file,
343
+ controlnet_seg_model_id,
344
+ controlnet_seg_prompt,
345
+ controlnet_seg_negative_prompt,
346
+ controlnet_seg_num_images_per_prompt,
347
+ controlnet_seg_guidance_scale,
348
+ controlnet_seg_num_inference_step,
349
+ controlnet_seg_scheduler,
350
+ controlnet_seg_seed_generator,
351
+ ],
352
+ outputs=[output_image],
353
+ )
diffusion_webui/diffusion_models/stable_diffusion/__init__.py ADDED
File without changes
diffusion_webui/diffusion_models/stable_diffusion/__pycache__/__init__.cpython-38.pyc ADDED
Binary file (215 Bytes). View file
 
diffusion_webui/diffusion_models/stable_diffusion/__pycache__/img2img_app.cpython-38.pyc ADDED
Binary file (3.48 kB). View file
 
diffusion_webui/diffusion_models/stable_diffusion/__pycache__/inpaint_app.cpython-38.pyc ADDED
Binary file (3.44 kB). View file
 
diffusion_webui/diffusion_models/stable_diffusion/__pycache__/text2img_app.cpython-38.pyc ADDED
Binary file (3.52 kB). View file
 
diffusion_webui/diffusion_models/stable_diffusion/img2img_app.py ADDED
@@ -0,0 +1,153 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from diffusers import StableDiffusionImg2ImgPipeline
4
+ from PIL import Image
5
+
6
+ from diffusion_webui.utils.model_list import stable_model_list
7
+ from diffusion_webui.utils.scheduler_list import (
8
+ SCHEDULER_LIST,
9
+ get_scheduler_list,
10
+ )
11
+
12
+
13
+ class StableDiffusionImage2ImageGenerator:
14
+ def __init__(self):
15
+ self.pipe = None
16
+
17
+ def load_model(self, model_path, scheduler):
18
+ if self.pipe is None:
19
+ self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
20
+ model_path, safety_checker=None, torch_dtype=torch.float16
21
+ )
22
+
23
+ self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
24
+ self.pipe.to("cuda")
25
+ self.pipe.enable_xformers_memory_efficient_attention()
26
+
27
+ return self.pipe
28
+
29
+ def generate_image(
30
+ self,
31
+ image_path: str,
32
+ model_path: str,
33
+ prompt: str,
34
+ negative_prompt: str,
35
+ num_images_per_prompt: int,
36
+ scheduler: str,
37
+ guidance_scale: int,
38
+ num_inference_step: int,
39
+ seed_generator=0,
40
+ ):
41
+ pipe = self.load_model(
42
+ model_path=model_path,
43
+ scheduler=scheduler,
44
+ )
45
+
46
+ if seed_generator == 0:
47
+ random_seed = torch.randint(0, 1000000, (1,))
48
+ generator = torch.manual_seed(random_seed)
49
+ else:
50
+ generator = torch.manual_seed(seed_generator)
51
+
52
+ image = Image.open(image_path)
53
+ images = pipe(
54
+ prompt,
55
+ image=image,
56
+ negative_prompt=negative_prompt,
57
+ num_images_per_prompt=num_images_per_prompt,
58
+ num_inference_steps=num_inference_step,
59
+ guidance_scale=guidance_scale,
60
+ generator=generator,
61
+ ).images
62
+
63
+ return images
64
+
65
+ def app():
66
+ with gr.Blocks():
67
+ with gr.Row():
68
+ with gr.Column():
69
+ image2image_image_file = gr.Image(
70
+ type="filepath", label="Image"
71
+ ).style(height=260)
72
+
73
+ image2image_prompt = gr.Textbox(
74
+ lines=1,
75
+ placeholder="Prompt",
76
+ show_label=False,
77
+ )
78
+
79
+ image2image_negative_prompt = gr.Textbox(
80
+ lines=1,
81
+ placeholder="Negative Prompt",
82
+ show_label=False,
83
+ )
84
+
85
+ with gr.Row():
86
+ with gr.Column():
87
+ image2image_model_path = gr.Dropdown(
88
+ choices=stable_model_list,
89
+ value=stable_model_list[0],
90
+ label="Stable Model Id",
91
+ )
92
+
93
+ image2image_guidance_scale = gr.Slider(
94
+ minimum=0.1,
95
+ maximum=15,
96
+ step=0.1,
97
+ value=7.5,
98
+ label="Guidance Scale",
99
+ )
100
+ image2image_num_inference_step = gr.Slider(
101
+ minimum=1,
102
+ maximum=100,
103
+ step=1,
104
+ value=50,
105
+ label="Num Inference Step",
106
+ )
107
+ with gr.Row():
108
+ with gr.Column():
109
+ image2image_scheduler = gr.Dropdown(
110
+ choices=SCHEDULER_LIST,
111
+ value=SCHEDULER_LIST[0],
112
+ label="Scheduler",
113
+ )
114
+ image2image_num_images_per_prompt = gr.Slider(
115
+ minimum=1,
116
+ maximum=30,
117
+ step=1,
118
+ value=1,
119
+ label="Number Of Images",
120
+ )
121
+
122
+ image2image_seed_generator = gr.Slider(
123
+ minimum=0,
124
+ maximum=1000000,
125
+ step=1,
126
+ value=0,
127
+ label="Seed(0 for random)",
128
+ )
129
+
130
+ image2image_predict_button = gr.Button(value="Generator")
131
+
132
+ with gr.Column():
133
+ output_image = gr.Gallery(
134
+ label="Generated images",
135
+ show_label=False,
136
+ elem_id="gallery",
137
+ ).style(grid=(1, 2))
138
+
139
+ image2image_predict_button.click(
140
+ fn=StableDiffusionImage2ImageGenerator().generate_image,
141
+ inputs=[
142
+ image2image_image_file,
143
+ image2image_model_path,
144
+ image2image_prompt,
145
+ image2image_negative_prompt,
146
+ image2image_num_images_per_prompt,
147
+ image2image_scheduler,
148
+ image2image_guidance_scale,
149
+ image2image_num_inference_step,
150
+ image2image_seed_generator,
151
+ ],
152
+ outputs=[output_image],
153
+ )
diffusion_webui/diffusion_models/stable_diffusion/inpaint_app.py ADDED
@@ -0,0 +1,148 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from diffusers import DiffusionPipeline
4
+
5
+ from diffusion_webui.utils.model_list import stable_inpiant_model_list
6
+
7
+
8
+ class StableDiffusionInpaintGenerator:
9
+ def __init__(self):
10
+ self.pipe = None
11
+
12
+ def load_model(self, model_path):
13
+ if self.pipe is None:
14
+ self.pipe = DiffusionPipeline.from_pretrained(
15
+ model_path, revision="fp16", torch_dtype=torch.float16
16
+ )
17
+
18
+ self.pipe.to("cuda")
19
+ self.pipe.enable_xformers_memory_efficient_attention()
20
+
21
+ return self.pipe
22
+
23
+ def generate_image(
24
+ self,
25
+ pil_image: str,
26
+ model_path: str,
27
+ prompt: str,
28
+ negative_prompt: str,
29
+ num_images_per_prompt: int,
30
+ guidance_scale: int,
31
+ num_inference_step: int,
32
+ seed_generator=0,
33
+ ):
34
+ image = pil_image["image"].convert("RGB").resize((512, 512))
35
+ mask_image = pil_image["mask"].convert("RGB").resize((512, 512))
36
+ pipe = self.load_model(model_path)
37
+
38
+ if seed_generator == 0:
39
+ random_seed = torch.randint(0, 1000000, (1,))
40
+ generator = torch.manual_seed(random_seed)
41
+ else:
42
+ generator = torch.manual_seed(seed_generator)
43
+
44
+ output = pipe(
45
+ prompt=prompt,
46
+ image=image,
47
+ mask_image=mask_image,
48
+ negative_prompt=negative_prompt,
49
+ num_images_per_prompt=num_images_per_prompt,
50
+ num_inference_steps=num_inference_step,
51
+ guidance_scale=guidance_scale,
52
+ generator=generator,
53
+ ).images
54
+
55
+ return output
56
+
57
+ def app():
58
+ with gr.Blocks():
59
+ with gr.Row():
60
+ with gr.Column():
61
+ stable_diffusion_inpaint_image_file = gr.Image(
62
+ source="upload",
63
+ tool="sketch",
64
+ elem_id="image_upload",
65
+ type="pil",
66
+ label="Upload",
67
+ ).style(height=260)
68
+
69
+ stable_diffusion_inpaint_prompt = gr.Textbox(
70
+ lines=1,
71
+ placeholder="Prompt",
72
+ show_label=False,
73
+ )
74
+
75
+ stable_diffusion_inpaint_negative_prompt = gr.Textbox(
76
+ lines=1,
77
+ placeholder="Negative Prompt",
78
+ show_label=False,
79
+ )
80
+ stable_diffusion_inpaint_model_id = gr.Dropdown(
81
+ choices=stable_inpiant_model_list,
82
+ value=stable_inpiant_model_list[0],
83
+ label="Inpaint Model Id",
84
+ )
85
+ with gr.Row():
86
+ with gr.Column():
87
+ stable_diffusion_inpaint_guidance_scale = gr.Slider(
88
+ minimum=0.1,
89
+ maximum=15,
90
+ step=0.1,
91
+ value=7.5,
92
+ label="Guidance Scale",
93
+ )
94
+
95
+ stable_diffusion_inpaint_num_inference_step = (
96
+ gr.Slider(
97
+ minimum=1,
98
+ maximum=100,
99
+ step=1,
100
+ value=50,
101
+ label="Num Inference Step",
102
+ )
103
+ )
104
+
105
+ with gr.Row():
106
+ with gr.Column():
107
+ stable_diffusion_inpiant_num_images_per_prompt = gr.Slider(
108
+ minimum=1,
109
+ maximum=10,
110
+ step=1,
111
+ value=1,
112
+ label="Number Of Images",
113
+ )
114
+ stable_diffusion_inpaint_seed_generator = (
115
+ gr.Slider(
116
+ minimum=0,
117
+ maximum=1000000,
118
+ step=1,
119
+ value=0,
120
+ label="Seed(0 for random)",
121
+ )
122
+ )
123
+
124
+ stable_diffusion_inpaint_predict = gr.Button(
125
+ value="Generator"
126
+ )
127
+
128
+ with gr.Column():
129
+ output_image = gr.Gallery(
130
+ label="Generated images",
131
+ show_label=False,
132
+ elem_id="gallery",
133
+ ).style(grid=(1, 2))
134
+
135
+ stable_diffusion_inpaint_predict.click(
136
+ fn=StableDiffusionInpaintGenerator().generate_image,
137
+ inputs=[
138
+ stable_diffusion_inpaint_image_file,
139
+ stable_diffusion_inpaint_model_id,
140
+ stable_diffusion_inpaint_prompt,
141
+ stable_diffusion_inpaint_negative_prompt,
142
+ stable_diffusion_inpiant_num_images_per_prompt,
143
+ stable_diffusion_inpaint_guidance_scale,
144
+ stable_diffusion_inpaint_num_inference_step,
145
+ stable_diffusion_inpaint_seed_generator,
146
+ ],
147
+ outputs=[output_image],
148
+ )
diffusion_webui/diffusion_models/stable_diffusion/text2img_app.py ADDED
@@ -0,0 +1,170 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from diffusers import StableDiffusionPipeline
4
+
5
+ from diffusion_webui.utils.model_list import stable_model_list
6
+ from diffusion_webui.utils.scheduler_list import get_scheduler_list
7
+
8
+
9
+ class StableDiffusionText2ImageGenerator:
10
+ def __init__(self):
11
+ self.pipe = None
12
+
13
+ def load_model(
14
+ self,
15
+ model_path,
16
+ scheduler,
17
+ ):
18
+ if self.pipe is None:
19
+ self.pipe = StableDiffusionPipeline.from_pretrained(
20
+ model_path, safety_checker=None, torch_dtype=torch.float16
21
+ )
22
+
23
+ self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
24
+ self.pipe.to("cuda")
25
+ self.pipe.enable_xformers_memory_efficient_attention()
26
+
27
+ return self.pipe
28
+
29
+ def generate_image(
30
+ self,
31
+ model_path: str,
32
+ prompt: str,
33
+ negative_prompt: str,
34
+ num_images_per_prompt: int,
35
+ scheduler: str,
36
+ guidance_scale: int,
37
+ num_inference_step: int,
38
+ height: int,
39
+ width: int,
40
+ seed_generator=0,
41
+ ):
42
+ pipe = self.load_model(
43
+ model_path=model_path,
44
+ scheduler=scheduler,
45
+ )
46
+ if seed_generator == 0:
47
+ random_seed = torch.randint(0, 1000000, (1,))
48
+ generator = torch.manual_seed(random_seed)
49
+ else:
50
+ generator = torch.manual_seed(seed_generator)
51
+
52
+ images = pipe(
53
+ prompt=prompt,
54
+ height=height,
55
+ width=width,
56
+ negative_prompt=negative_prompt,
57
+ num_images_per_prompt=num_images_per_prompt,
58
+ num_inference_steps=num_inference_step,
59
+ guidance_scale=guidance_scale,
60
+ generator=generator,
61
+ ).images
62
+
63
+ return images
64
+
65
+ def app():
66
+ with gr.Blocks():
67
+ with gr.Row():
68
+ with gr.Column():
69
+ text2image_prompt = gr.Textbox(
70
+ lines=1,
71
+ placeholder="Prompt",
72
+ show_label=False,
73
+ )
74
+
75
+ text2image_negative_prompt = gr.Textbox(
76
+ lines=1,
77
+ placeholder="Negative Prompt",
78
+ show_label=False,
79
+ )
80
+ with gr.Row():
81
+ with gr.Column():
82
+ text2image_model_path = gr.Dropdown(
83
+ choices=stable_model_list,
84
+ value=stable_model_list[0],
85
+ label="Text-Image Model Id",
86
+ )
87
+
88
+ text2image_guidance_scale = gr.Slider(
89
+ minimum=0.1,
90
+ maximum=15,
91
+ step=0.1,
92
+ value=7.5,
93
+ label="Guidance Scale",
94
+ )
95
+
96
+ text2image_num_inference_step = gr.Slider(
97
+ minimum=1,
98
+ maximum=100,
99
+ step=1,
100
+ value=50,
101
+ label="Num Inference Step",
102
+ )
103
+ text2image_num_images_per_prompt = gr.Slider(
104
+ minimum=1,
105
+ maximum=30,
106
+ step=1,
107
+ value=1,
108
+ label="Number Of Images",
109
+ )
110
+ with gr.Row():
111
+ with gr.Column():
112
+
113
+ text2image_scheduler = gr.Dropdown(
114
+ choices=[
115
+ "DDIM",
116
+ "EulerA",
117
+ "Euler",
118
+ "LMS",
119
+ "Heun",
120
+ ],
121
+ value="DDIM",
122
+ label="Scheduler",
123
+ )
124
+
125
+ text2image_height = gr.Slider(
126
+ minimum=128,
127
+ maximum=1280,
128
+ step=32,
129
+ value=512,
130
+ label="Image Height",
131
+ )
132
+
133
+ text2image_width = gr.Slider(
134
+ minimum=128,
135
+ maximum=1280,
136
+ step=32,
137
+ value=512,
138
+ label="Image Width",
139
+ )
140
+ text2image_seed_generator = gr.Slider(
141
+ label="Seed(0 for random)",
142
+ minimum=0,
143
+ maximum=1000000,
144
+ value=0,
145
+ )
146
+ text2image_predict = gr.Button(value="Generator")
147
+
148
+ with gr.Column():
149
+ output_image = gr.Gallery(
150
+ label="Generated images",
151
+ show_label=False,
152
+ elem_id="gallery",
153
+ ).style(grid=(1, 2), height=200)
154
+
155
+ text2image_predict.click(
156
+ fn=StableDiffusionText2ImageGenerator().generate_image,
157
+ inputs=[
158
+ text2image_model_path,
159
+ text2image_prompt,
160
+ text2image_negative_prompt,
161
+ text2image_num_images_per_prompt,
162
+ text2image_scheduler,
163
+ text2image_guidance_scale,
164
+ text2image_num_inference_step,
165
+ text2image_height,
166
+ text2image_width,
167
+ text2image_seed_generator,
168
+ ],
169
+ outputs=output_image,
170
+ )
diffusion_webui/helpers.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusion_webui.diffusion_models.controlnet.controlnet_canny import (
2
+ StableDiffusionControlNetCannyGenerator,
3
+ )
4
+ from diffusion_webui.diffusion_models.controlnet.controlnet_depth import (
5
+ StableDiffusionControlNetDepthGenerator,
6
+ )
7
+ from diffusion_webui.diffusion_models.controlnet.controlnet_hed import (
8
+ StableDiffusionControlNetHEDGenerator,
9
+ )
10
+ from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.controlnet_inpaint_app import (
11
+ StableDiffusionControlInpaintNetCannyGenerator,
12
+ )
13
+ from diffusion_webui.diffusion_models.controlnet.controlnet_mlsd import (
14
+ StableDiffusionControlNetMLSDGenerator,
15
+ )
16
+ from diffusion_webui.diffusion_models.controlnet.controlnet_pose import (
17
+ StableDiffusionControlNetPoseGenerator,
18
+ )
19
+ from diffusion_webui.diffusion_models.controlnet.controlnet_scribble import (
20
+ StableDiffusionControlNetScribbleGenerator,
21
+ )
22
+ from diffusion_webui.diffusion_models.controlnet.controlnet_seg import (
23
+ StableDiffusionControlNetSegGenerator,
24
+ )
25
+ from diffusion_webui.diffusion_models.stable_diffusion.img2img_app import (
26
+ StableDiffusionImage2ImageGenerator,
27
+ )
28
+ from diffusion_webui.diffusion_models.stable_diffusion.inpaint_app import (
29
+ StableDiffusionInpaintGenerator,
30
+ )
31
+ from diffusion_webui.diffusion_models.stable_diffusion.text2img_app import (
32
+ StableDiffusionText2ImageGenerator,
33
+ )
diffusion_webui/utils/__init__.py ADDED
File without changes
diffusion_webui/utils/model_list.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ stable_model_list = [
2
+ "runwayml/stable-diffusion-v1-5",
3
+ "stabilityai/stable-diffusion-2-1",
4
+ ]
5
+
6
+ controlnet_canny_model_list = [
7
+ "lllyasviel/sd-controlnet-canny",
8
+ "thibaud/controlnet-sd21-canny-diffusers",
9
+ ]
10
+
11
+ controlnet_depth_model_list = [
12
+ "lllyasviel/sd-controlnet-depth",
13
+ "thibaud/controlnet-sd21-depth-diffusers",
14
+ ]
15
+
16
+ controlnet_pose_model_list = [
17
+ "lllyasviel/sd-controlnet-openpose",
18
+ "thibaud/controlnet-sd21-openpose-diffusers",
19
+ ]
20
+
21
+ controlnet_hed_model_list = [
22
+ "lllyasviel/sd-controlnet-hed",
23
+ "thibaud/controlnet-sd21-hed-diffusers",
24
+ ]
25
+
26
+ controlnet_scribble_model_list = [
27
+ "lllyasviel/sd-controlnet-scribble",
28
+ "thibaud/controlnet-sd21-scribble-diffusers",
29
+ ]
30
+ stable_inpiant_model_list = [
31
+ "stabilityai/stable-diffusion-2-inpainting",
32
+ "runwayml/stable-diffusion-inpainting",
33
+ ]
diffusion_webui/utils/scheduler_list.py ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusers import (
2
+ DDIMScheduler,
3
+ EulerAncestralDiscreteScheduler,
4
+ EulerDiscreteScheduler,
5
+ HeunDiscreteScheduler,
6
+ LMSDiscreteScheduler,
7
+ UniPCMultistepScheduler,
8
+ )
9
+
10
+ SCHEDULER_LIST = [
11
+ "DDIM",
12
+ "EulerA",
13
+ "Euler",
14
+ "LMS",
15
+ "Heun",
16
+ "UniPC",
17
+ ]
18
+
19
+
20
+ def get_scheduler_list(pipe, scheduler):
21
+ if scheduler == SCHEDULER_LIST[0]:
22
+ pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
23
+
24
+ elif scheduler == SCHEDULER_LIST[1]:
25
+ pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(
26
+ pipe.scheduler.config
27
+ )
28
+
29
+ elif scheduler == SCHEDULER_LIST[2]:
30
+ pipe.scheduler = EulerDiscreteScheduler.from_config(
31
+ pipe.scheduler.config
32
+ )
33
+
34
+ elif scheduler == SCHEDULER_LIST[3]:
35
+ pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config)
36
+
37
+ elif scheduler == SCHEDULER_LIST[4]:
38
+ pipe.scheduler = HeunDiscreteScheduler.from_config(
39
+ pipe.scheduler.config
40
+ )
41
+
42
+ elif scheduler == SCHEDULER_LIST[5]:
43
+ pipe.scheduler = UniPCMultistepScheduler.from_config(
44
+ pipe.scheduler.config
45
+ )
46
+
47
+ return pipe