| | 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 |
| |
|
