From 19d0b563ac96f8a40deda655c70c427a246d1876 Mon Sep 17 00:00:00 2001 From: Bryce Date: Fri, 16 Sep 2022 09:24:24 -0700 Subject: [PATCH] style: lint fixes. remove unused code --- README.md | 3 +- imaginairy/api.py | 10 ++-- imaginairy/cmds.py | 12 +---- .../enhancers/face_restoration_codeformer.py | 2 +- .../enhancers/face_restoration_gfpgan.py | 50 ------------------- imaginairy/img_log.py | 4 +- imaginairy/schema.py | 4 +- imaginairy/utils.py | 24 +++++---- 8 files changed, 25 insertions(+), 84 deletions(-) delete mode 100644 imaginairy/enhancers/face_restoration_gfpgan.py diff --git a/README.md b/README.md index 6313c0d..736c2ea 100644 --- a/README.md +++ b/README.md @@ -171,9 +171,10 @@ docker run -it --gpus all -v $HOME/.cache/huggingface:/root/.cache/huggingface - - outpainting - inpainting - https://github.com/andreas128/RePaint - - img2img but keeps img stable + - img2img but keeps img stable - https://www.reddit.com/r/StableDiffusion/comments/xboy90/a_better_way_of_doing_img2img_by_finding_the/ - https://gist.github.com/trygvebw/c71334dd127d537a15e9d59790f7f5e1 + - https://github.com/pesser/stable-diffusion/commit/bbb52981460707963e2a62160890d7ecbce00e79 - CPU support - img2img for plms? - images as actual prompts instead of just init images diff --git a/imaginairy/api.py b/imaginairy/api.py index 4f8c06c..92743c5 100755 --- a/imaginairy/api.py +++ b/imaginairy/api.py @@ -165,7 +165,6 @@ def imagine( prompts = [ImaginePrompt(prompts)] if isinstance(prompts, str) else prompts prompts = [prompts] if isinstance(prompts, ImaginePrompt) else prompts _img_callback = None - step_count = 0 precision_scope = ( autocast @@ -185,11 +184,8 @@ def imagine( uc = model.get_learned_conditioning(1 * [""]) total_weight = sum(wp.weight for wp in prompt.prompts) c = sum( - [ - model.get_learned_conditioning(wp.text) - * (wp.weight / total_weight) - for wp in prompt.prompts - ] + model.get_learned_conditioning(wp.text) * (wp.weight / total_weight) + for wp in prompt.prompts ) shape = [ @@ -205,7 +201,7 @@ def imagine( ddim_steps = int(prompt.steps / generation_strength) sampler.make_schedule(ddim_num_steps=ddim_steps, ddim_eta=ddim_eta) - init_image, w, h = pillow_img_to_torch_image( + init_image, w, h = pillow_img_to_torch_image( # noqa prompt.init_image, max_height=prompt.height, max_width=prompt.width, diff --git a/imaginairy/cmds.py b/imaginairy/cmds.py index 99954b4..5bb5e75 100644 --- a/imaginairy/cmds.py +++ b/imaginairy/cmds.py @@ -3,8 +3,9 @@ import logging.config import click from imaginairy import LazyLoadingImage -from imaginairy.api import load_model +from imaginairy.api import imagine_image_files, load_model from imaginairy.samplers.base import SAMPLER_TYPE_OPTIONS +from imaginairy.schema import ImaginePrompt from imaginairy.suppress_logs import suppress_annoying_logs_and_warnings logger = logging.getLogger(__name__) @@ -96,11 +97,7 @@ def configure_logging(level="INFO"): help="What seed to use for randomness. Allows reproducible image renders", ) @click.option("--upscale", is_flag=True) -@click.option( - "--upscale-method", default="realesrgan", type=click.Choice(["realesrgan"]) -) @click.option("--fix-faces", is_flag=True) -@click.option("--fix-faces-method", default="gfpgan", type=click.Choice(["gfpgan"])) @click.option( "--sampler-type", default="plms", @@ -138,9 +135,7 @@ def imagine_cmd( steps, seed, upscale, - upscale_method, fix_faces, - fix_faces_method, sampler_type, ddim_eta, log_level, @@ -151,9 +146,6 @@ def imagine_cmd( suppress_annoying_logs_and_warnings() configure_logging(log_level) - from imaginairy.api import imagine_image_files - from imaginairy.schema import ImaginePrompt - total_image_count = len(prompt_texts) * repeats logger.info( f"🤖🧠 imaginAIry received {len(prompt_texts)} prompt(s) and will repeat them {repeats} times to create {total_image_count} images." diff --git a/imaginairy/enhancers/face_restoration_codeformer.py b/imaginairy/enhancers/face_restoration_codeformer.py index 2a38196..dcc4314 100644 --- a/imaginairy/enhancers/face_restoration_codeformer.py +++ b/imaginairy/enhancers/face_restoration_codeformer.py @@ -60,7 +60,7 @@ def enhance_faces(img, fidelity=0): face_helper.align_warp_face() # face restoration for each cropped face - for idx, cropped_face in enumerate(face_helper.cropped_faces): + for cropped_face in face_helper.cropped_faces: # prepare data cropped_face_t = img2tensor(cropped_face / 255.0, bgr2rgb=True, float32=True) normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True) diff --git a/imaginairy/enhancers/face_restoration_gfpgan.py b/imaginairy/enhancers/face_restoration_gfpgan.py deleted file mode 100644 index c0ff5ca..0000000 --- a/imaginairy/enhancers/face_restoration_gfpgan.py +++ /dev/null @@ -1,50 +0,0 @@ -from functools import lru_cache - -import numpy as np -import torch -from PIL import Image - -from imaginairy.utils import get_cached_url_path, get_device - - -@lru_cache() -def face_enhance_model(model_type="codeformer"): - from gfpgan import GFPGANer - - if model_type == "gfpgan": - arch = "clean" - url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth" - elif model_type == "codeformer": - arch = "CodeFormer" - url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/CodeFormer.pth" - else: - raise ValueError("model_type must be one of gfpgan, codeformer") - - model_path = get_cached_url_path(url) - - if get_device() == "cuda": - device = "cuda" - else: - device = "cpu" - - return GFPGANer( - model_path=model_path, - upscale=1, - arch=arch, - channel_multiplier=2, - bg_upsampler=None, - device=device, - ) - - -def fix_faces_gfpgan(image, model_type): - image = image.convert("RGB") - np_img = np.array(image, dtype=np.uint8) - cropped_faces, restored_faces, restored_img = face_enhance_model( - model_type - ).enhance( - np_img, has_aligned=False, only_center_face=False, paste_back=True, weight=0 - ) - res = Image.fromarray(restored_img) - - return res diff --git a/imaginairy/img_log.py b/imaginairy/img_log.py index ad2a3ca..7610f3b 100644 --- a/imaginairy/img_log.py +++ b/imaginairy/img_log.py @@ -31,12 +31,12 @@ class LatentLoggingContext: self.img_callback = img_callback def __enter__(self): - global _CURRENT_LOGGING_CONTEXT + global _CURRENT_LOGGING_CONTEXT # noqa _CURRENT_LOGGING_CONTEXT = self return self def __exit__(self, exc_type, exc_val, exc_tb): - global _CURRENT_LOGGING_CONTEXT + global _CURRENT_LOGGING_CONTEXT # noqa _CURRENT_LOGGING_CONTEXT = None def log_latents(self, latents, description): diff --git a/imaginairy/schema.py b/imaginairy/schema.py index 8cacb32..dea5cca 100644 --- a/imaginairy/schema.py +++ b/imaginairy/schema.py @@ -36,7 +36,7 @@ class LazyLoadingImage: try: parsed_url = parse_url(url) except LocationParseError: - raise InvalidUrlError(f"Invalid url: {url}") + raise InvalidUrlError(f"Invalid url: {url}") # noqa if parsed_url.scheme not in {"http", "https"} or not parsed_url.host: raise InvalidUrlError(f"Invalid url: {url}") @@ -55,7 +55,7 @@ class LazyLoadingImage: f"Loaded input 🖼 of size {self._img.size} from {self._lazy_filepath}" ) elif self._lazy_url: - self._img = Image.open(requests.get(self._lazy_url, stream=True).raw) + self._img = Image.open(requests.get(self._lazy_url, stream=True, timeout=60).raw) logger.info( f"Loaded input 🖼 of size {self._img.size} from {self._lazy_url}" ) diff --git a/imaginairy/utils.py b/imaginairy/utils.py index ab60915..f9bd064 100644 --- a/imaginairy/utils.py +++ b/imaginairy/utils.py @@ -11,6 +11,7 @@ import requests import torch from PIL import Image from torch import Tensor +from torch.nn import functional from torch.overrides import handle_torch_function, has_torch_function_variadic from transformers import cached_path @@ -21,10 +22,11 @@ logger = logging.getLogger(__name__) def get_device(): if torch.cuda.is_available(): return "cuda" - elif torch.backends.mps.is_available(): + + if torch.backends.mps.is_available(): return "mps" - else: - return "cpu" + + return "cpu" @lru_cache() @@ -40,13 +42,13 @@ def log_params(model): def instantiate_from_config(config): - if not "target" in config: + if "target" not in config: if config == "__is_first_stage__": return None - elif config == "__is_unconditional__": + if config == "__is_unconditional__": return None raise KeyError("Expected key `target` to instantiate.") - return get_obj_from_str(config["target"])(**config.get("params", dict())) + return get_obj_from_str(config["target"])(**config.get("params", {})) def get_obj_from_str(string, reload=False): @@ -58,13 +60,15 @@ def get_obj_from_str(string, reload=False): def _fixed_layer_norm( - input: Tensor, + input: Tensor, # noqa normalized_shape: List[int], weight: Optional[Tensor] = None, bias: Optional[Tensor] = None, eps: float = 1e-5, ) -> Tensor: - r"""Applies Layer Normalization for last certain number of dimensions. + """ + Applies Layer Normalization for last certain number of dimensions. + See :class:`~torch.nn.LayerNorm` for details. """ if has_torch_function_variadic(input, weight, bias): @@ -90,8 +94,6 @@ def _fixed_layer_norm( @contextmanager def fix_torch_nn_layer_norm(): """https://github.com/CompVis/stable-diffusion/issues/25#issuecomment-1221416526""" - from torch.nn import functional - orig_function = functional.layer_norm functional.layer_norm = _fixed_layer_norm try: @@ -143,7 +145,7 @@ def get_cached_url_path(url): dest_path = os.path.join(dest, filename) if os.path.exists(dest_path): return dest_path - r = requests.get(url) + r = requests.get(url) # noqa with open(dest_path, "wb") as f: f.write(r.content)