Spaces:
Sleeping
Sleeping
Delete diffusers_patches.py
Browse files- diffusers_patches.py +0 -541
diffusers_patches.py
DELETED
@@ -1,541 +0,0 @@
|
|
1 |
-
import torch
|
2 |
-
from diffusers import ImagePipelineOutput, PixArtAlphaPipeline, AutoencoderKL, Transformer2DModel, \
|
3 |
-
DPMSolverMultistepScheduler
|
4 |
-
from diffusers.image_processor import VaeImageProcessor
|
5 |
-
from diffusers.models.attention import BasicTransformerBlock
|
6 |
-
from diffusers.models.embeddings import PixArtAlphaTextProjection, PatchEmbed
|
7 |
-
from diffusers.models.normalization import AdaLayerNormSingle
|
8 |
-
from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha import retrieve_timesteps
|
9 |
-
from typing import Callable, List, Optional, Tuple, Union
|
10 |
-
|
11 |
-
from diffusers.utils import deprecate
|
12 |
-
from torch import nn
|
13 |
-
from transformers import T5Tokenizer, T5EncoderModel
|
14 |
-
|
15 |
-
ASPECT_RATIO_2048_BIN = {
|
16 |
-
"0.25": [1024.0, 4096.0],
|
17 |
-
"0.26": [1024.0, 3968.0],
|
18 |
-
"0.27": [1024.0, 3840.0],
|
19 |
-
"0.28": [1024.0, 3712.0],
|
20 |
-
"0.32": [1152.0, 3584.0],
|
21 |
-
"0.33": [1152.0, 3456.0],
|
22 |
-
"0.35": [1152.0, 3328.0],
|
23 |
-
"0.4": [1280.0, 3200.0],
|
24 |
-
"0.42": [1280.0, 3072.0],
|
25 |
-
"0.48": [1408.0, 2944.0],
|
26 |
-
"0.5": [1408.0, 2816.0],
|
27 |
-
"0.52": [1408.0, 2688.0],
|
28 |
-
"0.57": [1536.0, 2688.0],
|
29 |
-
"0.6": [1536.0, 2560.0],
|
30 |
-
"0.68": [1664.0, 2432.0],
|
31 |
-
"0.72": [1664.0, 2304.0],
|
32 |
-
"0.78": [1792.0, 2304.0],
|
33 |
-
"0.82": [1792.0, 2176.0],
|
34 |
-
"0.88": [1920.0, 2176.0],
|
35 |
-
"0.94": [1920.0, 2048.0],
|
36 |
-
"1.0": [2048.0, 2048.0],
|
37 |
-
"1.07": [2048.0, 1920.0],
|
38 |
-
"1.13": [2176.0, 1920.0],
|
39 |
-
"1.21": [2176.0, 1792.0],
|
40 |
-
"1.29": [2304.0, 1792.0],
|
41 |
-
"1.38": [2304.0, 1664.0],
|
42 |
-
"1.46": [2432.0, 1664.0],
|
43 |
-
"1.67": [2560.0, 1536.0],
|
44 |
-
"1.75": [2688.0, 1536.0],
|
45 |
-
"2.0": [2816.0, 1408.0],
|
46 |
-
"2.09": [2944.0, 1408.0],
|
47 |
-
"2.4": [3072.0, 1280.0],
|
48 |
-
"2.5": [3200.0, 1280.0],
|
49 |
-
"2.89": [3328.0, 1152.0],
|
50 |
-
"3.0": [3456.0, 1152.0],
|
51 |
-
"3.11": [3584.0, 1152.0],
|
52 |
-
"3.62": [3712.0, 1024.0],
|
53 |
-
"3.75": [3840.0, 1024.0],
|
54 |
-
"3.88": [3968.0, 1024.0],
|
55 |
-
"4.0": [4096.0, 1024.0]
|
56 |
-
}
|
57 |
-
|
58 |
-
ASPECT_RATIO_256_BIN = {
|
59 |
-
"0.25": [128.0, 512.0],
|
60 |
-
"0.28": [128.0, 464.0],
|
61 |
-
"0.32": [144.0, 448.0],
|
62 |
-
"0.33": [144.0, 432.0],
|
63 |
-
"0.35": [144.0, 416.0],
|
64 |
-
"0.4": [160.0, 400.0],
|
65 |
-
"0.42": [160.0, 384.0],
|
66 |
-
"0.48": [176.0, 368.0],
|
67 |
-
"0.5": [176.0, 352.0],
|
68 |
-
"0.52": [176.0, 336.0],
|
69 |
-
"0.57": [192.0, 336.0],
|
70 |
-
"0.6": [192.0, 320.0],
|
71 |
-
"0.68": [208.0, 304.0],
|
72 |
-
"0.72": [208.0, 288.0],
|
73 |
-
"0.78": [224.0, 288.0],
|
74 |
-
"0.82": [224.0, 272.0],
|
75 |
-
"0.88": [240.0, 272.0],
|
76 |
-
"0.94": [240.0, 256.0],
|
77 |
-
"1.0": [256.0, 256.0],
|
78 |
-
"1.07": [256.0, 240.0],
|
79 |
-
"1.13": [272.0, 240.0],
|
80 |
-
"1.21": [272.0, 224.0],
|
81 |
-
"1.29": [288.0, 224.0],
|
82 |
-
"1.38": [288.0, 208.0],
|
83 |
-
"1.46": [304.0, 208.0],
|
84 |
-
"1.67": [320.0, 192.0],
|
85 |
-
"1.75": [336.0, 192.0],
|
86 |
-
"2.0": [352.0, 176.0],
|
87 |
-
"2.09": [368.0, 176.0],
|
88 |
-
"2.4": [384.0, 160.0],
|
89 |
-
"2.5": [400.0, 160.0],
|
90 |
-
"3.0": [432.0, 144.0],
|
91 |
-
"4.0": [512.0, 128.0]
|
92 |
-
}
|
93 |
-
|
94 |
-
ASPECT_RATIO_1024_BIN = {
|
95 |
-
"0.25": [512.0, 2048.0],
|
96 |
-
"0.28": [512.0, 1856.0],
|
97 |
-
"0.32": [576.0, 1792.0],
|
98 |
-
"0.33": [576.0, 1728.0],
|
99 |
-
"0.35": [576.0, 1664.0],
|
100 |
-
"0.4": [640.0, 1600.0],
|
101 |
-
"0.42": [640.0, 1536.0],
|
102 |
-
"0.48": [704.0, 1472.0],
|
103 |
-
"0.5": [704.0, 1408.0],
|
104 |
-
"0.52": [704.0, 1344.0],
|
105 |
-
"0.57": [768.0, 1344.0],
|
106 |
-
"0.6": [768.0, 1280.0],
|
107 |
-
"0.68": [832.0, 1216.0],
|
108 |
-
"0.72": [832.0, 1152.0],
|
109 |
-
"0.78": [896.0, 1152.0],
|
110 |
-
"0.82": [896.0, 1088.0],
|
111 |
-
"0.88": [960.0, 1088.0],
|
112 |
-
"0.94": [960.0, 1024.0],
|
113 |
-
"1.0": [1024.0, 1024.0],
|
114 |
-
"1.07": [1024.0, 960.0],
|
115 |
-
"1.13": [1088.0, 960.0],
|
116 |
-
"1.21": [1088.0, 896.0],
|
117 |
-
"1.29": [1152.0, 896.0],
|
118 |
-
"1.38": [1152.0, 832.0],
|
119 |
-
"1.46": [1216.0, 832.0],
|
120 |
-
"1.67": [1280.0, 768.0],
|
121 |
-
"1.75": [1344.0, 768.0],
|
122 |
-
"2.0": [1408.0, 704.0],
|
123 |
-
"2.09": [1472.0, 704.0],
|
124 |
-
"2.4": [1536.0, 640.0],
|
125 |
-
"2.5": [1600.0, 640.0],
|
126 |
-
"3.0": [1728.0, 576.0],
|
127 |
-
"4.0": [2048.0, 512.0],
|
128 |
-
}
|
129 |
-
|
130 |
-
ASPECT_RATIO_512_BIN = {
|
131 |
-
"0.25": [256.0, 1024.0],
|
132 |
-
"0.28": [256.0, 928.0],
|
133 |
-
"0.32": [288.0, 896.0],
|
134 |
-
"0.33": [288.0, 864.0],
|
135 |
-
"0.35": [288.0, 832.0],
|
136 |
-
"0.4": [320.0, 800.0],
|
137 |
-
"0.42": [320.0, 768.0],
|
138 |
-
"0.48": [352.0, 736.0],
|
139 |
-
"0.5": [352.0, 704.0],
|
140 |
-
"0.52": [352.0, 672.0],
|
141 |
-
"0.57": [384.0, 672.0],
|
142 |
-
"0.6": [384.0, 640.0],
|
143 |
-
"0.68": [416.0, 608.0],
|
144 |
-
"0.72": [416.0, 576.0],
|
145 |
-
"0.78": [448.0, 576.0],
|
146 |
-
"0.82": [448.0, 544.0],
|
147 |
-
"0.88": [480.0, 544.0],
|
148 |
-
"0.94": [480.0, 512.0],
|
149 |
-
"1.0": [512.0, 512.0],
|
150 |
-
"1.07": [512.0, 480.0],
|
151 |
-
"1.13": [544.0, 480.0],
|
152 |
-
"1.21": [544.0, 448.0],
|
153 |
-
"1.29": [576.0, 448.0],
|
154 |
-
"1.38": [576.0, 416.0],
|
155 |
-
"1.46": [608.0, 416.0],
|
156 |
-
"1.67": [640.0, 384.0],
|
157 |
-
"1.75": [672.0, 384.0],
|
158 |
-
"2.0": [704.0, 352.0],
|
159 |
-
"2.09": [736.0, 352.0],
|
160 |
-
"2.4": [768.0, 320.0],
|
161 |
-
"2.5": [800.0, 320.0],
|
162 |
-
"3.0": [864.0, 288.0],
|
163 |
-
"4.0": [1024.0, 256.0],
|
164 |
-
}
|
165 |
-
|
166 |
-
|
167 |
-
def pipeline_pixart_alpha_call(
|
168 |
-
self,
|
169 |
-
prompt: Union[str, List[str]] = None,
|
170 |
-
negative_prompt: str = "",
|
171 |
-
num_inference_steps: int = 20,
|
172 |
-
timesteps: List[int] = None,
|
173 |
-
guidance_scale: float = 4.5,
|
174 |
-
num_images_per_prompt: Optional[int] = 1,
|
175 |
-
height: Optional[int] = None,
|
176 |
-
width: Optional[int] = None,
|
177 |
-
eta: float = 0.0,
|
178 |
-
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
179 |
-
latents: Optional[torch.FloatTensor] = None,
|
180 |
-
prompt_embeds: Optional[torch.FloatTensor] = None,
|
181 |
-
prompt_attention_mask: Optional[torch.FloatTensor] = None,
|
182 |
-
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
|
183 |
-
negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
|
184 |
-
output_type: Optional[str] = "pil",
|
185 |
-
return_dict: bool = True,
|
186 |
-
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
|
187 |
-
callback_steps: int = 1,
|
188 |
-
clean_caption: bool = True,
|
189 |
-
use_resolution_binning: bool = True,
|
190 |
-
max_sequence_length: int = 120,
|
191 |
-
**kwargs,
|
192 |
-
) -> Union[ImagePipelineOutput, Tuple]:
|
193 |
-
"""
|
194 |
-
Function invoked when calling the pipeline for generation.
|
195 |
-
|
196 |
-
Args:
|
197 |
-
prompt (`str` or `List[str]`, *optional*):
|
198 |
-
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
|
199 |
-
instead.
|
200 |
-
negative_prompt (`str` or `List[str]`, *optional*):
|
201 |
-
The prompt or prompts not to guide the image generation. If not defined, one has to pass
|
202 |
-
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
|
203 |
-
less than `1`).
|
204 |
-
num_inference_steps (`int`, *optional*, defaults to 100):
|
205 |
-
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
|
206 |
-
expense of slower inference.
|
207 |
-
timesteps (`List[int]`, *optional*):
|
208 |
-
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
|
209 |
-
timesteps are used. Must be in descending order.
|
210 |
-
guidance_scale (`float`, *optional*, defaults to 4.5):
|
211 |
-
Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
|
212 |
-
`guidance_scale` is defined as `w` of equation 2. of [Imagen
|
213 |
-
Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
|
214 |
-
1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
|
215 |
-
usually at the expense of lower image quality.
|
216 |
-
num_images_per_prompt (`int`, *optional*, defaults to 1):
|
217 |
-
The number of images to generate per prompt.
|
218 |
-
height (`int`, *optional*, defaults to self.unet.config.sample_size):
|
219 |
-
The height in pixels of the generated image.
|
220 |
-
width (`int`, *optional*, defaults to self.unet.config.sample_size):
|
221 |
-
The width in pixels of the generated image.
|
222 |
-
eta (`float`, *optional*, defaults to 0.0):
|
223 |
-
Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
|
224 |
-
[`schedulers.DDIMScheduler`], will be ignored for others.
|
225 |
-
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
|
226 |
-
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
|
227 |
-
to make generation deterministic.
|
228 |
-
latents (`torch.FloatTensor`, *optional*):
|
229 |
-
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
|
230 |
-
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
|
231 |
-
tensor will ge generated by sampling using the supplied random `generator`.
|
232 |
-
prompt_embeds (`torch.FloatTensor`, *optional*):
|
233 |
-
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
|
234 |
-
provided, text embeddings will be generated from `prompt` input argument.
|
235 |
-
prompt_attention_mask (`torch.FloatTensor`, *optional*): Pre-generated attention mask for text embeddings.
|
236 |
-
negative_prompt_embeds (`torch.FloatTensor`, *optional*):
|
237 |
-
Pre-generated negative text embeddings. For PixArt-Alpha this negative prompt should be "". If not
|
238 |
-
provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
|
239 |
-
negative_prompt_attention_mask (`torch.FloatTensor`, *optional*):
|
240 |
-
Pre-generated attention mask for negative text embeddings.
|
241 |
-
output_type (`str`, *optional*, defaults to `"pil"`):
|
242 |
-
The output format of the generate image. Choose between
|
243 |
-
[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
|
244 |
-
return_dict (`bool`, *optional*, defaults to `True`):
|
245 |
-
Whether or not to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple.
|
246 |
-
callback (`Callable`, *optional*):
|
247 |
-
A function that will be called every `callback_steps` steps during inference. The function will be
|
248 |
-
called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
|
249 |
-
callback_steps (`int`, *optional*, defaults to 1):
|
250 |
-
The frequency at which the `callback` function will be called. If not specified, the callback will be
|
251 |
-
called at every step.
|
252 |
-
clean_caption (`bool`, *optional*, defaults to `True`):
|
253 |
-
Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to
|
254 |
-
be installed. If the dependencies are not installed, the embeddings will be created from the raw
|
255 |
-
prompt.
|
256 |
-
use_resolution_binning (`bool` defaults to `True`):
|
257 |
-
If set to `True`, the requested height and width are first mapped to the closest resolutions using
|
258 |
-
`ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to
|
259 |
-
the requested resolution. Useful for generating non-square images.
|
260 |
-
|
261 |
-
Examples:
|
262 |
-
|
263 |
-
Returns:
|
264 |
-
[`~pipelines.ImagePipelineOutput`] or `tuple`:
|
265 |
-
If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is
|
266 |
-
returned where the first element is a list with the generated images
|
267 |
-
"""
|
268 |
-
if "mask_feature" in kwargs:
|
269 |
-
deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version."
|
270 |
-
deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False)
|
271 |
-
# 1. Check inputs. Raise error if not correct
|
272 |
-
height = height or self.transformer.config.sample_size * self.vae_scale_factor
|
273 |
-
width = width or self.transformer.config.sample_size * self.vae_scale_factor
|
274 |
-
if use_resolution_binning:
|
275 |
-
if self.transformer.config.sample_size == 32:
|
276 |
-
aspect_ratio_bin = ASPECT_RATIO_256_BIN
|
277 |
-
elif self.transformer.config.sample_size == 64:
|
278 |
-
aspect_ratio_bin = ASPECT_RATIO_512_BIN
|
279 |
-
elif self.transformer.config.sample_size == 128:
|
280 |
-
aspect_ratio_bin = ASPECT_RATIO_1024_BIN
|
281 |
-
elif self.transformer.config.sample_size == 256:
|
282 |
-
aspect_ratio_bin = ASPECT_RATIO_2048_BIN
|
283 |
-
else:
|
284 |
-
raise ValueError("Invalid sample size")
|
285 |
-
orig_height, orig_width = height, width
|
286 |
-
height, width = self.classify_height_width_bin(height, width, ratios=aspect_ratio_bin)
|
287 |
-
|
288 |
-
self.check_inputs(
|
289 |
-
prompt,
|
290 |
-
height,
|
291 |
-
width,
|
292 |
-
negative_prompt,
|
293 |
-
callback_steps,
|
294 |
-
prompt_embeds,
|
295 |
-
negative_prompt_embeds,
|
296 |
-
prompt_attention_mask,
|
297 |
-
negative_prompt_attention_mask,
|
298 |
-
)
|
299 |
-
|
300 |
-
# 2. Default height and width to transformer
|
301 |
-
if prompt is not None and isinstance(prompt, str):
|
302 |
-
batch_size = 1
|
303 |
-
elif prompt is not None and isinstance(prompt, list):
|
304 |
-
batch_size = len(prompt)
|
305 |
-
else:
|
306 |
-
batch_size = prompt_embeds.shape[0]
|
307 |
-
|
308 |
-
device = self._execution_device
|
309 |
-
|
310 |
-
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
|
311 |
-
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
|
312 |
-
# corresponds to doing no classifier free guidance.
|
313 |
-
do_classifier_free_guidance = guidance_scale > 1.0
|
314 |
-
|
315 |
-
# 3. Encode input prompt
|
316 |
-
(
|
317 |
-
prompt_embeds,
|
318 |
-
prompt_attention_mask,
|
319 |
-
negative_prompt_embeds,
|
320 |
-
negative_prompt_attention_mask,
|
321 |
-
) = self.encode_prompt(
|
322 |
-
prompt,
|
323 |
-
do_classifier_free_guidance,
|
324 |
-
negative_prompt=negative_prompt,
|
325 |
-
num_images_per_prompt=num_images_per_prompt,
|
326 |
-
device=device,
|
327 |
-
prompt_embeds=prompt_embeds,
|
328 |
-
negative_prompt_embeds=negative_prompt_embeds,
|
329 |
-
prompt_attention_mask=prompt_attention_mask,
|
330 |
-
negative_prompt_attention_mask=negative_prompt_attention_mask,
|
331 |
-
clean_caption=clean_caption,
|
332 |
-
max_sequence_length=max_sequence_length,
|
333 |
-
)
|
334 |
-
if do_classifier_free_guidance:
|
335 |
-
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
|
336 |
-
prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
|
337 |
-
|
338 |
-
# 4. Prepare timesteps
|
339 |
-
timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
|
340 |
-
|
341 |
-
# 5. Prepare latents.
|
342 |
-
latent_channels = self.transformer.config.in_channels
|
343 |
-
latents = self.prepare_latents(
|
344 |
-
batch_size * num_images_per_prompt,
|
345 |
-
latent_channels,
|
346 |
-
height,
|
347 |
-
width,
|
348 |
-
prompt_embeds.dtype,
|
349 |
-
device,
|
350 |
-
generator,
|
351 |
-
latents,
|
352 |
-
)
|
353 |
-
|
354 |
-
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
|
355 |
-
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
|
356 |
-
|
357 |
-
# 6.1 Prepare micro-conditions.
|
358 |
-
added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
|
359 |
-
if self.transformer.config.sample_size == 128:
|
360 |
-
resolution = torch.tensor([height, width]).repeat(batch_size * num_images_per_prompt, 1)
|
361 |
-
aspect_ratio = torch.tensor([float(height / width)]).repeat(batch_size * num_images_per_prompt, 1)
|
362 |
-
resolution = resolution.to(dtype=prompt_embeds.dtype, device=device)
|
363 |
-
aspect_ratio = aspect_ratio.to(dtype=prompt_embeds.dtype, device=device)
|
364 |
-
|
365 |
-
if do_classifier_free_guidance:
|
366 |
-
resolution = torch.cat([resolution, resolution], dim=0)
|
367 |
-
aspect_ratio = torch.cat([aspect_ratio, aspect_ratio], dim=0)
|
368 |
-
|
369 |
-
added_cond_kwargs = {"resolution": resolution, "aspect_ratio": aspect_ratio}
|
370 |
-
|
371 |
-
# 7. Denoising loop
|
372 |
-
num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
|
373 |
-
|
374 |
-
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
375 |
-
for i, t in enumerate(timesteps):
|
376 |
-
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
|
377 |
-
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
|
378 |
-
|
379 |
-
current_timestep = t
|
380 |
-
if not torch.is_tensor(current_timestep):
|
381 |
-
# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
|
382 |
-
# This would be a good case for the `match` statement (Python 3.10+)
|
383 |
-
is_mps = latent_model_input.device.type == "mps"
|
384 |
-
if isinstance(current_timestep, float):
|
385 |
-
dtype = torch.float32 if is_mps else torch.float64
|
386 |
-
else:
|
387 |
-
dtype = torch.int32 if is_mps else torch.int64
|
388 |
-
current_timestep = torch.tensor([current_timestep], dtype=dtype, device=latent_model_input.device)
|
389 |
-
elif len(current_timestep.shape) == 0:
|
390 |
-
current_timestep = current_timestep[None].to(latent_model_input.device)
|
391 |
-
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
392 |
-
current_timestep = current_timestep.expand(latent_model_input.shape[0])
|
393 |
-
|
394 |
-
# predict noise model_output
|
395 |
-
noise_pred = self.transformer(
|
396 |
-
latent_model_input,
|
397 |
-
encoder_hidden_states=prompt_embeds,
|
398 |
-
encoder_attention_mask=prompt_attention_mask,
|
399 |
-
timestep=current_timestep,
|
400 |
-
added_cond_kwargs=added_cond_kwargs,
|
401 |
-
return_dict=False,
|
402 |
-
)[0]
|
403 |
-
|
404 |
-
# perform guidance
|
405 |
-
if do_classifier_free_guidance:
|
406 |
-
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
|
407 |
-
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
|
408 |
-
|
409 |
-
# learned sigma
|
410 |
-
if self.transformer.config.out_channels // 2 == latent_channels:
|
411 |
-
noise_pred = noise_pred.chunk(2, dim=1)[0]
|
412 |
-
else:
|
413 |
-
noise_pred = noise_pred
|
414 |
-
|
415 |
-
# compute previous image: x_t -> x_t-1
|
416 |
-
if num_inference_steps == 1:
|
417 |
-
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).pred_original_sample
|
418 |
-
else:
|
419 |
-
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
|
420 |
-
|
421 |
-
# call the callback, if provided
|
422 |
-
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
|
423 |
-
progress_bar.update()
|
424 |
-
if callback is not None and i % callback_steps == 0:
|
425 |
-
step_idx = i // getattr(self.scheduler, "order", 1)
|
426 |
-
callback(step_idx, t, latents)
|
427 |
-
|
428 |
-
if not output_type == "latent":
|
429 |
-
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
|
430 |
-
if use_resolution_binning:
|
431 |
-
image = self.resize_and_crop_tensor(image, orig_width, orig_height)
|
432 |
-
else:
|
433 |
-
image = latents
|
434 |
-
|
435 |
-
if not output_type == "latent":
|
436 |
-
image = self.image_processor.postprocess(image, output_type=output_type)
|
437 |
-
|
438 |
-
# Offload all models
|
439 |
-
self.maybe_free_model_hooks()
|
440 |
-
|
441 |
-
if not return_dict:
|
442 |
-
return (image,)
|
443 |
-
|
444 |
-
return ImagePipelineOutput(images=image)
|
445 |
-
|
446 |
-
|
447 |
-
class PixArtSigmaPipeline(PixArtAlphaPipeline):
|
448 |
-
r"""
|
449 |
-
tmp Pipeline for text-to-image generation using PixArt-Sigma.
|
450 |
-
"""
|
451 |
-
|
452 |
-
def __init__(
|
453 |
-
self,
|
454 |
-
tokenizer: T5Tokenizer,
|
455 |
-
text_encoder: T5EncoderModel,
|
456 |
-
vae: AutoencoderKL,
|
457 |
-
transformer: Transformer2DModel,
|
458 |
-
scheduler: DPMSolverMultistepScheduler,
|
459 |
-
):
|
460 |
-
super().__init__(tokenizer, text_encoder, vae, transformer, scheduler)
|
461 |
-
|
462 |
-
self.register_modules(
|
463 |
-
tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler
|
464 |
-
)
|
465 |
-
|
466 |
-
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
|
467 |
-
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
|
468 |
-
|
469 |
-
|
470 |
-
def pixart_sigma_init_patched_inputs(self, norm_type):
|
471 |
-
assert self.config.sample_size is not None, "Transformer2DModel over patched input must provide sample_size"
|
472 |
-
|
473 |
-
self.height = self.config.sample_size
|
474 |
-
self.width = self.config.sample_size
|
475 |
-
|
476 |
-
self.patch_size = self.config.patch_size
|
477 |
-
interpolation_scale = (
|
478 |
-
self.config.interpolation_scale
|
479 |
-
if self.config.interpolation_scale is not None
|
480 |
-
else max(self.config.sample_size // 64, 1)
|
481 |
-
)
|
482 |
-
self.pos_embed = PatchEmbed(
|
483 |
-
height=self.config.sample_size,
|
484 |
-
width=self.config.sample_size,
|
485 |
-
patch_size=self.config.patch_size,
|
486 |
-
in_channels=self.in_channels,
|
487 |
-
embed_dim=self.inner_dim,
|
488 |
-
interpolation_scale=interpolation_scale,
|
489 |
-
)
|
490 |
-
|
491 |
-
self.transformer_blocks = nn.ModuleList(
|
492 |
-
[
|
493 |
-
BasicTransformerBlock(
|
494 |
-
self.inner_dim,
|
495 |
-
self.config.num_attention_heads,
|
496 |
-
self.config.attention_head_dim,
|
497 |
-
dropout=self.config.dropout,
|
498 |
-
cross_attention_dim=self.config.cross_attention_dim,
|
499 |
-
activation_fn=self.config.activation_fn,
|
500 |
-
num_embeds_ada_norm=self.config.num_embeds_ada_norm,
|
501 |
-
attention_bias=self.config.attention_bias,
|
502 |
-
only_cross_attention=self.config.only_cross_attention,
|
503 |
-
double_self_attention=self.config.double_self_attention,
|
504 |
-
upcast_attention=self.config.upcast_attention,
|
505 |
-
norm_type=norm_type,
|
506 |
-
norm_elementwise_affine=self.config.norm_elementwise_affine,
|
507 |
-
norm_eps=self.config.norm_eps,
|
508 |
-
attention_type=self.config.attention_type,
|
509 |
-
)
|
510 |
-
for _ in range(self.config.num_layers)
|
511 |
-
]
|
512 |
-
)
|
513 |
-
|
514 |
-
if self.config.norm_type != "ada_norm_single":
|
515 |
-
self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
|
516 |
-
self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim)
|
517 |
-
self.proj_out_2 = nn.Linear(
|
518 |
-
self.inner_dim, self.config.patch_size * self.config.patch_size * self.out_channels
|
519 |
-
)
|
520 |
-
elif self.config.norm_type == "ada_norm_single":
|
521 |
-
self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
|
522 |
-
self.scale_shift_table = nn.Parameter(torch.randn(2, self.inner_dim) / self.inner_dim ** 0.5)
|
523 |
-
self.proj_out = nn.Linear(
|
524 |
-
self.inner_dim, self.config.patch_size * self.config.patch_size * self.out_channels
|
525 |
-
)
|
526 |
-
|
527 |
-
# PixArt-Sigma blocks.
|
528 |
-
self.adaln_single = None
|
529 |
-
self.use_additional_conditions = False
|
530 |
-
if self.config.norm_type == "ada_norm_single":
|
531 |
-
# TODO(Sayak, PVP) clean this, PixArt-Sigma doesn't use additional_conditions anymore
|
532 |
-
# additional conditions until we find better name
|
533 |
-
self.adaln_single = AdaLayerNormSingle(
|
534 |
-
self.inner_dim, use_additional_conditions=self.use_additional_conditions
|
535 |
-
)
|
536 |
-
|
537 |
-
self.caption_projection = None
|
538 |
-
if self.caption_channels is not None:
|
539 |
-
self.caption_projection = PixArtAlphaTextProjection(
|
540 |
-
in_features=self.caption_channels, hidden_size=self.inner_dim
|
541 |
-
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|