133 lines
4.9 KiB
Python
133 lines
4.9 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import argparse
|
|
import os
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
import torch.cuda.amp as amp
|
|
|
|
from scepter.modules.solver.registry import SOLVERS
|
|
from scepter.modules.utils.config import Config
|
|
from scepter.modules.utils.data import transfer_data_to_cuda
|
|
from scepter.modules.utils.distribute import we
|
|
from scepter.modules.utils.file_system import FS
|
|
from scepter.modules.utils.logger import get_logger
|
|
|
|
|
|
def run_task(cfg):
|
|
std_logger = get_logger(name='scepter')
|
|
solver = SOLVERS.build(cfg.SOLVER, logger=std_logger)
|
|
solver.set_up()
|
|
if not cfg.args.pretrained_model == '':
|
|
with FS.get_from(cfg.args.pretrained_model,
|
|
wait_finish=True) as local_path:
|
|
solver.model.load_state_dict(
|
|
torch.load(local_path, map_location='cuda')['model'])
|
|
solver.test_mode()
|
|
num_samples = cfg.args.num_samples
|
|
prompt = [cfg.args.prompt] * num_samples
|
|
n_prompt = [cfg.args.n_prompt] * num_samples
|
|
sampler = cfg.args.sampler
|
|
sample_steps = cfg.args.sample_steps
|
|
seed = cfg.args.seed
|
|
guide_scale = cfg.args.guide_scale
|
|
guide_rescale = cfg.args.guide_rescale
|
|
image_size = cfg.args.image_size
|
|
if image_size is not None:
|
|
if ',' in image_size:
|
|
h, w = image_size.split(',')
|
|
image_size = [int(h), int(w)]
|
|
else:
|
|
image_size = [int(image_size), int(image_size)]
|
|
|
|
batch_data = {}
|
|
if solver.sample_args:
|
|
batch_data.update(solver.sample_args.get_lowercase_dict())
|
|
if image_size is not None:
|
|
batch_data.update({'image_size': image_size})
|
|
batch_data.update({
|
|
'prompt': prompt,
|
|
'n_prompt': n_prompt,
|
|
'sampler': sampler,
|
|
'sample_steps': sample_steps,
|
|
'seed': seed,
|
|
'guide_scale': guide_scale,
|
|
'guide_rescale': guide_rescale,
|
|
})
|
|
|
|
dtype = getattr(torch, cfg.SOLVER.DTYPE)
|
|
with amp.autocast(enabled=True, dtype=dtype):
|
|
batch_data = transfer_data_to_cuda(batch_data)
|
|
ret = solver.run_step_test(batch_data)
|
|
save_folder = os.path.join(solver.work_dir, cfg.args.save_folder)
|
|
for idx, out in enumerate(ret):
|
|
img = out['image']
|
|
img = img.permute(1, 2, 0).cpu().numpy()
|
|
img = (img * 255).astype(np.uint8)
|
|
filename = '{}_{}.png'.format('inference', idx)
|
|
save_file = os.path.join(save_folder, filename)
|
|
with FS.put_to(save_file) as local_path:
|
|
image = img.copy()
|
|
cv2.cvtColor(image, cv2.COLOR_RGB2BGR, image)
|
|
cv2.imwrite(local_path, image)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(description='Argparser for Scepter:\n')
|
|
parser.add_argument(
|
|
'--prompt',
|
|
dest='prompt',
|
|
help='Prompt sentence!',
|
|
default='a woman is walking on the street in a rainy day.')
|
|
parser.add_argument('--n_prompt',
|
|
dest='n_prompt',
|
|
help='Add Prompt sentence!',
|
|
default='')
|
|
parser.add_argument('--num_samples',
|
|
dest='num_samples',
|
|
help="Output image's number!",
|
|
default=4,
|
|
type=int)
|
|
parser.add_argument('--sampler',
|
|
dest='sampler',
|
|
help='Sampler method!',
|
|
default='ddim',
|
|
type=str)
|
|
parser.add_argument('--sample_steps',
|
|
dest='sample_steps',
|
|
help='Sample steps!',
|
|
default=50,
|
|
type=int)
|
|
parser.add_argument('--seed',
|
|
dest='seed',
|
|
help='Random seed!',
|
|
default=2023,
|
|
type=int)
|
|
parser.add_argument('--guide_scale',
|
|
dest='guide_scale',
|
|
help='Guidance scale!',
|
|
default=7.5,
|
|
type=float)
|
|
parser.add_argument('--guide_rescale',
|
|
dest='guide_rescale',
|
|
help='Guidance rescale!',
|
|
default=0.5,
|
|
type=float)
|
|
parser.add_argument('--image_size',
|
|
dest='image_size',
|
|
help='Output image size! (h, w)',
|
|
default=None,
|
|
type=str)
|
|
parser.add_argument('--save_folder',
|
|
dest='save_folder',
|
|
help="Output image's save folder!",
|
|
default='test_images')
|
|
parser.add_argument('--pretrained_model',
|
|
dest='pretrained_model',
|
|
help='The pretrained model for our network!',
|
|
default='')
|
|
cfg = Config(load=True, parser_ins=parser)
|
|
we.init_env(cfg, logger=None, fn=run_task)
|