fix(ci): cache-dependency-glob (#24)

* drops redundant paths from GA workflow
* fixes `cache-dependency-glob`
* enables capturing std-out on CI 
* test on `3.10` and above

Author: SauravMaheshkar

.github/workflows/python.yml

@@ -5,46 +5,46 @@ on:
     branches: [main]
     paths:
       - "**.py"
-      - ".devcontainer/requirements.txt"
       - ".github/workflows/python.yml"
   pull_request:
     branches: [main]
     paths:
       - "**.py"
-      - ".devcontainer/requirements.txt"
       - ".github/workflows/python.yml"
   release:
     types: [created]
   schedule:
     - cron: "0 0 * * 0"
+  workflow_dispatch:
 
 jobs:
   build:
     runs-on: ${{ matrix.os }}
 
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12"]
         os: [ubuntu-latest, macos-latest]
 
     steps:
       - uses: actions/checkout@v4
 
-      - name: Install uv
-        uses: astral-sh/setup-uv@v3
+      - name: Setup Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
         with:
-          enable-cache: true
-          cache-dependency-glob: "uv.lock"
-
-      - name: Set up Python ${{ matrix.python-version }}
-        run: uv python install ${{ matrix.python-version }}
+          python-version: ${{ matrix.python-version }}
+          cache: "pip"
+          cache-dependency-path: "pyproject.toml"
 
       - name: Install dependencies
-        run: uv sync --all-extras --dev
-
-      - name: Ruff
         run: |
-          uv run ruff check jflux
+          pip install -U pip
+          pip install jax flax fire pytest ruff
+          pip install git+https://github.com/black-forest-labs/flux.git
+
+      - name: ruff
+        run: ruff check jflux
+      
       - name: Test with PyTest
         run: |
-          uv run pytest -v .
+          pytest -v -s .

jflux/cli.py

@@ -4,19 +4,19 @@
 from dataclasses import dataclass
 from glob import iglob
 
-import torch
 import jax
 import jax.numpy as jnp
 import numpy as np
+import torch
 from einops import rearrange
 from fire import Fire
 from flax import nnx
 from PIL import Image
 
 from jflux.sampling import denoise, get_noise, get_schedule, prepare, unpack
-from jflux.util import configs, load_ae, load_clip, load_flow_model, load_t5, torch2jax
+from jflux.util import configs, load_ae, load_clip, load_flow_model, load_t5
+
 
-import os
 os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
 
 @dataclass

jflux/modules/layers.py

@@ -1,7 +1,6 @@
 import math
 from dataclasses import dataclass
 
-import jax
 import jax.numpy as jnp
 from chex import Array
 from einops import rearrange