monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Make sure cleanup happens even if tests fail

Browse Source
This commit is contained in:
filipstrand 2024-10-09 23:38:56 +02:00
parent a9461cacf0
commit 78e2b726fb
2 changed files with 62 additions and 58 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

64
tests/helpers/image_generation_controlnet_test_helper.py
View File

@ -27,37 +27,39 @@ class ImageGeneratorControlnetTestHelper:
controlnet_image_path = str(ImageGeneratorTestHelper.resolve_path(controlnet_image_path)) controlnet_image_path = str(ImageGeneratorTestHelper.resolve_path(controlnet_image_path))
lora_paths = [str(ImageGeneratorTestHelper.resolve_path(p)) for p in lora_paths] if lora_paths else None lora_paths = [str(ImageGeneratorTestHelper.resolve_path(p)) for p in lora_paths] if lora_paths else None
# given try:
flux = Flux1Controlnet( # given
model_config=model_config, flux = Flux1Controlnet(
quantize=8, model_config=model_config,
lora_paths=lora_paths, quantize=8,
lora_scales=lora_scales, lora_paths=lora_paths,
) lora_scales=lora_scales,
)
# when # when
image = flux.generate_image( image = flux.generate_image(
seed=seed, seed=seed,
prompt=prompt, prompt=prompt,
output=str(output_image_path), output=str(output_image_path),
controlnet_image_path=controlnet_image_path, controlnet_image_path=controlnet_image_path,
controlnet_save_canny=False, controlnet_save_canny=False,
config=ConfigControlnet( config=ConfigControlnet(
num_inference_steps=steps, num_inference_steps=steps,
height=768, height=768,
width=493, width=493,
controlnet_strength=controlnet_strength, controlnet_strength=controlnet_strength,
), ),
) )
image.save(path=output_image_path) image.save(path=output_image_path)
# then # then
np.testing.assert_array_equal( np.testing.assert_array_equal(
np.array(Image.open(output_image_path)), np.array(Image.open(output_image_path)),
np.array(Image.open(reference_image_path)), np.array(Image.open(reference_image_path)),
err_msg="Generated image doesn't match reference image", err_msg="Generated image doesn't match reference image",
) )
# cleanup finally:
if os.path.exists(output_image_path): # cleanup
os.remove(output_image_path) if os.path.exists(output_image_path):
os.remove(output_image_path)

56
tests/helpers/image_generation_test_helper.py
View File

@ -24,36 +24,38 @@ class ImageGeneratorTestHelper:
output_image_path = ImageGeneratorTestHelper.resolve_path(output_image_path) output_image_path = ImageGeneratorTestHelper.resolve_path(output_image_path)
lora_paths = [str(ImageGeneratorTestHelper.resolve_path(p)) for p in lora_paths] if lora_paths else None lora_paths = [str(ImageGeneratorTestHelper.resolve_path(p)) for p in lora_paths] if lora_paths else None
# given try:
flux = Flux1( # given
model_config=model_config, flux = Flux1(
quantize=8, model_config=model_config,
lora_paths=lora_paths, quantize=8,
lora_scales=lora_scales lora_paths=lora_paths,
) # fmt: off lora_scales=lora_scales
) # fmt: off
# when # when
image = flux.generate_image( image = flux.generate_image(
seed=seed, seed=seed,
prompt=prompt, prompt=prompt,
config=Config( config=Config(
num_inference_steps=steps, num_inference_steps=steps,
height=341, height=341,
width=768, width=768,
), ),
) )
image.save(path=output_image_path) image.save(path=output_image_path)
# then # then
np.testing.assert_array_equal( np.testing.assert_array_equal(
np.array(Image.open(output_image_path)), np.array(Image.open(output_image_path)),
np.array(Image.open(reference_image_path)), np.array(Image.open(reference_image_path)),
err_msg="Generated image doesn't match reference image", err_msg="Generated image doesn't match reference image",
) )
# cleanup finally:
if os.path.exists(output_image_path): # cleanup
os.remove(output_image_path) if os.path.exists(output_image_path):
os.remove(output_image_path)
@staticmethod @staticmethod
def resolve_path(path) -> Path: def resolve_path(path) -> Path: