Views
No views yet
1# `pip install scepter>0.0.4` or
2# clone newest SCEPTER and run `PYTHONPATH=./ python <this_script>` at the main branch root.
3import os
4import unittest
5
6from PIL import Image
7from torchvision.utils import save_image
8
9from scepter.modules.inference.stylebooth_inference import StyleboothInference
10from scepter.modules.utils.config import Config
11from scepter.modules.utils.file_system import FS
12from scepter.modules.utils.logger import get_logger
13
14
15class DiffusionInferenceTest(unittest.TestCase):
16 def setUp(self):
17 print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
18 self.logger = get_logger(name='scepter')
19 config_file = 'scepter/methods/studio/scepter_ui.yaml'
20 cfg = Config(cfg_file=config_file)
21 if 'FILE_SYSTEM' in cfg:
22 for fs_info in cfg['FILE_SYSTEM']:
23 FS.init_fs_client(fs_info)
24 self.tmp_dir = './cache/save_data/diffusion_inference'
25 if not os.path.exists(self.tmp_dir):
26 os.makedirs(self.tmp_dir)
27
28 def tearDown(self):
29 super().tearDown()
30
31 # uncomment this line to skip this module.
32 # @unittest.skip('')
33 def test_stylebooth(self):
34 config_file = 'scepter/methods/studio/inference/edit/stylebooth_tb_pro.yaml'
35 cfg = Config(cfg_file=config_file)
36 diff_infer = StyleboothInference(logger=self.logger)
37 diff_infer.init_from_cfg(cfg)
38
39 output = diff_infer({'prompt': 'Let this image be in the style of sai-lowpoly'},
40 style_edit_image=Image.open('asset/images/inpainting_text_ref/ex4_scene_im.jpg'),
41 style_guide_scale_text=7.5,
42 style_guide_scale_image=1.5,
43 stylebooth_state=True)
44 save_path = os.path.join(self.tmp_dir,
45 'stylebooth_test_lowpoly_cute_dog.png')
46 save_image(output['images'], save_path)
47
48
49if __name__ == '__main__':
50 unittest.main()