feat: initial nnx port (#3)

Fixes #2 --------- Co-authored-by: Saurav Maheshkar <[email protected]>
ml-gde · Sep 29, 2024 · 0e8491f · 0e8491f
1 parent 7b744dd
commit 0e8491f
Show file tree

Hide file tree

Showing 23 changed files with 4,194 additions and 107 deletions.
diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
diff --git a/.devcontainer/requirements.txt b/.devcontainer/requirements.txt
diff --git a/.github/README.md b/.github/README.md
@@ -1,3 +1,20 @@
 # FL(U/A)X
 
 JAX Implementation of Black Forest Labs' Flux.1 family of models
+
+
+## Installation
+
+```shell
+$ uv sync
+```
+
+## Running
+
+```shell
+$ uv jflux
+```
+
+## References
+
+* Original Implementation: [black-forest-labs/flux](https://github.com/black-forest-labs/flux)
diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml
@@ -29,24 +29,22 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-      - name: Setup Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v2
         with:
-          python-version: ${{ matrix.python-version }}
-          cache: "pip"
-          cache-dependency-path: ".devcontainer/requirements.txt"
+          enable-cache: true
+          cache-dependency-glob: "uv.lock"
+
+      - name: Set up Python ${{ matrix.python-version }}
+        run: uv python install ${{ matrix.python-version }}
 
       - name: Install dependencies
-        run: |
-          python -m venv .venv && export PATH=".venv/bin:$PATH"
-          python -m pip install uv
-          python -m uv pip install --upgrade wheel setuptools
-          python -m uv pip install -r .devcontainer/requirements.txt
+        run: uv sync --all-extras --dev
+
       - name: Ruff
         run: |
-          python -m venv .venv && export PATH=".venv/bin:$PATH"
-          python -m ruff check src
+          uv run ruff check jflux
       - name: Test with PyTest
         run: |
-          python -m venv .venv && export PATH=".venv/bin:$PATH"
-          python -m pytest -v .
+          uv run pytest -v .
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -14,10 +14,6 @@ repos:
     rev: 24.8.0
     hooks:
       - id: black
-  - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.11.2
-    hooks:
-      - id: mypy
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.6.2
     hooks:

diff --git a/src/__init__.py → jflux/__init__.py b/src/__init__.py → jflux/__init__.py
diff --git a/jflux/__main__.py b/jflux/__main__.py
@@ -0,0 +1,4 @@
+from jflux.cli import app
+
+if __name__ == "__main__":
+    app()