| | import json |
| | import logging |
| | import os |
| | from collections import defaultdict |
| | from pathlib import Path |
| |
|
| | from huggingface_hub import HfApi |
| |
|
| | import diffusers |
| |
|
| |
|
| | PATH_TO_REPO = Path(__file__).parent.parent.resolve() |
| | ALWAYS_TEST_PIPELINE_MODULES = [ |
| | "controlnet", |
| | "controlnet_flux", |
| | "controlnet_sd3", |
| | "stable_diffusion", |
| | "stable_diffusion_2", |
| | "stable_diffusion_3", |
| | "stable_diffusion_xl", |
| | "ip_adapters", |
| | "flux", |
| | ] |
| | PIPELINE_USAGE_CUTOFF = int(os.getenv("PIPELINE_USAGE_CUTOFF", 50000)) |
| |
|
| | logger = logging.getLogger(__name__) |
| | api = HfApi() |
| |
|
| |
|
| | def filter_pipelines(usage_dict, usage_cutoff=10000): |
| | output = [] |
| | for diffusers_object, usage in usage_dict.items(): |
| | if usage < usage_cutoff: |
| | continue |
| |
|
| | is_diffusers_pipeline = hasattr(diffusers.pipelines, diffusers_object) |
| | if not is_diffusers_pipeline: |
| | continue |
| |
|
| | output.append(diffusers_object) |
| |
|
| | return output |
| |
|
| |
|
| | def fetch_pipeline_objects(): |
| | models = api.list_models(library="diffusers") |
| | downloads = defaultdict(int) |
| |
|
| | for model in models: |
| | is_counted = False |
| | for tag in model.tags: |
| | if tag.startswith("diffusers:"): |
| | is_counted = True |
| | downloads[tag[len("diffusers:") :]] += model.downloads |
| |
|
| | if not is_counted: |
| | downloads["other"] += model.downloads |
| |
|
| | |
| | downloads = {k: v for k, v in downloads.items() if v > 0} |
| | pipeline_objects = filter_pipelines(downloads, PIPELINE_USAGE_CUTOFF) |
| |
|
| | return pipeline_objects |
| |
|
| |
|
| | def fetch_pipeline_modules_to_test(): |
| | try: |
| | pipeline_objects = fetch_pipeline_objects() |
| | except Exception as e: |
| | logger.error(e) |
| | raise RuntimeError("Unable to fetch model list from HuggingFace Hub.") |
| |
|
| | test_modules = [] |
| | for pipeline_name in pipeline_objects: |
| | module = getattr(diffusers, pipeline_name) |
| |
|
| | test_module = module.__module__.split(".")[-2].strip() |
| | test_modules.append(test_module) |
| |
|
| | return test_modules |
| |
|
| |
|
| | def main(): |
| | test_modules = fetch_pipeline_modules_to_test() |
| | test_modules.extend(ALWAYS_TEST_PIPELINE_MODULES) |
| |
|
| | |
| | test_modules = sorted(set(test_modules)) |
| | print(json.dumps(test_modules)) |
| |
|
| | save_path = f"{PATH_TO_REPO}/reports" |
| | os.makedirs(save_path, exist_ok=True) |
| |
|
| | with open(f"{save_path}/test-pipelines.json", "w") as f: |
| | json.dump({"pipeline_test_modules": test_modules}, f) |
| |
|
| |
|
| | if __name__ == "__main__": |
| | main() |
| |
|
