From f88f3cbbb36e24236a23f0d041ce6b306a35127f Mon Sep 17 00:00:00 2001 From: Qing Date: Fri, 5 Jan 2024 16:38:55 +0800 Subject: [PATCH] fix import issue --- iopaint/const.py | 4 ++++ iopaint/download.py | 2 +- iopaint/model/instruct_pix2pix.py | 3 ++- iopaint/model/kandinsky.py | 4 ++-- iopaint/model/power_paint/power_paint.py | 3 ++- iopaint/model_info.py | 20 +++++++++++--------- requirements.txt | 2 +- 7 files changed, 23 insertions(+), 15 deletions(-) diff --git a/iopaint/const.py b/iopaint/const.py index a28beba..11b1241 100644 --- a/iopaint/const.py +++ b/iopaint/const.py @@ -3,6 +3,10 @@ import os from enum import Enum from pydantic import BaseModel +INSTRUCT_PIX2PIX_NAME = "timbrooks/instruct-pix2pix" +KANDINSKY22_NAME = "kandinsky-community/kandinsky-2-2-decoder-inpaint" +POWERPAINT_NAME = "Sanster/PowerPaint-V1-stable-diffusion-inpainting" + DIFFUSERS_SD_CLASS_NAME = "StableDiffusionPipeline" DIFFUSERS_SD_INPAINT_CLASS_NAME = "StableDiffusionInpaintPipeline" diff --git a/iopaint/download.py b/iopaint/download.py index 8afb44e..f8fc7c9 100644 --- a/iopaint/download.py +++ b/iopaint/download.py @@ -13,12 +13,12 @@ from iopaint.const import ( DIFFUSERS_SDXL_CLASS_NAME, DIFFUSERS_SDXL_INPAINT_CLASS_NAME, ) -from iopaint.model.utils import handle_from_pretrained_exceptions from iopaint.model_info import ModelInfo, ModelType def cli_download_model(model: str): from iopaint.model import models + from iopaint.model.utils import handle_from_pretrained_exceptions if model in models and models[model].is_erase_model: logger.info(f"Downloading {model}...") diff --git a/iopaint/model/instruct_pix2pix.py b/iopaint/model/instruct_pix2pix.py index 56bae64..1c43d46 100644 --- a/iopaint/model/instruct_pix2pix.py +++ b/iopaint/model/instruct_pix2pix.py @@ -3,12 +3,13 @@ import cv2 import torch from loguru import logger +from iopaint.const import INSTRUCT_PIX2PIX_NAME from iopaint.model.base import DiffusionInpaintModel from iopaint.schema import InpaintRequest class InstructPix2Pix(DiffusionInpaintModel): - name = "timbrooks/instruct-pix2pix" + name = INSTRUCT_PIX2PIX_NAME pad_mod = 8 min_size = 512 diff --git a/iopaint/model/kandinsky.py b/iopaint/model/kandinsky.py index 481a72e..cd2a665 100644 --- a/iopaint/model/kandinsky.py +++ b/iopaint/model/kandinsky.py @@ -3,8 +3,8 @@ import cv2 import numpy as np import torch +from iopaint.const import KANDINSKY22_NAME from iopaint.model.base import DiffusionInpaintModel -from iopaint.model.utils import get_scheduler from iopaint.schema import InpaintRequest @@ -62,4 +62,4 @@ class Kandinsky(DiffusionInpaintModel): class Kandinsky22(Kandinsky): - name = "kandinsky-community/kandinsky-2-2-decoder-inpaint" + name = KANDINSKY22_NAME diff --git a/iopaint/model/power_paint/power_paint.py b/iopaint/model/power_paint/power_paint.py index 174dbac..980cf93 100644 --- a/iopaint/model/power_paint/power_paint.py +++ b/iopaint/model/power_paint/power_paint.py @@ -9,10 +9,11 @@ from iopaint.model.helper.cpu_text_encoder import CPUTextEncoderWrapper from iopaint.model.utils import handle_from_pretrained_exceptions from iopaint.schema import InpaintRequest from .powerpaint_tokenizer import add_task_to_prompt +from ...const import POWERPAINT_NAME class PowerPaint(DiffusionInpaintModel): - name = "Sanster/PowerPaint-V1-stable-diffusion-inpainting" + name = POWERPAINT_NAME pad_mod = 8 min_size = 512 lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5" diff --git a/iopaint/model_info.py b/iopaint/model_info.py index d3f6af9..ddf5c60 100644 --- a/iopaint/model_info.py +++ b/iopaint/model_info.py @@ -6,8 +6,10 @@ from iopaint.const import ( SDXL_CONTROLNET_CHOICES, SD2_CONTROLNET_CHOICES, SD_CONTROLNET_CHOICES, + INSTRUCT_PIX2PIX_NAME, + KANDINSKY22_NAME, + POWERPAINT_NAME, ) -from iopaint.model import InstructPix2Pix, Kandinsky22, PowerPaint, SD2 from iopaint.schema import ModelType @@ -26,9 +28,9 @@ class ModelInfo(BaseModel): ModelType.DIFFUSERS_SD_INPAINT, ModelType.DIFFUSERS_SDXL_INPAINT, ] or self.name in [ - InstructPix2Pix.name, - Kandinsky22.name, - PowerPaint.name, + INSTRUCT_PIX2PIX_NAME, + KANDINSKY22_NAME, + POWERPAINT_NAME, ] @computed_field @@ -40,11 +42,11 @@ class ModelInfo(BaseModel): ]: return SDXL_CONTROLNET_CHOICES if self.model_type in [ModelType.DIFFUSERS_SD, ModelType.DIFFUSERS_SD_INPAINT]: - if self.name in [SD2.name]: + if "sd2" in self.name.lower(): return SD2_CONTROLNET_CHOICES else: return SD_CONTROLNET_CHOICES - if self.name == PowerPaint.name: + if self.name == POWERPAINT_NAME: return SD_CONTROLNET_CHOICES return [] @@ -66,7 +68,7 @@ class ModelInfo(BaseModel): ModelType.DIFFUSERS_SDXL, ModelType.DIFFUSERS_SD_INPAINT, ModelType.DIFFUSERS_SDXL_INPAINT, - ] or self.name in [Kandinsky22.name, PowerPaint.name] + ] or self.name in [KANDINSKY22_NAME, POWERPAINT_NAME] @computed_field @property @@ -86,7 +88,7 @@ class ModelInfo(BaseModel): ModelType.DIFFUSERS_SDXL, ModelType.DIFFUSERS_SD_INPAINT, ModelType.DIFFUSERS_SDXL_INPAINT, - ] or self.name in [PowerPaint.name] + ] or self.name in [POWERPAINT_NAME] @computed_field @property @@ -96,4 +98,4 @@ class ModelInfo(BaseModel): ModelType.DIFFUSERS_SDXL, ModelType.DIFFUSERS_SD_INPAINT, ModelType.DIFFUSERS_SDXL_INPAINT, - ] or self.name in [InstructPix2Pix.name] + ] or self.name in [INSTRUCT_PIX2PIX_NAME] diff --git a/requirements.txt b/requirements.txt index 649af09..f27a99f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,7 +7,7 @@ controlnet-aux==0.0.3 fastapi==0.108.0 python-socketio==5.7.2 typer -pydantic +pydantic>=2.5.2 rich loguru yacs