add tests

config.json

@@ -56,7 +56,7 @@
     "training_image_repeats_per_epoch": 1,
     "training_write_csv_every": 10.0,
     "training_xattention_optimizations": false,
-    "sd_model_checkpoint": "sd-v15-runwayml.ckpt [81761151]",
+    "sd_model_checkpoint": "sd21-hot-fantasy.safetensors [651c7c98]",
     "sd_checkpoint_cache": 0,
     "sd_vae": "vae-ft-mse-840000-ema-pruned",
     "sd_vae_as_default": false,

modules/sd_models.py

@@ -49,6 +49,9 @@ def checkpoint_tiles():
 
 
 def find_checkpoint_config(info):
+    if info is None:
+        return shared.cmd_opts.config
+
     config = os.path.splitext(info.filename)[0] + ".yaml"
     if os.path.exists(config):
         return config
@@ -174,7 +177,7 @@ def read_state_dict(checkpoint_file, print_global_state=False, map_location=None
     if device is None:
         device = devices.get_cuda_device_string() if torch.cuda.is_available() else "cpu"
     if extension.lower() == ".safetensors":
-        pl_sd = safetensors.torch.load_file(checkpoint_file, device=device)
+        pl_sd = safetensors.torch.load_file(checkpoint_file, device="cpu") # safetensors cannot load into gpu directly
     else:
         pl_sd = torch.load(checkpoint_file, map_location=device)
 
@@ -345,14 +348,16 @@ def reload_model_weights(sd_model=None, info=None):
 
     if not sd_model:
         sd_model = shared.sd_model
+    if sd_model is None: # previous model load failed
+        current_checkpoint_info = None
+    else:
+        current_checkpoint_info = sd_model.sd_checkpoint_info
+        if sd_model.sd_model_checkpoint == checkpoint_info.filename:
+            return
 
-    current_checkpoint_info = sd_model.sd_checkpoint_info
     checkpoint_config = find_checkpoint_config(current_checkpoint_info)
 
-    if sd_model.sd_model_checkpoint == checkpoint_info.filename:
-        return
-
-    if checkpoint_config != find_checkpoint_config(checkpoint_info) or should_hijack_inpainting(checkpoint_info) != should_hijack_inpainting(sd_model.sd_checkpoint_info):
+    if current_checkpoint_info is None or checkpoint_config != find_checkpoint_config(checkpoint_info) or should_hijack_inpainting(checkpoint_info) != should_hijack_inpainting(sd_model.sd_checkpoint_info):
         del sd_model
         checkpoints_loaded.clear()
         load_model(checkpoint_info)

test/advanced_features/extras_test.py

@@ -0,0 +1,29 @@
+import unittest
+
+
+class TestExtrasWorking(unittest.TestCase):
+    def setUp(self):
+        self.url_img2img = "http://localhost:7860/sdapi/v1/extra-single-image"
+        self.simple_extras = {
+            "resize_mode": 0,
+            "show_extras_results": True,
+            "gfpgan_visibility": 0,
+            "codeformer_visibility": 0,
+            "codeformer_weight": 0,
+            "upscaling_resize": 2,
+            "upscaling_resize_w": 128,
+            "upscaling_resize_h": 128,
+            "upscaling_crop": True,
+            "upscaler_1": "None",
+            "upscaler_2": "None",
+            "extras_upscaler_2_visibility": 0,
+            "image": ""
+            }
+
+
+class TestExtrasCorrectness(unittest.TestCase):
+    pass
+
+
+if __name__ == "__main__":
+    unittest.main()

test/advanced_features/txt2img_test.py

@@ -0,0 +1,47 @@
+import unittest
+import requests
+
+
+class TestTxt2ImgWorking(unittest.TestCase):
+    def setUp(self):
+        self.url_txt2img = "http://localhost:7860/sdapi/v1/txt2img"
+        self.simple_txt2img = {
+            "enable_hr": False,
+            "denoising_strength": 0,
+            "firstphase_width": 0,
+            "firstphase_height": 0,
+            "prompt": "example prompt",
+            "styles": [],
+            "seed": -1,
+            "subseed": -1,
+            "subseed_strength": 0,
+            "seed_resize_from_h": -1,
+            "seed_resize_from_w": -1,
+            "batch_size": 1,
+            "n_iter": 1,
+            "steps": 3,
+            "cfg_scale": 7,
+            "width": 64,
+            "height": 64,
+            "restore_faces": False,
+            "tiling": False,
+            "negative_prompt": "",
+            "eta": 0,
+            "s_churn": 0,
+            "s_tmax": 0,
+            "s_tmin": 0,
+            "s_noise": 1,
+            "sampler_index": "Euler a"
+        }
+
+    def test_txt2img_with_restore_faces_performed(self):
+        self.simple_txt2img["restore_faces"] = True
+        self.assertEqual(requests.post(self.url_txt2img, json=self.simple_txt2img).status_code, 200)
+
+
+class TestTxt2ImgCorrectness(unittest.TestCase):
+    pass
+
+
+if __name__ == "__main__":
+    unittest.main()

test/basic_features/img2img_test.py

@@ -0,0 +1,55 @@
+import unittest
+import requests
+from gradio.processing_utils import encode_pil_to_base64
+from PIL import Image
+
+
+class TestImg2ImgWorking(unittest.TestCase):
+    def setUp(self):
+        self.url_img2img = "http://localhost:7860/sdapi/v1/img2img"
+        self.simple_img2img = {
+            "init_images": [encode_pil_to_base64(Image.open(r"test/test_files/img2img_basic.png"))],
+            "resize_mode": 0,
+            "denoising_strength": 0.75,
+            "mask": None,
+            "mask_blur": 4,
+            "inpainting_fill": 0,
+            "inpaint_full_res": False,
+            "inpaint_full_res_padding": 0,
+            "inpainting_mask_invert": 0,
+            "prompt": "example prompt",
+            "styles": [],
+            "seed": -1,
+            "subseed": -1,
+            "subseed_strength": 0,
+            "seed_resize_from_h": -1,
+            "seed_resize_from_w": -1,
+            "batch_size": 1,
+            "n_iter": 1,
+            "steps": 3,
+            "cfg_scale": 7,
+            "width": 64,
+            "height": 64,
+            "restore_faces": False,
+            "tiling": False,
+            "negative_prompt": "",
+            "eta": 0,
+            "s_churn": 0,
+            "s_tmax": 0,
+            "s_tmin": 0,
+            "s_noise": 1,
+            "override_settings": {},
+            "sampler_index": "Euler a",
+            "include_init_images": False
+            }
+
+    def test_img2img_simple_performed(self):
+        self.assertEqual(requests.post(self.url_img2img, json=self.simple_img2img).status_code, 200)
+
+    def test_inpainting_masked_performed(self):
+        self.simple_img2img["mask"] = encode_pil_to_base64(Image.open(r"test/test_files/mask_basic.png"))
+        self.assertEqual(requests.post(self.url_img2img, json=self.simple_img2img).status_code, 200)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/basic_features/txt2img_test.py

@@ -0,0 +1,68 @@
+import unittest
+import requests
+
+
+class TestTxt2ImgWorking(unittest.TestCase):
+    def setUp(self):
+        self.url_txt2img = "http://localhost:7860/sdapi/v1/txt2img"
+        self.simple_txt2img = {
+            "enable_hr": False,
+            "denoising_strength": 0,
+            "firstphase_width": 0,
+            "firstphase_height": 0,
+            "prompt": "example prompt",
+            "styles": [],
+            "seed": -1,
+            "subseed": -1,
+            "subseed_strength": 0,
+            "seed_resize_from_h": -1,
+            "seed_resize_from_w": -1,
+            "batch_size": 1,
+            "n_iter": 1,
+            "steps": 3,
+            "cfg_scale": 7,
+            "width": 64,
+            "height": 64,
+            "restore_faces": False,
+            "tiling": False,
+            "negative_prompt": "",
+            "eta": 0,
+            "s_churn": 0,
+            "s_tmax": 0,
+            "s_tmin": 0,
+            "s_noise": 1,
+            "sampler_index": "Euler a"
+        }
+
+    def test_txt2img_simple_performed(self):
+        self.assertEqual(requests.post(self.url_txt2img, json=self.simple_txt2img).status_code, 200)
+
+    def test_txt2img_with_negative_prompt_performed(self):
+        self.simple_txt2img["negative_prompt"] = "example negative prompt"
+        self.assertEqual(requests.post(self.url_txt2img, json=self.simple_txt2img).status_code, 200)
+
+    def test_txt2img_not_square_image_performed(self):
+        self.simple_txt2img["height"] = 128
+        self.assertEqual(requests.post(self.url_txt2img, json=self.simple_txt2img).status_code, 200)
+
+    def test_txt2img_with_hrfix_performed(self):
+        self.simple_txt2img["enable_hr"] = True
+        self.assertEqual(requests.post(self.url_txt2img, json=self.simple_txt2img).status_code, 200)
+
+    def test_txt2img_with_tiling_performed(self):
+        self.simple_txt2img["tiling"] = True
+        self.assertEqual(requests.post(self.url_txt2img, json=self.simple_txt2img).status_code, 200)
+
+    def test_txt2img_with_vanilla_sampler_performed(self):
+        self.simple_txt2img["sampler_index"] = "PLMS"
+        self.assertEqual(requests.post(self.url_txt2img, json=self.simple_txt2img).status_code, 200)
+        self.simple_txt2img["sampler_index"] = "DDIM"
+        self.assertEqual(requests.post(self.url_txt2img, json=self.simple_txt2img).status_code, 200)
+
+    def test_txt2img_multiple_batches_performed(self):
+        self.simple_txt2img["n_iter"] = 2
+        self.assertEqual(requests.post(self.url_txt2img, json=self.simple_txt2img).status_code, 200)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/basic_features/utils_test.py

@@ -0,0 +1,53 @@
+import unittest
+import requests
+
+class UtilsTests(unittest.TestCase):
+  def setUp(self):
+    self.url_options = "http://localhost:7860/sdapi/v1/options"
+    self.url_cmd_flags = "http://localhost:7860/sdapi/v1/cmd-flags"
+    self.url_samplers = "http://localhost:7860/sdapi/v1/samplers"
+    self.url_upscalers = "http://localhost:7860/sdapi/v1/upscalers"
+    self.url_sd_models = "http://localhost:7860/sdapi/v1/sd-models"
+    self.url_hypernetworks = "http://localhost:7860/sdapi/v1/hypernetworks"
+    self.url_face_restorers = "http://localhost:7860/sdapi/v1/face-restorers"
+    self.url_realesrgan_models = "http://localhost:7860/sdapi/v1/realesrgan-models"
+    self.url_prompt_styles = "http://localhost:7860/sdapi/v1/prompt-styles"
+    self.url_artist_categories = "http://localhost:7860/sdapi/v1/artist-categories"
+    self.url_artists = "http://localhost:7860/sdapi/v1/artists"
+
+  def test_options_get(self):
+    self.assertEqual(requests.get(self.url_options).status_code, 200)
+
+  def test_cmd_flags(self):
+    self.assertEqual(requests.get(self.url_cmd_flags).status_code, 200)
+
+  def test_samplers(self):
+    self.assertEqual(requests.get(self.url_samplers).status_code, 200)
+
+  def test_upscalers(self):
+    self.assertEqual(requests.get(self.url_upscalers).status_code, 200)
+
+  def test_sd_models(self):
+    self.assertEqual(requests.get(self.url_sd_models).status_code, 200)
+
+  def test_hypernetworks(self):
+    self.assertEqual(requests.get(self.url_hypernetworks).status_code, 200)
+
+  def test_face_restorers(self):
+    self.assertEqual(requests.get(self.url_face_restorers).status_code, 200)
+  
+  def test_realesrgan_models(self):
+    self.assertEqual(requests.get(self.url_realesrgan_models).status_code, 200)
+  
+  def test_prompt_styles(self):
+    self.assertEqual(requests.get(self.url_prompt_styles).status_code, 200)
+  
+  def test_artist_categories(self):
+    self.assertEqual(requests.get(self.url_artist_categories).status_code, 200)
+
+  def test_artists(self):
+    self.assertEqual(requests.get(self.url_artists).status_code, 200)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/server_poll.py

@@ -0,0 +1,24 @@
+import unittest
+import requests
+import time
+
+
+def run_tests(proc, test_dir):
+    timeout_threshold = 240
+    start_time = time.time()
+    while time.time()-start_time < timeout_threshold:
+        try:
+            requests.head("http://localhost:7860/")
+            break
+        except requests.exceptions.ConnectionError:
+            if proc.poll() is not None:
+                break
+    if proc.poll() is None:
+        if test_dir is None:
+            test_dir = ""
+        suite = unittest.TestLoader().discover(test_dir, pattern="*_test.py", top_level_dir="test")
+        result = unittest.TextTestRunner(verbosity=2).run(suite)
+        return len(result.failures) + len(result.errors)
+    else:
+        print("Launch unsuccessful")
+        return 1