Spaces:
Running
on
Zero
Running
on
Zero
zhiweili
commited on
Commit
•
85f4bc7
1
Parent(s):
ff1a4e2
fix mask
Browse files
pipelines/masked_stable_diffusion_img2img.py
ADDED
@@ -0,0 +1,262 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from typing import Any, Callable, Dict, List, Optional, Union
|
2 |
+
|
3 |
+
import numpy as np
|
4 |
+
import PIL.Image
|
5 |
+
import torch
|
6 |
+
|
7 |
+
from diffusers import StableDiffusionImg2ImgPipeline
|
8 |
+
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
|
9 |
+
|
10 |
+
|
11 |
+
class MaskedStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):
|
12 |
+
debug_save = False
|
13 |
+
|
14 |
+
@torch.no_grad()
|
15 |
+
def __call__(
|
16 |
+
self,
|
17 |
+
prompt: Union[str, List[str]] = None,
|
18 |
+
image: Union[
|
19 |
+
torch.Tensor,
|
20 |
+
PIL.Image.Image,
|
21 |
+
np.ndarray,
|
22 |
+
List[torch.Tensor],
|
23 |
+
List[PIL.Image.Image],
|
24 |
+
List[np.ndarray],
|
25 |
+
] = None,
|
26 |
+
strength: float = 0.8,
|
27 |
+
num_inference_steps: Optional[int] = 50,
|
28 |
+
guidance_scale: Optional[float] = 7.5,
|
29 |
+
negative_prompt: Optional[Union[str, List[str]]] = None,
|
30 |
+
num_images_per_prompt: Optional[int] = 1,
|
31 |
+
eta: Optional[float] = 0.0,
|
32 |
+
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
33 |
+
prompt_embeds: Optional[torch.Tensor] = None,
|
34 |
+
negative_prompt_embeds: Optional[torch.Tensor] = None,
|
35 |
+
output_type: Optional[str] = "pil",
|
36 |
+
return_dict: bool = True,
|
37 |
+
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
|
38 |
+
callback_steps: int = 1,
|
39 |
+
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
|
40 |
+
mask: Union[
|
41 |
+
torch.Tensor,
|
42 |
+
PIL.Image.Image,
|
43 |
+
np.ndarray,
|
44 |
+
List[torch.Tensor],
|
45 |
+
List[PIL.Image.Image],
|
46 |
+
List[np.ndarray],
|
47 |
+
] = None,
|
48 |
+
):
|
49 |
+
r"""
|
50 |
+
The call function to the pipeline for generation.
|
51 |
+
|
52 |
+
Args:
|
53 |
+
prompt (`str` or `List[str]`, *optional*):
|
54 |
+
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
|
55 |
+
image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
|
56 |
+
`Image` or tensor representing an image batch to be used as the starting point. Can also accept image
|
57 |
+
latents as `image`, but if passing latents directly it is not encoded again.
|
58 |
+
strength (`float`, *optional*, defaults to 0.8):
|
59 |
+
Indicates extent to transform the reference `image`. Must be between 0 and 1. `image` is used as a
|
60 |
+
starting point and more noise is added the higher the `strength`. The number of denoising steps depends
|
61 |
+
on the amount of noise initially added. When `strength` is 1, added noise is maximum and the denoising
|
62 |
+
process runs for the full number of iterations specified in `num_inference_steps`. A value of 1
|
63 |
+
essentially ignores `image`.
|
64 |
+
num_inference_steps (`int`, *optional*, defaults to 50):
|
65 |
+
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
|
66 |
+
expense of slower inference. This parameter is modulated by `strength`.
|
67 |
+
guidance_scale (`float`, *optional*, defaults to 7.5):
|
68 |
+
A higher guidance scale value encourages the model to generate images closely linked to the text
|
69 |
+
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
|
70 |
+
negative_prompt (`str` or `List[str]`, *optional*):
|
71 |
+
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
|
72 |
+
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
|
73 |
+
num_images_per_prompt (`int`, *optional*, defaults to 1):
|
74 |
+
The number of images to generate per prompt.
|
75 |
+
eta (`float`, *optional*, defaults to 0.0):
|
76 |
+
Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
|
77 |
+
to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
|
78 |
+
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
|
79 |
+
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
|
80 |
+
generation deterministic.
|
81 |
+
prompt_embeds (`torch.Tensor`, *optional*):
|
82 |
+
Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
|
83 |
+
provided, text embeddings are generated from the `prompt` input argument.
|
84 |
+
negative_prompt_embeds (`torch.Tensor`, *optional*):
|
85 |
+
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
|
86 |
+
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
|
87 |
+
output_type (`str`, *optional*, defaults to `"pil"`):
|
88 |
+
The output format of the generated image. Choose between `PIL.Image` or `np.array`.
|
89 |
+
return_dict (`bool`, *optional*, defaults to `True`):
|
90 |
+
Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
|
91 |
+
plain tuple.
|
92 |
+
callback (`Callable`, *optional*):
|
93 |
+
A function that calls every `callback_steps` steps during inference. The function is called with the
|
94 |
+
following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.
|
95 |
+
callback_steps (`int`, *optional*, defaults to 1):
|
96 |
+
The frequency at which the `callback` function is called. If not specified, the callback is called at
|
97 |
+
every step.
|
98 |
+
cross_attention_kwargs (`dict`, *optional*):
|
99 |
+
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
|
100 |
+
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
|
101 |
+
mask (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`, *optional*):
|
102 |
+
A mask with non-zero elements for the area to be inpainted. If not specified, no mask is applied.
|
103 |
+
Examples:
|
104 |
+
|
105 |
+
Returns:
|
106 |
+
[`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
|
107 |
+
If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,
|
108 |
+
otherwise a `tuple` is returned where the first element is a list with the generated images and the
|
109 |
+
second element is a list of `bool`s indicating whether the corresponding generated image contains
|
110 |
+
"not-safe-for-work" (nsfw) content.
|
111 |
+
"""
|
112 |
+
# code adapted from parent class StableDiffusionImg2ImgPipeline
|
113 |
+
|
114 |
+
# 0. Check inputs. Raise error if not correct
|
115 |
+
self.check_inputs(prompt, strength, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds)
|
116 |
+
|
117 |
+
# 1. Define call parameters
|
118 |
+
if prompt is not None and isinstance(prompt, str):
|
119 |
+
batch_size = 1
|
120 |
+
elif prompt is not None and isinstance(prompt, list):
|
121 |
+
batch_size = len(prompt)
|
122 |
+
else:
|
123 |
+
batch_size = prompt_embeds.shape[0]
|
124 |
+
device = self._execution_device
|
125 |
+
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
|
126 |
+
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
|
127 |
+
# corresponds to doing no classifier free guidance.
|
128 |
+
do_classifier_free_guidance = guidance_scale > 1.0
|
129 |
+
|
130 |
+
# 2. Encode input prompt
|
131 |
+
text_encoder_lora_scale = (
|
132 |
+
cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
|
133 |
+
)
|
134 |
+
prompt_embeds = self._encode_prompt(
|
135 |
+
prompt,
|
136 |
+
device,
|
137 |
+
num_images_per_prompt,
|
138 |
+
do_classifier_free_guidance,
|
139 |
+
negative_prompt,
|
140 |
+
prompt_embeds=prompt_embeds,
|
141 |
+
negative_prompt_embeds=negative_prompt_embeds,
|
142 |
+
lora_scale=text_encoder_lora_scale,
|
143 |
+
)
|
144 |
+
|
145 |
+
# 3. Preprocess image
|
146 |
+
image = self.image_processor.preprocess(image)
|
147 |
+
|
148 |
+
# 4. set timesteps
|
149 |
+
self.scheduler.set_timesteps(num_inference_steps, device=device)
|
150 |
+
timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
|
151 |
+
latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
|
152 |
+
|
153 |
+
# 5. Prepare latent variables
|
154 |
+
# it is sampled from the latent distribution of the VAE
|
155 |
+
latents = self.prepare_latents(
|
156 |
+
image, latent_timestep, batch_size, num_images_per_prompt, prompt_embeds.dtype, device, generator
|
157 |
+
)
|
158 |
+
|
159 |
+
# mean of the latent distribution
|
160 |
+
init_latents = [
|
161 |
+
self.vae.encode(image.to(device=device, dtype=prompt_embeds.dtype)[i : i + 1]).latent_dist.mean
|
162 |
+
for i in range(batch_size)
|
163 |
+
]
|
164 |
+
init_latents = torch.cat(init_latents, dim=0)
|
165 |
+
|
166 |
+
# 6. create latent mask
|
167 |
+
latent_mask = self._make_latent_mask(latents, mask)
|
168 |
+
|
169 |
+
# 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
|
170 |
+
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
|
171 |
+
|
172 |
+
# 8. Denoising loop
|
173 |
+
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
|
174 |
+
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
175 |
+
for i, t in enumerate(timesteps):
|
176 |
+
# expand the latents if we are doing classifier free guidance
|
177 |
+
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
|
178 |
+
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
|
179 |
+
|
180 |
+
# predict the noise residual
|
181 |
+
noise_pred = self.unet(
|
182 |
+
latent_model_input,
|
183 |
+
t,
|
184 |
+
encoder_hidden_states=prompt_embeds,
|
185 |
+
cross_attention_kwargs=cross_attention_kwargs,
|
186 |
+
return_dict=False,
|
187 |
+
)[0]
|
188 |
+
|
189 |
+
# perform guidance
|
190 |
+
if do_classifier_free_guidance:
|
191 |
+
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
|
192 |
+
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
|
193 |
+
|
194 |
+
if latent_mask is not None:
|
195 |
+
latents = torch.lerp(init_latents * self.vae.config.scaling_factor, latents, latent_mask)
|
196 |
+
noise_pred = torch.lerp(torch.zeros_like(noise_pred), noise_pred, latent_mask)
|
197 |
+
|
198 |
+
# compute the previous noisy sample x_t -> x_t-1
|
199 |
+
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
|
200 |
+
|
201 |
+
# call the callback, if provided
|
202 |
+
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
|
203 |
+
progress_bar.update()
|
204 |
+
if callback is not None and i % callback_steps == 0:
|
205 |
+
step_idx = i // getattr(self.scheduler, "order", 1)
|
206 |
+
callback(step_idx, t, latents)
|
207 |
+
|
208 |
+
if not output_type == "latent":
|
209 |
+
scaled = latents / self.vae.config.scaling_factor
|
210 |
+
if latent_mask is not None:
|
211 |
+
# scaled = latents / self.vae.config.scaling_factor * latent_mask + init_latents * (1 - latent_mask)
|
212 |
+
scaled = torch.lerp(init_latents, scaled, latent_mask)
|
213 |
+
image = self.vae.decode(scaled, return_dict=False)[0]
|
214 |
+
if self.debug_save:
|
215 |
+
image_gen = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
|
216 |
+
image_gen = self.image_processor.postprocess(image_gen, output_type=output_type, do_denormalize=[True])
|
217 |
+
image_gen[0].save("from_latent.png")
|
218 |
+
image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
|
219 |
+
else:
|
220 |
+
image = latents
|
221 |
+
has_nsfw_concept = None
|
222 |
+
|
223 |
+
if has_nsfw_concept is None:
|
224 |
+
do_denormalize = [True] * image.shape[0]
|
225 |
+
else:
|
226 |
+
do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
|
227 |
+
|
228 |
+
image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
|
229 |
+
|
230 |
+
# Offload last model to CPU
|
231 |
+
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
|
232 |
+
self.final_offload_hook.offload()
|
233 |
+
|
234 |
+
if not return_dict:
|
235 |
+
return (image, has_nsfw_concept)
|
236 |
+
|
237 |
+
return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
|
238 |
+
|
239 |
+
def _make_latent_mask(self, latents, mask):
|
240 |
+
if mask is not None:
|
241 |
+
latent_mask = []
|
242 |
+
if not isinstance(mask, list):
|
243 |
+
tmp_mask = [mask]
|
244 |
+
else:
|
245 |
+
tmp_mask = mask
|
246 |
+
_, l_channels, l_height, l_width = latents.shape
|
247 |
+
for m in tmp_mask:
|
248 |
+
if not isinstance(m, PIL.Image.Image):
|
249 |
+
if len(m.shape) == 2:
|
250 |
+
m = m[..., np.newaxis]
|
251 |
+
if m.max() > 1:
|
252 |
+
m = m / 255.0
|
253 |
+
m = self.image_processor.numpy_to_pil(m)[0]
|
254 |
+
if m.mode != "L":
|
255 |
+
m = m.convert("L")
|
256 |
+
resized = self.image_processor.resize(m, l_height, l_width)
|
257 |
+
if self.debug_save:
|
258 |
+
resized.save("latent_mask.png")
|
259 |
+
latent_mask.append(np.repeat(np.array(resized)[np.newaxis, :, :], l_channels, axis=0))
|
260 |
+
latent_mask = torch.as_tensor(np.stack(latent_mask)).to(latents)
|
261 |
+
latent_mask = latent_mask / latent_mask.max()
|
262 |
+
return latent_mask
|
pipelines/masked_stable_diffusion_xl_img2img.py
CHANGED
@@ -208,8 +208,10 @@ class MaskedStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
|
|
208 |
|
209 |
if not isinstance(mask, Image.Image):
|
210 |
pil_mask = Image.fromarray(mask)
|
211 |
-
|
212 |
-
|
|
|
|
|
213 |
mask_blur = self.blur_mask(pil_mask, blur)
|
214 |
mask_compose = self.blur_mask(pil_mask, blur_compose)
|
215 |
if original_image is None:
|
|
|
208 |
|
209 |
if not isinstance(mask, Image.Image):
|
210 |
pil_mask = Image.fromarray(mask)
|
211 |
+
else:
|
212 |
+
pil_mask = mask
|
213 |
+
if pil_mask.mode != "L":
|
214 |
+
pil_mask = pil_mask.convert("L")
|
215 |
mask_blur = self.blur_mask(pil_mask, blur)
|
216 |
mask_compose = self.blur_mask(pil_mask, blur_compose)
|
217 |
if original_image is None:
|