feat: chex-ify testsuite

.github/README.md

@@ -12,7 +12,7 @@ $ uv sync
 ## Running
 
 ```shell
-$ uv jflux
+$ uv run jflux
 ```
 
 ## References

.vscode/settings.json

@@ -0,0 +1,15 @@
+{
+    "python.testing.pytestArgs": [
+        "tests"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true,
+    "[python]": {
+        "editor.formatOnSave": true,
+        "editor.codeActionsOnSave": {
+            "source.fixAll": "explicit",
+            "source.organizeImports": "explicit"
+          },
+        "editor.defaultFormatter": "charliermarsh.ruff",
+    },
+}

jflux/cli.py

@@ -6,10 +6,9 @@
 
 import jax
 import jax.numpy as jnp
+import numpy as np
 from einops import rearrange
 from fire import Fire
-from flax import nnx
-from jax.typing import DTypeLike
 from PIL import Image
 
 from jflux.sampling import denoise, get_noise, get_schedule, prepare, unpack
@@ -124,7 +123,8 @@ def main(
             by the index of the sample
         prompt: Prompt used for sampling
         device: Pytorch device
-        num_steps: number of sampling steps (default 4 for schnell, 50 for guidance distilled)
+        num_steps: number of sampling steps
+            (default 4 for schnell, 50 for guidance distilled)
         loop: start an interactive session and sample multiple times
         guidance: guidance value used for guidance distillation
         add_sampling_metadata: Add the prompt to the image Exif metadata
@@ -216,7 +216,12 @@ def main(
         x = x.clip(-1, 1)
         x = rearrange(x[0], "c h w -> h w c")
 
-        img = Image.fromarray((127.5 * (x + 1.0)))
+        x = 127.5 * (x + 1.0)
+        x_numpy = np.array(x.astype(jnp.uint8))
+        img = Image.fromarray(x_numpy)
+
+        img.save(fn, quality=95, subsampling=0)
+        idx += 1
 
         if loop:
             print("-" * 80)

jflux/model.py

@@ -49,7 +49,7 @@ def __init__(self, params: FluxParams):
         self.out_channels = self.in_channels
         if params.hidden_size % params.num_heads != 0:
             raise ValueError(
-                f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
+                f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"  # noqa: E501
             )
         pe_dim = params.hidden_size // params.num_heads
         if sum(params.axes_dim) != pe_dim:

jflux/port.py

@@ -1,5 +1,6 @@
 from einops import rearrange
 
+
 ##############################################################################################
 # AUTOENCODER MODEL PORTING
 ##############################################################################################
@@ -481,3 +482,5 @@ def port_flux(flux, tensors):
         tensors=tensors,
         prefix="final_layer",
     )
+
+    return flux

jflux/util.py

@@ -1,12 +1,10 @@
 import os
 from dataclasses import dataclass
 
-import jax
 import torch  # need for t5 and clip
 from flax import nnx
 from huggingface_hub import hf_hub_download
 from jax import numpy as jnp
-from jax.typing import DTypeLike
 from safetensors import safe_open
 
 from jflux.model import Flux, FluxParams

pyproject.toml

@@ -9,7 +9,6 @@ dependencies = [
     "einops>=0.8.0",
     "fire>=0.6.0",
     "flax>=0.9.0",
-    "jflux",
     # FIXME: Allow for local installation without GPUs as well `jax[cuda12]`
     "jax>=0.4.31",
     "mypy>=1.11.2",
@@ -22,6 +21,7 @@ dependencies = [
 jflux = "jflux.cli:app"
 
 [tool.uv]
+package = true
 dev-dependencies = [
     "flux",
     "pytest>=8.3.3",
@@ -32,7 +32,10 @@ jflux = { workspace = true }
 flux = { git = "https://github.com/black-forest-labs/flux.git" }
 
 [tool.ruff.lint]
-select = ["I001"]
+select = ["E", "F", "I001", "W"]
+
+[tool.ruff.lint.isort]
+lines-after-imports = 2
 
 [tool.ruff.lint.pydocstyle]
 convention = "google"