add control tests

2026-09-19 17:24:32 +02:00 · 2023-12-21 13:36:42 -05:00
parent 9b40bdcd91
commit d7b0d59143
3 changed files with 122 additions and 9 deletions
@@ -1,10 +1,10 @@
 import math
-from PIL import Image
-from modules.control import processors # patrickvonplaten controlnet_aux
-from modules import shared
+from PIL import Image, ImageChops
+from modules import shared, errors


 def test_processors(image):
+    from modules.control import processors
    if image is None:
        shared.log.error('Image not loaded')
        return None, None, None
@@ -15,11 +15,16 @@ def test_processors(image):
            continue
        shared.log.info(f'Testing processor: {processor_id}')
        processor = processors.Processor(processor_id)
+        output = image
        if processor is None:
            shared.log.error(f'Processor load failed: id="{processor_id}"')
-            continue
-        output = processor(image)
-        processor.reset()
+            processor_id = f'{processor_id} error'
+        else:
+            output = processor(image)
+            processor.reset()
+        diff = ImageChops.difference(image, output)
+        if not diff.getbbox():
+            processor_id = f'{processor_id} null'
        draw = ImageDraw.Draw(output)
        font = ImageFont.truetype('DejaVuSansMono', 48)
        draw.text((10, 10), processor_id, (0,0,0), font=font)
@@ -40,3 +45,107 @@ def test_processors(image):
        grid.paste(thumb, box=(x, y))
    yield None, grid, None, images
    return None, grid, None, images # preview_process, output_image, output_video, output_gallery
+
+
+def test_controlnets(prompt, negative, image):
+    from modules import devices, sd_models
+    from modules.control import controlnets
+    if image is None:
+        shared.log.error('Image not loaded')
+        return None, None, None
+    from PIL import ImageDraw, ImageFont
+    images = []
+    for model_id in controlnets.list_models():
+        if model_id is None:
+            model_id = 'None'
+        if shared.state.interrupted:
+            continue
+        output = image
+        if model_id != 'None':
+            controlnet = controlnets.ControlNet(model_id=model_id, device=devices.device, dtype=devices.dtype)
+            if controlnet is None:
+                shared.log.error(f'ControlNet load failed: id="{model_id}"')
+                continue
+            shared.log.info(f'Testing ControlNet: {model_id}')
+            pipe = controlnets.ControlNetPipeline(controlnet=controlnet.model, pipeline=shared.sd_model)
+            pipe.pipeline.to(device=devices.device, dtype=devices.dtype)
+            sd_models.set_diffuser_options(pipe)
+            try:
+                res = pipe.pipeline(prompt=prompt, negative_prompt=negative, image=image, num_inference_steps=10, output_type='pil')
+                output = res.images[0]
+            except Exception as e:
+                errors.display(e, f'ControlNet {model_id} inference')
+                model_id = f'{model_id} error'
+            pipe.restore()
+        draw = ImageDraw.Draw(output)
+        font = ImageFont.truetype('DejaVuSansMono', 48)
+        draw.text((10, 10), model_id, (0,0,0), font=font)
+        draw.text((8, 8), model_id, (255,255,255), font=font)
+        images.append(output)
+        yield output, None, None, images
+    rows = round(math.sqrt(len(images)))
+    cols = math.ceil(len(images) / rows)
+    w, h = 256, 256
+    size = (cols * w + cols, rows * h + rows)
+    grid = Image.new('RGB', size=size, color='black')
+    shared.log.info(f'Test ControlNets: images={len(images)} grid={grid}')
+    for i, image in enumerate(images):
+        x = (i % cols * w) + (i % cols)
+        y = (i // cols * h) + (i // cols)
+        thumb = image.copy().convert('RGB')
+        thumb.thumbnail((w, h), Image.Resampling.HAMMING)
+        grid.paste(thumb, box=(x, y))
+    yield None, grid, None, images
+    return None, grid, None, images # preview_process, output_image, output_video, output_gallery
+
+
+def test_adapters(prompt, negative, image):
+    from modules import devices, sd_models
+    from modules.control import adapters
+    if image is None:
+        shared.log.error('Image not loaded')
+        return None, None, None
+    from PIL import ImageDraw, ImageFont
+    images = []
+    for model_id in adapters.list_models():
+        if model_id is None:
+            model_id = 'None'
+        if shared.state.interrupted:
+            continue
+        output = image
+        if model_id != 'None':
+            adapter = adapters.Adapter(model_id=model_id, device=devices.device, dtype=devices.dtype)
+            if adapter is None:
+                shared.log.error(f'Adapter load failed: id="{model_id}"')
+                continue
+            shared.log.info(f'Testing Adapter: {model_id}')
+            pipe = adapters.AdapterPipeline(adapter=adapter.model, pipeline=shared.sd_model)
+            pipe.pipeline.to(device=devices.device, dtype=devices.dtype)
+            sd_models.set_diffuser_options(pipe)
+            try:
+                res = pipe.pipeline(prompt=prompt, negative_prompt=negative, image=image, num_inference_steps=10, output_type='pil')
+                output = res.images[0]
+            except Exception as e:
+                errors.display(e, f'Adapter {model_id} inference')
+                model_id = f'{model_id} error'
+            pipe.restore()
+        draw = ImageDraw.Draw(output)
+        font = ImageFont.truetype('DejaVuSansMono', 48)
+        draw.text((10, 10), model_id, (0,0,0), font=font)
+        draw.text((8, 8), model_id, (255,255,255), font=font)
+        images.append(output)
+        yield output, None, None, images
+    rows = round(math.sqrt(len(images)))
+    cols = math.ceil(len(images) / rows)
+    w, h = 256, 256
+    size = (cols * w + cols, rows * h + rows)
+    grid = Image.new('RGB', size=size, color='black')
+    shared.log.info(f'Test Adapters: images={len(images)} grid={grid}')
+    for i, image in enumerate(images):
+        x = (i % cols * w) + (i % cols)
+        y = (i // cols * h) + (i // cols)
+        thumb = image.copy().convert('RGB')
+        thumb.thumbnail((w, h), Image.Resampling.HAMMING)
+        grid.paste(thumb, box=(x, y))
+    yield None, grid, None, images
+    return None, grid, None, images # preview_process, output_image, output_video, output_gallery
@@ -477,10 +477,14 @@ def create_ui(_blocks: gr.Blocks=None):
                generation_parameters_copypaste.register_paste_params_button(bindings)

                if debug:
-                    from modules.control.test import test_processors
+                    from modules.control.test import test_processors, test_controlnets, test_adapters
                    gr.HTML('<br><h1>Debug</h1><br>')
                    with gr.Row():
-                        run_test_processors_btn = gr.Button(value="Test all processors", variant='primary', elem_classes=['control-button'])
+                        run_test_processors_btn = gr.Button(value="Test:Processors", variant='primary', elem_classes=['control-button'])
+                        run_test_controlnets_btn = gr.Button(value="Test:ControlNets", variant='primary', elem_classes=['control-button'])
+                        run_test_adapters_btn = gr.Button(value="Test:Adapters", variant='primary', elem_classes=['control-button'])
                        run_test_processors_btn.click(fn=test_processors, inputs=[input_image], outputs=[preview_process, output_image, output_video, output_gallery])
+                        run_test_controlnets_btn.click(fn=test_controlnets, inputs=[prompt, negative, input_image], outputs=[preview_process, output_image, output_video, output_gallery])
+                        run_test_adapters_btn.click(fn=test_adapters, inputs=[prompt, negative, input_image], outputs=[preview_process, output_image, output_video, output_gallery])

    return [(control_ui, 'Control', 'control')]