suvadityamuk
commited on
Commit
·
73472c7
1
Parent(s):
2a034bd
Delete pipeline.py
Browse files- pipeline.py +0 -365
pipeline.py
DELETED
@@ -1,365 +0,0 @@
|
|
1 |
-
from typing import Any, Callable, Dict, List, Optional, Union
|
2 |
-
|
3 |
-
import torch
|
4 |
-
|
5 |
-
from diffusers import DiffusionPipeline, StableDiffusionPipeline
|
6 |
-
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
|
7 |
-
|
8 |
-
|
9 |
-
pipe1_model_id = "CompVis/stable-diffusion-v1-1"
|
10 |
-
pipe2_model_id = "CompVis/stable-diffusion-v1-2"
|
11 |
-
pipe3_model_id = "CompVis/stable-diffusion-v1-3"
|
12 |
-
pipe4_model_id = "CompVis/stable-diffusion-v1-4"
|
13 |
-
|
14 |
-
|
15 |
-
class StableDiffusionComparisonPipeline(DiffusionPipeline):
|
16 |
-
r"""
|
17 |
-
Pipeline for parallel comparison of Stable Diffusion v1-v4
|
18 |
-
This pipeline inherits from DiffusionPipeline and depends on the use of an Auth Token for
|
19 |
-
downloading pre-trained checkpoints from Hugging Face Hub.
|
20 |
-
Args:
|
21 |
-
pipe1 ('StableDiffusionPipeline' or 'str', optional):
|
22 |
-
A Stable Diffusion Pipeline prepared from the SD1.1 Checkpoints on Hugging Face Hub
|
23 |
-
pipe2 ('StableDiffusionPipeline' or 'str', optional):
|
24 |
-
A Stable Diffusion Pipeline prepared from the SD1.2 Checkpoints on Hugging Face Hub
|
25 |
-
pipe3 ('StableDiffusionPipeline' or 'str', optional):
|
26 |
-
A Stable Diffusion Pipeline prepared from the SD1.3 Checkpoints on Hugging Face Hub
|
27 |
-
pipe4 ('StableDiffusionPipeline' or 'str', optional):
|
28 |
-
A Stable Diffusion Pipeline prepared from the SD1.4 Checkpoints on Hugging Face Hub
|
29 |
-
"""
|
30 |
-
|
31 |
-
def _init_(
|
32 |
-
self,
|
33 |
-
sd1_1: Union[StableDiffusionPipeline, str],
|
34 |
-
sd1_2: Union[StableDiffusionPipeline, str],
|
35 |
-
sd1_3: Union[StableDiffusionPipeline, str],
|
36 |
-
sd1_4: Union[StableDiffusionPipeline, str],
|
37 |
-
):
|
38 |
-
super()._init_()
|
39 |
-
|
40 |
-
if not isinstance(sd1_1, StableDiffusionPipeline):
|
41 |
-
self.pipe1 = StableDiffusionPipeline.from_pretrained(
|
42 |
-
pipe1_model_id, torch_dtype=torch.float16, revision="fp16", use_auth_token=True
|
43 |
-
)
|
44 |
-
else:
|
45 |
-
self.pipe1 = sd1_1
|
46 |
-
if not isinstance(sd1_2, StableDiffusionPipeline):
|
47 |
-
self.pipe2 = StableDiffusionPipeline.from_pretrained(
|
48 |
-
pipe2_model_id, torch_dtype=torch.float16, revision="fp16", use_auth_token=True
|
49 |
-
)
|
50 |
-
else:
|
51 |
-
self.pipe2 = sd1_2
|
52 |
-
if not isinstance(sd1_3, StableDiffusionPipeline):
|
53 |
-
self.pipe3 = StableDiffusionPipeline.from_pretrained(
|
54 |
-
pipe3_model_id, torch_dtype=torch.float16, revision="fp16", use_auth_token=True
|
55 |
-
)
|
56 |
-
else:
|
57 |
-
self.pipe3 = sd1_3
|
58 |
-
if not isinstance(sd1_4, StableDiffusionPipeline):
|
59 |
-
self.pipe4 = StableDiffusionPipeline.from_pretrained(
|
60 |
-
pipe4_model_id, torch_dtype=torch.float16, revision="fp16", use_auth_token=True
|
61 |
-
)
|
62 |
-
else:
|
63 |
-
self.pipe4 = sd1_4
|
64 |
-
|
65 |
-
self.register_modules(pipeline1=self.pipe1, pipeline2=self.pipe2, pipeline3=self.pipe3, pipeline4=self.pipe4)
|
66 |
-
|
67 |
-
@property
|
68 |
-
def layers(self) -> Dict[str, Any]:
|
69 |
-
return {k: getattr(self, k) for k in self.config.keys() if not k.startswith("_")}
|
70 |
-
|
71 |
-
@torch.no_grad()
|
72 |
-
def text2img_sd1_1(
|
73 |
-
self,
|
74 |
-
prompt: Union[str, List[str]],
|
75 |
-
height: int = 512,
|
76 |
-
width: int = 512,
|
77 |
-
num_inference_steps: int = 50,
|
78 |
-
guidance_scale: float = 7.5,
|
79 |
-
negative_prompt: Optional[Union[str, List[str]]] = None,
|
80 |
-
num_images_per_prompt: Optional[int] = 1,
|
81 |
-
eta: float = 0.0,
|
82 |
-
generator: Optional[torch.Generator] = None,
|
83 |
-
latents: Optional[torch.FloatTensor] = None,
|
84 |
-
output_type: Optional[str] = "pil",
|
85 |
-
return_dict: bool = True,
|
86 |
-
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
|
87 |
-
callback_steps: Optional[int] = 1,
|
88 |
-
**kwargs,
|
89 |
-
):
|
90 |
-
return self.pipe1(
|
91 |
-
prompt=prompt,
|
92 |
-
height=height,
|
93 |
-
width=width,
|
94 |
-
num_inference_steps=num_inference_steps,
|
95 |
-
guidance_scale=guidance_scale,
|
96 |
-
negative_prompt=negative_prompt,
|
97 |
-
num_images_per_prompt=num_images_per_prompt,
|
98 |
-
eta=eta,
|
99 |
-
generator=generator,
|
100 |
-
latents=latents,
|
101 |
-
output_type=output_type,
|
102 |
-
return_dict=return_dict,
|
103 |
-
callback=callback,
|
104 |
-
callback_steps=callback_steps,
|
105 |
-
**kwargs,
|
106 |
-
)
|
107 |
-
|
108 |
-
@torch.no_grad()
|
109 |
-
def text2img_sd1_2(
|
110 |
-
self,
|
111 |
-
prompt: Union[str, List[str]],
|
112 |
-
height: int = 512,
|
113 |
-
width: int = 512,
|
114 |
-
num_inference_steps: int = 50,
|
115 |
-
guidance_scale: float = 7.5,
|
116 |
-
negative_prompt: Optional[Union[str, List[str]]] = None,
|
117 |
-
num_images_per_prompt: Optional[int] = 1,
|
118 |
-
eta: float = 0.0,
|
119 |
-
generator: Optional[torch.Generator] = None,
|
120 |
-
latents: Optional[torch.FloatTensor] = None,
|
121 |
-
output_type: Optional[str] = "pil",
|
122 |
-
return_dict: bool = True,
|
123 |
-
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
|
124 |
-
callback_steps: Optional[int] = 1,
|
125 |
-
**kwargs,
|
126 |
-
):
|
127 |
-
return self.pipe2(
|
128 |
-
prompt=prompt,
|
129 |
-
height=height,
|
130 |
-
width=width,
|
131 |
-
num_inference_steps=num_inference_steps,
|
132 |
-
guidance_scale=guidance_scale,
|
133 |
-
negative_prompt=negative_prompt,
|
134 |
-
num_images_per_prompt=num_images_per_prompt,
|
135 |
-
eta=eta,
|
136 |
-
generator=generator,
|
137 |
-
latents=latents,
|
138 |
-
output_type=output_type,
|
139 |
-
return_dict=return_dict,
|
140 |
-
callback=callback,
|
141 |
-
callback_steps=callback_steps,
|
142 |
-
**kwargs,
|
143 |
-
)
|
144 |
-
|
145 |
-
@torch.no_grad()
|
146 |
-
def text2img_sd1_3(
|
147 |
-
self,
|
148 |
-
prompt: Union[str, List[str]],
|
149 |
-
height: int = 512,
|
150 |
-
width: int = 512,
|
151 |
-
num_inference_steps: int = 50,
|
152 |
-
guidance_scale: float = 7.5,
|
153 |
-
negative_prompt: Optional[Union[str, List[str]]] = None,
|
154 |
-
num_images_per_prompt: Optional[int] = 1,
|
155 |
-
eta: float = 0.0,
|
156 |
-
generator: Optional[torch.Generator] = None,
|
157 |
-
latents: Optional[torch.FloatTensor] = None,
|
158 |
-
output_type: Optional[str] = "pil",
|
159 |
-
return_dict: bool = True,
|
160 |
-
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
|
161 |
-
callback_steps: Optional[int] = 1,
|
162 |
-
**kwargs,
|
163 |
-
):
|
164 |
-
return self.pipe3(
|
165 |
-
prompt=prompt,
|
166 |
-
height=height,
|
167 |
-
width=width,
|
168 |
-
num_inference_steps=num_inference_steps,
|
169 |
-
guidance_scale=guidance_scale,
|
170 |
-
negative_prompt=negative_prompt,
|
171 |
-
num_images_per_prompt=num_images_per_prompt,
|
172 |
-
eta=eta,
|
173 |
-
generator=generator,
|
174 |
-
latents=latents,
|
175 |
-
output_type=output_type,
|
176 |
-
return_dict=return_dict,
|
177 |
-
callback=callback,
|
178 |
-
callback_steps=callback_steps,
|
179 |
-
**kwargs,
|
180 |
-
)
|
181 |
-
|
182 |
-
@torch.no_grad()
|
183 |
-
def text2img_sd1_4(
|
184 |
-
self,
|
185 |
-
prompt: Union[str, List[str]],
|
186 |
-
height: int = 512,
|
187 |
-
width: int = 512,
|
188 |
-
num_inference_steps: int = 50,
|
189 |
-
guidance_scale: float = 7.5,
|
190 |
-
negative_prompt: Optional[Union[str, List[str]]] = None,
|
191 |
-
num_images_per_prompt: Optional[int] = 1,
|
192 |
-
eta: float = 0.0,
|
193 |
-
generator: Optional[torch.Generator] = None,
|
194 |
-
latents: Optional[torch.FloatTensor] = None,
|
195 |
-
output_type: Optional[str] = "pil",
|
196 |
-
return_dict: bool = True,
|
197 |
-
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
|
198 |
-
callback_steps: Optional[int] = 1,
|
199 |
-
**kwargs,
|
200 |
-
):
|
201 |
-
return self.pipe4(
|
202 |
-
prompt=prompt,
|
203 |
-
height=height,
|
204 |
-
width=width,
|
205 |
-
num_inference_steps=num_inference_steps,
|
206 |
-
guidance_scale=guidance_scale,
|
207 |
-
negative_prompt=negative_prompt,
|
208 |
-
num_images_per_prompt=num_images_per_prompt,
|
209 |
-
eta=eta,
|
210 |
-
generator=generator,
|
211 |
-
latents=latents,
|
212 |
-
output_type=output_type,
|
213 |
-
return_dict=return_dict,
|
214 |
-
callback=callback,
|
215 |
-
callback_steps=callback_steps,
|
216 |
-
**kwargs,
|
217 |
-
)
|
218 |
-
|
219 |
-
@torch.no_grad()
|
220 |
-
def _call_(
|
221 |
-
self,
|
222 |
-
prompt: Union[str, List[str]],
|
223 |
-
height: int = 512,
|
224 |
-
width: int = 512,
|
225 |
-
num_inference_steps: int = 50,
|
226 |
-
guidance_scale: float = 7.5,
|
227 |
-
negative_prompt: Optional[Union[str, List[str]]] = None,
|
228 |
-
num_images_per_prompt: Optional[int] = 1,
|
229 |
-
eta: float = 0.0,
|
230 |
-
generator: Optional[torch.Generator] = None,
|
231 |
-
latents: Optional[torch.FloatTensor] = None,
|
232 |
-
output_type: Optional[str] = "pil",
|
233 |
-
return_dict: bool = True,
|
234 |
-
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
|
235 |
-
callback_steps: Optional[int] = 1,
|
236 |
-
**kwargs,
|
237 |
-
):
|
238 |
-
r"""
|
239 |
-
Function invoked when calling the pipeline for generation. This function will generate 4 results as part
|
240 |
-
of running all the 4 pipelines for SD1.1-1.4 together in a serial-processing, parallel-invocation fashion.
|
241 |
-
Args:
|
242 |
-
prompt (`str` or `List[str]`):
|
243 |
-
The prompt or prompts to guide the image generation.
|
244 |
-
height (`int`, optional, defaults to 512):
|
245 |
-
The height in pixels of the generated image.
|
246 |
-
width (`int`, optional, defaults to 512):
|
247 |
-
The width in pixels of the generated image.
|
248 |
-
num_inference_steps (`int`, optional, defaults to 50):
|
249 |
-
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
|
250 |
-
expense of slower inference.
|
251 |
-
guidance_scale (`float`, optional, defaults to 7.5):
|
252 |
-
Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
|
253 |
-
`guidance_scale` is defined as `w` of equation 2. of [Imagen
|
254 |
-
Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
|
255 |
-
1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
|
256 |
-
usually at the expense of lower image quality.
|
257 |
-
eta (`float`, optional, defaults to 0.0):
|
258 |
-
Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
|
259 |
-
[`schedulers.DDIMScheduler`], will be ignored for others.
|
260 |
-
generator (`torch.Generator`, optional):
|
261 |
-
A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation
|
262 |
-
deterministic.
|
263 |
-
latents (`torch.FloatTensor`, optional):
|
264 |
-
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
|
265 |
-
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
|
266 |
-
tensor will ge generated by sampling using the supplied random `generator`.
|
267 |
-
output_type (`str`, optional, defaults to `"pil"`):
|
268 |
-
The output format of the generate image. Choose between
|
269 |
-
[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
|
270 |
-
return_dict (`bool`, optional, defaults to `True`):
|
271 |
-
Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
|
272 |
-
plain tuple.
|
273 |
-
Returns:
|
274 |
-
[`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
|
275 |
-
[`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
|
276 |
-
When returning a tuple, the first element is a list with the generated images, and the second element is a
|
277 |
-
list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
|
278 |
-
(nsfw) content, according to the `safety_checker`.
|
279 |
-
"""
|
280 |
-
|
281 |
-
device = "cuda" if torch.cuda.is_available() else "cpu"
|
282 |
-
self.to(device)
|
283 |
-
|
284 |
-
# Checks if the height and width are divisible by 8 or not
|
285 |
-
if height % 8 != 0 or width % 8 != 0:
|
286 |
-
raise ValueError(f"`height` and `width` must be divisible by 8 but are {height} and {width}.")
|
287 |
-
|
288 |
-
# Get first result from Stable Diffusion Checkpoint v1.1
|
289 |
-
res1 = self.text2img_sd1_1(
|
290 |
-
prompt=prompt,
|
291 |
-
height=height,
|
292 |
-
width=width,
|
293 |
-
num_inference_steps=num_inference_steps,
|
294 |
-
guidance_scale=guidance_scale,
|
295 |
-
negative_prompt=negative_prompt,
|
296 |
-
num_images_per_prompt=num_images_per_prompt,
|
297 |
-
eta=eta,
|
298 |
-
generator=generator,
|
299 |
-
latents=latents,
|
300 |
-
output_type=output_type,
|
301 |
-
return_dict=return_dict,
|
302 |
-
callback=callback,
|
303 |
-
callback_steps=callback_steps,
|
304 |
-
**kwargs,
|
305 |
-
)
|
306 |
-
|
307 |
-
# Get first result from Stable Diffusion Checkpoint v1.2
|
308 |
-
res2 = self.text2img_sd1_2(
|
309 |
-
prompt=prompt,
|
310 |
-
height=height,
|
311 |
-
width=width,
|
312 |
-
num_inference_steps=num_inference_steps,
|
313 |
-
guidance_scale=guidance_scale,
|
314 |
-
negative_prompt=negative_prompt,
|
315 |
-
num_images_per_prompt=num_images_per_prompt,
|
316 |
-
eta=eta,
|
317 |
-
generator=generator,
|
318 |
-
latents=latents,
|
319 |
-
output_type=output_type,
|
320 |
-
return_dict=return_dict,
|
321 |
-
callback=callback,
|
322 |
-
callback_steps=callback_steps,
|
323 |
-
**kwargs,
|
324 |
-
)
|
325 |
-
|
326 |
-
# Get first result from Stable Diffusion Checkpoint v1.3
|
327 |
-
res3 = self.text2img_sd1_3(
|
328 |
-
prompt=prompt,
|
329 |
-
height=height,
|
330 |
-
width=width,
|
331 |
-
num_inference_steps=num_inference_steps,
|
332 |
-
guidance_scale=guidance_scale,
|
333 |
-
negative_prompt=negative_prompt,
|
334 |
-
num_images_per_prompt=num_images_per_prompt,
|
335 |
-
eta=eta,
|
336 |
-
generator=generator,
|
337 |
-
latents=latents,
|
338 |
-
output_type=output_type,
|
339 |
-
return_dict=return_dict,
|
340 |
-
callback=callback,
|
341 |
-
callback_steps=callback_steps,
|
342 |
-
**kwargs,
|
343 |
-
)
|
344 |
-
|
345 |
-
# Get first result from Stable Diffusion Checkpoint v1.4
|
346 |
-
res4 = self.text2img_sd1_4(
|
347 |
-
prompt=prompt,
|
348 |
-
height=height,
|
349 |
-
width=width,
|
350 |
-
num_inference_steps=num_inference_steps,
|
351 |
-
guidance_scale=guidance_scale,
|
352 |
-
negative_prompt=negative_prompt,
|
353 |
-
num_images_per_prompt=num_images_per_prompt,
|
354 |
-
eta=eta,
|
355 |
-
generator=generator,
|
356 |
-
latents=latents,
|
357 |
-
output_type=output_type,
|
358 |
-
return_dict=return_dict,
|
359 |
-
callback=callback,
|
360 |
-
callback_steps=callback_steps,
|
361 |
-
**kwargs,
|
362 |
-
)
|
363 |
-
|
364 |
-
# Get all result images into a single list and pass it via StableDiffusionPipelineOutput for final result
|
365 |
-
return StableDiffusionPipelineOutput([res1[0], res2[0], res3[0], res4[0]])
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|