| import gc |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| from flow.flow_utils import flow_warp |
|
|
| |
|
|
|
|
| def calc_mean_std(feat, eps=1e-5): |
| |
| size = feat.size() |
| assert (len(size) == 4) |
| N, C = size[:2] |
| feat_var = feat.view(N, C, -1).var(dim=2) + eps |
| feat_std = feat_var.sqrt().view(N, C, 1, 1) |
| feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1) |
| return feat_mean, feat_std |
|
|
|
|
| class AttentionControl(): |
|
|
| def __init__(self, inner_strength, mask_period, cross_period, ada_period, |
| warp_period): |
| self.step_store = self.get_empty_store() |
| self.cur_step = 0 |
| self.total_step = 0 |
| self.cur_index = 0 |
| self.init_store = False |
| self.restore = False |
| self.update = False |
| self.flow = None |
| self.mask = None |
| self.restorex0 = False |
| self.updatex0 = False |
| self.inner_strength = inner_strength |
| self.cross_period = cross_period |
| self.mask_period = mask_period |
| self.ada_period = ada_period |
| self.warp_period = warp_period |
|
|
| @staticmethod |
| def get_empty_store(): |
| return { |
| 'first': [], |
| 'previous': [], |
| 'x0_previous': [], |
| 'first_ada': [] |
| } |
|
|
| def forward(self, context, is_cross: bool, place_in_unet: str): |
| cross_period = (self.total_step * self.cross_period[0], |
| self.total_step * self.cross_period[1]) |
| if not is_cross and place_in_unet == 'up': |
| if self.init_store: |
| self.step_store['first'].append(context.detach()) |
| self.step_store['previous'].append(context.detach()) |
| if self.update: |
| tmp = context.clone().detach() |
| if self.restore and self.cur_step >= cross_period[0] and \ |
| self.cur_step <= cross_period[1]: |
| context = torch.cat( |
| (self.step_store['first'][self.cur_index], |
| self.step_store['previous'][self.cur_index]), |
| dim=1).clone() |
| if self.update: |
| self.step_store['previous'][self.cur_index] = tmp |
| self.cur_index += 1 |
| return context |
|
|
| def update_x0(self, x0): |
| if self.init_store: |
| self.step_store['x0_previous'].append(x0.detach()) |
| style_mean, style_std = calc_mean_std(x0.detach()) |
| self.step_store['first_ada'].append(style_mean.detach()) |
| self.step_store['first_ada'].append(style_std.detach()) |
| if self.updatex0: |
| tmp = x0.clone().detach() |
| if self.restorex0: |
| if self.cur_step >= self.total_step * self.ada_period[ |
| 0] and self.cur_step <= self.total_step * self.ada_period[ |
| 1]: |
| x0 = F.instance_norm(x0) * self.step_store['first_ada'][ |
| 2 * self.cur_step + |
| 1] + self.step_store['first_ada'][2 * self.cur_step] |
| if self.cur_step >= self.total_step * self.warp_period[ |
| 0] and self.cur_step <= self.total_step * self.warp_period[ |
| 1]: |
| pre = self.step_store['x0_previous'][self.cur_step] |
| x0 = flow_warp(pre, self.flow, mode='nearest') * self.mask + ( |
| 1 - self.mask) * x0 |
| if self.updatex0: |
| self.step_store['x0_previous'][self.cur_step] = tmp |
| return x0 |
|
|
| def set_warp(self, flow, mask): |
| self.flow = flow.clone() |
| self.mask = mask.clone() |
|
|
| def __call__(self, context, is_cross: bool, place_in_unet: str): |
| context = self.forward(context, is_cross, place_in_unet) |
| return context |
|
|
| def set_step(self, step): |
| self.cur_step = step |
|
|
| def set_total_step(self, total_step): |
| self.total_step = total_step |
| self.cur_index = 0 |
|
|
| def clear_store(self): |
| del self.step_store |
| torch.cuda.empty_cache() |
| gc.collect() |
| self.step_store = self.get_empty_store() |
|
|
| def set_task(self, task, restore_step=1.0): |
| self.init_store = False |
| self.restore = False |
| self.update = False |
| self.cur_index = 0 |
| self.restore_step = restore_step |
| self.updatex0 = False |
| self.restorex0 = False |
| if 'initfirst' in task: |
| self.init_store = True |
| self.clear_store() |
| if 'updatestyle' in task: |
| self.update = True |
| if 'keepstyle' in task: |
| self.restore = True |
| if 'updatex0' in task: |
| self.updatex0 = True |
| if 'keepx0' in task: |
| self.restorex0 = True |
|
|
