Function main

utils/generate_model_tests.py:509–597

()

def main():
    parser = argparse.ArgumentParser(description="Generate test suite for a diffusers model class")
    parser.add_argument(
        "model_filepath",
        type=str,
        help="Path to the model file (e.g., src/diffusers/models/transformers/transformer_flux.py)",
    )
    parser.add_argument(
        "--output", "-o", type=str, default=None, help="Output file path (default: auto-generated based on model path)"
    )
    parser.add_argument(
        "--include",
        "-i",
        type=str,
        nargs="*",
        default=[],
        choices=[
            "bnb",
            "quanto",
            "torchao",
            "gguf",
            "modelopt",
            "bnb_compile",
            "quanto_compile",
            "torchao_compile",
            "gguf_compile",
            "modelopt_compile",
            "pab_cache",
            "fbc_cache",
            "faster_cache",
            "single_file",
            "ip_adapter",
            "cp_attn",
            "all",
        ],
        help="Optional testers to include",
    )
    parser.add_argument(
        "--class-name",
        "-c",
        type=str,
        default=None,
        help="Specific model class to generate tests for (default: first model class found)",
    )
    parser.add_argument("--dry-run", action="store_true", help="Print generated code without writing to file")

    args = parser.parse_args()

    if not Path(args.model_filepath).exists():
        print(f"Error: File not found: {args.model_filepath}", file=sys.stderr)
        sys.exit(1)

    model_classes, imports = analyze_model_file(args.model_filepath)

    if not model_classes:
        print(f"Error: No model classes found in {args.model_filepath}", file=sys.stderr)
        sys.exit(1)