update benchmark scripts (#425)

Co-authored-by: Yao Chi <later@usopp.net>

.github/workflows/examples.yml

@@ -123,8 +123,6 @@ jobs:
         run: docker exec ${{ env.CONTAINER_NAME }} python3 examples/text_to_image_sdxl.py --base /share_nfs/hf_models/stable-diffusion-xl-base-1.0 --compile_unet False --height 512 --width 512 --use_multiple_resolutions True
       - if: matrix.test-suite == 'diffusers_examples'
         run: docker exec ${{ env.CONTAINER_NAME }} python3 examples/text_to_image_controlnet.py --base=/share_nfs/hf_models/stable-diffusion-v1-5 --controlnet=/share_nfs/hf_models/sd-controlnet-canny --input_image=/share_nfs/hf_models/input_image_vermeer.png
-      - if: matrix.test-suite == 'diffusers_examples'
-        run: docker exec ${{ env.CONTAINER_NAME }} python3 benchmarks/stable_diffusion_2_unet.py --model_id=/share_nfs/hf_models/stable-diffusion-2-1
       - if: matrix.test-suite == 'diffusers_examples'
         run: docker exec ${{ env.CONTAINER_NAME }} bash examples/unet_save_and_load.sh --model_id=/share_nfs/hf_models/stable-diffusion-2-1
 

benchmarks/README.md

@@ -0,0 +1,69 @@
+# Bench OneDiff
+
+## Build docker image
+
+```bash
+python3 -m pip install -r requirements.txt
+python3 docker/main.py --yaml ./docker/config/community-default.yaml
+```
+
+## Prepare models
+
+Download models from [here](#About-the-models). If you have downloaded the models before, please skip it.
+
+## Run OneDiff Benchmark
+
+Start docker container and run the benchmark by the following command.
+
+```bash
+export BENCHMARK_MODEL_PATH=./benchmark_model
+docker compose -f ./docker-compose.onediff:benchmark-community-default.yaml up
+```
+
+Wait for a while, you will see the following logs,
+
+```bash
+onediff-benchmark-community-default  | Run SD1.5(FP16) 1024x1024...
+onediff-benchmark-community-default  | + python3 ./text_to_image.py --model /benchmark_model/stable-diffusion-v1-5 --warmup 5 --height 1024 --width 1024
+Loading pipeline components...:  43% 3/7 [00:00<00:00, 20.94it/s]`text_config_dict` is provided which will be used to initialize `CLIPTextConfig`. The value `text_config["id2label"]` will be overriden.
+Loading pipeline components...: 100% 7/7 [00:00<00:00, 12.79it/s]
+100% 30/30 [00:43<00:00,  1.45s/it]  |
+100% 30/30 [00:03<00:00,  7.76it/s]  |
+100% 30/30 [00:03<00:00,  7.74it/s]  |
+100% 30/30 [00:03<00:00,  7.74it/s]  |
+100% 30/30 [00:03<00:00,  7.72it/s]  |
+100% 30/30 [00:03<00:00,  7.72it/s]  |
+onediff-benchmark-community-default  | e2e (30 steps) elapsed: 4.1393163204193115 s, cuda memory usage: 7226.875 MiB
+......
+```
+
+## About the models
+
+The structure of `/benchmark_model` should follow this hierarchy:
+
+```text
+benchmark_model
+├── stable-diffusion-2-1
+├── stable-diffusion-v1-5
+├── stable-diffusion-xl-base-1.0
+├── stable-diffusion-xl-base-1.0-int8
+```
+
+You can obtain the models from [HuggingFace](https://huggingface.co) (excluding the int8 model) or download them from OSS (including the int8 model). The OSS download method is as follows:
+
+- Obtain ossutil by executing the following command:
+
+  ```bash
+  wget http://gosspublic.alicdn.com/ossutil/1.7.3/ossutil64  && chmod u+x ossutil64
+  ```
+
+- Configure ossutil by referring to [the official example](https://www.alibabacloud.com/help/en/oss/developer-reference/configure-ossutil?spm=a2c63.p38356.0.0.337f374a4pcwa4)
+  ```bash
+  ossutil64 config
+  ```
+
+- Download the benchmark models finally
+
+  ```bash
+  ./ossutil64 cp -r oss://oneflow-pro/onediff_benchmark_model/  benchmark_model  --update 
+  ```

benchmarks/docker/.gitignore

@@ -0,0 +1,6 @@
+config/
+Dockerfile-*
+docker-compose.*
+onediff/
+ComfyUI/
+diffusers_quant/

benchmarks/docker/_logger.py

@@ -0,0 +1,10 @@
+import logging
+
+logger = logging.getLogger("onediff-benchmark")
+logger.setLevel(logging.INFO)
+formatter = logging.Formatter(
+    "%(asctime)s - %(levelname)s - %(message)s", "%Y-%m-%d %H:%M:%S"
+)
+console_handler = logging.StreamHandler()
+console_handler.setFormatter(formatter)
+logger.addHandler(console_handler)

benchmarks/docker/_utils.py

@@ -0,0 +1,184 @@
+import hashlib
+import os
+import subprocess
+import yaml
+
+from git import Repo
+
+from _logger import logger
+
+
+def load_yaml(*, file):
+    if not os.path.exists(file):
+        raise RuntimeError(f"config file not existed: {file}")
+
+    with open(file, "r") as file:
+        yaml_content = yaml.safe_load(file)
+
+    return yaml_content
+
+
+def calculate_sha256(file_path):
+    sha256_hash = hashlib.sha256()
+
+    with open(file_path, "rb") as file:
+        for chunk in iter(lambda: file.read(4096), b""):
+            sha256_hash.update(chunk)
+    return sha256_hash.hexdigest()
+
+
+def setup_repo(repo_item: dict):
+    if len(repo_item.keys()) != 1:
+        raise RuntimeError(f"Only one key required, but got {repo_item.keys()}")
+    else:
+        for key in repo_item.keys():
+            repo_name = key
+            repo_item = repo_item.pop(repo_name)
+            break
+
+    repo_url = repo_item.pop("repo_url")
+    branch = repo_item.pop("branch")
+    commit = repo_item.pop("commit", None)
+    cmds = repo_item.pop("cmds", None)
+
+    repo_path = os.path.join(".", repo_name)
+    if not os.path.exists(repo_path):
+        logger.info(f"git clone {repo_url} to {repo_name}, branch: {branch}")
+        git_repo = Repo.clone_from(repo_url, repo_path, branch=branch)
+    else:
+        logger.info(f"git repository {repo_name} has existed, use it")
+        git_repo = Repo(repo_path)
+    if commit is not None:
+        git_repo.git.checkout(commit)
+        logger.info(f"checkout {repo_name} to {commit}")
+    docker_commands = f"COPY {repo_name} /app/{repo_name}"
+    extra_cmds = ""
+    if cmds is not None:
+        extra_cmds = ["RUN ", " && \\\n".join(cmds)]
+        extra_cmds = " ".join(extra_cmds)
+        extra_cmds = "\n".join([f"WORKDIR /app/{repo_name}", extra_cmds])
+    docker_commands = "\n".join([docker_commands, extra_cmds])
+
+    return docker_commands
+
+
+def generate_docker_file(yaml_file, file_hash, output_dir, **kwargs):
+    image_config = kwargs
+
+    base_image = image_config.pop("base_image", None)
+    context_path = image_config.pop("context_path", None)
+    oneflow_pip_index = image_config.pop("oneflow_pip_index", None)
+    repos = image_config.pop("repos", None)
+    proxy = image_config.pop("proxy", None)
+    set_pip_mirror = image_config.pop("set_pip_mirror", None)
+
+    origin_file_info = f"""#==== Generated from {yaml_file} ====
+#     yaml file SHA256: {file_hash}
+"""
+
+    dockerfile_head = f"""
+#==== Docker Base Image ====
+FROM {base_image}
+"""
+
+    if set_pip_mirror is not None:
+        dockerfile_set_pip_mirror = f"RUN {set_pip_mirror}"
+    else:
+        dockerfile_set_pip_mirror = ""
+
+    dockerfile_oneflow = f"""
+#==== Install the OneFlow ====
+RUN pip install -f {oneflow_pip_index} oneflow
+"""
+
+    repos_cmds = []
+    for repo in repos:
+        repo = setup_repo(repo)
+        repos_cmds.append(repo)
+    repos_cmds = "\n\n".join(repos_cmds)
+    dockerfile_repos = f"""
+#==== Download and set up the repos ====
+{repos_cmds}
+"""
+
+    docker_post_cmds = f"""
+#==== Post setting
+WORKDIR /app
+"""
+
+    dockerfile_content = "\n".join(
+        [
+            origin_file_info,
+            dockerfile_head,
+            dockerfile_set_pip_mirror,
+            dockerfile_oneflow,
+            dockerfile_repos,
+            docker_post_cmds,
+        ]
+    )
+
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+    dockerfile_name = os.path.join(output_dir, f"Dockerfile-{file_hash[0:8]}")
+    logger.info(f"write Dockerfile to {dockerfile_name}")
+    with open(dockerfile_name, "w") as f:
+        f.write(dockerfile_content)
+    return dockerfile_name
+
+
+def build_image(docker_file, imagename, context):
+    command = ["docker", "build", "-f", docker_file, "-t", imagename, context]
+    try:
+        process = subprocess.Popen(command, stdout=subprocess.PIPE, text=True)
+        for line in iter(process.stdout.readline, ""):
+            print(line, end="")  # Print each line
+        process.wait()
+    except subprocess.CalledProcessError as e:
+        print(f"Command execution failed: {e}")
+
+
+def gen_docker_compose_yaml(container_name, image, envs, volumes, output_dir):
+    from collections import OrderedDict
+
+    onediff_benchmark_service = {
+        "container_name": None,
+        "image": None,
+        # "entrypoint": "sleep infinity",
+        "tty": True,
+        "stdin_open": True,
+        "privileged": True,
+        "shm_size": "8g",
+        "network_mode": "host",
+        "pids_limit": 2000,
+        "cap_add": ["SYS_PTRACE"],
+        "security_opt": ["seccomp=unconfined"],
+        "environment": ["HF_HUB_OFFLINE=1"],
+        "volumes": [],
+        "working_dir": "/app",
+        "restart": "no",
+        "command": "/bin/bash -c \"cd /app/onediff/benchmarks && bash run_benchmark.sh /benchmark_model\"",
+    }
+    onediff_benchmark_service["container_name"] = container_name
+    onediff_benchmark_service["image"] = image
+    onediff_benchmark_service["environment"].extend(envs)
+    onediff_benchmark_service["volumes"].extend(volumes)
+    docker_compose_dict = {
+        "version": "3.8",
+        "services": {"onediff-benchmark": onediff_benchmark_service},
+    }
+
+    yaml_string = yaml.dump(docker_compose_dict)
+
+    docker_compose_file = os.path.join(output_dir, f"docker-compose.{image}.yaml")
+
+    docker_compose_readme = f"""#========
+# run the OneDiff benchmark container by:
+#    docker compose -f {docker_compose_file} up
+#========
+
+"""
+    run_command = f"docker compose -f {docker_compose_file} up"
+    with open(docker_compose_file, "w") as f:
+        content = [docker_compose_readme, yaml_string]
+        f.write("\n".join(content))
+    return docker_compose_file, run_command

benchmarks/docker/config/community-default.yaml

@@ -0,0 +1,20 @@
+base_image: nvcr.io/nvidia/pytorch:23.08-py3
+set_pip_mirror: "pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple"
+oneflow_pip_index: "https://oneflow-pro.oss-cn-beijing.aliyuncs.com/branch/community/cu122"
+repos:
+  - onediff:
+      repo_url: https://github.com/Oneflow-Inc/onediff.git
+      branch: dev_update_benchmark_scripts
+      cmds:
+        - "python3 -m pip install transformers==4.27.1 diffusers[torch]==0.19.3"
+        - "python3 -m pip install -e ."
+  # - ComfyUI:
+  #     repo_url: "https://github.com/comfyanonymous/ComfyUI.git"
+  #     branch: master
+  #     commit: "6c5990f7dba2d5d0ad04c7ed5a702b067926cbe2"
+  #     cmds:
+  #       - "python3 -m pip install -r requirements.txt"
+proxy: ""
+volumes:
+  - '$BENCHMARK_MODEL_PATH:/benchmark_model:ro'
+envs: []