28 lines
1.2 KiB
Python
28 lines
1.2 KiB
Python
from ....pl_module_params_controlnet import AttentionMaskParams
|
|
import torch
|
|
|
|
|
|
class MaskGenerator():
|
|
|
|
def __init__(self, params: AttentionMaskParams, num_frame_conditioning, num_frames):
|
|
self.params = params
|
|
self.num_frame_conditioning = num_frame_conditioning
|
|
self.num_frames = num_frames
|
|
def get_mask(self, precision, device):
|
|
|
|
params = self.params
|
|
if params.temporal_self_attention_only_on_conditioning:
|
|
with torch.no_grad():
|
|
attention_mask = torch.zeros((1, self.num_frames, self.num_frames), dtype=torch.float16 if precision.startswith(
|
|
"16") else torch.float32, device=device)
|
|
for frame in range(self.num_frame_conditioning, self.num_frames):
|
|
attention_mask[:, frame,
|
|
self.num_frame_conditioning:] = float("-inf")
|
|
if params.temporal_self_attention_mask_included_itself:
|
|
attention_mask[:, frame, frame] = 0
|
|
if params.temp_attend_on_uncond_include_past:
|
|
attention_mask[:, frame, :frame] = 0
|
|
else:
|
|
attention_mask = None
|
|
return attention_mask
|