182 lines
8.5 KiB
Python
182 lines
8.5 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import numpy as np
|
|
import torch
|
|
from scepter.modules.solver import LatentDiffusionSolver
|
|
from scepter.modules.solver.registry import SOLVERS
|
|
from scepter.modules.utils.data import transfer_data_to_cuda
|
|
from scepter.modules.utils.distribute import we
|
|
from scepter.modules.utils.probe import ProbeData
|
|
from tqdm import tqdm
|
|
@SOLVERS.register_class()
|
|
class FormalACEPlusSolver(LatentDiffusionSolver):
|
|
def __init__(self, cfg, logger=None):
|
|
super().__init__(cfg, logger=logger)
|
|
self.probe_prompt = cfg.get("PROBE_PROMPT", None)
|
|
self.probe_hw = cfg.get("PROBE_HW", [])
|
|
|
|
@torch.no_grad()
|
|
def run_eval(self):
|
|
self.eval_mode()
|
|
self.before_all_iter(self.hooks_dict[self._mode])
|
|
all_results = []
|
|
for batch_idx, batch_data in tqdm(
|
|
enumerate(self.datas[self._mode].dataloader)):
|
|
self.before_iter(self.hooks_dict[self._mode])
|
|
if self.sample_args:
|
|
batch_data.update(self.sample_args.get_lowercase_dict())
|
|
with torch.autocast(device_type='cuda',
|
|
enabled=self.use_amp,
|
|
dtype=self.dtype):
|
|
results = self.run_step_eval(transfer_data_to_cuda(batch_data),
|
|
batch_idx,
|
|
step=self.total_iter,
|
|
rank=we.rank)
|
|
all_results.extend(results)
|
|
self.after_iter(self.hooks_dict[self._mode])
|
|
log_data, log_label = self.save_results(all_results)
|
|
self.register_probe({'eval_label': log_label})
|
|
self.register_probe({
|
|
'eval_image':
|
|
ProbeData(log_data,
|
|
is_image=True,
|
|
build_html=True,
|
|
build_label=log_label)
|
|
})
|
|
self.after_all_iter(self.hooks_dict[self._mode])
|
|
|
|
@torch.no_grad()
|
|
def run_test(self):
|
|
self.test_mode()
|
|
self.before_all_iter(self.hooks_dict[self._mode])
|
|
all_results = []
|
|
for batch_idx, batch_data in tqdm(
|
|
enumerate(self.datas[self._mode].dataloader)):
|
|
self.before_iter(self.hooks_dict[self._mode])
|
|
if self.sample_args:
|
|
batch_data.update(self.sample_args.get_lowercase_dict())
|
|
with torch.autocast(device_type='cuda',
|
|
enabled=self.use_amp,
|
|
dtype=self.dtype):
|
|
results = self.run_step_eval(transfer_data_to_cuda(batch_data),
|
|
batch_idx,
|
|
step=self.total_iter,
|
|
rank=we.rank)
|
|
all_results.extend(results)
|
|
self.after_iter(self.hooks_dict[self._mode])
|
|
log_data, log_label = self.save_results(all_results)
|
|
self.register_probe({'test_label': log_label})
|
|
self.register_probe({
|
|
'test_image':
|
|
ProbeData(log_data,
|
|
is_image=True,
|
|
build_html=True,
|
|
build_label=log_label)
|
|
})
|
|
|
|
self.after_all_iter(self.hooks_dict[self._mode])
|
|
|
|
def run_step_val(self, batch_data, batch_idx=0, step=None, rank=None):
|
|
sample_id_list = batch_data['sample_id']
|
|
loss_dict = {}
|
|
with torch.autocast(device_type='cuda',
|
|
enabled=self.use_amp,
|
|
dtype=self.dtype):
|
|
results = self.model.forward_train(**batch_data)
|
|
loss = results['loss']
|
|
for sample_id in sample_id_list:
|
|
loss_dict[sample_id] = loss.detach().cpu().numpy()
|
|
return loss_dict
|
|
|
|
def save_results(self, results):
|
|
log_data, log_label = [], []
|
|
for result in results:
|
|
ret_images, ret_labels = [], []
|
|
edit_image = result.get('edit_image', None)
|
|
modify_image = result.get('modify_image', None)
|
|
edit_mask = result.get('edit_mask', None)
|
|
if edit_image is not None:
|
|
for i, edit_img in enumerate(result['edit_image']):
|
|
if edit_img is None:
|
|
continue
|
|
ret_images.append((edit_img.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
|
|
ret_labels.append(f'edit_image{i}; ')
|
|
ret_images.append((modify_image[i].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
|
|
ret_labels.append(f'modify_image{i}; ')
|
|
if edit_mask is not None:
|
|
ret_images.append((edit_mask[i].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
|
|
ret_labels.append(f'edit_mask{i}; ')
|
|
|
|
target_image = result.get('target_image', None)
|
|
target_mask = result.get('target_mask', None)
|
|
if target_image is not None:
|
|
ret_images.append((target_image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
|
|
ret_labels.append(f'target_image; ')
|
|
if target_mask is not None:
|
|
ret_images.append((target_mask.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
|
|
ret_labels.append(f'target_mask; ')
|
|
teacher_image = result.get('image', None)
|
|
if teacher_image is not None:
|
|
ret_images.append((teacher_image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
|
|
ret_labels.append(f"teacher_image")
|
|
reconstruct_image = result.get('reconstruct_image', None)
|
|
if reconstruct_image is not None:
|
|
ret_images.append((reconstruct_image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
|
|
ret_labels.append(f"{result['instruction']}")
|
|
log_data.append(ret_images)
|
|
log_label.append(ret_labels)
|
|
return log_data, log_label
|
|
@property
|
|
def probe_data(self):
|
|
if not we.debug and self.mode == 'train':
|
|
batch_data = transfer_data_to_cuda(self.current_batch_data[self.mode])
|
|
self.eval_mode()
|
|
with torch.autocast(device_type='cuda',
|
|
enabled=self.use_amp,
|
|
dtype=self.dtype):
|
|
batch_data['log_num'] = self.log_train_num
|
|
batch_data.update(self.sample_args.get_lowercase_dict())
|
|
results = self.run_step_eval(batch_data)
|
|
self.train_mode()
|
|
log_data, log_label = self.save_results(results)
|
|
self.register_probe({
|
|
'train_image':
|
|
ProbeData(log_data,
|
|
is_image=True,
|
|
build_html=True,
|
|
build_label=log_label)
|
|
})
|
|
self.register_probe({'train_label': log_label})
|
|
if self.probe_prompt:
|
|
self.eval_mode()
|
|
all_results = []
|
|
for prompt in self.probe_prompt:
|
|
with torch.autocast(device_type='cuda',
|
|
enabled=self.use_amp,
|
|
dtype=self.dtype):
|
|
batch_data = {
|
|
"prompt": [[prompt]],
|
|
"image": [torch.zeros(3, self.probe_hw[0], self.probe_hw[1])],
|
|
"image_mask": [torch.ones(1, self.probe_hw[0], self.probe_hw[1])],
|
|
"src_image_list": [[]],
|
|
"modify_image_list": [[]],
|
|
"src_mask_list": [[]],
|
|
"edit_id": [[]],
|
|
"height": self.probe_hw[0],
|
|
"width": self.probe_hw[1]
|
|
}
|
|
batch_data.update(self.sample_args.get_lowercase_dict())
|
|
results = self.run_step_eval(batch_data)
|
|
all_results.extend(results)
|
|
self.train_mode()
|
|
log_data, log_label = self.save_results(all_results)
|
|
self.register_probe({
|
|
'probe_image':
|
|
ProbeData(log_data,
|
|
is_image=True,
|
|
build_html=True,
|
|
build_label=log_label)
|
|
})
|
|
|
|
return super(LatentDiffusionSolver, self).probe_data
|