v0.0.2

tests/tools/test_annotators.py

@@ -0,0 +1,209 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import os
+import unittest
+
+import numpy as np
+from PIL import Image
+
+from scepter.modules.annotator.registry import ANNOTATORS
+from scepter.modules.utils.config import Config
+from scepter.modules.utils.distribute import we
+from scepter.modules.utils.file_system import FS
+
+
+class AnnotatorTest(unittest.TestCase):
+    def setUp(self):
+        print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
+        _ = FS.init_fs_client(Config(cfg_dict={
+            'NAME': 'ModelscopeFs',
+            'TEMP_DIR': './cache/data'
+        },
+                                     load=False),
+                              overwrite=False)
+        image_path = 'asset/images/sunflower.jpeg'
+        image = Image.open(image_path)
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        self.image = np.array(image)
+
+        self.save_dir = './cache/save_data/images'
+        if not os.path.exists(self.save_dir):
+            os.makedirs(self.save_dir)
+
+    def tearDown(self):
+        super().tearDown()
+
+    @unittest.skip('')
+    def test_annotator_canny(self):
+        # canny
+        canny_dict = {
+            'NAME': 'CannyAnnotator',
+            'LOW_THRESHOLD': 100,
+            'HIGH_THRESHOLD': 200
+        }
+        canny_anno = Config(cfg_dict=canny_dict, load=False)
+        canny_ins = ANNOTATORS.build(canny_anno).to(we.device_id)
+        canny_image = canny_ins(self.image)
+        print("canny's shape:", canny_image.shape)
+        Image.fromarray(canny_image).save(
+            os.path.join(self.save_dir, 'sunflower_canny.png'))
+
+    @unittest.skip('')
+    def test_annotator_hed(self):
+        # hed
+        hed_dict = {
+            'NAME':
+            'HedAnnotator',
+            'PRETRAINED_MODEL':
+            'ms://damo/scepter_scedit@annotator/ckpts/ControlNetHED.pth'
+        }
+        hed_anno = Config(cfg_dict=hed_dict, load=False)
+        hed_ins = ANNOTATORS.build(hed_anno).to(we.device_id)
+        hed_image = hed_ins(self.image)
+        print("hed's shape:", hed_image.shape)
+        Image.fromarray(hed_image).save(
+            os.path.join(self.save_dir, 'sunflower_hed.png'))
+
+    @unittest.skip('')
+    def test_annotator_openpose(self):
+        # openpose
+        openpose_dict = {
+            'NAME':
+            'OpenposeAnnotator',
+            'BODY_MODEL_PATH':
+            'ms://damo/scepter_scedit@annotator/ckpts/body_pose_model.pth',
+            'HAND_MODEL_PATH':
+            'ms://damo/scepter_scedit@annotator/ckpts/hand_pose_model.pth'
+        }
+        openpose_anno = Config(cfg_dict=openpose_dict, load=False)
+        openpose_ins = ANNOTATORS.build(openpose_anno).to(we.device_id)
+        openpose_image = openpose_ins(self.image)
+        print("openpose's shape:", openpose_image.shape)
+        Image.fromarray(openpose_image).save(
+            os.path.join(self.save_dir, 'sunflower_openpose.png'))
+
+    @unittest.skip('')
+    def test_annotator_midas(self):
+        # midas
+        midas_dict = {
+            'NAME': 'MidasDetector',
+            'PRETRAINED_MODEL':
+            'ms://damo/scepter_scedit@annotator/ckpts/dpt_hybrid-midas-501f0c75.pt',
+            'A': 6.2,
+            'BG_TH': 0.1
+        }
+        midas_anno = Config(cfg_dict=midas_dict, load=False)
+        midas_ins = ANNOTATORS.build(midas_anno).to(we.device_id)
+        midas_image = midas_ins(self.image)
+        print("midas's shape:", midas_image.shape)
+        Image.fromarray(midas_image).save(
+            os.path.join(self.save_dir, 'sunflower_midas.png'))
+
+    @unittest.skip('')
+    def test_annotator_mlsd(self):
+        # mlsd
+        mlsd_dict = {
+            'NAME': 'MLSDdetector',
+            'PRETRAINED_MODEL':
+            'ms://damo/scepter_scedit@annotator/ckpts/mlsd_large_512_fp32.pth',
+            'THR_V': 0.1,
+            'THR_D': 0.1
+        }
+        mlsd_anno = Config(cfg_dict=mlsd_dict, load=False)
+        mlsd_ins = ANNOTATORS.build(mlsd_anno).to(we.device_id)
+        mlsd_image = mlsd_ins(self.image)
+        print("mlsd's shape:", mlsd_image.shape)
+        Image.fromarray(mlsd_image).save(
+            os.path.join(self.save_dir, 'sunflower_mlsd.png'))
+
+    @unittest.skip('')
+    def test_annotator_color(self):
+        # color
+        color_dict = {'NAME': 'ColorAnnotator', 'RATIO': 64}
+        color_anno = Config(cfg_dict=color_dict, load=False)
+        color_ins = ANNOTATORS.build(color_anno).to(we.device_id)
+        color_image = color_ins(self.image)
+        print("color's shape:", color_image.shape)
+        Image.fromarray(color_image).save(
+            os.path.join(self.save_dir, 'sunflower_color.png'))
+
+    @unittest.skip('')
+    def test_annotator_multi(self):
+        # multi annotators
+        canny_dict = {
+            'NAME': 'CannyAnnotator',
+            'LOW_THRESHOLD': 100,
+            'HIGH_THRESHOLD': 200,
+            'INPUT_KEYS': ['img'],
+            'OUTPUT_KEYS': ['canny_img']
+        }
+        hed_dict = {
+            'NAME': 'HedAnnotator',
+            'PRETRAINED_MODEL':
+            'ms://damo/scepter_scedit@annotator/ckpts/ControlNetHED.pth',
+            'INPUT_KEYS': ['img'],
+            'OUTPUT_KEYS': ['hed_img']
+        }
+        openpose_dict = {
+            'NAME': 'OpenposeAnnotator',
+            'BODY_MODEL_PATH':
+            'ms://damo/scepter_scedit@annotator/ckpts/body_pose_model.pth',
+            'HAND_MODEL_PATH':
+            'ms://damo/scepter_scedit@annotator/ckpts/hand_pose_model.pth',
+            'INPUT_KEYS': ['img'],
+            'OUTPUT_KEYS': ['openpose_img']
+        }
+        midas_dict = {
+            'NAME': 'MidasDetector',
+            'PRETRAINED_MODEL':
+            'ms://damo/scepter_scedit@annotator/ckpts/dpt_hybrid-midas-501f0c75.pt',
+            'INPUT_KEYS': ['img'],
+            'OUTPUT_KEYS': ['midas_img']
+        }
+        mlsd_dict = {
+            'NAME': 'MLSDdetector',
+            'PRETRAINED_MODEL':
+            'ms://damo/scepter_scedit@annotator/ckpts/mlsd_large_512_fp32.pth',
+            'INPUT_KEYS': ['img'],
+            'OUTPUT_KEYS': ['mlsd_img']
+        }
+        color_dict = {'NAME': 'ColorAnnotator', 'RATIO': 64}
+        general_dict = {
+            'NAME':
+            'GeneralAnnotator',
+            'ANNOTATORS': [
+                canny_dict, hed_dict, openpose_dict, midas_dict, mlsd_dict,
+                color_dict
+            ]
+        }
+        general_anno = Config(cfg_dict=general_dict, load=False)
+        general_ins = ANNOTATORS.build(general_anno).to(we.device_id)
+        output_image = general_ins({'img': self.image})
+        for key, save_image in output_image.items():
+            Image.fromarray(save_image).save(
+                os.path.join(self.save_dir, f'sunflower_multi_{key}.png'))
+
+    # @unittest.skip('')
+    def test_annotator_processor(self):
+        from scepter.modules.annotator.utils import AnnotatorProcessor
+        anno_processor = AnnotatorProcessor(anno_type='hed')
+        output_image = anno_processor.run(self.image, 'hed')
+        Image.fromarray(output_image).save(
+            os.path.join(self.save_dir, 'sunflower_processor_hed.png'))
+
+        anno_processor = AnnotatorProcessor(
+            anno_type=['canny', 'color', 'depth'])
+        output_image = anno_processor.run(self.image, 'color')
+        Image.fromarray(output_image).save(
+            os.path.join(self.save_dir, 'sunflower_processor_color.png'))
+
+        output_image = anno_processor.run(self.image, ['canny', 'depth'])
+        for key, save_image in output_image.items():
+            Image.fromarray(save_image).save(
+                os.path.join(self.save_dir, f'sunflower_processor_{key}.png'))
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/tools/test_inference.py

@@ -12,53 +12,171 @@ class InferenceTest(unittest.TestCase):
     def tearDown(self):
         super().tearDown()
 
-    # @unittest.skip('')
+    @unittest.skip('')
     def test_infer_args(self):
-        os.system("python scepter/tools/run_inference.py "
-          "--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml "
-          "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_2023'")
-        os.system("python scepter/tools/run_inference.py "
-          "--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml "
-          "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_2024' --seed 2024")
-        os.system("python scepter/tools/run_inference.py "
-          "--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml "
-          "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_size768' "
-          "--image_size '768'")
-        os.system("python scepter/tools/run_inference.py "
-          "--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml "
-          "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_1280_720' "
-          "--image_size '1280,720'")
-        os.system("python scepter/tools/run_inference.py "
-          "--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml "
-          "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_720_1280_step10' "
-          "--image_size '720,1280' --sample_steps 10")
-        os.system("python scepter/tools/run_inference.py "
-          "--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml "
-          "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_num2' --num_samples 2")
-        os.system("python scepter/tools/run_inference.py "
-          "--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml "
-          "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_dpmpp_2s_ancestral' "
-          "--sampler 'dpmpp_2s_ancestral'")
-        os.system("python scepter/tools/run_inference.py "
-          "--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml "
-          "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_scale5' "
-          "--guide_scale 5.0")
-        os.system("python scepter/tools/run_inference.py "
-          "--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml "
-          "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_rescale0_1' "
-          "--guide_rescale 0.1")
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml '
+            "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_2023'"
+        )
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml '
+            "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_2024' --seed 2024"
+        )
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml '
+            "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_size768' "
+            "--image_size '768'")
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml '
+            "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_1280_720' "
+            "--image_size '1280,720'")
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml '
+            "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_720_1280_step10' "
+            "--image_size '720,1280' --sample_steps 10")
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml '
+            "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_num2' --num_samples 2"
+        )
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml '
+            "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_dpmpp_2s_ancestral' "
+            "--sampler 'dpmpp_2s_ancestral'")
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml '
+            "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_scale5' "
+            '--guide_scale 5.0')
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml '
+            "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog_rescale0_1' "
+            '--guide_rescale 0.1')
+
+    @unittest.skip('')
+    def test_example_infer(self):
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml '
+            "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog'")
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_2.1_768.yaml '
+            "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog'")
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_xl_1024.yaml '
+            "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog'")
+
+    @unittest.skip('')
+    def test_trained_infer(self):
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml '
+            "--pretrained_model 'cache/save_data/sd15_512_full/checkpoints/ldm_step-100.pth' "
+            "--prompt 'A close up of a small rabbit wearing a hat and scarf' "
+            "--save_folder 'trained_test_prompt_rabbit' ")
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_2.1_768.yaml '
+            "--pretrained_model 'cache/save_data/sd21_768_full/checkpoints/ldm_step-100.pth' "
+            "--prompt 'A close up of a small rabbit wearing a hat and scarf' "
+            "--save_folder 'trained_test_prompt_rabbit' ")
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_xl_1024.yaml --guide_scale 5.0 '
+            "--pretrained_model 'cache/save_data/sdxl_1024_full/checkpoints/ldm_step-100.pth' "
+            "--prompt 'A close up of a small rabbit wearing a hat and scarf' "
+            "--save_folder 'trained_test_prompt_rabbit' ")
+
+    @unittest.skip('')
+    def test_trained_tuning_infer(self):
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512_lora.yaml '
+            "--pretrained_model 'cache/save_data/sd15_512_lora/checkpoints/ldm_step-100.pth' "
+            "--prompt 'A close up of a small rabbit wearing a hat and scarf' "
+            "--save_folder 'trained_test_prompt_rabbit' ")
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_2.1_768_lora.yaml '
+            "--pretrained_model 'cache/save_data/sd21_768_lora/checkpoints/ldm_step-100.pth' "
+            "--prompt 'A close up of a small rabbit wearing a hat and scarf' "
+            "--save_folder 'trained_test_prompt_rabbit' ")
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_xl_1024_lora.yaml '
+            "--pretrained_model 'cache/save_data/sdxl_1024_lora/checkpoints/ldm_step-100.pth' "
+            "--prompt 'A close up of a small rabbit wearing a hat and scarf' "
+            "--save_folder 'trained_test_prompt_rabbit' ")
+
+    @unittest.skip('')
+    def test_trained_scedit_infer(self):
+        # swift
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/scedit/t2i/sd15_512_sce_t2i_swift.yaml '
+            "--pretrained_model 'cache/save_data/sd15_512_sce_t2i_swift/checkpoints/ldm_step-100.pth' "
+            "--prompt 'A close up of a small rabbit wearing a hat and scarf' "
+            "--save_folder 'trained_test_prompt_rabbit' ")
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/scedit/t2i/sd21_768_sce_t2i_swift.yaml '
+            "--pretrained_model 'cache/save_data/sd21_768_sce_t2i_swift/checkpoints/ldm_step-100.pth' "
+            "--prompt 'A close up of a small rabbit wearing a hat and scarf' "
+            "--save_folder 'trained_test_prompt_rabbit' ")
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/scedit/t2i/sdxl_1024_sce_t2i_swift.yaml '
+            "--pretrained_model 'cache/save_data/sdxl_1024_sce_t2i_swift/checkpoints/ldm_step-100.pth' "
+            "--prompt 'A close up of a small rabbit wearing a hat and scarf' "
+            "--save_folder 'trained_test_prompt_rabbit' ")
+        # original
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/scedit/t2i/sd15_512_sce_t2i.yaml '
+            "--pretrained_model 'cache/save_data/sd15_512_sce_t2i/checkpoints/ldm_step-100.pth' "
+            "--prompt 'A close up of a small rabbit wearing a hat and scarf' "
+            "--save_folder 'trained_test_prompt_rabbit' ")
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/scedit/t2i/sd21_768_sce_t2i.yaml '
+            "--pretrained_model 'cache/save_data/sd21_768_sce_t2i/checkpoints/ldm_step-100.pth' "
+            "--prompt 'A close up of a small rabbit wearing a hat and scarf' "
+            "--save_folder 'trained_test_prompt_rabbit' ")
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/scedit/t2i/sdxl_1024_sce_t2i.yaml '
+            "--pretrained_model 'cache/save_data/sdxl_1024_sce_t2i/checkpoints/ldm_step-100.pth' "
+            "--prompt 'A close up of a small rabbit wearing a hat and scarf' "
+            "--save_folder 'trained_test_prompt_rabbit' ")
 
     # @unittest.skip('')
-    def test_example_infer(self):
-          os.system("python scepter/tools/run_inference.py "
-            "--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml "
-            "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog'")
-          os.system("python scepter/tools/run_inference.py "
-            "--cfg scepter/methods/examples/generation/stable_diffusion_2.1_768.yaml "
-            "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog'")
-          os.system("python scepter/tools/run_inference.py "
-            "--cfg scepter/methods/examples/generation/stable_diffusion_xl_1024.yaml "
-            "--prompt 'a cute dog' --save_folder 'test_prompt_a_cute_dog'")
+    def test_pretrained_scedit_control_infer(self):
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/scedit/ctr/sd21_768_sce_ctr_canny.yaml --num_samples 1 '
+            "--prompt 'a single flower is shown in front of a tree' --save_folder 'test_flower_canny' "
+            "--image_size 768 --task control --image 'asset/images/flower.jpg' --control_mode canny "
+            '--pretrained_model '
+            'ms://damo/scepter_scedit@controllable_model/SD2.1/canny_control/0_SwiftSCETuning/pytorch_model.bin'
+        )
+
+        os.system(
+            'python scepter/tools/run_inference.py '
+            '--cfg scepter/methods/scedit/ctr/sd21_768_sce_ctr_pose.yaml --num_samples 1 '
+            "--prompt 'super mario' --save_folder 'test_mario_pose' "
+            "--image_size 768 --task control --image 'asset/images/pose_source.png' --control_mode source "
+            '--pretrained_model '
+            'ms://damo/scepter_scedit@controllable_model/SD2.1/pose_control/0_SwiftSCETuning/pytorch_model.bin'
+        )
 
 
 if __name__ == '__main__':

tests/tools/test_train.py

@@ -8,37 +8,178 @@ import unittest
 class TrainTest(unittest.TestCase):
     def setUp(self):
         print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
+        self.tmp_dir = './cache/save_data'
 
     def tearDown(self):
         super().tearDown()
 
-    # @unittest.skip('')
+    @unittest.skip('')
     def test_generation_example_full(self):
-        os.system("python scepter/tools/run_train.py "
-          "--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml ")
-        os.system("python scepter/tools/run_train.py "
-          "--cfg scepter/methods/examples/generation/stable_diffusion_2.1_768.yaml ")
-        os.system("python scepter/tools/run_train.py "
-          "--cfg scepter/methods/examples/generation/stable_diffusion_xl_1024.yaml ")
-    
-    # @unittest.skip('')
+        os.system(
+            'python scepter/tools/run_train.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512.yaml '
+            '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir, 'sd15_512_full/checkpoints')))
+
+        os.system(
+            'python scepter/tools/run_train.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_2.1_768.yaml '
+            '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir, 'sd21_768_full/checkpoints')))
+
+        os.system(
+            'python scepter/tools/run_train.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_xl_1024.yaml '
+            '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir, 'sdxl_1024_full/checkpoints')))
+
+    @unittest.skip('')
     def test_generation_example_lora(self):
-        os.system("python scepter/tools/run_train.py "
-          "--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512_lora.yaml ")
-        os.system("python scepter/tools/run_train.py "
-          "--cfg scepter/methods/examples/generation/stable_diffusion_2.1_768_lora.yaml ")
-        os.system("python scepter/tools/run_train.py "
-          "--cfg scepter/methods/examples/generation/stable_diffusion_xl_1024_lora.yaml ")
-        
+        os.system(
+            'python scepter/tools/run_train.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_1.5_512_lora.yaml '
+            '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir, 'sd15_512_lora/checkpoints')))
+
+        os.system(
+            'python scepter/tools/run_train.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_2.1_768_lora.yaml '
+            '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir, 'sd21_768_lora/checkpoints')))
+
+        os.system(
+            'python scepter/tools/run_train.py '
+            '--cfg scepter/methods/examples/generation/stable_diffusion_xl_1024_lora.yaml '
+            '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir, 'sdxl_1024_lora/checkpoints')))
+
+    @unittest.skip('')
+    def test_generation_example_scedit_t2i_swift(self):
+        os.system(
+            'python scepter/tools/run_train.py '
+            '--cfg scepter/methods/scedit/t2i/sd15_512_sce_t2i_swift.yaml '
+            '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir,
+                             'sd15_512_sce_t2i_swift/checkpoints')))
+
+        os.system(
+            'python scepter/tools/run_train.py '
+            '--cfg scepter/methods/scedit/t2i/sd21_768_sce_t2i_swift.yaml '
+            '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir,
+                             'sd21_768_sce_t2i_swift/checkpoints')))
+
+        os.system(
+            'python scepter/tools/run_train.py '
+            '--cfg scepter/methods/scedit/t2i/sdxl_1024_sce_t2i_swift.yaml '
+            '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir,
+                             'sdxl_1024_sce_t2i_swift/checkpoints')))
+
+    @unittest.skip('')
+    def test_generation_example_scedit_t2i(self):
+        os.system('python scepter/tools/run_train.py '
+                  '--cfg scepter/methods/scedit/t2i/sd15_512_sce_t2i.yaml '
+                  '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir, 'sd15_512_sce_t2i/checkpoints')))
+
+        os.system('python scepter/tools/run_train.py '
+                  '--cfg scepter/methods/scedit/t2i/sd21_768_sce_t2i.yaml '
+                  '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir, 'sd21_768_sce_t2i/checkpoints')))
+
+        os.system('python scepter/tools/run_train.py '
+                  '--cfg scepter/methods/scedit/t2i/sdxl_1024_sce_t2i.yaml '
+                  '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir, 'sdxl_1024_sce_t2i/checkpoints')))
+
+    @unittest.skip('')
+    def test_generation_example_scedit_ctr(self):
+        os.system('python scepter/tools/run_train.py '
+                  '--cfg scepter/methods/scedit/ctr/sd15_512_sce_ctr_hed.yaml '
+                  '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir,
+                             'sd15_512_sce_ctr_hed/checkpoints')))
+
+        os.system(
+            'python scepter/tools/run_train.py '
+            '--cfg scepter/methods/scedit/ctr/sd21_768_sce_ctr_canny.yaml '
+            '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir,
+                             'sd21_768_sce_ctr_canny/checkpoints')))
+
+        os.system(
+            'python scepter/tools/run_train.py '
+            '--cfg scepter/methods/scedit/ctr/sd21_768_sce_ctr_pose.yaml '
+            '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir,
+                             'sd21_768_sce_ctr_pose/checkpoints')))
+
+        os.system(
+            'python scepter/tools/run_train.py '
+            '--cfg scepter/methods/scedit/ctr/sdxl_1024_sce_ctr_depth.yaml '
+            '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir,
+                             'sdxl_1024_sce_ctr_depth/checkpoints')))
+
+        os.system(
+            'python scepter/tools/run_train.py '
+            '--cfg scepter/methods/scedit/ctr/sdxl_1024_sce_ctr_color.yaml '
+            '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir,
+                             'sdxl_1024_sce_ctr_color/checkpoints')))
+
     # @unittest.skip('')
-    def test_generation_example_scedit(self):
-        os.system("python scepter/tools/run_train.py "
-          "--cfg scepter/methods/SCEdit/t2i_sd15_512_sce.yaml ")
-        os.system("python scepter/tools/run_train.py "
-          "--cfg scepter/methods/SCEdit/t2i_sd21_768_sce.yaml ")
-        os.system("python scepter/tools/run_train.py "
-          "--cfg scepter/methods/SCEdit/t2i_sdxl_1024_sce.yaml ")
-        
+    def test_generation_example_datatxt(self):
+        # os.system(
+        #     'python scepter/tools/run_train.py '
+        #     '--cfg scepter/methods/scedit/t2i/sdxl_1024_sce_t2i_datatxt.yaml '
+        #     '--max_steps 100'
+        # )
+        # self.assertTrue(os.path.exists(os.path.join(self.tmp_dir, 'sdxl_1024_sce_t2i_datatxt/checkpoints')))
+
+        os.system(
+            'python scepter/tools/run_train.py '
+            '--cfg scepter/methods/scedit/ctr/sdxl_1024_sce_ctr_color_datatxt.yaml '
+            '--max_steps 100')
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(self.tmp_dir,
+                             'sdxl_1024_sce_ctr_color_datatxt/checkpoints')))
 
 
 if __name__ == '__main__':

tests/utils/test_fs.py

@@ -45,6 +45,11 @@ class FSTest(unittest.TestCase):
             print(f'Download from {path} to {local_path}')
             self.assertTrue(os.path.exists(local_path))
 
+        path = 'ms://AI-ModelScope/clip-vit-large-patch14'
+        with FS.get_dir_to_local_dir(path, wait_finish=True) as local_path:
+            print(f'Download from {path} to {local_path}')
+            self.assertTrue(os.path.exists(local_path))
+
 
 if __name__ == '__main__':
     unittest.main()