code/tests/test_inference_smoke.py

"""Optional end-to-end GPU smoke tests for the maintained demo cases.

Required environment variables:

- MING_GENERATION_MODEL
- MING_LAYER_MODEL

The model variables may be local directories or Hugging Face Hub IDs.

Optional environment variables:

- MING_SMOKE_OUTPUT_DIR: keep outputs under this directory (created when
  missing) instead of a temporary directory.
- MING_SMOKE_MODE: "short" (default) runs the two-step probes only;
  "full" additionally runs the default-parameter acceptance cases.
- MING_SMOKE_DEVICE, MING_SMOKE_DEVICE_MAP, MING_SMOKE_NUM_GPUS,
  MING_SMOKE_DTYPE and MING_SMOKE_ATTN_IMPLEMENTATION are forwarded to
  infer.py. The smoke defaults to the validated single-GPU layout in BF16
  and FlashAttention 2; set MING_SMOKE_NUM_GPUS (e.g. 8) for a multi-GPU
  run.

Text-to-image and layer decomposition use the same assets displayed in the
public README. The image-edit compatibility probe keeps its established input
and prompt.
"""

import json
import os
from pathlib import Path
import subprocess
import sys
import tempfile
import unittest


REPOSITORY = Path(__file__).resolve().parents[1]
INFER = REPOSITORY / "infer.py"
T2I_PROMPT = REPOSITORY / "assets" / "t2i_four_seasons_cabin_prompt.json"
IMAGE_EDIT_INPUT = REPOSITORY / "tests" / "assets" / "smoke_input.png"
IMAGE_EDIT_PROMPT = "Change the background to blue"
LAYER_INPUT = REPOSITORY / "assets" / "layer_samples" / "card_making_input.png"
LAYER_PROMPT = REPOSITORY / "assets" / "layer_samples" / "card_making_prompt.txt"
LAYER_COUNT = 6


class SmokeAssetContractTest(unittest.TestCase):
    def test_smoke_assets_are_present_and_well_formed(self):
        for path in (T2I_PROMPT, IMAGE_EDIT_INPUT, LAYER_INPUT, LAYER_PROMPT):
            self.assertTrue(path.is_file(), path)

        with T2I_PROMPT.open(encoding="utf-8") as handle:
            structured_prompt = json.load(handle)
        self.assertEqual(set(structured_prompt), {"canvas_settings", "layers"})

        layer_prompt = LAYER_PROMPT.read_text(encoding="utf-8")
        self.assertIn(f"Number of layers: {LAYER_COUNT}", layer_prompt)


def _smoke_output_root():
    value = os.environ.get("MING_SMOKE_OUTPUT_DIR")
    return Path(value).expanduser().resolve() if value else None


@unittest.skipUnless(
    os.environ.get("MING_GENERATION_MODEL")
    and os.environ.get("MING_LAYER_MODEL"),
    "set MING_GENERATION_MODEL and MING_LAYER_MODEL",
)
class InferenceSmokeTest(unittest.TestCase):
    def _run(self, *arguments, steps=2, resolution=None):
        command = [
            sys.executable,
            str(INFER),
            *map(str, arguments),
            "--seed",
            "42",
            "--device",
            os.environ.get("MING_SMOKE_DEVICE", "cuda:0"),
            "--device-map",
            os.environ.get("MING_SMOKE_DEVICE_MAP", "balanced"),
            "--num-gpus",
            os.environ.get("MING_SMOKE_NUM_GPUS", "1"),
            "--dtype",
            os.environ.get("MING_SMOKE_DTYPE", "bfloat16"),
            "--attn-implementation",
            os.environ.get("MING_SMOKE_ATTN_IMPLEMENTATION", "flash_attention_2"),
        ]
        if resolution is not None:
            command += ["--resolution", str(resolution)]
        if steps is not None:
            command += ["--steps", str(steps)]
        subprocess.run(command, cwd=REPOSITORY, check=True)

    def _output_root(self, test_case_dir):
        configured = _smoke_output_root()
        if configured is not None:
            return configured / test_case_dir
        return Path(self._temporary_directory.name) / test_case_dir

    def setUp(self):
        self._temporary_directory = (
            tempfile.TemporaryDirectory() if _smoke_output_root() is None else None
        )
        self.generation_model = os.environ["MING_GENERATION_MODEL"]
        self.layer_model = os.environ["MING_LAYER_MODEL"]
        for path in (T2I_PROMPT, IMAGE_EDIT_INPUT, LAYER_INPUT, LAYER_PROMPT):
            self.assertTrue(path.is_file(), path)

    def tearDown(self):
        if self._temporary_directory is not None:
            self._temporary_directory.cleanup()

    def _assert_png_set(self, directory, count, mode=None):
        from PIL import Image

        paths = sorted(Path(directory).glob("*.png"))
        self.assertEqual(len(paths), count, f"{directory}: {paths}")
        pixels = []
        for path in paths:
            with Image.open(path) as image:
                if mode is not None:
                    self.assertEqual(image.mode, mode, path)
                sample = image.convert("RGBA")
                pixels.append(sample.tobytes())
                extrema = sample.getextrema()
                for channel, (low, high) in enumerate(extrema):
                    self.assertGreater(
                        high, low, f"{path} channel {channel} is constant"
                    )
        return paths, pixels

    def _assert_showcase_layers(self, directory):
        from PIL import Image

        paths, pixels = self._assert_png_set(directory, LAYER_COUNT, mode="RGBA")
        self.assertEqual(
            len(set(pixels)),
            LAYER_COUNT,
            "showcase layer output must contain six distinct images",
        )
        alpha_ranges = []
        for path in paths:
            with Image.open(path) as image:
                alpha_ranges.append(image.getchannel("A").getextrema())
        self.assertTrue(
            any(low < high for low, high in alpha_ranges),
            f"no alpha variation in showcase layer outputs: {alpha_ranges}",
        )

    def _run_showcase_cases(self, root, *, steps, t2i_resolution, layer_resolution):
        self._run(
            "--model", self.generation_model,
            "--task", "text-to-image",
            "--prompt", T2I_PROMPT,
            "--output-dir", root / "text-to-image",
            steps=steps,
            resolution=t2i_resolution,
        )
        self._run(
            "--model", self.generation_model,
            "--task", "image-edit",
            "--input-image", IMAGE_EDIT_INPUT,
            "--prompt", IMAGE_EDIT_PROMPT,
            "--output-dir", root / "image-edit",
            steps=steps,
            resolution=1024,
        )
        self._run(
            "--model", self.layer_model,
            "--task", "layer-decompose",
            "--input-image", LAYER_INPUT,
            "--prompt", LAYER_PROMPT,
            "--output-dir", root / "layer-decompose",
            steps=steps,
            resolution=layer_resolution,
        )

        self._assert_png_set(root / "text-to-image", 1)
        self._assert_png_set(root / "image-edit", 1)
        self._assert_showcase_layers(root / "layer-decompose")

    def test_two_step_showcase_smoke(self):
        """Two-step probes for both showcase cases and image-edit regression."""
        self._run_showcase_cases(
            self._output_root("short"),
            steps=2,
            t2i_resolution=1024,
            layer_resolution=512,
        )

    @unittest.skipUnless(
        os.environ.get("MING_SMOKE_MODE") == "full",
        "set MING_SMOKE_MODE=full to run default-parameter acceptance",
    )
    def test_default_parameter_acceptance(self):
        """Default steps, CFG and recommended showcase resolutions; seed 42."""
        self._run_showcase_cases(
            self._output_root("default"),
            steps=None,
            t2i_resolution=2048,
            layer_resolution=1024,
        )


if __name__ == "__main__":
    unittest.main()