From 341e55fb266790b7e2391c26759093b6045ce0b5 Mon Sep 17 00:00:00 2001 From: DN6 Date: Wed, 26 Aug 2026 10:34:36 +0530 Subject: [PATCH 1/4] [CI] Rewrite `tests_fetcher.py` as an AST-based import-graph selector with feature-split matrix output --- .github/workflows/pr_test_fetcher.yml | 32 +- utils/tests_fetcher.py | 1387 ++++++++----------------- 2 files changed, 424 insertions(+), 995 deletions(-) diff --git a/.github/workflows/pr_test_fetcher.yml b/.github/workflows/pr_test_fetcher.yml index 17789ec8a9cd..aef711f01cb2 100644 --- a/.github/workflows/pr_test_fetcher.yml +++ b/.github/workflows/pr_test_fetcher.yml @@ -28,7 +28,6 @@ jobs: shell: bash outputs: matrix: ${{ steps.set_matrix.outputs.matrix }} - test_map: ${{ steps.set_matrix.outputs.test_map }} steps: - name: Checkout diffusers uses: actions/checkout@v6 @@ -51,31 +50,26 @@ jobs: path: test_preparation.txt - id: set_matrix name: Create Test Matrix - # The `keys` is used as GitHub actions matrix for jobs, i.e. `models`, `pipelines`, etc. - # The `test_map` is used to get the actual identified test files under each key. - # If no test to run (so no `test_map.json` file), create a dummy map (empty matrix will fail) + # `test_map.json` is a list of matrix entries `{"name", "paths", "markers"}`, one job each. + # If no test to run (so no `test_map.json` file), emit an empty list and skip the test job. run: | if [ -f test_map.json ]; then - keys=$(python3 -c 'import json; fp = open("test_map.json"); test_map = json.load(fp); fp.close(); d = list(test_map.keys()); print(json.dumps(d))') - test_map=$(python3 -c 'import json; fp = open("test_map.json"); test_map = json.load(fp); fp.close(); print(json.dumps(test_map))') + matrix=$(python3 -c 'import json; print(json.dumps(json.load(open("test_map.json"))))') else - keys=$(python3 -c 'keys = ["dummy"]; print(keys)') - test_map=$(python3 -c 'test_map = {"dummy": []}; print(test_map)') + matrix='[]' fi - echo $keys - echo $test_map - echo "matrix=$keys" >> $GITHUB_OUTPUT - echo "test_map=$test_map" >> $GITHUB_OUTPUT + echo "matrix=$matrix" >> $GITHUB_OUTPUT + echo "Matrix: $matrix" run_pr_tests: - name: Run PR Tests + name: ${{ matrix.name }} needs: setup_pr_tests - if: contains(fromJson(needs.setup_pr_tests.outputs.matrix), 'dummy') != true + if: needs.setup_pr_tests.outputs.matrix != '[]' strategy: fail-fast: false max-parallel: 2 matrix: - modules: ${{ fromJson(needs.setup_pr_tests.outputs.matrix) }} + include: ${{ fromJson(needs.setup_pr_tests.outputs.matrix) }} runs-on: group: aws-general-8-plus container: @@ -101,20 +95,20 @@ jobs: - name: Run all selected tests on CPU run: | - pytest -n 2 --dist=loadfile -v --make-reports=${{ matrix.modules }}_tests_cpu ${{ fromJson(needs.setup_pr_tests.outputs.test_map)[matrix.modules] }} + pytest -n 2 --dist=loadfile -v -m "${{ matrix.markers }}" --make-reports=${{ matrix.name }}_tests_cpu ${{ matrix.paths }} - name: Failure short reports if: ${{ failure() }} continue-on-error: true run: | - cat reports/${{ matrix.modules }}_tests_cpu_stats.txt - cat reports/${{ matrix.modules }}_tests_cpu_failures_short.txt + cat reports/${{ matrix.name }}_tests_cpu_stats.txt + cat reports/${{ matrix.name }}_tests_cpu_failures_short.txt - name: Test suite reports artifacts if: ${{ always() }} uses: actions/upload-artifact@v6 with: - name: ${{ matrix.modules }}_test_reports + name: ${{ matrix.name }}_test_reports path: reports run_staging_tests: diff --git a/utils/tests_fetcher.py b/utils/tests_fetcher.py index d487efb40518..8afa2994fa44 100644 --- a/utils/tests_fetcher.py +++ b/utils/tests_fetcher.py @@ -14,1115 +14,550 @@ # limitations under the License. """ -Welcome to tests_fetcher V2. - -This util is designed to fetch tests to run on a PR so that only the tests impacted by the modifications are run, and -when too many models are being impacted, only run the tests of a subset of core models. It works like this. - -Stage 1: Identify the modified files. For jobs that run on the main branch, it's just the diff with the last commit. -On a PR, this takes all the files from the branching point to the current commit (so all modifications in a PR, not -just the last commit) but excludes modifications that are on docstrings or comments only. - -Stage 2: Extract the tests to run. This is done by looking at the imports in each module and test file: if module A -imports module B, then changing module B impacts module A, so the tests using module A should be run. We thus get the -dependencies of each model and then recursively builds the 'reverse' map of dependencies to get all modules and tests -impacted by a given file. We then only keep the tests (and only the core models tests if there are too many modules). - -Caveats: - - This module only filters tests by files (not individual tests) so it's better to have tests for different things - in different files. - - This module assumes inits are just importing things, not really building objects, so it's better to structure - them this way and move objects building in separate submodules. +Diffusers tests_fetcher (graph-based). + +For each PR, walk the AST of every modified Python file to extract diffusers-internal imports, build a +forward dependency graph for the repo, invert it to a reverse map (file → tests transitively depending on +it), and select the impacted tests. + +There is no automatic full-suite trigger. If a change is in territory the import graph can't see +correctly (dynamic dispatch via auto-mappings, lazy `_import_structure`, etc.), apply the `run-all-tests` +PR label or pass `--force_full_suite` to bypass selection. + +Pipeline-specific note: diffusers' `__init__.py` files use the `_import_structure = {...}` lazy-loading +pattern paired with an `if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:` block containing real +`from .submodule import Class` statements. AST extraction sees the TYPE_CHECKING imports (since +`ast.walk` descends into `If` / `Try` blocks regardless of runtime conditions), so the import graph +mirrors the actual public API. + +Stage 1 — diff: list modified Python files (vs. merge-base with main, or the previous commit on main). + Docstring/comment-only changes are filtered out by content comparison. +Stage 2 — graph: parse every .py under `src/diffusers/` and `tests/` with `ast`, build the forward + dependency map, transitively close it, then invert to the reverse map. +Stage 3 — select: for each modified file, look up `reverse_map[file]` to get impacted tests. +Stage 4 — bucket: group tests by top-level `tests/` folder for the CI matrix. `tests/models` and + `tests/pipelines` are split further into one job per feature mixin (via the pytest markers the + mixins carry) so a large selection fans out instead of serialising in one job. Usage: -Base use to fetch the tests in a pull request - ```bash -python utils/tests_fetcher.py -``` - -Base use to fetch the tests on a the main branch (with diff from the last commit): - -```bash -python utils/tests_fetcher.py --diff_with_last_commit +python utils/tests_fetcher.py # PR mode: diff against main +python utils/tests_fetcher.py --diff_with_last_commit # main mode: diff against last commit +python utils/tests_fetcher.py --force_full_suite # bypass selection, run everything ``` """ import argparse +import ast import collections import json -import os -import re from contextlib import contextmanager from pathlib import Path -from typing import Dict, List, Tuple, Union +from typing import Dict, List, Optional, Tuple from git import Repo PATH_TO_REPO = Path(__file__).parent.parent.resolve() -PATH_TO_EXAMPLES = PATH_TO_REPO / "examples" PATH_TO_DIFFUSERS = PATH_TO_REPO / "src/diffusers" PATH_TO_TESTS = PATH_TO_REPO / "tests" -# Ignore fixtures in tests folder -# Ignore lora since they are always tested -MODULES_TO_IGNORE = ["fixtures", "lora"] - -IMPORTANT_PIPELINES = [ - "controlnet", - "stable_diffusion", - "stable_diffusion_2", - "stable_diffusion_xl", - "stable_video_diffusion", - "deepfloyd_if", - "kandinsky", - "kandinsky2_2", - "text_to_video_synthesis", - "wuerstchen", -] +# `tests/models` and `tests/pipelines` compose one test class per feature mixin, and each mixin carries a +# pytest marker (see the `is_*` decorators in tests/testing_utils.py). Each group below becomes its own +# matrix job, selected with `pytest -m `, but only when a selected file composes a mixin carrying +# one of the group's markers; `core` is everything no group and no CPU-skipped marker claims. +SPLIT_BY_FEATURE = {"models", "pipelines"} +FEATURE_GROUPS = { + "lora": ["lora"], + "attention": ["attention"], + "memory": ["memory", "cpu_offload", "group_offload"], + "cache": ["cache"], + "ip_adapter": ["ip_adapter"], +} +# Gated on an accelerator, multi-GPU, or an HF token, so every test skips on the CPU runner. Excluded +# from `core` and given no job rather than spinning up a runner to skip everything. +CPU_SKIPPED_MARKERS = ["quantization", "compile", "single_file", "training", "context_parallel", "tensor_parallel"] +FEATURE_MARKERS = [m for markers in FEATURE_GROUPS.values() for m in markers] +CORE_MARKERS = "not (" + " or ".join(FEATURE_MARKERS + CPU_SKIPPED_MARKERS) + ")" +# ============================================================ +# Generic helpers +# ============================================================ @contextmanager def checkout_commit(repo: Repo, commit_id: str): - """ - Context manager that checks out a given commit when entered, but gets back to the reference it was at on exit. - - Args: - repo (`git.Repo`): A git repository (for instance the Transformers repo). - commit_id (`str`): The commit reference to checkout inside the context manager. - """ + """Check out `commit_id` for the duration of the block, restoring the prior HEAD on exit.""" current_head = repo.head.commit if repo.head.is_detached else repo.head.ref - try: repo.git.checkout(commit_id) yield - finally: repo.git.checkout(current_head) -def clean_code(content: str) -> str: - """ - Remove docstrings, empty line or comments from some code (used to detect if a diff is real or only concern - comments or docstrings). - - Args: - content (`str`): The code to clean - - Returns: - `str`: The cleaned code. - """ - # We need to deactivate autoformatting here to write escaped triple quotes (we cannot use real triple quotes or - # this would mess up the result if this function applied to this particular file). - # fmt: off - # Remove docstrings by splitting on triple " then triple ': - splits = content.split('\"\"\"') - content = "".join(splits[::2]) - splits = content.split("\'\'\'") - # fmt: on - content = "".join(splits[::2]) - - # Remove empty lines and comments - lines_to_keep = [] - for line in content.split("\n"): - # remove anything that is after a # sign. - line = re.sub("#.*$", "", line) - # remove white lines - if len(line) != 0 and not line.isspace(): - lines_to_keep.append(line) - return "\n".join(lines_to_keep) - - -def keep_doc_examples_only(content: str) -> str: - """ - Remove everything from the code content except the doc examples (used to determined if a diff should trigger doc - tests or not). - - Args: - content (`str`): The code to clean - - Returns: - `str`: The cleaned code. - """ - # Keep doc examples only by splitting on triple "`" - splits = content.split("```") - # Add leading and trailing "```" so the navigation is easier when compared to the original input `content` - content = "```" + "```".join(splits[1::2]) + "```" - - # Remove empty lines and comments - lines_to_keep = [] - for line in content.split("\n"): - # remove anything that is after a # sign. - line = re.sub("#.*$", "", line) - # remove white lines - if len(line) != 0 and not line.isspace(): - lines_to_keep.append(line) - return "\n".join(lines_to_keep) - - -def get_all_tests() -> List[str]: - """ - Walks the `tests` folder to return a list of files/subfolders. This is used to split the tests to run when using - parallelism. The split is: +# ============================================================ +# Diff detection +# ============================================================ - - folders under `tests`: (`tokenization`, `pipelines`, etc) except the subfolder `models` is excluded. - - folders under `tests/models`: `bert`, `gpt2`, etc. - - test files under `tests`: `test_modeling_common.py`, `test_tokenization_common.py`, etc. - """ - - # test folders/files directly under `tests` folder - tests = os.listdir(PATH_TO_TESTS) - tests = [f"tests/{f}" for f in tests if "__pycache__" not in f] - tests = sorted([f for f in tests if (PATH_TO_REPO / f).is_dir() or f.startswith("tests/test_")]) - - return tests - - -def diff_is_docstring_only(repo: Repo, branching_point: str, filename: str) -> bool: - """ - Check if the diff is only in docstrings (or comments and whitespace) in a filename. - - Args: - repo (`git.Repo`): A git repository (for instance the Transformers repo). - branching_point (`str`): The commit reference of where to compare for the diff. - filename (`str`): The filename where we want to know if the diff isonly in docstrings/comments. - - Returns: - `bool`: Whether the diff is docstring/comments only or not. - """ - folder = Path(repo.working_dir) - with checkout_commit(repo, branching_point): - with open(folder / filename, "r", encoding="utf-8") as f: - old_content = f.read() - - with open(folder / filename, "r", encoding="utf-8") as f: - new_content = f.read() - - old_content_clean = clean_code(old_content) - new_content_clean = clean_code(new_content) - - return old_content_clean == new_content_clean - - -def diff_contains_doc_examples(repo: Repo, branching_point: str, filename: str) -> bool: - """ - Check if the diff is only in code examples of the doc in a filename. - Args: - repo (`git.Repo`): A git repository (for instance the Transformers repo). - branching_point (`str`): The commit reference of where to compare for the diff. - filename (`str`): The filename where we want to know if the diff is only in codes examples. +def _strip_comments_and_docstrings(source: str) -> str: + """Return source with all docstrings and comments removed via AST round-trip. - Returns: - `bool`: Whether the diff is only in code examples of the doc or not. + Used by `diff_is_docstring_only` to detect diffs that are purely cosmetic. """ - folder = Path(repo.working_dir) + try: + tree = ast.parse(source) + except SyntaxError: + return source + + for node in ast.walk(tree): + # Strip module/class/function docstrings. + if isinstance(node, (ast.Module, ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + body = node.body + if ( + body + and isinstance(body[0], ast.Expr) + and isinstance(body[0].value, ast.Constant) + and isinstance(body[0].value.value, str) + ): + body.pop(0) + return ast.unparse(tree) + + +def diff_is_docstring_only(repo: Repo, branching_point, filename: str) -> bool: + """True if the diff in `filename` between `branching_point` and HEAD only changes docstrings/comments.""" with checkout_commit(repo, branching_point): - with open(folder / filename, "r", encoding="utf-8") as f: - old_content = f.read() - - with open(folder / filename, "r", encoding="utf-8") as f: - new_content = f.read() + old_content = (PATH_TO_REPO / filename).read_text(encoding="utf-8") + new_content = (PATH_TO_REPO / filename).read_text(encoding="utf-8") + return _strip_comments_and_docstrings(old_content) == _strip_comments_and_docstrings(new_content) - old_content_clean = keep_doc_examples_only(old_content) - new_content_clean = keep_doc_examples_only(new_content) - return old_content_clean != new_content_clean - - -def get_diff(repo: Repo, base_commit: str, commits: List[str]) -> List[str]: - """ - Get the diff between a base commit and one or several commits. - - Args: - repo (`git.Repo`): - A git repository (for instance the Transformers repo). - base_commit (`str`): - The commit reference of where to compare for the diff. This is the current commit, not the branching point! - commits (`List[str]`): - The list of commits with which to compare the repo at `base_commit` (so the branching point). - - Returns: - `List[str]`: The list of Python files with a diff (files added, renamed or deleted are always returned, files - modified are returned if the diff in the file is not only in docstrings or comments, see - `diff_is_docstring_only`). - """ - print("\n### DIFF ###\n") +def get_diff(repo: Repo, base_commit, commits) -> List[str]: + """Return Python files changed between `commits` (branching point) and `base_commit` (HEAD).""" code_diff = [] for commit in commits: - for diff_obj in commit.diff(base_commit): - # We always add new python files - if diff_obj.change_type == "A" and diff_obj.b_path.endswith(".py"): - code_diff.append(diff_obj.b_path) - # We check that deleted python files won't break corresponding tests. - elif diff_obj.change_type == "D" and diff_obj.a_path.endswith(".py"): - code_diff.append(diff_obj.a_path) - # Now for modified files - elif diff_obj.change_type in ["M", "R"] and diff_obj.b_path.endswith(".py"): - # In case of renames, we'll look at the tests using both the old and new name. - if diff_obj.a_path != diff_obj.b_path: - code_diff.extend([diff_obj.a_path, diff_obj.b_path]) - else: - # Otherwise, we check modifications are in code and not docstrings. - if diff_is_docstring_only(repo, commit, diff_obj.b_path): - print(f"Ignoring diff in {diff_obj.b_path} as it only concerns docstrings or comments.") - else: - code_diff.append(diff_obj.a_path) - + for d in commit.diff(base_commit): + paths = [p for p in (d.a_path, d.b_path) if p and p.endswith(".py")] + if not paths: + continue + # Add/delete/rename: keep every changed path verbatim. Pure modification: skip if the diff is + # docstring/comment-only. + if d.change_type in ("A", "D") or d.a_path != d.b_path: + code_diff.extend(paths) + elif not diff_is_docstring_only(repo, commit, d.b_path): + code_diff.append(d.b_path) return code_diff def get_modified_python_files(diff_with_last_commit: bool = False) -> List[str]: - """ - Return a list of python files that have been modified between: - - - the current head and the main branch if `diff_with_last_commit=False` (default) - - the current head and its parent commit otherwise. - - Returns: - `List[str]`: The list of Python files with a diff (files added, renamed or deleted are always returned, files - modified are returned if the diff in the file is not only in docstrings or comments, see - `diff_is_docstring_only`). - """ + """List Python files modified between HEAD and either main (default) or the previous commit.""" repo = Repo(PATH_TO_REPO) - - if not diff_with_last_commit: - # Need to fetch refs for main using remotes when running with github actions. + if diff_with_last_commit: + base_label = "previous commit" + commits = repo.head.commit.parents + else: upstream_main = repo.remotes.origin.refs.main + base_label = f"merge-base with main ({upstream_main.commit})" + commits = repo.merge_base(upstream_main, repo.head) + print(f"Diffing HEAD ({repo.head.commit}) against {base_label}: {[str(c) for c in commits]}") + return get_diff(repo, repo.head.commit, commits) - print(f"main is at {upstream_main.commit}") - print(f"Current head is at {repo.head.commit}") - branching_commits = repo.merge_base(upstream_main, repo.head) - for commit in branching_commits: - print(f"Branching commit: {commit}") - return get_diff(repo, repo.head.commit, branching_commits) - else: - print(f"main is at {repo.head.commit}") - parent_commits = repo.head.commit.parents - for commit in parent_commits: - print(f"Parent commit: {commit}") - return get_diff(repo, repo.head.commit, parent_commits) +def get_all_tests() -> List[str]: + """Top-level entries under `tests/` (folders + `tests/test_*.py`), used to expand a full-suite selection.""" + return sorted( + f"tests/{p.name}" + for p in PATH_TO_TESTS.iterdir() + if "__pycache__" not in p.name and (p.is_dir() or p.name.startswith("test_")) + ) -def get_diff_for_doctesting(repo: Repo, base_commit: str, commits: List[str]) -> List[str]: - """ - Get the diff in doc examples between a base commit and one or several commits. +# ============================================================ +# AST-based import extraction +# ============================================================ - Args: - repo (`git.Repo`): - A git repository (for instance the Transformers repo). - base_commit (`str`): - The commit reference of where to compare for the diff. This is the current commit, not the branching point! - commits (`List[str]`): - The list of commits with which to compare the repo at `base_commit` (so the branching point). - Returns: - `List[str]`: The list of Python and Markdown files with a diff (files added or renamed are always returned, files - modified are returned if the diff in the file is only in doctest examples). - """ - print("\n### DIFF ###\n") - code_diff = [] - for commit in commits: - for diff_obj in commit.diff(base_commit): - # We only consider Python files and doc files. - if not diff_obj.b_path.endswith(".py") and not diff_obj.b_path.endswith(".md"): - continue - # We always add new python/md files - if diff_obj.change_type in ["A"]: - code_diff.append(diff_obj.b_path) - # Now for modified files - elif diff_obj.change_type in ["M", "R"]: - # In case of renames, we'll look at the tests using both the old and new name. - if diff_obj.a_path != diff_obj.b_path: - code_diff.extend([diff_obj.a_path, diff_obj.b_path]) - else: - # Otherwise, we check modifications contain some doc example(s). - if diff_contains_doc_examples(repo, commit, diff_obj.b_path): - code_diff.append(diff_obj.a_path) - else: - print(f"Ignoring diff in {diff_obj.b_path} as it doesn't contain any doc example.") +def _resolve_import(module: Optional[str], level: int, importer_pkg: List[str]) -> Optional[List[str]]: + """Resolve an `ImportFrom` node to a list of repo-rooted path parts. - return code_diff + Args: + module: the `X.Y` part of `from X.Y import Z` (None for `from . import Z`). + level: number of leading dots (0 for absolute, 1+ for relative). + importer_pkg: parts of the importing module's *package* (parent dir parts), e.g. + `["src", "diffusers", "pipelines", "flux"]` for `pipelines/flux/pipeline_flux.py`. + Returns: + Path parts like `["src", "diffusers", "pipelines", "flux", "pipeline_flux"]` (no extension), + or None if the import is external or can't be resolved. + """ + if level == 0: + if module is None or not (module == "diffusers" or module.startswith("diffusers.")): + return None + sub = module.split(".")[1:] + return ["src", "diffusers", *sub] + + if level > len(importer_pkg): + return None + base = importer_pkg[: len(importer_pkg) - level + 1] + if module: + return [*base, *module.split(".")] + return base + + +def _to_module_file(path_parts: List[str]) -> Optional[str]: + """Resolve `path_parts` to either `.py` or `/__init__.py`. Returns repo-relative path.""" + candidate = PATH_TO_REPO.joinpath(*path_parts).with_suffix(".py") + if candidate.is_file(): + return str(candidate.relative_to(PATH_TO_REPO)) + init = PATH_TO_REPO.joinpath(*path_parts) / "__init__.py" + if init.is_file(): + return str(init.relative_to(PATH_TO_REPO)) + return None + + +def _iter_module_level_imports(node): + """Yield `ImportFrom` nodes that execute at module load. + + Recurses into `If` / `Try` / `ClassDef` bodies (those run at import time) but stops at + `FunctionDef` / `AsyncFunctionDef` / `Lambda` boundaries — imports inside function bodies are + deferred runtime imports (e.g. lazy stubs in deprecation shims) and shouldn't count as dependencies. + """ + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)): + return + if isinstance(node, ast.ImportFrom): + yield node + for child in ast.iter_child_nodes(node): + yield from _iter_module_level_imports(child) -def get_all_doctest_files() -> List[str]: - """ - Return the complete list of python and Markdown files on which we run doctest. - At this moment, we restrict this to only take files from `src/` or `docs/source/en/` that are not in `utils/not_doctested.txt`. +def _extract_imports(module_file: str) -> List[Tuple[str, List[str]]]: + """Parse `module_file` and return [(target_file, [imported_symbols]), ...] for diffusers-internal imports. - Returns: - `List[str]`: The complete list of Python and Markdown files on which we run doctest. + Only module-level `from X import ...` statements are considered. Bare `import X`, `from X import *`, + external imports (transformers, torch, stdlib), and imports inside function bodies are skipped. """ - py_files = [str(x.relative_to(PATH_TO_REPO)) for x in PATH_TO_REPO.glob("**/*.py")] - md_files = [str(x.relative_to(PATH_TO_REPO)) for x in PATH_TO_REPO.glob("**/*.md")] - test_files_to_run = py_files + md_files - - # only include files in `src` or `docs/source/en/` - test_files_to_run = [x for x in test_files_to_run if x.startswith(("src/", "docs/source/en/"))] - # not include init files - test_files_to_run = [x for x in test_files_to_run if not x.endswith(("__init__.py",))] + abs_path = PATH_TO_REPO / module_file + try: + source = abs_path.read_text(encoding="utf-8") + tree = ast.parse(source) + except (SyntaxError, UnicodeDecodeError, FileNotFoundError): + return [] - # These are files not doctested yet. - with open("utils/not_doctested.txt") as fp: - not_doctested = {x.split(" ")[0] for x in fp.read().strip().split("\n")} + importer_pkg = list(Path(module_file).parts[:-1]) - # So far we don't have 100% coverage for doctest. This line will be removed once we achieve 100%. - test_files_to_run = [x for x in test_files_to_run if x not in not_doctested] + results: List[Tuple[str, List[str]]] = [] + for node in _iter_module_level_imports(tree): + names = [alias.name for alias in node.names if alias.name != "*"] + if not names: + continue + target_parts = _resolve_import(node.module, node.level, importer_pkg) + if target_parts is None: + continue + target_file = _to_module_file(target_parts) + if target_file is None: + continue + results.append((target_file, names)) - return sorted(test_files_to_run) + return results -def get_new_doctest_files(repo, base_commit, branching_commit) -> List[str]: - """ - Get the list of files that were removed from "utils/not_doctested.txt", between `base_commit` and - `branching_commit`. +# ============================================================ +# Dependency graph +# ============================================================ - Returns: - `List[str]`: List of files that were removed from "utils/not_doctested.txt". - """ - for diff_obj in branching_commit.diff(base_commit): - # Ignores all but the "utils/not_doctested.txt" file. - if diff_obj.a_path != "utils/not_doctested.txt": - continue - # Loads the two versions - folder = Path(repo.working_dir) - with checkout_commit(repo, branching_commit): - with open(folder / "utils/not_doctested.txt", "r", encoding="utf-8") as f: - old_content = f.read() - with open(folder / "utils/not_doctested.txt", "r", encoding="utf-8") as f: - new_content = f.read() - # Compute the removed lines and return them - removed_content = {x.split(" ")[0] for x in old_content.split("\n")} - { - x.split(" ")[0] for x in new_content.split("\n") - } - return sorted(removed_content) - return [] - - -def get_doctest_files(diff_with_last_commit: bool = False) -> List[str]: - """ - Return a list of python and Markdown files where doc example have been modified between: - - the current head and the main branch if `diff_with_last_commit=False` (default) - - the current head and its parent commit otherwise. +def get_module_dependencies(module_file: str, cache: Dict[str, List[Tuple[str, List[str]]]]) -> List[str]: + """Return source files `module_file` truly depends on, traversing inits to find the defining file. - Returns: - `List[str]`: The list of Python and Markdown files with a diff (files added or renamed are always returned, files - modified are returned if the diff in the file is only in doctest examples). + When an import lands on an `__init__.py`, walk its imports too, matching by symbol name to find the + actual submodule that re-exports each requested symbol. This collapses + `from diffusers import StableDiffusionPipeline` to `pipelines/stable_diffusion/pipeline_stable_diffusion.py` + instead of the root init (which would over-select to almost every test). """ - repo = Repo(PATH_TO_REPO) + if module_file not in cache: + cache[module_file] = _extract_imports(module_file) - test_files_to_run = [] # noqa - if not diff_with_last_commit: - upstream_main = repo.remotes.origin.refs.main - print(f"main is at {upstream_main.commit}") - print(f"Current head is at {repo.head.commit}") + dependencies: List[str] = [] + queue: List[Tuple[str, List[str]]] = list(cache[module_file]) + seen_inits: set = set() - branching_commits = repo.merge_base(upstream_main, repo.head) - for commit in branching_commits: - print(f"Branching commit: {commit}") - test_files_to_run = get_diff_for_doctesting(repo, repo.head.commit, branching_commits) - else: - print(f"main is at {repo.head.commit}") - parent_commits = repo.head.commit.parents - for commit in parent_commits: - print(f"Parent commit: {commit}") - test_files_to_run = get_diff_for_doctesting(repo, repo.head.commit, parent_commits) - - all_test_files_to_run = get_all_doctest_files() - - # Add to the test files to run any removed entry from "utils/not_doctested.txt". - new_test_files = get_new_doctest_files(repo, repo.head.commit, upstream_main.commit) - test_files_to_run = list(set(test_files_to_run + new_test_files)) - - # Do not run slow doctest tests on CircleCI - with open("utils/slow_documentation_tests.txt") as fp: - slow_documentation_tests = set(fp.read().strip().split("\n")) - test_files_to_run = [ - x for x in test_files_to_run if x in all_test_files_to_run and x not in slow_documentation_tests - ] + while queue: + target, symbols = queue.pop(0) - # Make sure we did not end up with a test file that was removed - test_files_to_run = [f for f in test_files_to_run if (PATH_TO_REPO / f).exists()] - - return sorted(test_files_to_run) - - -# (:?^|\n) -> Non-catching group for the beginning of the doc or a new line. -# \s*from\s+(\.+\S+)\s+import\s+([^\n]+) -> Line only contains from .xxx import yyy and we catch .xxx and yyy -# (?=\n) -> Look-ahead to a new line. We can't just put \n here or using find_all on this re will only catch every -# other import. -_re_single_line_relative_imports = re.compile(r"(?:^|\n)\s*from\s+(\.+\S+)\s+import\s+([^\n]+)(?=\n)") -# (:?^|\n) -> Non-catching group for the beginning of the doc or a new line. -# \s*from\s+(\.+\S+)\s+import\s+\(([^\)]+)\) -> Line continues with from .xxx import (yyy) and we catch .xxx and yyy -# yyy will take multiple lines otherwise there wouldn't be parenthesis. -_re_multi_line_relative_imports = re.compile(r"(?:^|\n)\s*from\s+(\.+\S+)\s+import\s+\(([^\)]+)\)") -# (:?^|\n) -> Non-catching group for the beginning of the doc or a new line. -# \s*from\s+transformers(\S*)\s+import\s+([^\n]+) -> Line only contains from transformers.xxx import yyy and we catch -# .xxx and yyy -# (?=\n) -> Look-ahead to a new line. We can't just put \n here or using find_all on this re will only catch every -# other import. -_re_single_line_direct_imports = re.compile(r"(?:^|\n)\s*from\s+diffusers(\S*)\s+import\s+([^\n]+)(?=\n)") -# (:?^|\n) -> Non-catching group for the beginning of the doc or a new line. -# \s*from\s+transformers(\S*)\s+import\s+\(([^\)]+)\) -> Line continues with from transformers.xxx import (yyy) and we -# catch .xxx and yyy. yyy will take multiple lines otherwise there wouldn't be parenthesis. -_re_multi_line_direct_imports = re.compile(r"(?:^|\n)\s*from\s+diffusers(\S*)\s+import\s+\(([^\)]+)\)") - - -def extract_imports(module_fname: str, cache: Dict[str, List[str]] = None) -> List[str]: - """ - Get the imports a given module makes. + if not target.endswith("__init__.py"): + dependencies.append(target) + continue - Args: - module_fname (`str`): - The name of the file of the module where we want to look at the imports (given relative to the root of - the repo). - cache (Dictionary `str` to `List[str]`, *optional*): - To speed up this function if it was previously called on `module_fname`, the cache of all previously - computed results. + # Avoid cycles through inits importing each other. + if target in seen_inits: + dependencies.append(target) + continue + seen_inits.add(target) - Returns: - `List[str]`: The list of module filenames imported in the input `module_fname` (a submodule we import from that - is a subfolder will give its init file). - """ - if cache is not None and module_fname in cache: - return cache[module_fname] - - with open(PATH_TO_REPO / module_fname, "r", encoding="utf-8") as f: - content = f.read() - - # Filter out all docstrings to not get imports in code examples. As before we need to deactivate formatting to - # keep this as escaped quotes and avoid this function failing on this file. - # fmt: off - splits = content.split('\"\"\"') - # fmt: on - content = "".join(splits[::2]) - - module_parts = str(module_fname).split(os.path.sep) - imported_modules = [] - - # Let's start with relative imports - relative_imports = _re_single_line_relative_imports.findall(content) - relative_imports = [ - (mod, imp) for mod, imp in relative_imports if "# tests_ignore" not in imp and imp.strip() != "(" - ] - multiline_relative_imports = _re_multi_line_relative_imports.findall(content) - relative_imports += [(mod, imp) for mod, imp in multiline_relative_imports if "# tests_ignore" not in imp] - - # We need to remove parts of the module name depending on the depth of the relative imports. - for module, imports in relative_imports: - level = 0 - while module.startswith("."): - module = module[1:] - level += 1 - - if len(module) > 0: - dep_parts = module_parts[: len(module_parts) - level] + module.split(".") - else: - dep_parts = module_parts[: len(module_parts) - level] - imported_module = os.path.sep.join(dep_parts) - imported_modules.append((imported_module, [imp.strip() for imp in imports.split(",")])) - - # Let's continue with direct imports - direct_imports = _re_single_line_direct_imports.findall(content) - direct_imports = [(mod, imp) for mod, imp in direct_imports if "# tests_ignore" not in imp and imp.strip() != "("] - multiline_direct_imports = _re_multi_line_direct_imports.findall(content) - direct_imports += [(mod, imp) for mod, imp in multiline_direct_imports if "# tests_ignore" not in imp] - - # We need to find the relative path of those imports. - for module, imports in direct_imports: - import_parts = module.split(".")[1:] # ignore the name of the repo since we add it below. - dep_parts = ["src", "diffusers"] + import_parts - imported_module = os.path.sep.join(dep_parts) - imported_modules.append((imported_module, [imp.strip() for imp in imports.split(",")])) - - result = [] - # Double check we get proper modules (either a python file or a folder with an init). - for module_file, imports in imported_modules: - if (PATH_TO_REPO / f"{module_file}.py").is_file(): - module_file = f"{module_file}.py" - elif (PATH_TO_REPO / module_file).is_dir() and (PATH_TO_REPO / module_file / "__init__.py").is_file(): - module_file = os.path.sep.join([module_file, "__init__.py"]) - imports = [imp for imp in imports if len(imp) > 0 and re.match("^[A-Za-z0-9_]*$", imp)] - if len(imports) > 0: - result.append((module_file, imports)) - - if cache is not None: - cache[module_fname] = result - - return result - - -def get_module_dependencies(module_fname: str, cache: Dict[str, List[str]] = None) -> List[str]: - """ - Refines the result of `extract_imports` to remove subfolders and get a proper list of module filenames: if a file - as an import `from utils import Foo, Bar`, with `utils` being a subfolder containing many files, this will traverse - the `utils` init file to check where those dependencies come from: for instance the files utils/foo.py and utils/bar.py. + if target not in cache: + cache[target] = _extract_imports(target) + init_imports = cache[target] - Warning: This presupposes that all intermediate inits are properly built (with imports from the respective - submodules) and work better if objects are defined in submodules and not the intermediate init (otherwise the - intermediate init is added, and inits usually have a lot of dependencies). + unresolved = list(symbols) + for sub_target, sub_names in init_imports: + matched = [s for s in unresolved if s in sub_names] + if matched: + queue.append((sub_target, matched)) + unresolved = [s for s in unresolved if s not in matched] - Args: - module_fname (`str`): - The name of the file of the module where we want to look at the imports (given relative to the root of - the repo). - cache (Dictionary `str` to `List[str]`, *optional*): - To speed up this function if it was previously called on `module_fname`, the cache of all previously - computed results. + if unresolved: + # Symbol(s) couldn't be resolved through the init's TYPE_CHECKING imports — likely lazy-loaded + # via `_import_structure` or defined directly in the init. Keep the init as the dep (coarse but + # correct: changes to the init will trigger this module). + dependencies.append(target) - Returns: - `List[str]`: The list of module filenames imported in the input `module_fname` (with submodule imports refined). - """ - dependencies = [] - imported_modules = extract_imports(module_fname, cache=cache) - # The while loop is to recursively traverse all inits we may encounter: we will add things as we go. - while len(imported_modules) > 0: - new_modules = [] - for module, imports in imported_modules: - # If we end up in an __init__ we are often not actually importing from this init (except in the case where - # the object is fully defined in the __init__) - if module.endswith("__init__.py"): - # So we get the imports from that init then try to find where our objects come from. - new_imported_modules = extract_imports(module, cache=cache) - for new_module, new_imports in new_imported_modules: - if any(i in new_imports for i in imports): - if new_module not in dependencies: - new_modules.append((new_module, [i for i in new_imports if i in imports])) - imports = [i for i in imports if i not in new_imports] - if len(imports) > 0: - # If there are any objects lefts, they may be a submodule - path_to_module = PATH_TO_REPO / module.replace("__init__.py", "") - dependencies.extend( - [ - os.path.join(module.replace("__init__.py", ""), f"{i}.py") - for i in imports - if (path_to_module / f"{i}.py").is_file() - ] - ) - imports = [i for i in imports if not (path_to_module / f"{i}.py").is_file()] - if len(imports) > 0: - # Then if there are still objects left, they are fully defined in the init, so we keep it as a - # dependency. - dependencies.append(module) - else: - dependencies.append(module) - - imported_modules = new_modules - - return dependencies - - -def create_reverse_dependency_tree() -> List[Tuple[str, str]]: - """ - Create a list of all edges (a, b) which mean that modifying a impacts b with a going over all module and test files. - """ - cache = {} - all_modules = list(PATH_TO_DIFFUSERS.glob("**/*.py")) + list(PATH_TO_TESTS.glob("**/*.py")) - all_modules = [str(mod.relative_to(PATH_TO_REPO)) for mod in all_modules] - edges = [(dep, mod) for mod in all_modules for dep in get_module_dependencies(mod, cache=cache)] + return list(set(dependencies)) - return list(set(edges)) +def _merged_nested_deps(m: str, direct_deps: Dict[str, List[str]]) -> bool: + """Pull each of m's deps' deps into m. Returns True if m grew. -def get_tree_starting_at(module: str, edges: List[Tuple[str, str]]) -> List[Union[str, List[str]]]: + Skips `__init__.py` targets — they re-export the entire package surface, so expanding through + them would pull in every diffusers symbol via the root init. """ - Returns the tree starting at a given module following all edges. - - Args: - module (`str`): The module that will be the root of the subtree we want. - edges (`List[Tuple[str, str]]`): The list of all edges of the tree. + merged = False + for d in list(direct_deps[m]): + if d.endswith("__init__.py"): + continue + new_deps = set(direct_deps[d]) - set(direct_deps[m]) + if new_deps: + direct_deps[m].extend(new_deps) + merged = True + return merged - Returns: - `List[Union[str, List[str]]]`: The tree to print in the following format: [module, [list of edges - starting at module], [list of edges starting at the preceding level], ...] - """ - vertices_seen = [module] - new_edges = [edge for edge in edges if edge[0] == module and edge[1] != module and "__init__.py" not in edge[1]] - tree = [module] - while len(new_edges) > 0: - tree.append(new_edges) - final_vertices = list({edge[1] for edge in new_edges}) - vertices_seen.extend(final_vertices) - new_edges = [ - edge - for edge in edges - if edge[0] in final_vertices and edge[1] not in vertices_seen and "__init__.py" not in edge[1] - ] - - return tree - - -def print_tree_deps_of(module, all_edges=None): - """ - Prints the tree of modules depending on a given module. - Args: - module (`str`): The module that will be the root of the subtree we want. - all_edges (`List[Tuple[str, str]]`, *optional*): - The list of all edges of the tree. Will be set to `create_reverse_dependency_tree()` if not passed. - """ - if all_edges is None: - all_edges = create_reverse_dependency_tree() - tree = get_tree_starting_at(module, all_edges) - - # The list of lines is a list of tuples (line_to_be_printed, module) - # Keeping the modules lets us know where to insert each new lines in the list. - lines = [(tree[0], tree[0])] - for index in range(1, len(tree)): - edges = tree[index] - start_edges = {edge[0] for edge in edges} - - for start in start_edges: - end_edges = {edge[1] for edge in edges if edge[0] == start} - # We will insert all those edges just after the line showing start. - pos = 0 - while lines[pos][1] != start: - pos += 1 - lines = lines[: pos + 1] + [(" " * (2 * index) + end, end) for end in end_edges] + lines[pos + 1 :] - - for line in lines: - # We don't print the refs that where just here to help build lines. - print(line[0]) - - -def init_test_examples_dependencies() -> Tuple[Dict[str, List[str]], List[str]]: - """ - The test examples do not import from the examples (which are just scripts, not modules) so we need some extra - care initializing the dependency map, which is the goal of this function. It initializes the dependency map for - example files by linking each example to the example test file for the example framework. +def create_reverse_dependency_map() -> Dict[str, List[str]]: + """Build the reverse dependency map: file → list of files that transitively depend on it. - Returns: - `Tuple[Dict[str, List[str]], List[str]]`: A tuple with two elements: the initialized dependency map which is a - dict test example file to list of example files potentially tested by that test file, and the list of all - example files (to avoid recomputing it later). + 1. Compute direct deps for every .py under `src/diffusers/` and `tests/`. + 2. Transitively close (skipping inits during recursion to avoid pulling in the universe via the root init). + 3. Invert. """ - test_example_deps = {} - all_examples = [] - for framework in ["pytorch", "tensorflow"]: - test_files = list((PATH_TO_EXAMPLES / framework).glob("test_*.py")) - all_examples.extend(test_files) - # Remove the files at the root of examples/framework since they are not proper examples (they are either utils - # or example test files). - examples = [ - f for f in (PATH_TO_EXAMPLES / framework).glob("**/*.py") if f.parent != PATH_TO_EXAMPLES / framework - ] - all_examples.extend(examples) - for test_file in test_files: - with open(test_file, "r", encoding="utf-8") as f: - content = f.read() - # Map all examples to the test files found in examples/framework. - test_example_deps[str(test_file.relative_to(PATH_TO_REPO))] = [ - str(e.relative_to(PATH_TO_REPO)) for e in examples if e.name in content - ] - # Also map the test files to themselves. - test_example_deps[str(test_file.relative_to(PATH_TO_REPO))].append( - str(test_file.relative_to(PATH_TO_REPO)) - ) - return test_example_deps, all_examples - - -def create_reverse_dependency_map() -> dict[str, List[str]]: - """ - Create the dependency map from module/test filename to the list of modules/tests that depend on it recursively. + cache: Dict[str, List[Tuple[str, List[str]]]] = {} + all_modules = [ + str(p.relative_to(PATH_TO_REPO)) + for p in list(PATH_TO_DIFFUSERS.glob("**/*.py")) + list(PATH_TO_TESTS.glob("**/*.py")) + ] + direct_deps: Dict[str, List[str]] = {m: get_module_dependencies(m, cache) for m in all_modules} - Returns: - `Dict[str, List[str]]`: The reverse dependency map as a dictionary mapping filenames to all the filenames - depending on it recursively. This way the tests impacted by a change in file A are the test files in the list - corresponding to key A in this result. - """ - cache = {} - # Start from the example deps init. - example_deps, examples = init_test_examples_dependencies() - # Add all modules and all tests to all examples - all_modules = list(PATH_TO_DIFFUSERS.glob("**/*.py")) + list(PATH_TO_TESTS.glob("**/*.py")) + examples - all_modules = [str(mod.relative_to(PATH_TO_REPO)) for mod in all_modules] - # Compute the direct dependencies of all modules. - direct_deps = {m: get_module_dependencies(m, cache=cache) for m in all_modules} - direct_deps.update(example_deps) - - # This recurses the dependencies - something_changed = True - while something_changed: - something_changed = False + # Each pass propagates dependency info one level deeper. Loop until a full pass adds nothing. + changed = True + while changed: + changed = False for m in all_modules: - for d in direct_deps[m]: - # We stop recursing at an init (cause we always end up in the main init and we don't want to add all - # files which the main init imports) - if d.endswith("__init__.py"): - continue - if d not in direct_deps: - raise ValueError(f"KeyError:{d}. From {m}") - new_deps = set(direct_deps[d]) - set(direct_deps[m]) - if len(new_deps) > 0: - direct_deps[m].extend(list(new_deps)) - something_changed = True - - # Finally we can build the reverse map. - reverse_map = collections.defaultdict(list) + if _merged_nested_deps(m, direct_deps): + changed = True + + reverse_map: Dict[str, List[str]] = collections.defaultdict(list) for m in all_modules: for d in direct_deps[m]: reverse_map[d].append(m) - # For inits, we don't do the reverse deps but the direct deps: if modifying an init, we want to make sure we test - # all the modules impacted by that init. - for m in [f for f in all_modules if f.endswith("__init__.py")]: - direct_deps = get_module_dependencies(m, cache=cache) - deps = sum([reverse_map[d] for d in direct_deps if not d.endswith("__init__.py")], direct_deps) - reverse_map[m] = list(set(deps) - {m}) - - return reverse_map - - -def create_module_to_test_map(reverse_map: Dict[str, List[str]] = None) -> dict[str, List[str]]: - """ - Extract the tests from the reverse_dependency_map and potentially filters the model tests. - - Args: - reverse_map (`Dict[str, List[str]]`, *optional*): - The reverse dependency map as created by `create_reverse_dependency_map`. Will default to the result of - that function if not provided. - filter_pipelines (`bool`, *optional*, defaults to `False`): - Whether or not to filter pipeline tests to only include core pipelines if a file impacts a lot of models. - - Returns: - `Dict[str, List[str]]`: A dictionary that maps each file to the tests to execute if that file was modified. - """ - if reverse_map is None: - reverse_map = create_reverse_dependency_map() - - # Utility that tells us if a given file is a test (taking test examples into account) - def is_test(fname): - if fname.startswith("tests"): - return True - if fname.startswith("examples") and fname.split(os.path.sep)[-1].startswith("test"): - return True - return False - - # Build the test map - test_map = {module: [f for f in deps if is_test(f)] for module, deps in reverse_map.items()} - - return test_map - - -def check_imports_all_exist(): - """ - Isn't used per se by the test fetcher but might be used later as a quality check. Putting this here for now so the - code is not lost. This checks all imports in a given file do exist. - """ - cache = {} - all_modules = list(PATH_TO_DIFFUSERS.glob("**/*.py")) + list(PATH_TO_TESTS.glob("**/*.py")) - all_modules = [str(mod.relative_to(PATH_TO_REPO)) for mod in all_modules] - direct_deps = {m: get_module_dependencies(m, cache=cache) for m in all_modules} - - for module, deps in direct_deps.items(): - for dep in deps: - if not (PATH_TO_REPO / dep).is_file(): - print(f"{module} has dependency on {dep} which does not exist.") - - -def _print_list(l) -> str: - """ - Pretty print a list of elements with one line per element and a - starting each line. - """ - return "\n".join([f"- {f}" for f in l]) - + # For inits, do the forward direction: editing an init impacts everything it re-exports. + for init in [m for m in all_modules if m.endswith("__init__.py")]: + deps = get_module_dependencies(init, cache) + impacted = set(deps) + for d in deps: + if not d.endswith("__init__.py"): + impacted.update(reverse_map.get(d, [])) + reverse_map[init] = sorted(impacted - {init}) -def update_test_map_with_core_pipelines(json_output_file: str): - print(f"\n### ADD CORE PIPELINE TESTS ###\n{_print_list(IMPORTANT_PIPELINES)}") - with open(json_output_file, "rb") as fp: - test_map = json.load(fp) + return dict(reverse_map) - # Add core pipelines as their own test group - test_map["core_pipelines"] = " ".join( - sorted([str(PATH_TO_TESTS / f"pipelines/{pipe}") for pipe in IMPORTANT_PIPELINES]) - ) - # If there are no existing pipeline tests save the map - if "pipelines" not in test_map: - with open(json_output_file, "w", encoding="UTF-8") as fp: - json.dump(test_map, fp, ensure_ascii=False) +# ============================================================ +# Test selection +# ============================================================ - pipeline_tests = test_map.pop("pipelines") - pipeline_tests = pipeline_tests.split(" ") - # Remove core pipeline tests from the fetched pipeline tests - updated_pipeline_tests = [] - for pipe in pipeline_tests: - if pipe == "tests/pipelines" or Path(pipe).parts[2] in IMPORTANT_PIPELINES: +def _bucket_for_matrix(test_paths: List[str]) -> Dict[str, List[str]]: + """Group test paths by top-level folder under `tests/`. Files directly under `tests/` go to `common`.""" + test_map: Dict[str, List[str]] = collections.defaultdict(list) + for p in test_paths: + parts = p.split("/") + if len(parts) < 2 or parts[0] != "tests": continue - updated_pipeline_tests.append(pipe) + bucket = "common" if len(parts) == 2 else parts[1] + test_map[bucket].append(p) + return {k: sorted(set(v)) for k, v in test_map.items()} - if len(updated_pipeline_tests) > 0: - test_map["pipelines"] = " ".join(sorted(updated_pipeline_tests)) - with open(json_output_file, "w", encoding="UTF-8") as fp: - json.dump(test_map, fp, ensure_ascii=False) - - -def create_json_map(test_files_to_run: List[str], json_output_file: str | None = None): - """ - Creates a map from a list of tests to run to easily split them by category, when running parallelism of slow tests. +def _base_names(node: ast.ClassDef) -> List[str]: + return [ + b.id if isinstance(b, ast.Name) else b.attr for b in node.bases if isinstance(b, (ast.Name, ast.Attribute)) + ] - Args: - test_files_to_run (`List[str]`): The list of tests to run. - json_output_file (`str`): The path where to store the built json map. - """ - if json_output_file is None: - return - test_map = {} - for test_file in test_files_to_run: - # `test_file` is a path to a test folder/file, starting with `tests/`. For example, - # - `tests/models/bert/test_modeling_bert.py` or `tests/models/bert` - # - `tests/trainer/test_trainer.py` or `tests/trainer` - # - `tests/test_modeling_common.py` - names = test_file.split(os.path.sep) - module = names[1] - if module in MODULES_TO_IGNORE: +def _marker_decorators() -> Dict[str, str]: + """Map decorator function name → pytest marker, from `def is_x(test_case): return pytest.mark.(test_case)`.""" + tree = ast.parse((PATH_TO_TESTS / "testing_utils.py").read_text(encoding="utf-8")) + decorators = {} + for node in ast.walk(tree): + if not isinstance(node, ast.FunctionDef) or not node.name.startswith("is_"): continue + for ret in ast.walk(node): + if isinstance(ret, ast.Return) and isinstance(ret.value, ast.Call): + func = ret.value.func + if ( + isinstance(func, ast.Attribute) + and isinstance(func.value, ast.Attribute) + and func.value.attr == "mark" + ): + decorators[node.name] = func.attr + return decorators + + +def _mixin_markers(bucket: str, decorators: Dict[str, str]) -> Dict[str, List[str]]: + """Map mixin class name → markers it carries, from the `@is_*` decorators on classes in + `tests//testing_utils/`. Markers are inherited through pytest, so a mixin also carries those of + any base mixin. Scoped per bucket because class names repeat across packages (the pipelines' legacy, + unmarked `IPAdapterTesterMixin` vs. the models' marked one).""" + bases: Dict[str, List[str]] = {} + markers: Dict[str, set] = {} + for module in (PATH_TO_TESTS / bucket / "testing_utils").glob("*.py"): + for node in ast.walk(ast.parse(module.read_text(encoding="utf-8"))): + if not isinstance(node, ast.ClassDef): + continue + bases[node.name] = _base_names(node) + markers[node.name] = { + decorators[d.id] for d in node.decorator_list if isinstance(d, ast.Name) and d.id in decorators + } + + changed = True + while changed: + changed = False + for cls, cls_bases in bases.items(): + inherited = set().union(*(markers.get(b, set()) for b in cls_bases)) + if not inherited <= markers[cls]: + markers[cls] |= inherited + changed = True + return {cls: sorted(m) for cls, m in markers.items() if m} + + +def _class_markers(test_file: str, mixin_markers: Dict[str, List[str]]) -> List[set]: + """Markers of every test class in `test_file`, derived from the marked mixins it composes.""" + tree = ast.parse((PATH_TO_REPO / test_file).read_text(encoding="utf-8")) + return [ + {m for base in _base_names(node) for m in mixin_markers.get(base, [])} + for node in ast.walk(tree) + if isinstance(node, ast.ClassDef) + ] - if len(names) > 2 or not test_file.endswith(".py"): - # test folders under `tests` or python files under them - # take the part like tokenization, `pipeline`, etc. for other test categories - key = os.path.sep.join(names[1:2]) - else: - # common test files directly under `tests/` - key = "common" - if key not in test_map: - test_map[key] = [] - test_map[key].append(test_file) +def _feature_groups_for(paths: List[str], mixin_markers: Dict[str, List[str]]) -> List[str]: + """Names of the feature groups (including `core`) that at least one test class in `paths` would land in.""" + class_markers = [m for p in paths for m in _class_markers(p, mixin_markers)] + non_core = set(FEATURE_MARKERS + CPU_SKIPPED_MARKERS) + groups = [] + if any(not m & non_core for m in class_markers): + groups.append("core") + for group, markers in FEATURE_GROUPS.items(): + if any(m & set(markers) for m in class_markers): + groups.append(group) + return groups + + +def _matrix_entries(test_map: Dict[str, List[str]]) -> List[Dict[str, str]]: + """Expand buckets into CI matrix entries: `{"name", "paths", "markers"}`, one job each.""" + decorators = _marker_decorators() + entries = [] + for bucket, paths in sorted(test_map.items()): + joined = " ".join(paths) + if bucket not in SPLIT_BY_FEATURE: + entries.append({"name": bucket, "paths": joined, "markers": ""}) + continue + for group in _feature_groups_for(paths, _mixin_markers(bucket, decorators)): + expr = CORE_MARKERS if group == "core" else " or ".join(FEATURE_GROUPS[group]) + entries.append({"name": f"{bucket}-{group}", "paths": joined, "markers": expr}) + return entries - # sort the keys & values - keys = sorted(test_map.keys()) - test_map = {k: " ".join(sorted(test_map[k])) for k in keys} +def _write_matrix(json_output_file: str, test_files: List[str]): with open(json_output_file, "w", encoding="UTF-8") as fp: - json.dump(test_map, fp, ensure_ascii=False) + json.dump(_matrix_entries(_bucket_for_matrix(test_files)), fp, ensure_ascii=False) -def infer_tests_to_run( - output_file: str, - diff_with_last_commit: bool = False, - json_output_file: str | None = None, -): - """ - The main function called by the test fetcher. Determines the tests to run from the diff. +def _is_test_file(path: str) -> bool: + """True if `path` is a `tests/.../test_*.py` file (the kind pytest collects).""" + return path.startswith("tests/") and Path(path).name.startswith("test_") - Args: - output_file (`str`): - The path where to store the summary of the test fetcher analysis. Other files will be stored in the same - folder: - - - examples_test_list.txt: The list of examples tests to run. - - test_repo_utils.txt: Will indicate if the repo utils tests should be run or not. - - doctest_list.txt: The list of doctests to run. - - diff_with_last_commit (`bool`, *optional*, defaults to `False`): - Whether to analyze the diff with the last commit (for use on the main branch after a PR is merged) or with - the branching point from main (for use on each PR). - filter_models (`bool`, *optional*, defaults to `True`): - Whether or not to filter the tests to core models only, when a file modified results in a lot of model - tests. - json_output_file (`str`, *optional*): - The path where to store the json file mapping categories of tests to tests to run (used for parallelism or - the slow tests). - """ + +def fetch_tests_to_run(json_output_file: str, diff_with_last_commit: bool): + """Determine the tests to run from the diff and write `test_map.json`.""" modified_files = get_modified_python_files(diff_with_last_commit=diff_with_last_commit) - print(f"\n### MODIFIED FILES ###\n{_print_list(modified_files)}") - # Create the map that will give us all impacted modules. reverse_map = create_reverse_dependency_map() - impacted_files = modified_files.copy() - for f in modified_files: - if f in reverse_map: - impacted_files.extend(reverse_map[f]) - # Remove duplicates - impacted_files = sorted(set(impacted_files)) - print(f"\n### IMPACTED FILES ###\n{_print_list(impacted_files)}") - - # Grab the corresponding test files: - if any(x in modified_files for x in ["setup.py"]): - test_files_to_run = ["tests", "examples"] - - # in order to trigger pipeline tests even if no code change at all - if "tests/utils/tiny_model_summary.json" in modified_files: - test_files_to_run = ["tests"] - any(f.split(os.path.sep)[0] == "utils" for f in modified_files) - else: - # All modified tests need to be run. - test_files_to_run = [ - f for f in modified_files if f.startswith("tests") and f.split(os.path.sep)[-1].startswith("test") - ] - # Then we grab the corresponding test files. - test_map = create_module_to_test_map(reverse_map=reverse_map) - for f in modified_files: - if f in test_map: - test_files_to_run.extend(test_map[f]) - test_files_to_run = sorted(set(test_files_to_run)) - # Make sure we did not end up with a test file that was removed - test_files_to_run = [f for f in test_files_to_run if (PATH_TO_REPO / f).exists()] - - any(f.split(os.path.sep)[0] == "utils" for f in modified_files) - - examples_tests_to_run = [f for f in test_files_to_run if f.startswith("examples")] - test_files_to_run = [f for f in test_files_to_run if not f.startswith("examples")] - print(f"\n### TEST TO RUN ###\n{_print_list(test_files_to_run)}") - if len(test_files_to_run) > 0: - with open(output_file, "w", encoding="utf-8") as f: - f.write(" ".join(test_files_to_run)) - - # Create a map that maps test categories to test files, i.e. `models/bert` -> [...test_modeling_bert.py, ...] - - # Get all test directories (and some common test files) under `tests` and `tests/models` if `test_files_to_run` - # contains `tests` (i.e. when `setup.py` is changed). - if "tests" in test_files_to_run: - test_files_to_run = get_all_tests() - - create_json_map(test_files_to_run, json_output_file) - - print(f"\n### EXAMPLES TEST TO RUN ###\n{_print_list(examples_tests_to_run)}") - if len(examples_tests_to_run) > 0: - # We use `all` in the case `commit_flags["test_all"]` as well as in `create_circleci_config.py` for processing - if examples_tests_to_run == ["examples"]: - examples_tests_to_run = ["all"] - example_file = Path(output_file).parent / "examples_test_list.txt" - with open(example_file, "w", encoding="utf-8") as f: - f.write(" ".join(examples_tests_to_run)) - - -def filter_tests(output_file: str, filters: List[str]): - """ - Reads the content of the output file and filters out all the tests in a list of given folders. - - Args: - output_file (`str` or `os.PathLike`): The path to the output file of the tests fetcher. - filters (`List[str]`): A list of folders to filter. - """ - if not os.path.isfile(output_file): - print("No test file found.") - return - with open(output_file, "r", encoding="utf-8") as f: - test_files = f.read().split(" ") + # Each modified file contributes itself (if it's a test) plus tests transitively impacted by it. + selected = set() + for f in modified_files: + if _is_test_file(f): + selected.add(f) + selected.update(t for t in reverse_map.get(f, []) if _is_test_file(t)) - if len(test_files) == 0 or test_files == [""]: - print("No tests to filter.") - return + test_files_to_run = sorted(p for p in selected if (PATH_TO_REPO / p).exists()) + _write_matrix(json_output_file, test_files_to_run) - if test_files == ["tests"]: - test_files = [os.path.join("tests", f) for f in os.listdir("tests") if f not in ["__init__.py"] + filters] - else: - test_files = [f for f in test_files if f.split(os.path.sep)[1] not in filters] - with open(output_file, "w", encoding="utf-8") as f: - f.write(" ".join(test_files)) +def _all_test_files() -> List[str]: + """Enumerate every `tests/.../test_*.py` (used for the full-suite path).""" + return sorted( + str(p.relative_to(PATH_TO_REPO)) for p in PATH_TO_TESTS.glob("**/test_*.py") if "__pycache__" not in p.parts + ) -def parse_commit_message(commit_message: str) -> dict[str, bool]: - """ - Parses the commit message to detect if a command is there to skip, force all or part of the CI. +def _write_full_suite(json_output_file: str): + """Schedule the entire test suite. Used by `--force_full_suite` and as exception fallback.""" + _write_matrix(json_output_file, _all_test_files()) - Args: - commit_message (`str`): The commit message of the current commit. - Returns: - `Dict[str, bool]`: A dictionary of strings to bools with keys the following keys: `"skip"`, - `"test_all_models"` and `"test_all"`. - """ - if commit_message is None: - return {"skip": False, "no_filter": False, "test_all": False} - - command_search = re.search(r"\[([^\]]*)\]", commit_message) - if command_search is not None: - command = command_search.groups()[0] - command = command.lower().replace("-", " ").replace("_", " ") - skip = command in ["ci skip", "skip ci", "circleci skip", "skip circleci"] - no_filter = set(command.split(" ")) == {"no", "filter"} - test_all = set(command.split(" ")) == {"test", "all"} - return {"skip": skip, "no_filter": no_filter, "test_all": test_all} - else: - return {"skip": False, "no_filter": False, "test_all": False} +# ============================================================ +# CLI +# ============================================================ if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument( - "--output_file", type=str, default="test_list.txt", help="Where to store the list of tests to run" - ) parser.add_argument( "--json_output_file", type=str, default="test_map.json", - help="Where to store the tests to run in a dictionary format mapping test categories to test files", + help="Where to store the list of matrix entries (name / paths / markers) consumed by CI.", ) parser.add_argument( "--diff_with_last_commit", action="store_true", - help="To fetch the tests between the current commit and the last commit", + help="Diff against the previous commit instead of main (use on main branch jobs)", ) parser.add_argument( - "--filter_tests", + "--force_full_suite", action="store_true", - help="Will filter the pipeline/repo utils tests outside of the generated list of tests.", - ) - parser.add_argument( - "--print_dependencies_of", - type=str, - help="Will only print the tree of modules depending on the file passed.", - default=None, - ) - parser.add_argument( - "--commit_message", - type=str, - help="The commit message (which could contain a command to force all tests or skip the CI).", - default=None, + help="Bypass selection and write outputs that schedule the entire test suite.", ) args = parser.parse_args() - if args.print_dependencies_of is not None: - print_tree_deps_of(args.print_dependencies_of) - else: - repo = Repo(PATH_TO_REPO) - commit_message = repo.head.commit.message - commit_flags = parse_commit_message(commit_message) - if commit_flags["skip"]: - print("Force-skipping the CI") - quit() - if commit_flags["no_filter"]: - print("Running all tests fetched without filtering.") - if commit_flags["test_all"]: - print("Force-launching all tests") - - diff_with_last_commit = args.diff_with_last_commit - if not diff_with_last_commit and not repo.head.is_detached and repo.head.ref == repo.refs.main: - print("main branch detected, fetching tests against last commit.") - diff_with_last_commit = True - - if not commit_flags["test_all"]: - try: - infer_tests_to_run( - args.output_file, - diff_with_last_commit=diff_with_last_commit, - json_output_file=args.json_output_file, - ) - filter_tests(args.output_file, ["repo_utils"]) - update_test_map_with_core_pipelines(json_output_file=args.json_output_file) - - except Exception as e: - print(f"\nError when trying to grab the relevant tests: {e}\n\nRunning all tests.") - commit_flags["test_all"] = True - - if commit_flags["test_all"]: - with open(args.output_file, "w", encoding="utf-8") as f: - f.write("tests") - example_file = Path(args.output_file).parent / "examples_test_list.txt" - with open(example_file, "w", encoding="utf-8") as f: - f.write("all") - - test_files_to_run = get_all_tests() - create_json_map(test_files_to_run, args.json_output_file) - update_test_map_with_core_pipelines(json_output_file=args.json_output_file) + + if args.force_full_suite: + print("Forcing full test suite.") + _write_full_suite(args.json_output_file) + raise SystemExit(0) + + repo = Repo(PATH_TO_REPO) + diff_with_last_commit = args.diff_with_last_commit + if not diff_with_last_commit and not repo.head.is_detached and repo.head.ref == repo.refs.main: + print("main branch detected, fetching tests against last commit.") + diff_with_last_commit = True + + try: + fetch_tests_to_run(args.json_output_file, diff_with_last_commit) + except Exception as e: + import traceback + + print(f"\nError when trying to grab the relevant tests: {e}\n") + traceback.print_exc() + print("\nRunning all tests.") + _write_full_suite(args.json_output_file) From c7c8cc41dcd1fcb7520ca5b57cac4a22baeb5bca Mon Sep 17 00:00:00 2001 From: DN6 Date: Wed, 26 Aug 2026 10:38:51 +0530 Subject: [PATCH 2/4] [tests] Import pipeline/model classes inside the tests that use them so the fetcher does not over-select through shared mixins --- tests/lora/test_lora_layers_sd.py | 8 ++- tests/models/test_modeling_common.py | 6 +- .../test_models_transformer_flux.py | 51 +--------------- tests/models/transformers/utils.py | 58 +++++++++++++++++++ tests/others/test_flashpack.py | 3 +- tests/pipelines/kandinsky3/test_kandinsky3.py | 6 +- .../kandinsky3/test_kandinsky3_img2img.py | 3 +- tests/pipelines/pag/test_pag_sd.py | 5 +- tests/pipelines/pag/test_pag_sd3_img2img.py | 5 +- tests/pipelines/pag/test_pag_sd_img2img.py | 5 +- tests/pipelines/pag/test_pag_sd_inpaint.py | 5 +- tests/pipelines/pag/test_pag_sdxl.py | 5 +- tests/pipelines/pag/test_pag_sdxl_img2img.py | 5 +- tests/pipelines/pag/test_pag_sdxl_inpaint.py | 5 +- tests/pipelines/test_pipelines_common.py | 2 +- 15 files changed, 107 insertions(+), 65 deletions(-) create mode 100644 tests/models/transformers/utils.py diff --git a/tests/lora/test_lora_layers_sd.py b/tests/lora/test_lora_layers_sd.py index 4662075cfd0a..51060abed97b 100644 --- a/tests/lora/test_lora_layers_sd.py +++ b/tests/lora/test_lora_layers_sd.py @@ -24,8 +24,6 @@ from transformers import CLIPTextModel, CLIPTokenizer from diffusers import ( - AutoPipelineForImage2Image, - AutoPipelineForText2Image, DDIMScheduler, DiffusionPipeline, LCMScheduler, @@ -641,6 +639,8 @@ def test_load_unload_load_kohya_lora(self): release_memory(pipe) def test_not_empty_state_dict(self): + from diffusers import AutoPipelineForText2Image + # Makes sure https://github.com/huggingface/diffusers/issues/7054 does not happen again pipe = AutoPipelineForText2Image.from_pretrained( "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16 @@ -655,6 +655,8 @@ def test_not_empty_state_dict(self): release_memory(pipe) def test_load_unload_load_state_dict(self): + from diffusers import AutoPipelineForText2Image + # Makes sure https://github.com/huggingface/diffusers/issues/7054 does not happen again pipe = AutoPipelineForText2Image.from_pretrained( "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16 @@ -705,6 +707,8 @@ def test_sdv1_5_lcm_lora(self): release_memory(pipe) def test_sdv1_5_lcm_lora_img2img(self): + from diffusers import AutoPipelineForImage2Image + pipe = AutoPipelineForImage2Image.from_pretrained( "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16 ) diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index 9968add19dd9..139111fe5c6f 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -26,7 +26,7 @@ from huggingface_hub import ModelCard, delete_repo, snapshot_download, try_to_load_from_cache from huggingface_hub.utils import HfHubHTTPError, is_jinja_available -from diffusers.models import FluxTransformer2DModel, SD3Transformer2DModel, UNet2DConditionModel +from diffusers.models import UNet2DConditionModel from ..others.test_utils import TOKEN, USER, is_staging_test from ..testing_utils import ( @@ -118,6 +118,8 @@ def test_cached_files_are_used_when_no_internet(self): assert False, "Parameters not the same!" def test_local_files_only_with_sharded_checkpoint(self): + from diffusers.models import FluxTransformer2DModel + repo_id = "hf-internal-testing/tiny-flux-sharded" error_response = mock.Mock( status_code=500, @@ -232,6 +234,8 @@ def test_keep_modules_in_fp32(self): A simple tests to check if the modules under `_keep_in_fp32_modules` are kept in fp32 when we load the model in fp16/bf16 Also ensures if inference works. """ + from diffusers.models import SD3Transformer2DModel + fp32_modules = SD3Transformer2DModel._keep_in_fp32_modules for torch_dtype in [torch.bfloat16, torch.float16]: diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index d60e34e2ea3f..27a80c4b9a02 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -23,7 +23,6 @@ import torch from diffusers import BitsAndBytesConfig, FluxTransformer2DModel, GGUFQuantizationConfig -from diffusers.models.embeddings import ImageProjection from diffusers.models.transformers.transformer_flux import FluxIPAdapterAttnProcessor from diffusers.utils.torch_utils import randn_tensor @@ -60,60 +59,12 @@ TorchCompileTesterMixin, TrainingTesterMixin, ) +from .utils import create_flux_ip_adapter_state_dict enable_full_determinism() -# TODO: This standalone function maintains backward compatibility with pipeline tests -# (tests/pipelines/test_pipelines_common.py) and will be refactored. -def create_flux_ip_adapter_state_dict(model) -> dict[str, dict[str, Any]]: - """Create a dummy IP Adapter state dict for Flux transformer testing.""" - ip_cross_attn_state_dict = {} - key_id = 0 - - for name in model.attn_processors.keys(): - if name.startswith("single_transformer_blocks"): - continue - - joint_attention_dim = model.config["joint_attention_dim"] - hidden_size = model.config["num_attention_heads"] * model.config["attention_head_dim"] - sd = FluxIPAdapterAttnProcessor( - hidden_size=hidden_size, cross_attention_dim=joint_attention_dim, scale=1.0 - ).state_dict() - ip_cross_attn_state_dict.update( - { - f"{key_id}.to_k_ip.weight": sd["to_k_ip.0.weight"], - f"{key_id}.to_v_ip.weight": sd["to_v_ip.0.weight"], - f"{key_id}.to_k_ip.bias": sd["to_k_ip.0.bias"], - f"{key_id}.to_v_ip.bias": sd["to_v_ip.0.bias"], - } - ) - key_id += 1 - - image_projection = ImageProjection( - cross_attention_dim=model.config["joint_attention_dim"], - image_embed_dim=( - model.config["pooled_projection_dim"] if "pooled_projection_dim" in model.config.keys() else 768 - ), - num_image_text_embeds=4, - ) - - ip_image_projection_state_dict = {} - sd = image_projection.state_dict() - ip_image_projection_state_dict.update( - { - "proj.weight": sd["image_embeds.weight"], - "proj.bias": sd["image_embeds.bias"], - "norm.weight": sd["norm.weight"], - "norm.bias": sd["norm.bias"], - } - ) - - del sd - return {"image_proj": ip_image_projection_state_dict, "ip_adapter": ip_cross_attn_state_dict} - - class FluxTransformerTesterConfig(BaseModelTesterConfig): @property def model_class(self): diff --git a/tests/models/transformers/utils.py b/tests/models/transformers/utils.py new file mode 100644 index 000000000000..2759b0550c2b --- /dev/null +++ b/tests/models/transformers/utils.py @@ -0,0 +1,58 @@ +"""Shared test helpers for transformer model tests. + +Diffusers imports inside helpers are deliberately deferred to function bodies so the test-fetcher's +module-level import graph doesn't propagate edits to specific transformer source files through this +file into every pipeline test that imports it. +""" + +from typing import Any + + +def create_flux_ip_adapter_state_dict(model) -> dict[str, dict[str, Any]]: + """Create a dummy IP Adapter state dict for Flux transformer testing.""" + from diffusers.models.embeddings import ImageProjection + from diffusers.models.transformers.transformer_flux import FluxIPAdapterAttnProcessor + + ip_cross_attn_state_dict = {} + key_id = 0 + + for name in model.attn_processors.keys(): + if name.startswith("single_transformer_blocks"): + continue + + joint_attention_dim = model.config["joint_attention_dim"] + hidden_size = model.config["num_attention_heads"] * model.config["attention_head_dim"] + sd = FluxIPAdapterAttnProcessor( + hidden_size=hidden_size, cross_attention_dim=joint_attention_dim, scale=1.0 + ).state_dict() + ip_cross_attn_state_dict.update( + { + f"{key_id}.to_k_ip.weight": sd["to_k_ip.0.weight"], + f"{key_id}.to_v_ip.weight": sd["to_v_ip.0.weight"], + f"{key_id}.to_k_ip.bias": sd["to_k_ip.0.bias"], + f"{key_id}.to_v_ip.bias": sd["to_v_ip.0.bias"], + } + ) + key_id += 1 + + image_projection = ImageProjection( + cross_attention_dim=model.config["joint_attention_dim"], + image_embed_dim=( + model.config["pooled_projection_dim"] if "pooled_projection_dim" in model.config.keys() else 768 + ), + num_image_text_embeds=4, + ) + + ip_image_projection_state_dict = {} + sd = image_projection.state_dict() + ip_image_projection_state_dict.update( + { + "proj.weight": sd["image_embeds.weight"], + "proj.bias": sd["image_embeds.bias"], + "norm.weight": sd["norm.weight"], + "norm.bias": sd["norm.bias"], + } + ) + + del sd + return {"image_proj": ip_image_projection_state_dict, "ip_adapter": ip_cross_attn_state_dict} diff --git a/tests/others/test_flashpack.py b/tests/others/test_flashpack.py index c20836df1c98..e7eefd6835ff 100644 --- a/tests/others/test_flashpack.py +++ b/tests/others/test_flashpack.py @@ -15,7 +15,6 @@ import pytest -from diffusers import AutoPipelineForText2Image from diffusers.models.auto_model import AutoModel from ..testing_utils import is_torch_available, require_flashpack, require_torch_gpu @@ -40,6 +39,8 @@ def test_save_load_model(self, tmp_path): @require_flashpack def test_save_load_pipeline(self, tmp_path): + from diffusers import AutoPipelineForText2Image + pipeline = AutoPipelineForText2Image.from_pretrained(self.model_id) pipeline.save_pretrained(tmp_path, use_flashpack=True) assert (tmp_path / "transformer" / "model.flashpack").exists() diff --git a/tests/pipelines/kandinsky3/test_kandinsky3.py b/tests/pipelines/kandinsky3/test_kandinsky3.py index 1a1aa4b9d9ca..8a7554a8430a 100644 --- a/tests/pipelines/kandinsky3/test_kandinsky3.py +++ b/tests/pipelines/kandinsky3/test_kandinsky3.py @@ -22,8 +22,6 @@ from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import ( - AutoPipelineForImage2Image, - AutoPipelineForText2Image, Kandinsky3Pipeline, Kandinsky3UNet, VQModel, @@ -185,6 +183,8 @@ def tearDown(self): backend_empty_cache(torch_device) def test_kandinskyV3(self): + from diffusers import AutoPipelineForText2Image + pipe = AutoPipelineForText2Image.from_pretrained( "kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16 ) @@ -211,6 +211,8 @@ def test_kandinskyV3(self): self.assertTrue(np.allclose(image_np, expected_image_np, atol=5e-2)) def test_kandinskyV3_img2img(self): + from diffusers import AutoPipelineForImage2Image + pipe = AutoPipelineForImage2Image.from_pretrained( "kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16 ) diff --git a/tests/pipelines/kandinsky3/test_kandinsky3_img2img.py b/tests/pipelines/kandinsky3/test_kandinsky3_img2img.py index d9f02326b1b5..d41e644a45ef 100644 --- a/tests/pipelines/kandinsky3/test_kandinsky3_img2img.py +++ b/tests/pipelines/kandinsky3/test_kandinsky3_img2img.py @@ -23,7 +23,6 @@ from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import ( - AutoPipelineForImage2Image, Kandinsky3Img2ImgPipeline, Kandinsky3UNet, VQModel, @@ -207,6 +206,8 @@ def tearDown(self): backend_empty_cache(torch_device) def test_kandinskyV3_img2img(self): + from diffusers import AutoPipelineForImage2Image + pipe = AutoPipelineForImage2Image.from_pretrained( "kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16 ) diff --git a/tests/pipelines/pag/test_pag_sd.py b/tests/pipelines/pag/test_pag_sd.py index 1dd3ef298fd0..707603d0d5ed 100644 --- a/tests/pipelines/pag/test_pag_sd.py +++ b/tests/pipelines/pag/test_pag_sd.py @@ -23,7 +23,6 @@ from diffusers import ( AutoencoderKL, - AutoPipelineForText2Image, DDIMScheduler, StableDiffusionPAGPipeline, StableDiffusionPipeline, @@ -315,6 +314,8 @@ def get_inputs(self, device, generator_device="cpu", seed=1, guidance_scale=7.0) return inputs def test_pag_cfg(self): + from diffusers import AutoPipelineForText2Image + pipeline = AutoPipelineForText2Image.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16) pipeline.enable_model_cpu_offload(device=torch_device) pipeline.set_progress_bar_config(disable=None) @@ -333,6 +334,8 @@ def test_pag_cfg(self): ) def test_pag_uncond(self): + from diffusers import AutoPipelineForText2Image + pipeline = AutoPipelineForText2Image.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16) pipeline.enable_model_cpu_offload(device=torch_device) pipeline.set_progress_bar_config(disable=None) diff --git a/tests/pipelines/pag/test_pag_sd3_img2img.py b/tests/pipelines/pag/test_pag_sd3_img2img.py index ede146915c55..be5c2e118850 100644 --- a/tests/pipelines/pag/test_pag_sd3_img2img.py +++ b/tests/pipelines/pag/test_pag_sd3_img2img.py @@ -16,7 +16,6 @@ from diffusers import ( AutoencoderKL, - AutoPipelineForImage2Image, FlowMatchEulerDiscreteScheduler, SD3Transformer2DModel, StableDiffusion3Img2ImgPipeline, @@ -240,6 +239,8 @@ def get_inputs( return inputs def test_pag_cfg(self): + from diffusers import AutoPipelineForImage2Image + pipeline = AutoPipelineForImage2Image.from_pretrained( self.repo_id, enable_pag=True, torch_dtype=torch.float16, pag_applied_layers=["blocks.17"] ) @@ -268,6 +269,8 @@ def test_pag_cfg(self): ) def test_pag_uncond(self): + from diffusers import AutoPipelineForImage2Image + pipeline = AutoPipelineForImage2Image.from_pretrained( self.repo_id, enable_pag=True, torch_dtype=torch.float16, pag_applied_layers=["blocks.(4|17)"] ) diff --git a/tests/pipelines/pag/test_pag_sd_img2img.py b/tests/pipelines/pag/test_pag_sd_img2img.py index 1e9b3c24c9ac..0afbd27c677e 100644 --- a/tests/pipelines/pag/test_pag_sd_img2img.py +++ b/tests/pipelines/pag/test_pag_sd_img2img.py @@ -25,7 +25,6 @@ from diffusers import ( AutoencoderKL, AutoencoderTiny, - AutoPipelineForImage2Image, EulerDiscreteScheduler, StableDiffusionImg2ImgPipeline, StableDiffusionPAGImg2ImgPipeline, @@ -254,6 +253,8 @@ def get_inputs(self, device, generator_device="cpu", dtype=torch.float32, seed=0 return inputs def test_pag_cfg(self): + from diffusers import AutoPipelineForImage2Image + pipeline = AutoPipelineForImage2Image.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16) pipeline.enable_model_cpu_offload(device=torch_device) pipeline.set_progress_bar_config(disable=None) @@ -272,6 +273,8 @@ def test_pag_cfg(self): ) def test_pag_uncond(self): + from diffusers import AutoPipelineForImage2Image + pipeline = AutoPipelineForImage2Image.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16) pipeline.enable_model_cpu_offload(device=torch_device) pipeline.set_progress_bar_config(disable=None) diff --git a/tests/pipelines/pag/test_pag_sd_inpaint.py b/tests/pipelines/pag/test_pag_sd_inpaint.py index 5a78ac6ade12..f351bd03b265 100644 --- a/tests/pipelines/pag/test_pag_sd_inpaint.py +++ b/tests/pipelines/pag/test_pag_sd_inpaint.py @@ -24,7 +24,6 @@ from diffusers import ( AutoencoderKL, - AutoPipelineForInpainting, PNDMScheduler, StableDiffusionPAGInpaintPipeline, UNet2DConditionModel, @@ -289,6 +288,8 @@ def get_inputs(self, device, generator_device="cpu", seed=0, guidance_scale=7.0) return inputs def test_pag_cfg(self): + from diffusers import AutoPipelineForInpainting + pipeline = AutoPipelineForInpainting.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16) pipeline.enable_model_cpu_offload(device=torch_device) pipeline.set_progress_bar_config(disable=None) @@ -307,6 +308,8 @@ def test_pag_cfg(self): ) def test_pag_uncond(self): + from diffusers import AutoPipelineForInpainting + pipeline = AutoPipelineForInpainting.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16) pipeline.enable_model_cpu_offload(device=torch_device) pipeline.set_progress_bar_config(disable=None) diff --git a/tests/pipelines/pag/test_pag_sdxl.py b/tests/pipelines/pag/test_pag_sdxl.py index dc412d9b341a..74969903a553 100644 --- a/tests/pipelines/pag/test_pag_sdxl.py +++ b/tests/pipelines/pag/test_pag_sdxl.py @@ -23,7 +23,6 @@ from diffusers import ( AutoencoderKL, - AutoPipelineForText2Image, EulerDiscreteScheduler, StableDiffusionXLPAGPipeline, StableDiffusionXLPipeline, @@ -320,6 +319,8 @@ def get_inputs(self, device, generator_device="cpu", seed=0, guidance_scale=7.0) return inputs def test_pag_cfg(self): + from diffusers import AutoPipelineForText2Image + pipeline = AutoPipelineForText2Image.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16) pipeline.enable_model_cpu_offload(device=torch_device) pipeline.set_progress_bar_config(disable=None) @@ -342,6 +343,8 @@ def test_pag_cfg(self): ) def test_pag_uncond(self): + from diffusers import AutoPipelineForText2Image + pipeline = AutoPipelineForText2Image.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16) pipeline.enable_model_cpu_offload(device=torch_device) pipeline.set_progress_bar_config(disable=None) diff --git a/tests/pipelines/pag/test_pag_sdxl_img2img.py b/tests/pipelines/pag/test_pag_sdxl_img2img.py index 3157e3a08e8d..b4a108cdab42 100644 --- a/tests/pipelines/pag/test_pag_sdxl_img2img.py +++ b/tests/pipelines/pag/test_pag_sdxl_img2img.py @@ -32,7 +32,6 @@ from diffusers import ( AutoencoderKL, - AutoPipelineForImage2Image, EulerDiscreteScheduler, StableDiffusionXLImg2ImgPipeline, StableDiffusionXLPAGImg2ImgPipeline, @@ -302,6 +301,8 @@ def get_inputs(self, device, generator_device="cpu", seed=0, guidance_scale=7.0) return inputs def test_pag_cfg(self): + from diffusers import AutoPipelineForImage2Image + pipeline = AutoPipelineForImage2Image.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16) pipeline.enable_model_cpu_offload(device=torch_device) pipeline.set_progress_bar_config(disable=None) @@ -319,6 +320,8 @@ def test_pag_cfg(self): ) def test_pag_uncond(self): + from diffusers import AutoPipelineForImage2Image + pipeline = AutoPipelineForImage2Image.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16) pipeline.enable_model_cpu_offload(device=torch_device) pipeline.set_progress_bar_config(disable=None) diff --git a/tests/pipelines/pag/test_pag_sdxl_inpaint.py b/tests/pipelines/pag/test_pag_sdxl_inpaint.py index 628da17c8dc3..3b6b23ed17db 100644 --- a/tests/pipelines/pag/test_pag_sdxl_inpaint.py +++ b/tests/pipelines/pag/test_pag_sdxl_inpaint.py @@ -33,7 +33,6 @@ from diffusers import ( AutoencoderKL, - AutoPipelineForInpainting, EulerDiscreteScheduler, StableDiffusionXLInpaintPipeline, StableDiffusionXLPAGInpaintPipeline, @@ -308,6 +307,8 @@ def get_inputs(self, device, generator_device="cpu", seed=0, guidance_scale=7.0) return inputs def test_pag_cfg(self): + from diffusers import AutoPipelineForInpainting + pipeline = AutoPipelineForInpainting.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16) pipeline.enable_model_cpu_offload(device=torch_device) pipeline.set_progress_bar_config(disable=None) @@ -325,6 +326,8 @@ def test_pag_cfg(self): ) def test_pag_uncond(self): + from diffusers import AutoPipelineForInpainting + pipeline = AutoPipelineForInpainting.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16) pipeline.enable_model_cpu_offload(device=torch_device) pipeline.set_progress_bar_config(disable=None) diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py index 106ba55cf149..71c92a10311d 100644 --- a/tests/pipelines/test_pipelines_common.py +++ b/tests/pipelines/test_pipelines_common.py @@ -53,7 +53,7 @@ get_autoencoder_tiny_config, get_consistency_vae_config, ) -from ..models.transformers.test_models_transformer_flux import create_flux_ip_adapter_state_dict +from ..models.transformers.utils import create_flux_ip_adapter_state_dict from ..models.unets.test_models_unet_2d_condition import ( create_ip_adapter_faceid_state_dict, create_ip_adapter_state_dict, From 4242566fdee9b6ac21a139a652c4f369198bae7a Mon Sep 17 00:00:00 2001 From: DN6 Date: Wed, 26 Aug 2026 11:42:08 +0530 Subject: [PATCH 3/4] [CI] Give `single_file` model tests their own fetcher job; they need an HF token, not an accelerator --- utils/tests_fetcher.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/utils/tests_fetcher.py b/utils/tests_fetcher.py index 8afa2994fa44..718e029319fb 100644 --- a/utils/tests_fetcher.py +++ b/utils/tests_fetcher.py @@ -74,10 +74,12 @@ "memory": ["memory", "cpu_offload", "group_offload"], "cache": ["cache"], "ip_adapter": ["ip_adapter"], + # Needs only an HF token (gated checkpoints), which the CPU job provides; skips on fork PRs. + "single_file": ["single_file"], } -# Gated on an accelerator, multi-GPU, or an HF token, so every test skips on the CPU runner. Excluded -# from `core` and given no job rather than spinning up a runner to skip everything. -CPU_SKIPPED_MARKERS = ["quantization", "compile", "single_file", "training", "context_parallel", "tensor_parallel"] +# Gated on an accelerator or multi-GPU, so every test skips on the CPU runner. Excluded from `core` and +# given no job rather than spinning up a runner to skip everything. +CPU_SKIPPED_MARKERS = ["quantization", "compile", "training", "context_parallel", "tensor_parallel"] FEATURE_MARKERS = [m for markers in FEATURE_GROUPS.values() for m in markers] CORE_MARKERS = "not (" + " or ".join(FEATURE_MARKERS + CPU_SKIPPED_MARKERS) + ")" # ============================================================ From bd8d8b8f4a6d7755fd9e26b3491192453c361a1e Mon Sep 17 00:00:00 2001 From: DN6 Date: Wed, 26 Aug 2026 14:57:37 +0530 Subject: [PATCH 4/4] [tests] Keep `create_flux_ip_adapter_state_dict` in `test_models_transformer_flux.py` instead of a separate `utils.py` --- .../test_models_transformer_flux.py | 51 +++++++++++++++- tests/models/transformers/utils.py | 58 ------------------- tests/pipelines/test_pipelines_common.py | 2 +- 3 files changed, 51 insertions(+), 60 deletions(-) delete mode 100644 tests/models/transformers/utils.py diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index 27a80c4b9a02..d60e34e2ea3f 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -23,6 +23,7 @@ import torch from diffusers import BitsAndBytesConfig, FluxTransformer2DModel, GGUFQuantizationConfig +from diffusers.models.embeddings import ImageProjection from diffusers.models.transformers.transformer_flux import FluxIPAdapterAttnProcessor from diffusers.utils.torch_utils import randn_tensor @@ -59,12 +60,60 @@ TorchCompileTesterMixin, TrainingTesterMixin, ) -from .utils import create_flux_ip_adapter_state_dict enable_full_determinism() +# TODO: This standalone function maintains backward compatibility with pipeline tests +# (tests/pipelines/test_pipelines_common.py) and will be refactored. +def create_flux_ip_adapter_state_dict(model) -> dict[str, dict[str, Any]]: + """Create a dummy IP Adapter state dict for Flux transformer testing.""" + ip_cross_attn_state_dict = {} + key_id = 0 + + for name in model.attn_processors.keys(): + if name.startswith("single_transformer_blocks"): + continue + + joint_attention_dim = model.config["joint_attention_dim"] + hidden_size = model.config["num_attention_heads"] * model.config["attention_head_dim"] + sd = FluxIPAdapterAttnProcessor( + hidden_size=hidden_size, cross_attention_dim=joint_attention_dim, scale=1.0 + ).state_dict() + ip_cross_attn_state_dict.update( + { + f"{key_id}.to_k_ip.weight": sd["to_k_ip.0.weight"], + f"{key_id}.to_v_ip.weight": sd["to_v_ip.0.weight"], + f"{key_id}.to_k_ip.bias": sd["to_k_ip.0.bias"], + f"{key_id}.to_v_ip.bias": sd["to_v_ip.0.bias"], + } + ) + key_id += 1 + + image_projection = ImageProjection( + cross_attention_dim=model.config["joint_attention_dim"], + image_embed_dim=( + model.config["pooled_projection_dim"] if "pooled_projection_dim" in model.config.keys() else 768 + ), + num_image_text_embeds=4, + ) + + ip_image_projection_state_dict = {} + sd = image_projection.state_dict() + ip_image_projection_state_dict.update( + { + "proj.weight": sd["image_embeds.weight"], + "proj.bias": sd["image_embeds.bias"], + "norm.weight": sd["norm.weight"], + "norm.bias": sd["norm.bias"], + } + ) + + del sd + return {"image_proj": ip_image_projection_state_dict, "ip_adapter": ip_cross_attn_state_dict} + + class FluxTransformerTesterConfig(BaseModelTesterConfig): @property def model_class(self): diff --git a/tests/models/transformers/utils.py b/tests/models/transformers/utils.py deleted file mode 100644 index 2759b0550c2b..000000000000 --- a/tests/models/transformers/utils.py +++ /dev/null @@ -1,58 +0,0 @@ -"""Shared test helpers for transformer model tests. - -Diffusers imports inside helpers are deliberately deferred to function bodies so the test-fetcher's -module-level import graph doesn't propagate edits to specific transformer source files through this -file into every pipeline test that imports it. -""" - -from typing import Any - - -def create_flux_ip_adapter_state_dict(model) -> dict[str, dict[str, Any]]: - """Create a dummy IP Adapter state dict for Flux transformer testing.""" - from diffusers.models.embeddings import ImageProjection - from diffusers.models.transformers.transformer_flux import FluxIPAdapterAttnProcessor - - ip_cross_attn_state_dict = {} - key_id = 0 - - for name in model.attn_processors.keys(): - if name.startswith("single_transformer_blocks"): - continue - - joint_attention_dim = model.config["joint_attention_dim"] - hidden_size = model.config["num_attention_heads"] * model.config["attention_head_dim"] - sd = FluxIPAdapterAttnProcessor( - hidden_size=hidden_size, cross_attention_dim=joint_attention_dim, scale=1.0 - ).state_dict() - ip_cross_attn_state_dict.update( - { - f"{key_id}.to_k_ip.weight": sd["to_k_ip.0.weight"], - f"{key_id}.to_v_ip.weight": sd["to_v_ip.0.weight"], - f"{key_id}.to_k_ip.bias": sd["to_k_ip.0.bias"], - f"{key_id}.to_v_ip.bias": sd["to_v_ip.0.bias"], - } - ) - key_id += 1 - - image_projection = ImageProjection( - cross_attention_dim=model.config["joint_attention_dim"], - image_embed_dim=( - model.config["pooled_projection_dim"] if "pooled_projection_dim" in model.config.keys() else 768 - ), - num_image_text_embeds=4, - ) - - ip_image_projection_state_dict = {} - sd = image_projection.state_dict() - ip_image_projection_state_dict.update( - { - "proj.weight": sd["image_embeds.weight"], - "proj.bias": sd["image_embeds.bias"], - "norm.weight": sd["norm.weight"], - "norm.bias": sd["norm.bias"], - } - ) - - del sd - return {"image_proj": ip_image_projection_state_dict, "ip_adapter": ip_cross_attn_state_dict} diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py index 71c92a10311d..106ba55cf149 100644 --- a/tests/pipelines/test_pipelines_common.py +++ b/tests/pipelines/test_pipelines_common.py @@ -53,7 +53,7 @@ get_autoencoder_tiny_config, get_consistency_vae_config, ) -from ..models.transformers.utils import create_flux_ip_adapter_state_dict +from ..models.transformers.test_models_transformer_flux import create_flux_ip_adapter_state_dict from ..models.unets.test_models_unet_2d_condition import ( create_ip_adapter_faceid_state_dict, create_ip_adapter_state_dict,