From 633243c98379e34e80420049cda095c5f748b220 Mon Sep 17 00:00:00 2001 From: enzymezoo-code Date: Mon, 4 Sep 2023 22:11:01 -0500 Subject: [PATCH] Assert images are generated --- tests/inference/test_inference.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/inference/test_inference.py b/tests/inference/test_inference.py index 4a6c19296..5610f6266 100644 --- a/tests/inference/test_inference.py +++ b/tests/inference/test_inference.py @@ -1,4 +1,5 @@ from copy import deepcopy +from io import BytesIO from urllib import request import numpy import os @@ -128,8 +129,8 @@ comfy_graph_list = [DEFAULT_COMFY_GRAPH] prompt_list = [ 'a painting of a cat', ] -sampler_list = KSampler.SAMPLERS[0:2] -scheduler_list = [KSampler.SCHEDULERS[0]] +sampler_list = KSampler.SAMPLERS +scheduler_list = KSampler.SCHEDULERS @pytest.mark.inference @pytest.mark.parametrize("sampler", sampler_list) @@ -208,5 +209,13 @@ class TestInference: comfy_graph.set_prompt(prompt) # Generate - client.get_images(comfy_graph.graph) + images = client.get_images(comfy_graph.graph) + + assert len(images) == 1, "No images generated" + # assert all images are not blank + for images_output in images.values(): + for image_data in images_output: + pil_image = Image.open(BytesIO(image_data)) + assert numpy.array(pil_image).any() != 0, "Image is blank" +