style: apply ruff

ml-gde · Oct 15, 2024 · 914bfdc · 914bfdc
1 parent 7c14c22
commit 914bfdc
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/jflux/cli.py b/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

diff --git a/jflux/modules/layers.py b/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