nunchaku-tech-ComfyUI-nunchaku/docs/source/developer/contribution_guide.rst

Contribution Guide
==================

Welcome to **ComfyUI-nunchaku**! We appreciate your interest in contributing.
This guide outlines how to set up your environment, run tests, and submit a Pull Request (PR).
Whether you're fixing a minor bug or implementing a major feature, we encourage you to
follow these steps for a smooth and efficient contribution process.

🚀 Setting Up Test Environment
------------------------------

1. Fork and Clone the Repository

   .. note::

      As a new contributor, you won't have write access to the repository.
      Please fork the repository to your own GitHub account, then clone your fork locally:

      .. code-block:: shell

         git clone https://github.com/<your_username>/ComfyUI-nunchaku.git

2. Set Up ComfyUI Test Environment

   After installing Nunchaku, set up the test environment:

   a. Install `uv <https://docs.astral.sh/uv/getting-started/installation/>`__

   b. Install test dependencies. The proper ``nunchaku`` binary will be selected based on your installed ``torch`` by default, or you can build `nunchaku` from source as an editable install.

      **Option 1: Install published wheels (recommended for typical testing):**

      .. code-block:: shell

         cd ComfyUI-nunchaku
         uv venv  # Create a new virtual environment if you don't already have one
         TORCH_VERSION=$(uv pip freeze | sed -n 's/^torch==\([0-9]\+\)\.\([0-9]\+\).*/torch\1\2/p')
         uv pip install --torch-backend=auto -e ".[${TORCH_VERSION},dev]"

      **Option 2: Build and install Nunchaku from source (for development):**

      See :ref:`Build from Source <nunchaku:build-from-source>` for detailed instructions on building and installing ``nunchaku`` from source for local development.

🧹 Code Formatting with Pre-Commit
----------------------------------

We enforce code style using `pre-commit <https://pre-commit.com/>`__ hooks. Please install and run them before submitting changes:

.. code-block:: shell

   pip install pre-commit
   pre-commit install
   pre-commit run --all-files

- Run ``pre-commit run --all-files`` until all checks pass.
- Ensure your code passes all checks before opening a pull request.
- Do **not** commit directly to the ``main`` branch. Always use a feature branch (e.g., ``feat/my-feature``) and open a PR from that branch.

🧪 Running Unit Tests & Integrating with CI
-------------------------------------------

We use ``pytest`` for unit testing. When adding features or bug fixes, include new test cases in the ``tests`` directory. **Do not modify existing tests.**

.. _running-tests:

Running Tests
~~~~~~~~~~~~~

.. code-block:: shell

   cd ComfyUI-nunchaku
   HF_TOKEN=$YOUR_HF_TOKEN pytest -v tests/ -x -vv

To run only your newly added test, use the ``-k`` flag with your workflow folder name:

.. code-block:: shell

   HF_TOKEN=$YOUR_HF_TOKEN pytest -v tests/ -x -vv -k "nunchaku-flux.1-schnell"

.. note::

   ``$YOUR_HF_TOKEN`` is your Hugging Face access token, required for downloading models and datasets. Create one at https://huggingface.co/settings/tokens. If you have already logged in with ``hf auth login``, you may omit this variable.

Writing Tests
~~~~~~~~~~~~~

When contributing new features or bug fixes, you must register a new test in the ``tests/workflows`` directory. **Do not alter existing tests.**

To add a test case:

1. **Create a Workflow Folder**

   Create a new folder in ``tests/workflows/`` with a descriptive name (e.g., ``nunchaku-flux.1-schnell``). This folder must contain four JSON files:

   - ``ref.json``: The reference workflow using BF16/FP8 models (for benchmarking)
   - ``workflow.json``: The corresponding Nunchaku version of the workflow
   - ``api.json``: API version of ``workflow.json`` (exported via ComfyUI's ``Export (API)`` option)
   - ``test_cases.json``: Test configurations with different parameters

   .. note::

      Both ``ref.json`` and ``workflow.json`` are for backup purposes, making it easier for future maintenance, development, testing, and debugging.

2. **Create the API Workflow**

   In ComfyUI, after designing your workflow, export it using ``Export (API)`` and save it as ``api.json`` (see example below).

   .. image:: https://huggingface.co/datasets/nunchaku-tech/cdn/resolve/main/ComfyUI-nunchaku/export_api.png
      :alt: ComfyUI Export API Example
      :align: center

3. **Configure Test Cases**

   Create ``test_cases.json`` to define test parameters. You can override variables in ``api.json`` using the ``inputs`` field. Here's an example:

   .. code-block:: json

      [
        {
          "ref_image_url": "https://github.com/user/repo/issues/123#issuecomment-456789",
          "expected_clip_iqa": {
            "int4-bf16": 0.98,
            "fp4-bf16": 0.99
          },
          "expected_lpips": {
            "int4-bf16": 0.23,
            "fp4-bf16": 0.22
          },
          "expected_psnr": {
            "int4-bf16": 19,
            "fp4-bf16": 19
          },
          "inputs": {
            "30,inputs,model_path": "svdq-{precision}_r32-flux.1-schnell.safetensors",
            "25,inputs,noise_seed": 778459239
          }
        }
      ]

   Each test case should include:

   - ``ref_image_url``: URL to the reference image generated by your BF16/FP8 workflow with the same parameters (remember to fix the seed). Upload the image to a GitHub PR comment to get a public URL.
   - ``expected_clip_iqa``, ``expected_lpips``, ``expected_psnr``: Image quality metrics. These keys use the format ``{precision}-{torch_dtype}``:

     - ``int4``/``fp4``: Nunchaku model precision
     - ``bf16``/``fp16``: Activation torch dtype (fp16 is typically used on RTX 20-series GPUs; others use bf16)

   - ``inputs``: Override parameters in ``api.json`` for testing different configurations

   **How to determine the expected values:**

   Run your test locally first (see :ref:`running-tests`). Use the local results as reference values. If you can only test one precision type (int4 or fp4), you can use the same reference values for both.

4. **Add Additional Test Data (if needed)**

   If your test requires additional input images or models:

   - Upload input images to a GitHub PR comment to get a public URL
   - Register the URLs in `test_data/inputs.yaml <https://github.com/nunchaku-tech/ComfyUI-nunchaku/blob/main/test_data/inputs.yaml>`__
   - If new models are required, update `scripts/download_models.py <https://github.com/nunchaku-tech/ComfyUI-nunchaku/blob/main/scripts/download_models.py>`__ and `test_data/models.yaml <https://github.com/nunchaku-tech/ComfyUI-nunchaku/blob/main/test_data/models.yaml>`__

5. **Install Additional Nodes (if needed)**

   If your test requires additional custom nodes, modify `.github/workflows/pr-test.yaml <https://github.com/nunchaku-tech/ComfyUI-nunchaku/blob/main/.github/workflows/pr-test.yaml>`__ to install the required nodes and their dependencies in the CI environment.