From 438c2868adde7de8a4be369d607a2ef6b2aac9f1 Mon Sep 17 00:00:00 2001 From: Bryce Date: Sun, 11 Sep 2022 13:58:14 -0700 Subject: [PATCH] refactor: run import sorter --- imaginairy/__init__.py | 2 +- imaginairy/api.py | 6 +++--- imaginairy/cmd_wrap.py | 1 - imaginairy/cmds.py | 1 + imaginairy/modules/attention.py | 7 ++++--- imaginairy/modules/autoencoder.py | 2 +- imaginairy/modules/clip_embedders.py | 2 +- imaginairy/modules/diffusion/ddim.py | 2 +- imaginairy/modules/diffusion/ddpm.py | 11 +++-------- imaginairy/modules/diffusion/model.py | 2 +- imaginairy/modules/diffusion/openaimodel.py | 8 ++++---- imaginairy/modules/distributions.py | 2 +- imaginairy/utils.py | 2 +- setup.py | 2 +- tests/test_imagine.py | 3 ++- 15 files changed, 25 insertions(+), 28 deletions(-) diff --git a/imaginairy/__init__.py b/imaginairy/__init__.py index bc5647c..db9b2ee 100644 --- a/imaginairy/__init__.py +++ b/imaginairy/__init__.py @@ -2,5 +2,5 @@ import os os.putenv("PYTORCH_ENABLE_MPS_FALLBACK", "1") -from .api import imagine_images, imagine_image_files # noqa +from .api import imagine_image_files, imagine_images # noqa from .schema import ImaginePrompt, ImagineResult, WeightedPrompt # noqa diff --git a/imaginairy/api.py b/imaginairy/api.py index da6a893..bc84cd2 100755 --- a/imaginairy/api.py +++ b/imaginairy/api.py @@ -5,13 +5,13 @@ import subprocess from contextlib import nullcontext from functools import lru_cache -import PIL import numpy as np +import PIL import torch import torch.nn -from PIL import Image from einops import rearrange from omegaconf import OmegaConf +from PIL import Image from pytorch_lightning import seed_everything from torch import autocast from transformers import cached_path @@ -21,9 +21,9 @@ from imaginairy.modules.diffusion.plms import PLMSSampler from imaginairy.safety import is_nsfw from imaginairy.schema import ImaginePrompt, ImagineResult from imaginairy.utils import ( + fix_torch_nn_layer_norm, get_device, instantiate_from_config, - fix_torch_nn_layer_norm, ) LIB_PATH = os.path.dirname(__file__) diff --git a/imaginairy/cmd_wrap.py b/imaginairy/cmd_wrap.py index 2c0b094..b49fa08 100644 --- a/imaginairy/cmd_wrap.py +++ b/imaginairy/cmd_wrap.py @@ -59,6 +59,5 @@ setup_env() from imaginairy.cmds import imagine_cmd # noqa - if __name__ == "__main__": imagine_cmd() # noqa diff --git a/imaginairy/cmds.py b/imaginairy/cmds.py index 574c058..980831b 100644 --- a/imaginairy/cmds.py +++ b/imaginairy/cmds.py @@ -1,6 +1,7 @@ import logging.config import click + from imaginairy.api import load_model logger = logging.getLogger(__name__) diff --git a/imaginairy/modules/attention.py b/imaginairy/modules/attention.py index 624a3bb..d130969 100644 --- a/imaginairy/modules/attention.py +++ b/imaginairy/modules/attention.py @@ -1,12 +1,13 @@ -from inspect import isfunction import math +from inspect import isfunction + import torch import torch.nn.functional as F -from torch import nn, einsum from einops import rearrange, repeat +from torch import einsum, nn from imaginairy.modules.diffusion.util import checkpoint -from imaginairy.utils import get_device_name, get_device +from imaginairy.utils import get_device, get_device_name def exists(val): diff --git a/imaginairy/modules/autoencoder.py b/imaginairy/modules/autoencoder.py index 34d2fd1..247b2a4 100644 --- a/imaginairy/modules/autoencoder.py +++ b/imaginairy/modules/autoencoder.py @@ -6,7 +6,7 @@ import torch import torch.nn as nn from einops import rearrange -from imaginairy.modules.diffusion.model import Encoder, Decoder +from imaginairy.modules.diffusion.model import Decoder, Encoder from imaginairy.modules.distributions import DiagonalGaussianDistribution from imaginairy.utils import instantiate_from_config diff --git a/imaginairy/modules/clip_embedders.py b/imaginairy/modules/clip_embedders.py index ff12938..9186ef0 100644 --- a/imaginairy/modules/clip_embedders.py +++ b/imaginairy/modules/clip_embedders.py @@ -3,7 +3,7 @@ import kornia import torch import torch.nn as nn from einops import repeat -from transformers import CLIPTokenizer, CLIPTextModel +from transformers import CLIPTextModel, CLIPTokenizer from imaginairy.utils import get_device diff --git a/imaginairy/modules/diffusion/ddim.py b/imaginairy/modules/diffusion/ddim.py index 6f6427d..c7bd1a5 100644 --- a/imaginairy/modules/diffusion/ddim.py +++ b/imaginairy/modules/diffusion/ddim.py @@ -6,10 +6,10 @@ import torch from tqdm import tqdm from imaginairy.modules.diffusion.util import ( + extract_into_tensor, make_ddim_sampling_parameters, make_ddim_timesteps, noise_like, - extract_into_tensor, ) from imaginairy.utils import get_device diff --git a/imaginairy/modules/diffusion/ddpm.py b/imaginairy/modules/diffusion/ddpm.py index 4cc8ff6..0f31af0 100644 --- a/imaginairy/modules/diffusion/ddpm.py +++ b/imaginairy/modules/diffusion/ddpm.py @@ -16,15 +16,10 @@ from einops import rearrange from torchvision.utils import make_grid from tqdm import tqdm -from imaginairy.modules.autoencoder import ( - VQModelInterface, -) -from imaginairy.modules.diffusion.util import ( - make_beta_schedule, - noise_like, -) +from imaginairy.modules.autoencoder import VQModelInterface +from imaginairy.modules.diffusion.util import make_beta_schedule, noise_like from imaginairy.modules.distributions import DiagonalGaussianDistribution -from imaginairy.utils import log_params, instantiate_from_config +from imaginairy.utils import instantiate_from_config, log_params logger = logging.getLogger(__name__) __conditioning_keys__ = {"concat": "c_concat", "crossattn": "c_crossattn", "adm": "y"} diff --git a/imaginairy/modules/diffusion/model.py b/imaginairy/modules/diffusion/model.py index 179d3e0..abc688d 100644 --- a/imaginairy/modules/diffusion/model.py +++ b/imaginairy/modules/diffusion/model.py @@ -10,7 +10,7 @@ from einops import rearrange from imaginairy.modules.attention import LinearAttention from imaginairy.modules.distributions import DiagonalGaussianDistribution -from imaginairy.utils import instantiate_from_config, get_device +from imaginairy.utils import get_device, instantiate_from_config logger = logging.getLogger(__name__) diff --git a/imaginairy/modules/diffusion/openaimodel.py b/imaginairy/modules/diffusion/openaimodel.py index bb6a261..2690771 100644 --- a/imaginairy/modules/diffusion/openaimodel.py +++ b/imaginairy/modules/diffusion/openaimodel.py @@ -1,21 +1,21 @@ -from abc import abstractmethod import math +from abc import abstractmethod import numpy as np import torch as th import torch.nn as nn import torch.nn.functional as F +from imaginairy.modules.attention import SpatialTransformer from imaginairy.modules.diffusion.util import ( + avg_pool_nd, checkpoint, conv_nd, linear, - avg_pool_nd, - zero_module, normalization, timestep_embedding, + zero_module, ) -from imaginairy.modules.attention import SpatialTransformer # dummy replace diff --git a/imaginairy/modules/distributions.py b/imaginairy/modules/distributions.py index 0b61f03..016be35 100644 --- a/imaginairy/modules/distributions.py +++ b/imaginairy/modules/distributions.py @@ -1,5 +1,5 @@ -import torch import numpy as np +import torch class AbstractDistribution: diff --git a/imaginairy/utils.py b/imaginairy/utils.py index 8c084fa..7bce237 100644 --- a/imaginairy/utils.py +++ b/imaginairy/utils.py @@ -51,7 +51,7 @@ def get_obj_from_str(string, reload=False): return getattr(importlib.import_module(module, package=None), cls) -from torch.overrides import has_torch_function_variadic, handle_torch_function +from torch.overrides import handle_torch_function, has_torch_function_variadic def _fixed_layer_norm( diff --git a/setup.py b/setup.py index eab654e..fe0b57c 100644 --- a/setup.py +++ b/setup.py @@ -1,4 +1,4 @@ -from setuptools import setup, find_packages +from setuptools import find_packages, setup setup( name="imaginairy", diff --git a/tests/test_imagine.py b/tests/test_imagine.py index e6ce7f5..bc56bf1 100644 --- a/tests/test_imagine.py +++ b/tests/test_imagine.py @@ -1,5 +1,6 @@ -from imaginairy.api import imagine_images, imagine_image_files +from imaginairy.api import imagine_image_files, imagine_images from imaginairy.schema import ImaginePrompt + from . import TESTS_FOLDER