declare-lab/tango2

1import inspect2import logging3import os4import random5import re6import tempfile7import unittest8import urllib.parse9from distutils.util import strtobool10from io import BytesIO, StringIO11from pathlib import Path12from typing import List, Optional, Union13 14import numpy as np15import PIL.Image16import PIL.ImageOps17import requests18from packaging import version19 20from .import_utils import (21    BACKENDS_MAPPING,22    is_compel_available,23    is_flax_available,24    is_note_seq_available,25    is_onnx_available,26    is_opencv_available,27    is_torch_available,28    is_torch_version,29)30from .logging import get_logger31 32 33global_rng = random.Random()34 35logger = get_logger(__name__)36 37if is_torch_available():38    import torch39 40    if "DIFFUSERS_TEST_DEVICE" in os.environ:41        torch_device = os.environ["DIFFUSERS_TEST_DEVICE"]42 43        available_backends = ["cuda", "cpu", "mps"]44        if torch_device not in available_backends:45            raise ValueError(46                f"unknown torch backend for diffusers tests: {torch_device}. Available backends are:"47                f" {available_backends}"48            )49        logger.info(f"torch_device overrode to {torch_device}")50    else:51        torch_device = "cuda" if torch.cuda.is_available() else "cpu"52        is_torch_higher_equal_than_1_12 = version.parse(53            version.parse(torch.__version__).base_version54        ) >= version.parse("1.12")55 56        if is_torch_higher_equal_than_1_12:57            # Some builds of torch 1.12 don't have the mps backend registered. See #892 for more details58            mps_backend_registered = hasattr(torch.backends, "mps")59            torch_device = "mps" if (mps_backend_registered and torch.backends.mps.is_available()) else torch_device60 61 62def torch_all_close(a, b, *args, **kwargs):63    if not is_torch_available():64        raise ValueError("PyTorch needs to be installed to use this function.")65    if not torch.allclose(a, b, *args, **kwargs):66        assert False, f"Max diff is absolute {(a - b).abs().max()}. Diff tensor is {(a - b).abs()}."67    return True68 69 70def print_tensor_test(tensor, filename="test_corrections.txt", expected_tensor_name="expected_slice"):71    test_name = os.environ.get("PYTEST_CURRENT_TEST")72    if not torch.is_tensor(tensor):73        tensor = torch.from_numpy(tensor)74 75    tensor_str = str(tensor.detach().cpu().flatten().to(torch.float32)).replace("\n", "")76    # format is usually:77    # expected_slice = np.array([-0.5713, -0.3018, -0.9814, 0.04663, -0.879, 0.76, -1.734, 0.1044, 1.161])78    output_str = tensor_str.replace("tensor", f"{expected_tensor_name} = np.array")79    test_file, test_class, test_fn = test_name.split("::")80    test_fn = test_fn.split()[0]81    with open(filename, "a") as f:82        print(";".join([test_file, test_class, test_fn, output_str]), file=f)83 84 85def get_tests_dir(append_path=None):86    """87    Args:88        append_path: optional path to append to the tests dir path89    Return:90        The full path to the `tests` dir, so that the tests can be invoked from anywhere. Optionally `append_path` is91        joined after the `tests` dir the former is provided.92    """93    # this function caller's __file__94    caller__file__ = inspect.stack()[1][1]95    tests_dir = os.path.abspath(os.path.dirname(caller__file__))96 97    while not tests_dir.endswith("tests"):98        tests_dir = os.path.dirname(tests_dir)99 100    if append_path:101        return os.path.join(tests_dir, append_path)102    else:103        return tests_dir104 105 106def parse_flag_from_env(key, default=False):107    try:108        value = os.environ[key]109    except KeyError:110        # KEY isn't set, default to `default`.111        _value = default112    else:113        # KEY is set, convert it to True or False.114        try:115            _value = strtobool(value)116        except ValueError:117            # More values are supported, but let's keep the message simple.118            raise ValueError(f"If set, {key} must be yes or no.")119    return _value120 121 122_run_slow_tests = parse_flag_from_env("RUN_SLOW", default=False)123_run_nightly_tests = parse_flag_from_env("RUN_NIGHTLY", default=False)124 125 126def floats_tensor(shape, scale=1.0, rng=None, name=None):127    """Creates a random float32 tensor"""128    if rng is None:129        rng = global_rng130 131    total_dims = 1132    for dim in shape:133        total_dims *= dim134 135    values = []136    for _ in range(total_dims):137        values.append(rng.random() * scale)138 139    return torch.tensor(data=values, dtype=torch.float).view(shape).contiguous()140 141 142def slow(test_case):143    """144    Decorator marking a test as slow.145 146    Slow tests are skipped by default. Set the RUN_SLOW environment variable to a truthy value to run them.147 148    """149    return unittest.skipUnless(_run_slow_tests, "test is slow")(test_case)150 151 152def nightly(test_case):153    """154    Decorator marking a test that runs nightly in the diffusers CI.155 156    Slow tests are skipped by default. Set the RUN_NIGHTLY environment variable to a truthy value to run them.157 158    """159    return unittest.skipUnless(_run_nightly_tests, "test is nightly")(test_case)160 161 162def require_torch(test_case):163    """164    Decorator marking a test that requires PyTorch. These tests are skipped when PyTorch isn't installed.165    """166    return unittest.skipUnless(is_torch_available(), "test requires PyTorch")(test_case)167 168 169def require_torch_2(test_case):170    """171    Decorator marking a test that requires PyTorch 2. These tests are skipped when it isn't installed.172    """173    return unittest.skipUnless(is_torch_available() and is_torch_version(">=", "2.0.0"), "test requires PyTorch 2")(174        test_case175    )176 177 178def require_torch_gpu(test_case):179    """Decorator marking a test that requires CUDA and PyTorch."""180    return unittest.skipUnless(is_torch_available() and torch_device == "cuda", "test requires PyTorch+CUDA")(181        test_case182    )183 184 185def skip_mps(test_case):186    """Decorator marking a test to skip if torch_device is 'mps'"""187    return unittest.skipUnless(torch_device != "mps", "test requires non 'mps' device")(test_case)188 189 190def require_flax(test_case):191    """192    Decorator marking a test that requires JAX & Flax. These tests are skipped when one / both are not installed193    """194    return unittest.skipUnless(is_flax_available(), "test requires JAX & Flax")(test_case)195 196 197def require_compel(test_case):198    """199    Decorator marking a test that requires compel: https://github.com/damian0815/compel. These tests are skipped when200    the library is not installed.201    """202    return unittest.skipUnless(is_compel_available(), "test requires compel")(test_case)203 204 205def require_onnxruntime(test_case):206    """207    Decorator marking a test that requires onnxruntime. These tests are skipped when onnxruntime isn't installed.208    """209    return unittest.skipUnless(is_onnx_available(), "test requires onnxruntime")(test_case)210 211 212def require_note_seq(test_case):213    """214    Decorator marking a test that requires note_seq. These tests are skipped when note_seq isn't installed.215    """216    return unittest.skipUnless(is_note_seq_available(), "test requires note_seq")(test_case)217 218 219def load_numpy(arry: Union[str, np.ndarray], local_path: Optional[str] = None) -> np.ndarray:220    if isinstance(arry, str):221        # local_path = "/home/patrick_huggingface_co/"222        if local_path is not None:223            # local_path can be passed to correct images of tests224            return os.path.join(local_path, "/".join([arry.split("/")[-5], arry.split("/")[-2], arry.split("/")[-1]]))225        elif arry.startswith("http://") or arry.startswith("https://"):226            response = requests.get(arry)227            response.raise_for_status()228            arry = np.load(BytesIO(response.content))229        elif os.path.isfile(arry):230            arry = np.load(arry)231        else:232            raise ValueError(233                f"Incorrect path or url, URLs must start with `http://` or `https://`, and {arry} is not a valid path"234            )235    elif isinstance(arry, np.ndarray):236        pass237    else:238        raise ValueError(239            "Incorrect format used for numpy ndarray. Should be an url linking to an image, a local path, or a"240            " ndarray."241        )242 243    return arry244 245 246def load_pt(url: str):247    response = requests.get(url)248    response.raise_for_status()249    arry = torch.load(BytesIO(response.content))250    return arry251 252 253def load_image(image: Union[str, PIL.Image.Image]) -> PIL.Image.Image:254    """255    Args:256    Loads `image` to a PIL Image.257        image (`str` or `PIL.Image.Image`):258            The image to convert to the PIL Image format.259    Returns:260        `PIL.Image.Image`: A PIL Image.261    """262    if isinstance(image, str):263        if image.startswith("http://") or image.startswith("https://"):264            image = PIL.Image.open(requests.get(image, stream=True).raw)265        elif os.path.isfile(image):266            image = PIL.Image.open(image)267        else:268            raise ValueError(269                f"Incorrect path or url, URLs must start with `http://` or `https://`, and {image} is not a valid path"270            )271    elif isinstance(image, PIL.Image.Image):272        image = image273    else:274        raise ValueError(275            "Incorrect format used for image. Should be an url linking to an image, a local path, or a PIL image."276        )277    image = PIL.ImageOps.exif_transpose(image)278    image = image.convert("RGB")279    return image280 281 282def export_to_video(video_frames: List[np.ndarray], output_video_path: str = None) -> str:283    if is_opencv_available():284        import cv2285    else:286        raise ImportError(BACKENDS_MAPPING["opencv"][1].format("export_to_video"))287    if output_video_path is None:288        output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4").name289 290    fourcc = cv2.VideoWriter_fourcc(*"mp4v")291    h, w, c = video_frames[0].shape292    video_writer = cv2.VideoWriter(output_video_path, fourcc, fps=8, frameSize=(w, h))293    for i in range(len(video_frames)):294        img = cv2.cvtColor(video_frames[i], cv2.COLOR_RGB2BGR)295        video_writer.write(img)296    return output_video_path297 298 299def load_hf_numpy(path) -> np.ndarray:300    if not path.startswith("http://") or path.startswith("https://"):301        path = os.path.join(302            "https://huggingface.co/datasets/fusing/diffusers-testing/resolve/main", urllib.parse.quote(path)303        )304 305    return load_numpy(path)306 307 308# --- pytest conf functions --- #309 310# to avoid multiple invocation from tests/conftest.py and examples/conftest.py - make sure it's called only once311pytest_opt_registered = {}312 313 314def pytest_addoption_shared(parser):315    """316    This function is to be called from `conftest.py` via `pytest_addoption` wrapper that has to be defined there.317 318    It allows loading both `conftest.py` files at once without causing a failure due to adding the same `pytest`319    option.320 321    """322    option = "--make-reports"323    if option not in pytest_opt_registered:324        parser.addoption(325            option,326            action="store",327            default=False,328            help="generate report files. The value of this option is used as a prefix to report names",329        )330        pytest_opt_registered[option] = 1331 332 333def pytest_terminal_summary_main(tr, id):334    """335    Generate multiple reports at the end of test suite run - each report goes into a dedicated file in the current336    directory. The report files are prefixed with the test suite name.337 338    This function emulates --duration and -rA pytest arguments.339 340    This function is to be called from `conftest.py` via `pytest_terminal_summary` wrapper that has to be defined341    there.342 343    Args:344    - tr: `terminalreporter` passed from `conftest.py`345    - id: unique id like `tests` or `examples` that will be incorporated into the final reports filenames - this is346      needed as some jobs have multiple runs of pytest, so we can't have them overwrite each other.347 348    NB: this functions taps into a private _pytest API and while unlikely, it could break should349    pytest do internal changes - also it calls default internal methods of terminalreporter which350    can be hijacked by various `pytest-` plugins and interfere.351 352    """353    from _pytest.config import create_terminal_writer354 355    if not len(id):356        id = "tests"357 358    config = tr.config359    orig_writer = config.get_terminal_writer()360    orig_tbstyle = config.option.tbstyle361    orig_reportchars = tr.reportchars362 363    dir = "reports"364    Path(dir).mkdir(parents=True, exist_ok=True)365    report_files = {366        k: f"{dir}/{id}_{k}.txt"367        for k in [368            "durations",369            "errors",370            "failures_long",371            "failures_short",372            "failures_line",373            "passes",374            "stats",375            "summary_short",376            "warnings",377        ]378    }379 380    # custom durations report381    # note: there is no need to call pytest --durations=XX to get this separate report382    # adapted from https://github.com/pytest-dev/pytest/blob/897f151e/src/_pytest/runner.py#L66383    dlist = []384    for replist in tr.stats.values():385        for rep in replist:386            if hasattr(rep, "duration"):387                dlist.append(rep)388    if dlist:389        dlist.sort(key=lambda x: x.duration, reverse=True)390        with open(report_files["durations"], "w") as f:391            durations_min = 0.05  # sec392            f.write("slowest durations\n")393            for i, rep in enumerate(dlist):394                if rep.duration < durations_min:395                    f.write(f"{len(dlist)-i} durations < {durations_min} secs were omitted")396                    break397                f.write(f"{rep.duration:02.2f}s {rep.when:<8} {rep.nodeid}\n")398 399    def summary_failures_short(tr):400        # expecting that the reports were --tb=long (default) so we chop them off here to the last frame401        reports = tr.getreports("failed")402        if not reports:403            return404        tr.write_sep("=", "FAILURES SHORT STACK")405        for rep in reports:406            msg = tr._getfailureheadline(rep)407            tr.write_sep("_", msg, red=True, bold=True)408            # chop off the optional leading extra frames, leaving only the last one409            longrepr = re.sub(r".*_ _ _ (_ ){10,}_ _ ", "", rep.longreprtext, 0, re.M | re.S)410            tr._tw.line(longrepr)411            # note: not printing out any rep.sections to keep the report short412 413    # use ready-made report funcs, we are just hijacking the filehandle to log to a dedicated file each414    # adapted from https://github.com/pytest-dev/pytest/blob/897f151e/src/_pytest/terminal.py#L814415    # note: some pytest plugins may interfere by hijacking the default `terminalreporter` (e.g.416    # pytest-instafail does that)417 418    # report failures with line/short/long styles419    config.option.tbstyle = "auto"  # full tb420    with open(report_files["failures_long"], "w") as f:421        tr._tw = create_terminal_writer(config, f)422        tr.summary_failures()423 424    # config.option.tbstyle = "short" # short tb425    with open(report_files["failures_short"], "w") as f:426        tr._tw = create_terminal_writer(config, f)427        summary_failures_short(tr)428 429    config.option.tbstyle = "line"  # one line per error430    with open(report_files["failures_line"], "w") as f:431        tr._tw = create_terminal_writer(config, f)432        tr.summary_failures()433 434    with open(report_files["errors"], "w") as f:435        tr._tw = create_terminal_writer(config, f)436        tr.summary_errors()437 438    with open(report_files["warnings"], "w") as f:439        tr._tw = create_terminal_writer(config, f)440        tr.summary_warnings()  # normal warnings441        tr.summary_warnings()  # final warnings442 443    tr.reportchars = "wPpsxXEf"  # emulate -rA (used in summary_passes() and short_test_summary())444    with open(report_files["passes"], "w") as f:445        tr._tw = create_terminal_writer(config, f)446        tr.summary_passes()447 448    with open(report_files["summary_short"], "w") as f:449        tr._tw = create_terminal_writer(config, f)450        tr.short_test_summary()451 452    with open(report_files["stats"], "w") as f:453        tr._tw = create_terminal_writer(config, f)454        tr.summary_stats()455 456    # restore:457    tr._tw = orig_writer458    tr.reportchars = orig_reportchars459    config.option.tbstyle = orig_tbstyle460 461 462class CaptureLogger:463    """464    Args:465    Context manager to capture `logging` streams466        logger: 'logging` logger object467    Returns:468        The captured output is available via `self.out`469    Example:470    ```python471    >>> from diffusers import logging472    >>> from diffusers.testing_utils import CaptureLogger473 474    >>> msg = "Testing 1, 2, 3"475    >>> logging.set_verbosity_info()476    >>> logger = logging.get_logger("diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.py")477    >>> with CaptureLogger(logger) as cl:478    ...     logger.info(msg)479    >>> assert cl.out, msg + "\n"480    ```481    """482 483    def __init__(self, logger):484        self.logger = logger485        self.io = StringIO()486        self.sh = logging.StreamHandler(self.io)487        self.out = ""488 489    def __enter__(self):490        self.logger.addHandler(self.sh)491        return self492 493    def __exit__(self, *exc):494        self.logger.removeHandler(self.sh)495        self.out = self.io.getvalue()496 497    def __repr__(self):498        return f"captured: {self.out}\n"499