| 185 | |
| 186 | |
| 187 | def attribute_edit(generator, data): |
| 188 | assert args.attribute in ['young', 'old', 'beard', 'lip'] |
| 189 | # Recommend factor |
| 190 | if args.attribute == 'young': |
| 191 | factor = -5.0 |
| 192 | elif args.attribute == 'old': |
| 193 | factor = 5.0 |
| 194 | elif args.attribute == 'beard': |
| 195 | factor = -20.0 |
| 196 | elif args.attribute == 'lip': |
| 197 | factor = 20.0 |
| 198 | |
| 199 | bs = args.batch_size |
| 200 | source_image = data['source_align'].unsqueeze(0).cuda() |
| 201 | per_wx, per_ix, per_res = generator.generator.encode(source_image) |
| 202 | |
| 203 | inv_data = inference_util.hfgi_inversion(generator, source_image, args=args, batch_size=bs) |
| 204 | source_image = source_image.repeat(bs, 1, 1, 1) |
| 205 | |
| 206 | num_batch = len(data['target_image']) // bs + 1 |
| 207 | gt_images, video_warp_images, audio_warp_images, fake_images = [], [], [], [] |
| 208 | source_3dmm = data['source_semantics'].unsqueeze(-1).repeat(1, 1, 27) # 1, 73, 27 |
| 209 | for _i in range(num_batch): |
| 210 | target_images = data['target_image'][_i * bs:(_i + 1) * bs] |
| 211 | if len(target_images) == 0 or _i * bs > args.frame_limit: |
| 212 | break |
| 213 | target_3dmm = data['target_semantics'][_i * bs:(_i + 1) * bs] |
| 214 | target_3dmm = torch.stack(target_3dmm).cuda() |
| 215 | |
| 216 | _len_3dmm = len(target_3dmm) |
| 217 | if _len_3dmm < bs: |
| 218 | # Last batch |
| 219 | ix, wx, fx, inversion_condition = inv_data |
| 220 | ix, wx, fx = ix[:_len_3dmm], wx[:_len_3dmm], fx[:_len_3dmm] |
| 221 | if args.inversion_option == 'encode': |
| 222 | inversion_condition = (inversion_condition[0][:_len_3dmm], inversion_condition[1][:_len_3dmm]) |
| 223 | inv_data = ix, wx, fx, inversion_condition |
| 224 | source_3dmm = source_3dmm[:_len_3dmm] |
| 225 | |
| 226 | with torch.no_grad(): |
| 227 | if args.edit_expression_only: |
| 228 | target_3dmm[:, 64:, :] = source_3dmm[:, 64:, :] |
| 229 | output = generator.forward(source_image, target_3dmm, inv_data=inv_data, imsize=1024) |
| 230 | |
| 231 | ix_edit, wx_edit, fx_edit, inversion_condition = generator. \ |
| 232 | generator.edit(x=None, factor=factor / num_batch * (_i + 1), choice=args.attribute, wx=per_wx, res=per_res) |
| 233 | inv_data = [ |
| 234 | ix_edit.expand(bs, 3, 256, 256), |
| 235 | wx_edit.expand(bs, 18, 512), |
| 236 | fx_edit.expand(bs, 512, 64, 64), |
| 237 | (inversion_condition[0].expand(bs, 512, 64, 64), |
| 238 | inversion_condition[1].expand(bs, 512, 64, 64)) |
| 239 | ] |
| 240 | |
| 241 | gt_images.append(target_images) |
| 242 | fake_images.append(output['fake_image'].cpu().clamp_(-1, 1)) |
| 243 | video_warp_images.append(output['video_warp_image'].cpu().clamp_(-1, 1)) |
| 244 | |