feat: initial nnx port (#3)

Fixes #2

---------

Co-authored-by: Saurav Maheshkar <[email protected]>

Author: ariG23498

.devcontainer/devcontainer.json

@@ -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)

.devcontainer/requirements.txt

@@ -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 .

.github/README.md

@@ -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:

.github/workflows/python.yml

@@ -0,0 +1,4 @@
+from jflux.cli import app
+
+if __name__ == "__main__":
+    app()