| 114 | |
| 115 | |
| 116 | def resize_mask(mask, latent, process_first_frame_only=True): |
| 117 | latent_size = latent.size() |
| 118 | batch_size, channels, num_frames, height, width = mask.shape |
| 119 | |
| 120 | if process_first_frame_only: |
| 121 | target_size = list(latent_size[2:]) |
| 122 | target_size[0] = 1 |
| 123 | first_frame_resized = F.interpolate( |
| 124 | mask[:, :, 0:1, :, :], |
| 125 | size=target_size, |
| 126 | mode='trilinear', |
| 127 | align_corners=False |
| 128 | ) |
| 129 | |
| 130 | target_size = list(latent_size[2:]) |
| 131 | target_size[0] = target_size[0] - 1 |
| 132 | if target_size[0] != 0: |
| 133 | remaining_frames_resized = F.interpolate( |
| 134 | mask[:, :, 1:, :, :], |
| 135 | size=target_size, |
| 136 | mode='trilinear', |
| 137 | align_corners=False |
| 138 | ) |
| 139 | resized_mask = torch.cat([first_frame_resized, remaining_frames_resized], dim=2) |
| 140 | else: |
| 141 | resized_mask = first_frame_resized |
| 142 | else: |
| 143 | target_size = list(latent_size[2:]) |
| 144 | resized_mask = F.interpolate( |
| 145 | mask, |
| 146 | size=target_size, |
| 147 | mode='trilinear', |
| 148 | align_corners=False |
| 149 | ) |
| 150 | return resized_mask |
| 151 | |
| 152 | |
| 153 | # Will error if the minimal version of diffusers is not installed. Remove at your own risks. |