wip

lama_cleaner/tests/test_model_switch.py

@@ -1,6 +1,6 @@
 import os
 
-from lama_cleaner.schema import Config
+from lama_cleaner.schema import InpaintRequest
 
 os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
 
@@ -38,7 +38,7 @@ def test_controlnet_switch_onoff(caplog):
     )
 
     model.switch_controlnet_method(
-        Config(
+        InpaintRequest(
             name=name,
             enable_controlnet=False,
         )
@@ -63,7 +63,7 @@ def test_switch_controlnet_method(caplog):
     )
 
     model.switch_controlnet_method(
-        Config(
+        InpaintRequest(
             name=name,
             enable_controlnet=True,
             controlnet_method=new_method,

lama_cleaner/tests/test_save_exif.py

@@ -1,4 +1,5 @@
 import io
+import tempfile
 from pathlib import Path
 from typing import List
 
@@ -18,9 +19,9 @@ def extra_info(img_p: Path):
     ext = img_p.suffix.strip(".")
     img_bytes = img_p.read_bytes()
     np_img, _, infos = load_img(img_bytes, False, True)
-    pil_bytes = pil_to_bytes(Image.fromarray(np_img), ext=ext, infos=infos)
-    res_img = Image.open(io.BytesIO(pil_bytes))
-    return infos, res_img.info
+    res_pil_bytes = pil_to_bytes(Image.fromarray(np_img), ext=ext, infos=infos)
+    res_img = Image.open(io.BytesIO(res_pil_bytes))
+    return infos, res_img.info, res_pil_bytes
 
 
 def assert_keys(keys: List[str], infos, res_infos):
@@ -30,23 +31,29 @@ def assert_keys(keys: List[str], infos, res_infos):
         assert infos[k] == res_infos[k]
 
 
+def run_test(file_path, keys):
+    infos, res_infos, res_pil_bytes = extra_info(file_path)
+    assert_keys(keys, infos, res_infos)
+    with tempfile.NamedTemporaryFile("wb", suffix=file_path.suffix) as temp_file:
+        temp_file.write(res_pil_bytes)
+        temp_file.flush()
+        infos, res_infos, res_pil_bytes = extra_info(Path(temp_file.name))
+        assert_keys(keys, infos, res_infos)
+
+
 def test_png_icc_profile_png():
-    infos, res_infos = extra_info(current_dir / "icc_profile_test.png")
-    assert_keys(["icc_profile", "exif"], infos, res_infos)
+    run_test(current_dir / "icc_profile_test.png", ["icc_profile", "exif"])
 
 
 def test_png_icc_profile_jpeg():
-    infos, res_infos = extra_info(current_dir / "icc_profile_test.jpg")
-    assert_keys(["icc_profile", "exif"], infos, res_infos)
+    run_test(current_dir / "icc_profile_test.jpg", ["icc_profile", "exif"])
 
 
 def test_jpeg():
     jpg_img_p = current_dir / "bunny.jpeg"
-    infos, res_infos = extra_info(jpg_img_p)
-    assert_keys(["dpi", "exif"], infos, res_infos)
+    run_test(jpg_img_p, ["dpi", "exif"])
 
 
 def test_png_parameter():
     jpg_img_p = current_dir / "png_parameter_test.png"
-    infos, res_infos = extra_info(jpg_img_p)
-    assert_keys(["parameters"], infos, res_infos)
+    run_test(jpg_img_p, ["parameters"])

lama_cleaner/tests/utils.py

@@ -3,7 +3,7 @@ import cv2
 import pytest
 import torch
 
-from lama_cleaner.schema import LDMSampler, HDStrategy, Config, SDSampler
+from lama_cleaner.schema import LDMSampler, HDStrategy, InpaintRequest, SDSampler
 
 current_dir = Path(__file__).parent.absolute().resolve()
 save_dir = current_dir / "result"
@@ -72,4 +72,4 @@ def get_config(**kwargs):
         hd_strategy_resize_limit=200,
     )
     data.update(**kwargs)
-    return Config(**data)
+    return InpaintRequest(**data)