diffusers/community-pipelines-mirror

1import math
2from typing import Dict, Optional
3
4import torch
5import torchvision.transforms.functional as FF
6from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
7
8from diffusers import StableDiffusionPipeline
9from diffusers.models import AutoencoderKL, UNet2DConditionModel
10from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
11from diffusers.schedulers import KarrasDiffusionSchedulers
12from diffusers.utils import USE_PEFT_BACKEND
13
14
15try:
16    from compel import Compel
17except ImportError:
18    Compel = None
19
20KCOMM = "ADDCOMM"
21KBRK = "BREAK"
22
23
24class RegionalPromptingStableDiffusionPipeline(StableDiffusionPipeline):
25    r"""
26    Args for Regional Prompting Pipeline:
27        rp_args:dict
28        Required
29            rp_args["mode"]: cols, rows, prompt, prompt-ex
30        for cols, rows mode
31            rp_args["div"]: ex) 1;1;1(Divide into 3 regions)
32        for prompt, prompt-ex mode
33            rp_args["th"]: ex) 0.5,0.5,0.6 (threshold for prompt mode)
34
35        Optional
36            rp_args["save_mask"]: True/False (save masks in prompt mode)
37
38    Pipeline for text-to-image generation using Stable Diffusion.
39
40    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
41    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
42
43    Args:
44        vae ([`AutoencoderKL`]):
45            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
46        text_encoder ([`CLIPTextModel`]):
47            Frozen text-encoder. Stable Diffusion uses the text portion of
48            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
49            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
50        tokenizer (`CLIPTokenizer`):
51            Tokenizer of class
52            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
53        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
54        scheduler ([`SchedulerMixin`]):
55            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
56            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
57        safety_checker ([`StableDiffusionSafetyChecker`]):
58            Classification module that estimates whether generated images could be considered offensive or harmful.
59            Please, refer to the [model card](https://huggingface.co/CompVis/stable-diffusion-v1-4) for details.
60        feature_extractor ([`CLIPImageProcessor`]):
61            Model that extracts features from generated images to be used as inputs for the `safety_checker`.
62    """
63
64    def __init__(
65        self,
66        vae: AutoencoderKL,
67        text_encoder: CLIPTextModel,
68        tokenizer: CLIPTokenizer,
69        unet: UNet2DConditionModel,
70        scheduler: KarrasDiffusionSchedulers,
71        safety_checker: StableDiffusionSafetyChecker,
72        feature_extractor: CLIPFeatureExtractor,
73        requires_safety_checker: bool = True,
74    ):
75        super().__init__(
76            vae,
77            text_encoder,
78            tokenizer,
79            unet,
80            scheduler,
81            safety_checker,
82            feature_extractor,
83            requires_safety_checker,
84        )
85        self.register_modules(
86            vae=vae,
87            text_encoder=text_encoder,
88            tokenizer=tokenizer,
89            unet=unet,
90            scheduler=scheduler,
91            safety_checker=safety_checker,
92            feature_extractor=feature_extractor,
93        )
94
95    @torch.no_grad()
96    def __call__(
97        self,
98        prompt: str,
99        height: int = 512,
100        width: int = 512,
101        num_inference_steps: int = 50,
102        guidance_scale: float = 7.5,
103        negative_prompt: str = None,
104        num_images_per_prompt: Optional[int] = 1,
105        eta: float = 0.0,
106        generator: Optional[torch.Generator] = None,
107        latents: Optional[torch.Tensor] = None,
108        output_type: Optional[str] = "pil",
109        return_dict: bool = True,
110        rp_args: Dict[str, str] = None,
111    ):
112        active = KBRK in prompt[0] if isinstance(prompt, list) else KBRK in prompt
113        if negative_prompt is None:
114            negative_prompt = "" if isinstance(prompt, str) else [""] * len(prompt)
115
116        device = self._execution_device
117        regions = 0
118
119        self.power = int(rp_args["power"]) if "power" in rp_args else 1
120
121        prompts = prompt if isinstance(prompt, list) else [prompt]
122        n_prompts = negative_prompt if isinstance(prompt, str) else [negative_prompt]
123        self.batch = batch = num_images_per_prompt * len(prompts)
124        all_prompts_cn, all_prompts_p = promptsmaker(prompts, num_images_per_prompt)
125        all_n_prompts_cn, _ = promptsmaker(n_prompts, num_images_per_prompt)
126
127        equal = len(all_prompts_cn) == len(all_n_prompts_cn)
128
129        if Compel:
130            compel = Compel(tokenizer=self.tokenizer, text_encoder=self.text_encoder)
131
132            def getcompelembs(prps):
133                embl = []
134                for prp in prps:
135                    embl.append(compel.build_conditioning_tensor(prp))
136                return torch.cat(embl)
137
138            conds = getcompelembs(all_prompts_cn)
139            unconds = getcompelembs(all_n_prompts_cn)
140            embs = getcompelembs(prompts)
141            n_embs = getcompelembs(n_prompts)
142            prompt = negative_prompt = None
143        else:
144            conds = self.encode_prompt(prompts, device, 1, True)[0]
145            unconds = (
146                self.encode_prompt(n_prompts, device, 1, True)[0]
147                if equal
148                else self.encode_prompt(all_n_prompts_cn, device, 1, True)[0]
149            )
150            embs = n_embs = None
151
152        if not active:
153            pcallback = None
154            mode = None
155        else:
156            if any(x in rp_args["mode"].upper() for x in ["COL", "ROW"]):
157                mode = "COL" if "COL" in rp_args["mode"].upper() else "ROW"
158                ocells, icells, regions = make_cells(rp_args["div"])
159
160            elif "PRO" in rp_args["mode"].upper():
161                regions = len(all_prompts_p[0])
162                mode = "PROMPT"
163                reset_attnmaps(self)
164                self.ex = "EX" in rp_args["mode"].upper()
165                self.target_tokens = target_tokens = tokendealer(self, all_prompts_p)
166                thresholds = [float(x) for x in rp_args["th"].split(",")]
167
168            orig_hw = (height, width)
169            revers = True
170
171            def pcallback(s_self, step: int, timestep: int, latents: torch.Tensor, selfs=None):
172                if "PRO" in mode:  # in Prompt mode, make masks from sum of attension maps
173                    self.step = step
174
175                    if len(self.attnmaps_sizes) > 3:
176                        self.history[step] = self.attnmaps.copy()
177                        for hw in self.attnmaps_sizes:
178                            allmasks = []
179                            basemasks = [None] * batch
180                            for tt, th in zip(target_tokens, thresholds):
181                                for b in range(batch):
182                                    key = f"{tt}-{b}"
183                                    _, mask, _ = makepmask(self, self.attnmaps[key], hw[0], hw[1], th, step)
184                                    mask = mask.unsqueeze(0).unsqueeze(-1)
185                                    if self.ex:
186                                        allmasks[b::batch] = [x - mask for x in allmasks[b::batch]]
187                                        allmasks[b::batch] = [torch.where(x > 0, 1, 0) for x in allmasks[b::batch]]
188                                    allmasks.append(mask)
189                                    basemasks[b] = mask if basemasks[b] is None else basemasks[b] + mask
190                            basemasks = [1 - mask for mask in basemasks]
191                            basemasks = [torch.where(x > 0, 1, 0) for x in basemasks]
192                            allmasks = basemasks + allmasks
193
194                            self.attnmasks[hw] = torch.cat(allmasks)
195                        self.maskready = True
196                return latents
197
198            def hook_forward(module):
199                # diffusers==0.23.2
200                def forward(
201                    hidden_states: torch.Tensor,
202                    encoder_hidden_states: Optional[torch.Tensor] = None,
203                    attention_mask: Optional[torch.Tensor] = None,
204                    temb: Optional[torch.Tensor] = None,
205                    scale: float = 1.0,
206                ) -> torch.Tensor:
207                    attn = module
208                    xshape = hidden_states.shape
209                    self.hw = (h, w) = split_dims(xshape[1], *orig_hw)
210
211                    if revers:
212                        nx, px = hidden_states.chunk(2)
213                    else:
214                        px, nx = hidden_states.chunk(2)
215
216                    if equal:
217                        hidden_states = torch.cat(
218                            [px for i in range(regions)] + [nx for i in range(regions)],
219                            0,
220                        )
221                        encoder_hidden_states = torch.cat([conds] + [unconds])
222                    else:
223                        hidden_states = torch.cat([px for i in range(regions)] + [nx], 0)
224                        encoder_hidden_states = torch.cat([conds] + [unconds])
225
226                    residual = hidden_states
227
228                    args = () if USE_PEFT_BACKEND else (scale,)
229
230                    if attn.spatial_norm is not None:
231                        hidden_states = attn.spatial_norm(hidden_states, temb)
232
233                    input_ndim = hidden_states.ndim
234
235                    if input_ndim == 4:
236                        batch_size, channel, height, width = hidden_states.shape
237                        hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
238
239                    batch_size, sequence_length, _ = (
240                        hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
241                    )
242
243                    if attention_mask is not None:
244                        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
245                        attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
246
247                    if attn.group_norm is not None:
248                        hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
249
250                    args = () if USE_PEFT_BACKEND else (scale,)
251                    query = attn.to_q(hidden_states, *args)
252
253                    if encoder_hidden_states is None:
254                        encoder_hidden_states = hidden_states
255                    elif attn.norm_cross:
256                        encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
257
258                    key = attn.to_k(encoder_hidden_states, *args)
259                    value = attn.to_v(encoder_hidden_states, *args)
260
261                    inner_dim = key.shape[-1]
262                    head_dim = inner_dim // attn.heads
263
264                    query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
265
266                    key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
267                    value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
268
269                    # the output of sdp = (batch, num_heads, seq_len, head_dim)
270                    # TODO: add support for attn.scale when we move to Torch 2.1
271                    hidden_states = scaled_dot_product_attention(
272                        self,
273                        query,
274                        key,
275                        value,
276                        attn_mask=attention_mask,
277                        dropout_p=0.0,
278                        is_causal=False,
279                        getattn="PRO" in mode,
280                    )
281
282                    hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
283                    hidden_states = hidden_states.to(query.dtype)
284
285                    # linear proj
286                    hidden_states = attn.to_out[0](hidden_states, *args)
287                    # dropout
288                    hidden_states = attn.to_out[1](hidden_states)
289
290                    if input_ndim == 4:
291                        hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
292
293                    if attn.residual_connection:
294                        hidden_states = hidden_states + residual
295
296                    hidden_states = hidden_states / attn.rescale_output_factor
297
298                    #### Regional Prompting Col/Row mode
299                    if any(x in mode for x in ["COL", "ROW"]):
300                        reshaped = hidden_states.reshape(hidden_states.size()[0], h, w, hidden_states.size()[2])
301                        center = reshaped.shape[0] // 2
302                        px = reshaped[0:center] if equal else reshaped[0:-batch]
303                        nx = reshaped[center:] if equal else reshaped[-batch:]
304                        outs = [px, nx] if equal else [px]
305                        for out in outs:
306                            c = 0
307                            for i, ocell in enumerate(ocells):
308                                for icell in icells[i]:
309                                    if "ROW" in mode:
310                                        out[
311                                            0:batch,
312                                            int(h * ocell[0]) : int(h * ocell[1]),
313                                            int(w * icell[0]) : int(w * icell[1]),
314                                            :,
315                                        ] = out[
316                                            c * batch : (c + 1) * batch,
317                                            int(h * ocell[0]) : int(h * ocell[1]),
318                                            int(w * icell[0]) : int(w * icell[1]),
319                                            :,
320                                        ]
321                                    else:
322                                        out[
323                                            0:batch,
324                                            int(h * icell[0]) : int(h * icell[1]),
325                                            int(w * ocell[0]) : int(w * ocell[1]),
326                                            :,
327                                        ] = out[
328                                            c * batch : (c + 1) * batch,
329                                            int(h * icell[0]) : int(h * icell[1]),
330                                            int(w * ocell[0]) : int(w * ocell[1]),
331                                            :,
332                                        ]
333                                    c += 1
334                        px, nx = (px[0:batch], nx[0:batch]) if equal else (px[0:batch], nx)
335                        hidden_states = torch.cat([nx, px], 0) if revers else torch.cat([px, nx], 0)
336                        hidden_states = hidden_states.reshape(xshape)
337
338                    #### Regional Prompting Prompt mode
339                    elif "PRO" in mode:
340                        px, nx = (
341                            torch.chunk(hidden_states) if equal else hidden_states[0:-batch],
342                            hidden_states[-batch:],
343                        )
344
345                        if (h, w) in self.attnmasks and self.maskready:
346
347                            def mask(input):
348                                out = torch.multiply(input, self.attnmasks[(h, w)])
349                                for b in range(batch):
350                                    for r in range(1, regions):
351                                        out[b] = out[b] + out[r * batch + b]
352                                return out
353
354                            px, nx = (mask(px), mask(nx)) if equal else (mask(px), nx)
355                        px, nx = (px[0:batch], nx[0:batch]) if equal else (px[0:batch], nx)
356                        hidden_states = torch.cat([nx, px], 0) if revers else torch.cat([px, nx], 0)
357                    return hidden_states
358
359                return forward
360
361            def hook_forwards(root_module: torch.nn.Module):
362                for name, module in root_module.named_modules():
363                    if "attn2" in name and module.__class__.__name__ == "Attention":
364                        module.forward = hook_forward(module)
365
366            hook_forwards(self.unet)
367
368        output = StableDiffusionPipeline(**self.components)(
369            prompt=prompt,
370            prompt_embeds=embs,
371            negative_prompt=negative_prompt,
372            negative_prompt_embeds=n_embs,
373            height=height,
374            width=width,
375            num_inference_steps=num_inference_steps,
376            guidance_scale=guidance_scale,
377            num_images_per_prompt=num_images_per_prompt,
378            eta=eta,
379            generator=generator,
380            latents=latents,
381            output_type=output_type,
382            return_dict=return_dict,
383            callback_on_step_end=pcallback,
384        )
385
386        if "save_mask" in rp_args:
387            save_mask = rp_args["save_mask"]
388        else:
389            save_mask = False
390
391        if mode == "PROMPT" and save_mask:
392            saveattnmaps(
393                self,
394                output,
395                height,
396                width,
397                thresholds,
398                num_inference_steps // 2,
399                regions,
400            )
401
402        return output
403
404
405### Make prompt list for each regions
406def promptsmaker(prompts, batch):
407    out_p = []
408    plen = len(prompts)
409    for prompt in prompts:
410        add = ""
411        if KCOMM in prompt:
412            add, prompt = prompt.split(KCOMM)
413            add = add + " "
414        prompts = prompt.split(KBRK)
415        out_p.append([add + p for p in prompts])
416    out = [None] * batch * len(out_p[0]) * len(out_p)
417    for p, prs in enumerate(out_p):  # inputs prompts
418        for r, pr in enumerate(prs):  # prompts for regions
419            start = (p + r * plen) * batch
420            out[start : start + batch] = [pr] * batch  # P1R1B1,P1R1B2...,P1R2B1,P1R2B2...,P2R1B1...
421    return out, out_p
422
423
424### make regions from ratios
425### ";" makes outercells, "," makes inner cells
426def make_cells(ratios):
427    if ";" not in ratios and "," in ratios:
428        ratios = ratios.replace(",", ";")
429    ratios = ratios.split(";")
430    ratios = [inratios.split(",") for inratios in ratios]
431
432    icells = []
433    ocells = []
434
435    def startend(cells, array):
436        current_start = 0
437        array = [float(x) for x in array]
438        for value in array:
439            end = current_start + (value / sum(array))
440            cells.append([current_start, end])
441            current_start = end
442
443    startend(ocells, [r[0] for r in ratios])
444
445    for inratios in ratios:
446        if 2 > len(inratios):
447            icells.append([[0, 1]])
448        else:
449            add = []
450            startend(add, inratios[1:])
451            icells.append(add)
452
453    return ocells, icells, sum(len(cell) for cell in icells)
454
455
456def make_emblist(self, prompts):
457    with torch.no_grad():
458        tokens = self.tokenizer(
459            prompts,
460            max_length=self.tokenizer.model_max_length,
461            padding=True,
462            truncation=True,
463            return_tensors="pt",
464        ).input_ids.to(self.device)
465        embs = self.text_encoder(tokens, output_hidden_states=True).last_hidden_state.to(self.device, dtype=self.dtype)
466    return embs
467
468
469def split_dims(xs, height, width):
470    xs = xs
471
472    def repeat_div(x, y):
473        while y > 0:
474            x = math.ceil(x / 2)
475            y = y - 1
476        return x
477
478    scale = math.ceil(math.log2(math.sqrt(height * width / xs)))
479    dsh = repeat_div(height, scale)
480    dsw = repeat_div(width, scale)
481    return dsh, dsw
482
483
484##### for prompt mode
485def get_attn_maps(self, attn):
486    height, width = self.hw
487    target_tokens = self.target_tokens
488    if (height, width) not in self.attnmaps_sizes:
489        self.attnmaps_sizes.append((height, width))
490
491    for b in range(self.batch):
492        for t in target_tokens:
493            power = self.power
494            add = attn[b, :, :, t[0] : t[0] + len(t)] ** (power) * (self.attnmaps_sizes.index((height, width)) + 1)
495            add = torch.sum(add, dim=2)
496            key = f"{t}-{b}"
497            if key not in self.attnmaps:
498                self.attnmaps[key] = add
499            else:
500                if self.attnmaps[key].shape[1] != add.shape[1]:
501                    add = add.view(8, height, width)
502                    add = FF.resize(add, self.attnmaps_sizes[0], antialias=None)
503                    add = add.reshape_as(self.attnmaps[key])
504
505                self.attnmaps[key] = self.attnmaps[key] + add
506
507
508def reset_attnmaps(self):  # init parameters in every batch
509    self.step = 0
510    self.attnmaps = {}  # maked from attention maps
511    self.attnmaps_sizes = []  # height,width set of u-net blocks
512    self.attnmasks = {}  # maked from attnmaps for regions
513    self.maskready = False
514    self.history = {}
515
516
517def saveattnmaps(self, output, h, w, th, step, regions):
518    masks = []
519    for i, mask in enumerate(self.history[step].values()):
520        img, _, mask = makepmask(self, mask, h, w, th[i % len(th)], step)
521        if self.ex:
522            masks = [x - mask for x in masks]
523            masks.append(mask)
524            if len(masks) == regions - 1:
525                output.images.extend([FF.to_pil_image(mask) for mask in masks])
526                masks = []
527        else:
528            output.images.append(img)
529
530
531def makepmask(
532    self, mask, h, w, th, step
533):  # make masks from attention cache return [for preview, for attention, for Latent]
534    th = th - step * 0.005
535    if 0.05 >= th:
536        th = 0.05
537    mask = torch.mean(mask, dim=0)
538    mask = mask / mask.max().item()
539    mask = torch.where(mask > th, 1, 0)
540    mask = mask.float()
541    mask = mask.view(1, *self.attnmaps_sizes[0])
542    img = FF.to_pil_image(mask)
543    img = img.resize((w, h))
544    mask = FF.resize(mask, (h, w), interpolation=FF.InterpolationMode.NEAREST, antialias=None)
545    lmask = mask
546    mask = mask.reshape(h * w)
547    mask = torch.where(mask > 0.1, 1, 0)
548    return img, mask, lmask
549
550
551def tokendealer(self, all_prompts):
552    for prompts in all_prompts:
553        targets = [p.split(",")[-1] for p in prompts[1:]]
554        tt = []
555
556        for target in targets:
557            ptokens = (
558                self.tokenizer(
559                    prompts,
560                    max_length=self.tokenizer.model_max_length,
561                    padding=True,
562                    truncation=True,
563                    return_tensors="pt",
564                ).input_ids
565            )[0]
566            ttokens = (
567                self.tokenizer(
568                    target,
569                    max_length=self.tokenizer.model_max_length,
570                    padding=True,
571                    truncation=True,
572                    return_tensors="pt",
573                ).input_ids
574            )[0]
575
576            tlist = []
577
578            for t in range(ttokens.shape[0] - 2):
579                for p in range(ptokens.shape[0]):
580                    if ttokens[t + 1] == ptokens[p]:
581                        tlist.append(p)
582            if tlist != []:
583                tt.append(tlist)
584
585    return tt
586
587
588def scaled_dot_product_attention(
589    self,
590    query,
591    key,
592    value,
593    attn_mask=None,
594    dropout_p=0.0,
595    is_causal=False,
596    scale=None,
597    getattn=False,
598) -> torch.Tensor:
599    # Efficient implementation equivalent to the following:
600    L, S = query.size(-2), key.size(-2)
601    scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
602    attn_bias = torch.zeros(L, S, dtype=query.dtype, device=self.device)
603    if is_causal:
604        assert attn_mask is None
605        temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
606        attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
607        attn_bias.to(query.dtype)
608
609    if attn_mask is not None:
610        if attn_mask.dtype == torch.bool:
611            attn_mask.masked_fill_(attn_mask.logical_not(), float("-inf"))
612        else:
613            attn_bias += attn_mask
614    attn_weight = query @ key.transpose(-2, -1) * scale_factor
615    attn_weight += attn_bias
616    attn_weight = torch.softmax(attn_weight, dim=-1)
617    if getattn:
618        get_attn_maps(self, attn_weight)
619    attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
620    return attn_weight @ value
621