From 914bfdcdd29c1c9dbfe095eebdaa4ce9170bf109 Mon Sep 17 00:00:00 2001 From: Saurav Maheshkar Date: Tue, 15 Oct 2024 18:02:08 +0100 Subject: [PATCH] style: apply ruff --- jflux/cli.py | 6 +++--- jflux/modules/layers.py | 1 - 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/jflux/cli.py b/jflux/cli.py index fcc5ae7..4f56127 100644 --- 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 index bc70357..0f3b18e 100644 --- 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