diffusers/community-pipelines-mirror
Community Pipeline Examples For more information about community pipelines, please have a look at this issue. Community pipeline examples consist pipelines that have been added by the community. Please have a look at the following tables to get an overview of all community examples. Click on the Code Example to get a copy-and-paste ready code example that you can try out. If a community pipeline doesn't work as expected, please open an issue and ping the author on it. Please… See the full description on the dataset page: https://huggingface.co/datasets/diffusers/community-pipelines-mirror.
922k
1# Copyright 2024 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 at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# 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 and13# limitations under the License.14 15import gc16import inspect17from typing import Any, Callable, Dict, List, Optional, Tuple, Union18 19import numpy as np20import PIL.Image21import torch22import torch.nn.functional as F23import torch.utils.model_zoo24from einops import rearrange, repeat25from gmflow.gmflow import GMFlow26from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection27 28from diffusers.image_processor import PipelineImageInput, VaeImageProcessor29from diffusers.loaders import LoraLoaderMixin, TextualInversionLoaderMixin30from diffusers.models import AutoencoderKL, ControlNetModel, ImageProjection, UNet2DConditionModel31from diffusers.models.attention_processor import AttnProcessor2_032from diffusers.models.lora import adjust_lora_scale_text_encoder33from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput34from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel35from diffusers.pipelines.controlnet.pipeline_controlnet_img2img import StableDiffusionControlNetImg2ImgPipeline36from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput37from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker38from diffusers.schedulers import KarrasDiffusionSchedulers39from diffusers.utils import (40 USE_PEFT_BACKEND,41 deprecate,42 logging,43 scale_lora_layers,44 unscale_lora_layers,45)46from diffusers.utils.torch_utils import is_compiled_module, randn_tensor47 48 49logger = logging.get_logger(__name__) # pylint: disable=invalid-name50 51 52def clear_cache():53 gc.collect()54 torch.cuda.empty_cache()55 56 57def coords_grid(b, h, w, homogeneous=False, device=None):58 y, x = torch.meshgrid(torch.arange(h), torch.arange(w)) # [H, W]59 60 stacks = [x, y]61 62 if homogeneous:63 ones = torch.ones_like(x) # [H, W]64 stacks.append(ones)65 66 grid = torch.stack(stacks, dim=0).float() # [2, H, W] or [3, H, W]67 68 grid = grid[None].repeat(b, 1, 1, 1) # [B, 2, H, W] or [B, 3, H, W]69 70 if device is not None:71 grid = grid.to(device)72 73 return grid74 75 76def bilinear_sample(img, sample_coords, mode="bilinear", padding_mode="zeros", return_mask=False):77 # img: [B, C, H, W]78 # sample_coords: [B, 2, H, W] in image scale79 if sample_coords.size(1) != 2: # [B, H, W, 2]80 sample_coords = sample_coords.permute(0, 3, 1, 2)81 82 b, _, h, w = sample_coords.shape83 84 # Normalize to [-1, 1]85 x_grid = 2 * sample_coords[:, 0] / (w - 1) - 186 y_grid = 2 * sample_coords[:, 1] / (h - 1) - 187 88 grid = torch.stack([x_grid, y_grid], dim=-1) # [B, H, W, 2]89 90 img = F.grid_sample(img, grid, mode=mode, padding_mode=padding_mode, align_corners=True)91 92 if return_mask:93 mask = (x_grid >= -1) & (y_grid >= -1) & (x_grid <= 1) & (y_grid <= 1) # [B, H, W]94 95 return img, mask96 97 return img98 99 100class Dilate:101 def __init__(self, kernel_size=7, channels=1, device="cpu"):102 self.kernel_size = kernel_size103 self.channels = channels104 gaussian_kernel = torch.ones(1, 1, self.kernel_size, self.kernel_size)105 gaussian_kernel = gaussian_kernel.repeat(self.channels, 1, 1, 1)106 self.mean = (self.kernel_size - 1) // 2107 gaussian_kernel = gaussian_kernel.to(device)108 self.gaussian_filter = gaussian_kernel109 110 def __call__(self, x):111 x = F.pad(x, (self.mean, self.mean, self.mean, self.mean), "replicate")112 return torch.clamp(F.conv2d(x, self.gaussian_filter, bias=None), 0, 1)113 114 115def flow_warp(feature, flow, mask=False, mode="bilinear", padding_mode="zeros"):116 b, c, h, w = feature.size()117 assert flow.size(1) == 2118 119 grid = coords_grid(b, h, w).to(flow.device) + flow # [B, 2, H, W]120 grid = grid.to(feature.dtype)121 return bilinear_sample(feature, grid, mode=mode, padding_mode=padding_mode, return_mask=mask)122 123 124def forward_backward_consistency_check(fwd_flow, bwd_flow, alpha=0.01, beta=0.5):125 # fwd_flow, bwd_flow: [B, 2, H, W]126 # alpha and beta values are following UnFlow127 # (https://arxiv.org/abs/1711.07837)128 assert fwd_flow.dim() == 4 and bwd_flow.dim() == 4129 assert fwd_flow.size(1) == 2 and bwd_flow.size(1) == 2130 flow_mag = torch.norm(fwd_flow, dim=1) + torch.norm(bwd_flow, dim=1) # [B, H, W]131 132 warped_bwd_flow = flow_warp(bwd_flow, fwd_flow) # [B, 2, H, W]133 warped_fwd_flow = flow_warp(fwd_flow, bwd_flow) # [B, 2, H, W]134 135 diff_fwd = torch.norm(fwd_flow + warped_bwd_flow, dim=1) # [B, H, W]136 diff_bwd = torch.norm(bwd_flow + warped_fwd_flow, dim=1)137 138 threshold = alpha * flow_mag + beta139 140 fwd_occ = (diff_fwd > threshold).float() # [B, H, W]141 bwd_occ = (diff_bwd > threshold).float()142 143 return fwd_occ, bwd_occ144 145 146def numpy2tensor(img):147 x0 = torch.from_numpy(img.copy()).float().cuda() / 255.0 * 2.0 - 1.0148 x0 = torch.stack([x0], dim=0)149 # einops.rearrange(x0, 'b h w c -> b c h w').clone()150 return x0.permute(0, 3, 1, 2)151 152 153def calc_mean_std(feat, eps=1e-5, chunk=1):154 size = feat.size()155 assert len(size) == 4156 if chunk == 2:157 feat = torch.cat(feat.chunk(2), dim=3)158 N, C = size[:2]159 feat_var = feat.view(N // chunk, C, -1).var(dim=2) + eps160 feat_std = feat_var.sqrt().view(N, C, 1, 1)161 feat_mean = feat.view(N // chunk, C, -1).mean(dim=2).view(N // chunk, C, 1, 1)162 return feat_mean.repeat(chunk, 1, 1, 1), feat_std.repeat(chunk, 1, 1, 1)163 164 165def adaptive_instance_normalization(content_feat, style_feat, chunk=1):166 assert content_feat.size()[:2] == style_feat.size()[:2]167 size = content_feat.size()168 style_mean, style_std = calc_mean_std(style_feat, chunk)169 content_mean, content_std = calc_mean_std(content_feat)170 171 normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)172 return normalized_feat * style_std.expand(size) + style_mean.expand(size)173 174 175def optimize_feature(176 sample, flows, occs, correlation_matrix=[], intra_weight=1e2, iters=20, unet_chunk_size=2, optimize_temporal=True177):178 """179 FRESO-guided latent feature optimization180 * optimize spatial correspondence (match correlation_matrix)181 * optimize temporal correspondence (match warped_image)182 """183 if (flows is None or occs is None or (not optimize_temporal)) and (184 intra_weight == 0 or len(correlation_matrix) == 0185 ):186 return sample187 # flows=[fwd_flows, bwd_flows]: (N-1)*2*H1*W1188 # occs=[fwd_occs, bwd_occs]: (N-1)*H1*W1189 # sample: 2N*C*H*W190 torch.cuda.empty_cache()191 video_length = sample.shape[0] // unet_chunk_size192 latent = rearrange(sample.to(torch.float32), "(b f) c h w -> b f c h w", f=video_length)193 194 cs = torch.nn.Parameter((latent.detach().clone()))195 optimizer = torch.optim.Adam([cs], lr=0.2)196 197 # unify resolution198 if flows is not None and occs is not None:199 scale = sample.shape[2] * 1.0 / flows[0].shape[2]200 kernel = int(1 / scale)201 bwd_flow_ = F.interpolate(flows[1] * scale, scale_factor=scale, mode="bilinear").repeat(202 unet_chunk_size, 1, 1, 1203 )204 bwd_occ_ = F.max_pool2d(occs[1].unsqueeze(1), kernel_size=kernel).repeat(205 unet_chunk_size, 1, 1, 1206 ) # 2(N-1)*1*H1*W1207 fwd_flow_ = F.interpolate(flows[0] * scale, scale_factor=scale, mode="bilinear").repeat(208 unet_chunk_size, 1, 1, 1209 )210 fwd_occ_ = F.max_pool2d(occs[0].unsqueeze(1), kernel_size=kernel).repeat(211 unet_chunk_size, 1, 1, 1212 ) # 2(N-1)*1*H1*W1213 # match frame 0,1,2,3 and frame 1,2,3,0214 reshuffle_list = list(range(1, video_length)) + [0]215 216 # attention_probs is the GRAM matrix of the normalized feature217 attention_probs = None218 for tmp in correlation_matrix:219 if sample.shape[2] * sample.shape[3] == tmp.shape[1]:220 attention_probs = tmp # 2N*HW*HW221 break222 223 n_iter = [0]224 while n_iter[0] < iters:225 226 def closure():227 optimizer.zero_grad()228 229 loss = 0230 231 # temporal consistency loss232 if optimize_temporal and flows is not None and occs is not None:233 c1 = rearrange(cs[:, :], "b f c h w -> (b f) c h w")234 c2 = rearrange(cs[:, reshuffle_list], "b f c h w -> (b f) c h w")235 warped_image1 = flow_warp(c1, bwd_flow_)236 warped_image2 = flow_warp(c2, fwd_flow_)237 loss = (238 abs((c2 - warped_image1) * (1 - bwd_occ_)) + abs((c1 - warped_image2) * (1 - fwd_occ_))239 ).mean() * 2240 241 # spatial consistency loss242 if attention_probs is not None and intra_weight > 0:243 cs_vector = rearrange(cs, "b f c h w -> (b f) (h w) c")244 # attention_scores = torch.bmm(cs_vector, cs_vector.transpose(-1, -2))245 # cs_attention_probs = attention_scores.softmax(dim=-1)246 cs_vector = cs_vector / ((cs_vector**2).sum(dim=2, keepdims=True) ** 0.5)247 cs_attention_probs = torch.bmm(cs_vector, cs_vector.transpose(-1, -2))248 tmp = F.l1_loss(cs_attention_probs, attention_probs) * intra_weight249 loss = tmp + loss250 251 loss.backward()252 n_iter[0] += 1253 254 return loss255 256 optimizer.step(closure)257 258 torch.cuda.empty_cache()259 return adaptive_instance_normalization(rearrange(cs.data.to(sample.dtype), "b f c h w -> (b f) c h w"), sample)260 261 262@torch.no_grad()263def warp_tensor(sample, flows, occs, saliency, unet_chunk_size):264 """265 Warp images or features based on optical flow266 Fuse the warped imges or features based on occusion masks and saliency map267 """268 scale = sample.shape[2] * 1.0 / flows[0].shape[2]269 kernel = int(1 / scale)270 bwd_flow_ = F.interpolate(flows[1] * scale, scale_factor=scale, mode="bilinear")271 bwd_occ_ = F.max_pool2d(occs[1].unsqueeze(1), kernel_size=kernel) # (N-1)*1*H1*W1272 if scale == 1:273 bwd_occ_ = Dilate(kernel_size=13, device=sample.device)(bwd_occ_)274 fwd_flow_ = F.interpolate(flows[0] * scale, scale_factor=scale, mode="bilinear")275 fwd_occ_ = F.max_pool2d(occs[0].unsqueeze(1), kernel_size=kernel) # (N-1)*1*H1*W1276 if scale == 1:277 fwd_occ_ = Dilate(kernel_size=13, device=sample.device)(fwd_occ_)278 scale2 = sample.shape[2] * 1.0 / saliency.shape[2]279 saliency = F.interpolate(saliency, scale_factor=scale2, mode="bilinear")280 latent = sample.to(torch.float32)281 video_length = sample.shape[0] // unet_chunk_size282 warp_saliency = flow_warp(saliency, bwd_flow_)283 warp_saliency_ = flow_warp(saliency[0:1], fwd_flow_[video_length - 1 : video_length])284 285 for j in range(unet_chunk_size):286 for ii in range(video_length - 1):287 i = video_length * j + ii288 warped_image = flow_warp(latent[i : i + 1], bwd_flow_[ii : ii + 1])289 mask = (1 - bwd_occ_[ii : ii + 1]) * saliency[ii + 1 : ii + 2] * warp_saliency[ii : ii + 1]290 latent[i + 1 : i + 2] = latent[i + 1 : i + 2] * (1 - mask) + warped_image * mask291 i = video_length * j292 ii = video_length - 1293 warped_image = flow_warp(latent[i : i + 1], fwd_flow_[ii : ii + 1])294 mask = (1 - fwd_occ_[ii : ii + 1]) * saliency[ii : ii + 1] * warp_saliency_295 latent[ii + i : ii + i + 1] = latent[ii + i : ii + i + 1] * (1 - mask) + warped_image * mask296 297 return latent.to(sample.dtype)298 299 300def my_forward(301 self,302 steps=[],303 layers=[0, 1, 2, 3],304 flows=None,305 occs=None,306 correlation_matrix=[],307 intra_weight=1e2,308 iters=20,309 optimize_temporal=True,310 saliency=None,311):312 """313 Hacked pipe.unet.forward()314 copied from https://github.com/huggingface/diffusers/blob/v0.19.3/src/diffusers/models/unet_2d_condition.py#L700315 if you are using a new version of diffusers, please copy the source code and modify it accordingly (find [HACK] in the code)316 * restore and return the decoder features317 * optimize the decoder features318 * perform background smoothing319 """320 321 def forward(322 sample: torch.FloatTensor,323 timestep: Union[torch.Tensor, float, int],324 encoder_hidden_states: torch.Tensor,325 class_labels: Optional[torch.Tensor] = None,326 timestep_cond: Optional[torch.Tensor] = None,327 attention_mask: Optional[torch.Tensor] = None,328 cross_attention_kwargs: Optional[Dict[str, Any]] = None,329 added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,330 down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,331 mid_block_additional_residual: Optional[torch.Tensor] = None,332 encoder_attention_mask: Optional[torch.Tensor] = None,333 return_dict: bool = True,334 ) -> Union[UNet2DConditionOutput, Tuple]:335 r"""336 The [`UNet2DConditionModel`] forward method.337 338 Args:339 sample (`torch.FloatTensor`):340 The noisy input tensor with the following shape `(batch, channel, height, width)`.341 timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input.342 encoder_hidden_states (`torch.FloatTensor`):343 The encoder hidden states with shape `(batch, sequence_length, feature_dim)`.344 encoder_attention_mask (`torch.Tensor`):345 A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If346 `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias,347 which adds large negative values to the attention scores corresponding to "discard" tokens.348 return_dict (`bool`, *optional*, defaults to `True`):349 Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain350 tuple.351 cross_attention_kwargs (`dict`, *optional*):352 A kwargs dictionary that if specified is passed along to the [`AttnProcessor`].353 added_cond_kwargs: (`dict`, *optional*):354 A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that355 are passed along to the UNet blocks.356 357 Returns:358 [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:359 If `return_dict` is True, an [`~models.unet_2d_condition.UNet2DConditionOutput`] is returned, otherwise360 a `tuple` is returned where the first element is the sample tensor.361 """362 # By default samples have to be AT least a multiple of the overall upsampling factor.363 # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).364 # However, the upsampling interpolation output size can be forced to fit any upsampling size365 # on the fly if necessary.366 default_overall_up_factor = 2**self.num_upsamplers367 368 # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`369 forward_upsample_size = False370 upsample_size = None371 372 if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):373 logger.info("Forward upsample size to force interpolation output size.")374 forward_upsample_size = True375 376 # ensure attention_mask is a bias, and give it a singleton query_tokens dimension377 # expects mask of shape:378 # [batch, key_tokens]379 # adds singleton query_tokens dimension:380 # [batch, 1, key_tokens]381 # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:382 # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)383 # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)384 if attention_mask is not None:385 # assume that mask is expressed as:386 # (1 = keep, 0 = discard)387 # convert mask into a bias that can be added to attention scores:388 # (keep = +0, discard = -10000.0)389 attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0390 attention_mask = attention_mask.unsqueeze(1)391 392 # convert encoder_attention_mask to a bias the same way we do for attention_mask393 if encoder_attention_mask is not None:394 encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0395 encoder_attention_mask = encoder_attention_mask.unsqueeze(1)396 397 # 0. center input if necessary398 if self.config.center_input_sample:399 sample = 2 * sample - 1.0400 401 # 1. time402 timesteps = timestep403 if not torch.is_tensor(timesteps):404 # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can405 # This would be a good case for the `match` statement (Python 3.10+)406 is_mps = sample.device.type == "mps"407 if isinstance(timestep, float):408 dtype = torch.float32 if is_mps else torch.float64409 else:410 dtype = torch.int32 if is_mps else torch.int64411 timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)412 elif len(timesteps.shape) == 0:413 timesteps = timesteps[None].to(sample.device)414 415 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML416 timesteps = timesteps.expand(sample.shape[0])417 418 t_emb = self.time_proj(timesteps)419 420 # `Timesteps` does not contain any weights and will always return f32 tensors421 # but time_embedding might actually be running in fp16. so we need to cast here.422 # there might be better ways to encapsulate this.423 t_emb = t_emb.to(dtype=sample.dtype)424 425 emb = self.time_embedding(t_emb, timestep_cond)426 aug_emb = None427 428 if self.class_embedding is not None:429 if class_labels is None:430 raise ValueError("class_labels should be provided when num_class_embeds > 0")431 432 if self.config.class_embed_type == "timestep":433 class_labels = self.time_proj(class_labels)434 435 # `Timesteps` does not contain any weights and will always return f32 tensors436 # there might be better ways to encapsulate this.437 class_labels = class_labels.to(dtype=sample.dtype)438 439 class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype)440 441 if self.config.class_embeddings_concat:442 emb = torch.cat([emb, class_emb], dim=-1)443 else:444 emb = emb + class_emb445 446 if self.config.addition_embed_type == "text":447 aug_emb = self.add_embedding(encoder_hidden_states)448 elif self.config.addition_embed_type == "text_image":449 # Kandinsky 2.1 - style450 if "image_embeds" not in added_cond_kwargs:451 raise ValueError(452 f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"453 )454 455 image_embs = added_cond_kwargs.get("image_embeds")456 text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states)457 aug_emb = self.add_embedding(text_embs, image_embs)458 elif self.config.addition_embed_type == "text_time":459 # SDXL - style460 if "text_embeds" not in added_cond_kwargs:461 raise ValueError(462 f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"463 )464 text_embeds = added_cond_kwargs.get("text_embeds")465 if "time_ids" not in added_cond_kwargs:466 raise ValueError(467 f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"468 )469 time_ids = added_cond_kwargs.get("time_ids")470 time_embeds = self.add_time_proj(time_ids.flatten())471 time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))472 473 add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)474 add_embeds = add_embeds.to(emb.dtype)475 aug_emb = self.add_embedding(add_embeds)476 elif self.config.addition_embed_type == "image":477 # Kandinsky 2.2 - style478 if "image_embeds" not in added_cond_kwargs:479 raise ValueError(480 f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"481 )482 image_embs = added_cond_kwargs.get("image_embeds")483 aug_emb = self.add_embedding(image_embs)484 elif self.config.addition_embed_type == "image_hint":485 # Kandinsky 2.2 - style486 if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs:487 raise ValueError(488 f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`"489 )490 image_embs = added_cond_kwargs.get("image_embeds")491 hint = added_cond_kwargs.get("hint")492 aug_emb, hint = self.add_embedding(image_embs, hint)493 sample = torch.cat([sample, hint], dim=1)494 495 emb = emb + aug_emb if aug_emb is not None else emb496 497 if self.time_embed_act is not None:498 emb = self.time_embed_act(emb)499 500 if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj":501 encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states)502 elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj":503 # Kadinsky 2.1 - style504 if "image_embeds" not in added_cond_kwargs:505 raise ValueError(506 f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`"507 )508 509 image_embeds = added_cond_kwargs.get("image_embeds")510 encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds)511 elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj":512 # Kandinsky 2.2 - style513 if "image_embeds" not in added_cond_kwargs:514 raise ValueError(515 f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`"516 )517 image_embeds = added_cond_kwargs.get("image_embeds")518 encoder_hidden_states = self.encoder_hid_proj(image_embeds)519 # 2. pre-process520 sample = self.conv_in(sample)521 522 # 3. down523 524 is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None525 is_adapter = mid_block_additional_residual is None and down_block_additional_residuals is not None526 527 down_block_res_samples = (sample,)528 for downsample_block in self.down_blocks:529 if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:530 # For t2i-adapter CrossAttnDownBlock2D531 additional_residuals = {}532 if is_adapter and len(down_block_additional_residuals) > 0:533 additional_residuals["additional_residuals"] = down_block_additional_residuals.pop(0)534 535 sample, res_samples = downsample_block(536 hidden_states=sample,537 temb=emb,538 encoder_hidden_states=encoder_hidden_states,539 attention_mask=attention_mask,540 cross_attention_kwargs=cross_attention_kwargs,541 encoder_attention_mask=encoder_attention_mask,542 **additional_residuals,543 )544 else:545 sample, res_samples = downsample_block(hidden_states=sample, temb=emb)546 547 if is_adapter and len(down_block_additional_residuals) > 0:548 sample += down_block_additional_residuals.pop(0)549 down_block_res_samples += res_samples550 551 if is_controlnet:552 new_down_block_res_samples = ()553 554 for down_block_res_sample, down_block_additional_residual in zip(555 down_block_res_samples, down_block_additional_residuals556 ):557 down_block_res_sample = down_block_res_sample + down_block_additional_residual558 new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)559 560 down_block_res_samples = new_down_block_res_samples561 562 # 4. mid563 if self.mid_block is not None:564 sample = self.mid_block(565 sample,566 emb,567 encoder_hidden_states=encoder_hidden_states,568 attention_mask=attention_mask,569 cross_attention_kwargs=cross_attention_kwargs,570 encoder_attention_mask=encoder_attention_mask,571 )572 573 if is_controlnet:574 sample = sample + mid_block_additional_residual575 576 # 5. up577 """578 [HACK] restore the decoder features in up_samples579 """580 up_samples = ()581 # down_samples = ()582 for i, upsample_block in enumerate(self.up_blocks):583 is_final_block = i == len(self.up_blocks) - 1584 585 res_samples = down_block_res_samples[-len(upsample_block.resnets) :]586 down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]587 588 """589 [HACK] restore the decoder features in up_samples590 [HACK] optimize the decoder features591 [HACK] perform background smoothing592 """593 if i in layers:594 up_samples += (sample,)595 if timestep in steps and i in layers:596 sample = optimize_feature(597 sample, flows, occs, correlation_matrix, intra_weight, iters, optimize_temporal=optimize_temporal598 )599 if saliency is not None:600 sample = warp_tensor(sample, flows, occs, saliency, 2)601 602 # if we have not reached the final block and need to forward the603 # upsample size, we do it here604 if not is_final_block and forward_upsample_size:605 upsample_size = down_block_res_samples[-1].shape[2:]606 607 if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:608 sample = upsample_block(609 hidden_states=sample,610 temb=emb,611 res_hidden_states_tuple=res_samples,612 encoder_hidden_states=encoder_hidden_states,613 cross_attention_kwargs=cross_attention_kwargs,614 upsample_size=upsample_size,615 attention_mask=attention_mask,616 encoder_attention_mask=encoder_attention_mask,617 )618 else:619 sample = upsample_block(620 hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size621 )622 623 # 6. post-process624 if self.conv_norm_out:625 sample = self.conv_norm_out(sample)626 sample = self.conv_act(sample)627 sample = self.conv_out(sample)628 629 """630 [HACK] return the output feature as well as the decoder features631 """632 if not return_dict:633 return (sample,) + up_samples634 635 return UNet2DConditionOutput(sample=sample)636 637 return forward638 639 640@torch.no_grad()641def get_single_mapping_ind(bwd_flow, bwd_occ, imgs, scale=1.0):642 """643 FLATTEN: Optical fLow-guided attention (Temoporal-guided attention)644 Find the correspondence between every pixels in a pair of frames645 646 [input]647 bwd_flow: 1*2*H*W648 bwd_occ: 1*H*W i.e., f2 = warp(f1, bwd_flow) * bwd_occ649 imgs: 2*3*H*W i.e., [f1,f2]650 651 [output]652 mapping_ind: pixel index correspondence653 unlinkedmask: indicate whether a pixel has no correspondence654 i.e., f2 = f1[mapping_ind] * unlinkedmask655 """656 flows = F.interpolate(bwd_flow, scale_factor=1.0 / scale, mode="bilinear")[0][[1, 0]] / scale # 2*H*W657 _, H, W = flows.shape658 masks = torch.logical_not(F.interpolate(bwd_occ[None], scale_factor=1.0 / scale, mode="bilinear") > 0.5)[659 0660 ] # 1*H*W661 frames = F.interpolate(imgs, scale_factor=1.0 / scale, mode="bilinear").view(2, 3, -1) # 2*3*HW662 grid = torch.stack(torch.meshgrid([torch.arange(H), torch.arange(W)]), dim=0).to(flows.device) # 2*H*W663 warp_grid = torch.round(grid + flows)664 mask = torch.logical_and(665 torch.logical_and(666 torch.logical_and(torch.logical_and(warp_grid[0] >= 0, warp_grid[0] < H), warp_grid[1] >= 0),667 warp_grid[1] < W,668 ),669 masks[0],670 ).view(-1) # HW671 warp_grid = warp_grid.view(2, -1) # 2*HW672 warp_ind = (warp_grid[0] * W + warp_grid[1]).to(torch.long) # HW673 mapping_ind = torch.zeros_like(warp_ind) - 1 # HW674 675 for f0ind, f1ind in enumerate(warp_ind):676 if mask[f0ind]:677 if mapping_ind[f1ind] == -1:678 mapping_ind[f1ind] = f0ind679 else:680 targetv = frames[0, :, f1ind]681 pref0ind = mapping_ind[f1ind]682 prev = frames[1, :, pref0ind]683 v = frames[1, :, f0ind]684 if ((prev - targetv) ** 2).mean() > ((v - targetv) ** 2).mean():685 mask[pref0ind] = False686 mapping_ind[f1ind] = f0ind687 else:688 mask[f0ind] = False689 690 unusedind = torch.arange(len(mask)).to(mask.device)[~mask]691 unlinkedmask = mapping_ind == -1692 mapping_ind[unlinkedmask] = unusedind693 return mapping_ind, unlinkedmask694 695 696@torch.no_grad()697def get_mapping_ind(bwd_flows, bwd_occs, imgs, scale=1.0):698 """699 FLATTEN: Optical fLow-guided attention (Temoporal-guided attention)700 Find pixel correspondence between every consecutive frames in a batch701 702 [input]703 bwd_flow: (N-1)*2*H*W704 bwd_occ: (N-1)*H*W705 imgs: N*3*H*W706 707 [output]708 fwd_mappings: N*1*HW709 bwd_mappings: N*1*HW710 flattn_mask: HW*1*N*N711 i.e., imgs[i,:,fwd_mappings[i]] corresponds to imgs[0]712 i.e., imgs[i,:,fwd_mappings[i]][:,bwd_mappings[i]] restore the original imgs[i]713 """714 N, H, W = imgs.shape[0], int(imgs.shape[2] // scale), int(imgs.shape[3] // scale)715 iterattn_mask = torch.ones(H * W, N, N, dtype=torch.bool).to(imgs.device)716 for i in range(len(imgs) - 1):717 one_mask = torch.ones(N, N, dtype=torch.bool).to(imgs.device)718 one_mask[: i + 1, i + 1 :] = False719 one_mask[i + 1 :, : i + 1] = False720 mapping_ind, unlinkedmask = get_single_mapping_ind(721 bwd_flows[i : i + 1], bwd_occs[i : i + 1], imgs[i : i + 2], scale722 )723 if i == 0:724 fwd_mapping = [torch.arange(len(mapping_ind)).to(mapping_ind.device)]725 bwd_mapping = [torch.arange(len(mapping_ind)).to(mapping_ind.device)]726 iterattn_mask[unlinkedmask[fwd_mapping[-1]]] = torch.logical_and(727 iterattn_mask[unlinkedmask[fwd_mapping[-1]]], one_mask728 )729 fwd_mapping += [mapping_ind[fwd_mapping[-1]]]730 bwd_mapping += [torch.sort(fwd_mapping[-1])[1]]731 fwd_mappings = torch.stack(fwd_mapping, dim=0).unsqueeze(1)732 bwd_mappings = torch.stack(bwd_mapping, dim=0).unsqueeze(1)733 return fwd_mappings, bwd_mappings, iterattn_mask.unsqueeze(1)734 735 736def apply_FRESCO_opt(737 pipe,738 steps=[],739 layers=[0, 1, 2, 3],740 flows=None,741 occs=None,742 correlation_matrix=[],743 intra_weight=1e2,744 iters=20,745 optimize_temporal=True,746 saliency=None,747):748 """749 Apply FRESCO-based optimization to a StableDiffusionPipeline750 """751 pipe.unet.forward = my_forward(752 pipe.unet, steps, layers, flows, occs, correlation_matrix, intra_weight, iters, optimize_temporal, saliency753 )754 755 756@torch.no_grad()757def get_intraframe_paras(pipe, imgs, frescoProc, prompt_embeds, do_classifier_free_guidance=True, generator=None):758 """759 Get parameters for spatial-guided attention and optimization760 * perform one step denoising761 * collect attention feature, stored in frescoProc.controller.stored_attn['decoder_attn']762 * compute the gram matrix of the normalized feature for spatial consistency loss763 """764 765 noise_scheduler = pipe.scheduler766 timestep = noise_scheduler.timesteps[-1]767 device = pipe._execution_device768 B, C, H, W = imgs.shape769 770 frescoProc.controller.disable_controller()771 apply_FRESCO_opt(pipe)772 frescoProc.controller.clear_store()773 frescoProc.controller.enable_store()774 775 latents = pipe.prepare_latents(776 imgs.to(pipe.unet.dtype), timestep, B, 1, prompt_embeds.dtype, device, generator=generator, repeat_noise=False777 )778 779 latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents780 model_output = pipe.unet(781 latent_model_input,782 timestep,783 encoder_hidden_states=prompt_embeds,784 cross_attention_kwargs=None,785 return_dict=False,786 )787 788 frescoProc.controller.disable_store()789 790 # gram matrix of the normalized feature for spatial consistency loss791 correlation_matrix = []792 for tmp in model_output[1:]:793 latent_vector = rearrange(tmp, "b c h w -> b (h w) c")794 latent_vector = latent_vector / ((latent_vector**2).sum(dim=2, keepdims=True) ** 0.5)795 attention_probs = torch.bmm(latent_vector, latent_vector.transpose(-1, -2))796 correlation_matrix += [attention_probs.detach().clone().to(torch.float32)]797 del attention_probs, latent_vector, tmp798 del model_output799 800 clear_cache()801 802 return correlation_matrix803 804 805@torch.no_grad()806def get_flow_and_interframe_paras(flow_model, imgs):807 """808 Get parameters for temporal-guided attention and optimization809 * predict optical flow and occlusion mask810 * compute pixel index correspondence for FLATTEN811 """812 images = torch.stack([torch.from_numpy(img).permute(2, 0, 1).float() for img in imgs], dim=0).cuda()813 imgs_torch = torch.cat([numpy2tensor(img) for img in imgs], dim=0)814 815 reshuffle_list = list(range(1, len(images))) + [0]816 817 results_dict = flow_model(818 images,819 images[reshuffle_list],820 attn_splits_list=[2],821 corr_radius_list=[-1],822 prop_radius_list=[-1],823 pred_bidir_flow=True,824 )825 flow_pr = results_dict["flow_preds"][-1] # [2*B, 2, H, W]826 fwd_flows, bwd_flows = flow_pr.chunk(2) # [B, 2, H, W]827 fwd_occs, bwd_occs = forward_backward_consistency_check(fwd_flows, bwd_flows) # [B, H, W]828 829 warped_image1 = flow_warp(images, bwd_flows)830 bwd_occs = torch.clamp(831 bwd_occs + (abs(images[reshuffle_list] - warped_image1).mean(dim=1) > 255 * 0.25).float(), 0, 1832 )833 834 warped_image2 = flow_warp(images[reshuffle_list], fwd_flows)835 fwd_occs = torch.clamp(fwd_occs + (abs(images - warped_image2).mean(dim=1) > 255 * 0.25).float(), 0, 1)836 837 attn_mask = []838 for scale in [8.0, 16.0, 32.0]:839 bwd_occs_ = F.interpolate(bwd_occs[:-1].unsqueeze(1), scale_factor=1.0 / scale, mode="bilinear")840 attn_mask += [841 torch.cat((bwd_occs_[0:1].reshape(1, -1) > -1, bwd_occs_.reshape(bwd_occs_.shape[0], -1) > 0.5), dim=0)842 ]843 844 fwd_mappings = []845 bwd_mappings = []846 interattn_masks = []847 for scale in [8.0, 16.0]:848 fwd_mapping, bwd_mapping, interattn_mask = get_mapping_ind(bwd_flows, bwd_occs, imgs_torch, scale=scale)849 fwd_mappings += [fwd_mapping]850 bwd_mappings += [bwd_mapping]851 interattn_masks += [interattn_mask]852 853 interattn_paras = {}854 interattn_paras["fwd_mappings"] = fwd_mappings855 interattn_paras["bwd_mappings"] = bwd_mappings856 interattn_paras["interattn_masks"] = interattn_masks857 858 clear_cache()859 860 return [fwd_flows, bwd_flows], [fwd_occs, bwd_occs], attn_mask, interattn_paras861 862 863class AttentionControl:864 """865 Control FRESCO-based attention866 * enable/diable spatial-guided attention867 * enable/diable temporal-guided attention868 * enable/diable cross-frame attention869 * collect intermediate attention feature (for spatial-guided attention)870 """871 872 def __init__(self):873 self.stored_attn = self.get_empty_store()874 self.store = False875 self.index = 0876 self.attn_mask = None877 self.interattn_paras = None878 self.use_interattn = False879 self.use_cfattn = False880 self.use_intraattn = False881 self.intraattn_bias = 0882 self.intraattn_scale_factor = 0.2883 self.interattn_scale_factor = 0.2884 885 @staticmethod886 def get_empty_store():887 return {888 "decoder_attn": [],889 }890 891 def clear_store(self):892 del self.stored_attn893 torch.cuda.empty_cache()894 gc.collect()895 self.stored_attn = self.get_empty_store()896 self.disable_intraattn()897 898 # store attention feature of the input frame for spatial-guided attention899 def enable_store(self):900 self.store = True901 902 def disable_store(self):903 self.store = False904 905 # spatial-guided attention906 def enable_intraattn(self):907 self.index = 0908 self.use_intraattn = True909 self.disable_store()910 if len(self.stored_attn["decoder_attn"]) == 0:911 self.use_intraattn = False912 913 def disable_intraattn(self):914 self.index = 0915 self.use_intraattn = False916 self.disable_store()917 918 def disable_cfattn(self):919 self.use_cfattn = False920 921 # cross frame attention922 def enable_cfattn(self, attn_mask=None):923 if attn_mask:924 if self.attn_mask:925 del self.attn_mask926 torch.cuda.empty_cache()927 self.attn_mask = attn_mask928 self.use_cfattn = True929 else:930 if self.attn_mask:931 self.use_cfattn = True932 else:933 print("Warning: no valid cross-frame attention parameters available!")934 self.disable_cfattn()935 936 def disable_interattn(self):937 self.use_interattn = False938 939 # temporal-guided attention940 def enable_interattn(self, interattn_paras=None):941 if interattn_paras:942 if self.interattn_paras:943 del self.interattn_paras944 torch.cuda.empty_cache()945 self.interattn_paras = interattn_paras946 self.use_interattn = True947 else:948 if self.interattn_paras:949 self.use_interattn = True950 else:951 print("Warning: no valid temporal-guided attention parameters available!")952 self.disable_interattn()953 954 def disable_controller(self):955 self.disable_intraattn()956 self.disable_interattn()957 self.disable_cfattn()958 959 def enable_controller(self, interattn_paras=None, attn_mask=None):960 self.enable_intraattn()961 self.enable_interattn(interattn_paras)962 self.enable_cfattn(attn_mask)963 964 def forward(self, context):965 if self.store:966 self.stored_attn["decoder_attn"].append(context.detach())967 if self.use_intraattn and len(self.stored_attn["decoder_attn"]) > 0:968 tmp = self.stored_attn["decoder_attn"][self.index]969 self.index = self.index + 1970 if self.index >= len(self.stored_attn["decoder_attn"]):971 self.index = 0972 self.disable_store()973 return tmp974 return context975 976 def __call__(self, context):977 context = self.forward(context)978 return context979 980 981class FRESCOAttnProcessor2_0:982 """983 Hack self attention to FRESCO-based attention984 * adding spatial-guided attention985 * adding temporal-guided attention986 * adding cross-frame attention987 988 Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).989 Usage990 frescoProc = FRESCOAttnProcessor2_0(2, attn_mask)991 attnProc = AttnProcessor2_0()992 993 attn_processor_dict = {}994 for k in pipe.unet.attn_processors.keys():995 if k.startswith("up_blocks.2") or k.startswith("up_blocks.3"):996 attn_processor_dict[k] = frescoProc997 else:998 attn_processor_dict[k] = attnProc999 pipe.unet.set_attn_processor(attn_processor_dict)1000 """1001 1002 def __init__(self, unet_chunk_size=2, controller=None):1003 if not hasattr(F, "scaled_dot_product_attention"):1004 raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")1005 self.unet_chunk_size = unet_chunk_size1006 self.controller = controller1007 1008 def __call__(1009 self,1010 attn,1011 hidden_states,1012 encoder_hidden_states=None,1013 attention_mask=None,1014 temb=None,1015 ):1016 residual = hidden_states1017 1018 if attn.spatial_norm is not None:1019 hidden_states = attn.spatial_norm(hidden_states, temb)1020 1021 input_ndim = hidden_states.ndim1022 1023 if input_ndim == 4:1024 batch_size, channel, height, width = hidden_states.shape1025 hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)1026 1027 batch_size, sequence_length, _ = (1028 hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape1029 )1030 1031 if attention_mask is not None:1032 attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)1033 # scaled_dot_product_attention expects attention_mask shape to be1034 # (batch, heads, source_length, target_length)1035 attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])1036 1037 if attn.group_norm is not None:1038 hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)1039 1040 query = attn.to_q(hidden_states)1041 1042 crossattn = False1043 if encoder_hidden_states is None:1044 encoder_hidden_states = hidden_states1045 if self.controller and self.controller.store:1046 self.controller(hidden_states.detach().clone())1047 else:1048 crossattn = True1049 if attn.norm_cross:1050 encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)1051 1052 # BC * HW * 8D1053 key = attn.to_k(encoder_hidden_states)1054 value = attn.to_v(encoder_hidden_states)1055 1056 query_raw, key_raw = None, None1057 if self.controller and self.controller.use_interattn and (not crossattn):1058 query_raw, key_raw = query.clone(), key.clone()1059 1060 inner_dim = key.shape[-1] # 8D1061 head_dim = inner_dim // attn.heads # D1062 1063 """for efficient cross-frame attention"""1064 if self.controller and self.controller.use_cfattn and (not crossattn):1065 video_length = key.size()[0] // self.unet_chunk_size1066 former_frame_index = [0] * video_length1067 attn_mask = None1068 if self.controller.attn_mask is not None:1069 for m in self.controller.attn_mask:1070 if m.shape[1] == key.shape[1]:1071 attn_mask = m1072 # BC * HW * 8D --> B * C * HW * 8D1073 key = rearrange(key, "(b f) d c -> b f d c", f=video_length)1074 # B * C * HW * 8D --> B * C * HW * 8D1075 if attn_mask is None:1076 key = key[:, former_frame_index]1077 else:1078 key = repeat(key[:, attn_mask], "b d c -> b f d c", f=video_length)1079 # B * C * HW * 8D --> BC * HW * 8D1080 key = rearrange(key, "b f d c -> (b f) d c").detach()1081 value = rearrange(value, "(b f) d c -> b f d c", f=video_length)1082 if attn_mask is None:1083 value = value[:, former_frame_index]1084 else:1085 value = repeat(value[:, attn_mask], "b d c -> b f d c", f=video_length)1086 value = rearrange(value, "b f d c -> (b f) d c").detach()1087 1088 # BC * HW * 8D --> BC * HW * 8 * D --> BC * 8 * HW * D1089 query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)1090 # BC * 8 * HW2 * D1091 key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)1092 # BC * 8 * HW2 * D21093 value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)1094 1095 """for spatial-guided intra-frame attention"""1096 if self.controller and self.controller.use_intraattn and (not crossattn):1097 ref_hidden_states = self.controller(None)1098 assert ref_hidden_states.shape == encoder_hidden_states.shape1099 query_ = attn.to_q(ref_hidden_states)1100 key_ = attn.to_k(ref_hidden_states)1101 1102 # BC * 8 * HW * D1103 query_ = query_.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)1104 key_ = key_.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)1105 query = F.scaled_dot_product_attention(1106 query_,1107 key_ * self.controller.intraattn_scale_factor,1108 query,1109 attn_mask=torch.eye(query_.size(-2), key_.size(-2), dtype=query.dtype, device=query.device)1110 * self.controller.intraattn_bias,1111 ).detach()1112 1113 del query_, key_1114 torch.cuda.empty_cache()1115 1116 # the output of sdp = (batch, num_heads, seq_len, head_dim)1117 # TODO: add support for attn.scale when we move to Torch 2.11118 # output: BC * 8 * HW * D21119 hidden_states = F.scaled_dot_product_attention(1120 query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False1121 )1122 1123 """for temporal-guided inter-frame attention (FLATTEN)"""1124 if self.controller and self.controller.use_interattn and (not crossattn):1125 del query, key, value1126 torch.cuda.empty_cache()1127 bwd_mapping = None1128 fwd_mapping = None1129 for i, f in enumerate(self.controller.interattn_paras["fwd_mappings"]):1130 if f.shape[2] == hidden_states.shape[2]:1131 fwd_mapping = f1132 bwd_mapping = self.controller.interattn_paras["bwd_mappings"][i]1133 interattn_mask = self.controller.interattn_paras["interattn_masks"][i]1134 video_length = key_raw.size()[0] // self.unet_chunk_size1135 # BC * HW * 8D --> C * 8BD * HW1136 key = rearrange(key_raw, "(b f) d c -> f (b c) d", f=video_length)1137 query = rearrange(query_raw, "(b f) d c -> f (b c) d", f=video_length)1138 # BC * 8 * HW * D --> C * 8BD * HW1139 # key = rearrange(hidden_states, "(b f) h d c -> f (b h c) d", f=video_length) ########1140 # query = rearrange(hidden_states, "(b f) h d c -> f (b h c) d", f=video_length) #######1141 1142 value = rearrange(hidden_states, "(b f) h d c -> f (b h c) d", f=video_length)1143 key = torch.gather(key, 2, fwd_mapping.expand(-1, key.shape[1], -1))1144 query = torch.gather(query, 2, fwd_mapping.expand(-1, query.shape[1], -1))1145 value = torch.gather(value, 2, fwd_mapping.expand(-1, value.shape[1], -1))1146 # C * 8BD * HW --> BHW, C, 8D1147 key = rearrange(key, "f (b c) d -> (b d) f c", b=self.unet_chunk_size)1148 query = rearrange(query, "f (b c) d -> (b d) f c", b=self.unet_chunk_size)1149 value = rearrange(value, "f (b c) d -> (b d) f c", b=self.unet_chunk_size)1150 # BHW * C * 8D --> BHW * C * 8 * D--> BHW * 8 * C * D1151 query = query.view(-1, video_length, attn.heads, head_dim).transpose(1, 2).detach()1152 key = key.view(-1, video_length, attn.heads, head_dim).transpose(1, 2).detach()1153 value = value.view(-1, video_length, attn.heads, head_dim).transpose(1, 2).detach()1154 hidden_states_ = F.scaled_dot_product_attention(1155 query,1156 key * self.controller.interattn_scale_factor,1157 value,1158 # .to(query.dtype)-1.0) * 1e6 -1159 attn_mask=(interattn_mask.repeat(self.unet_chunk_size, 1, 1, 1)),1160 # torch.eye(interattn_mask.shape[2]).to(query.device).to(query.dtype) * 1e4,1161 )1162 1163 # BHW * 8 * C * D --> C * 8BD * HW1164 hidden_states_ = rearrange(hidden_states_, "(b d) h f c -> f (b h c) d", b=self.unet_chunk_size)1165 hidden_states_ = torch.gather(1166 hidden_states_, 2, bwd_mapping.expand(-1, hidden_states_.shape[1], -1)1167 ).detach()1168 # C * 8BD * HW --> BC * 8 * HW * D1169 hidden_states = rearrange(1170 hidden_states_, "f (b h c) d -> (b f) h d c", b=self.unet_chunk_size, h=attn.heads1171 )1172 1173 # BC * 8 * HW * D --> BC * HW * 8D1174 hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)1175 hidden_states = hidden_states.to(query.dtype)1176 1177 # linear proj1178 hidden_states = attn.to_out[0](hidden_states)1179 # dropout1180 hidden_states = attn.to_out[1](hidden_states)1181 1182 if input_ndim == 4:1183 hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)1184 1185 if attn.residual_connection:1186 hidden_states = hidden_states + residual1187 1188 hidden_states = hidden_states / attn.rescale_output_factor1189 1190 return hidden_states1191 1192 1193def apply_FRESCO_attn(pipe):1194 """1195 Apply FRESCO-guided attention to a StableDiffusionPipeline1196 """1197 frescoProc = FRESCOAttnProcessor2_0(2, AttentionControl())1198 attnProc = AttnProcessor2_0()1199 attn_processor_dict = {}1200 for k in pipe.unet.attn_processors.keys():