fix import issue

This commit is contained in:
Qing 2024-01-05 16:38:55 +08:00
parent 87d9078fe3
commit f88f3cbbb3
7 changed files with 23 additions and 15 deletions

View File

@ -3,6 +3,10 @@ import os
from enum import Enum from enum import Enum
from pydantic import BaseModel 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_CLASS_NAME = "StableDiffusionPipeline"
DIFFUSERS_SD_INPAINT_CLASS_NAME = "StableDiffusionInpaintPipeline" DIFFUSERS_SD_INPAINT_CLASS_NAME = "StableDiffusionInpaintPipeline"

View File

@ -13,12 +13,12 @@ from iopaint.const import (
DIFFUSERS_SDXL_CLASS_NAME, DIFFUSERS_SDXL_CLASS_NAME,
DIFFUSERS_SDXL_INPAINT_CLASS_NAME, DIFFUSERS_SDXL_INPAINT_CLASS_NAME,
) )
from iopaint.model.utils import handle_from_pretrained_exceptions
from iopaint.model_info import ModelInfo, ModelType from iopaint.model_info import ModelInfo, ModelType
def cli_download_model(model: str): def cli_download_model(model: str):
from iopaint.model import models from iopaint.model import models
from iopaint.model.utils import handle_from_pretrained_exceptions
if model in models and models[model].is_erase_model: if model in models and models[model].is_erase_model:
logger.info(f"Downloading {model}...") logger.info(f"Downloading {model}...")

View File

@ -3,12 +3,13 @@ import cv2
import torch import torch
from loguru import logger from loguru import logger
from iopaint.const import INSTRUCT_PIX2PIX_NAME
from iopaint.model.base import DiffusionInpaintModel from iopaint.model.base import DiffusionInpaintModel
from iopaint.schema import InpaintRequest from iopaint.schema import InpaintRequest
class InstructPix2Pix(DiffusionInpaintModel): class InstructPix2Pix(DiffusionInpaintModel):
name = "timbrooks/instruct-pix2pix" name = INSTRUCT_PIX2PIX_NAME
pad_mod = 8 pad_mod = 8
min_size = 512 min_size = 512

View File

@ -3,8 +3,8 @@ import cv2
import numpy as np import numpy as np
import torch import torch
from iopaint.const import KANDINSKY22_NAME
from iopaint.model.base import DiffusionInpaintModel from iopaint.model.base import DiffusionInpaintModel
from iopaint.model.utils import get_scheduler
from iopaint.schema import InpaintRequest from iopaint.schema import InpaintRequest
@ -62,4 +62,4 @@ class Kandinsky(DiffusionInpaintModel):
class Kandinsky22(Kandinsky): class Kandinsky22(Kandinsky):
name = "kandinsky-community/kandinsky-2-2-decoder-inpaint" name = KANDINSKY22_NAME

View File

@ -9,10 +9,11 @@ from iopaint.model.helper.cpu_text_encoder import CPUTextEncoderWrapper
from iopaint.model.utils import handle_from_pretrained_exceptions from iopaint.model.utils import handle_from_pretrained_exceptions
from iopaint.schema import InpaintRequest from iopaint.schema import InpaintRequest
from .powerpaint_tokenizer import add_task_to_prompt from .powerpaint_tokenizer import add_task_to_prompt
from ...const import POWERPAINT_NAME
class PowerPaint(DiffusionInpaintModel): class PowerPaint(DiffusionInpaintModel):
name = "Sanster/PowerPaint-V1-stable-diffusion-inpainting" name = POWERPAINT_NAME
pad_mod = 8 pad_mod = 8
min_size = 512 min_size = 512
lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5" lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5"

View File

@ -6,8 +6,10 @@ from iopaint.const import (
SDXL_CONTROLNET_CHOICES, SDXL_CONTROLNET_CHOICES,
SD2_CONTROLNET_CHOICES, SD2_CONTROLNET_CHOICES,
SD_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 from iopaint.schema import ModelType
@ -26,9 +28,9 @@ class ModelInfo(BaseModel):
ModelType.DIFFUSERS_SD_INPAINT, ModelType.DIFFUSERS_SD_INPAINT,
ModelType.DIFFUSERS_SDXL_INPAINT, ModelType.DIFFUSERS_SDXL_INPAINT,
] or self.name in [ ] or self.name in [
InstructPix2Pix.name, INSTRUCT_PIX2PIX_NAME,
Kandinsky22.name, KANDINSKY22_NAME,
PowerPaint.name, POWERPAINT_NAME,
] ]
@computed_field @computed_field
@ -40,11 +42,11 @@ class ModelInfo(BaseModel):
]: ]:
return SDXL_CONTROLNET_CHOICES return SDXL_CONTROLNET_CHOICES
if self.model_type in [ModelType.DIFFUSERS_SD, ModelType.DIFFUSERS_SD_INPAINT]: 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 return SD2_CONTROLNET_CHOICES
else: else:
return SD_CONTROLNET_CHOICES return SD_CONTROLNET_CHOICES
if self.name == PowerPaint.name: if self.name == POWERPAINT_NAME:
return SD_CONTROLNET_CHOICES return SD_CONTROLNET_CHOICES
return [] return []
@ -66,7 +68,7 @@ class ModelInfo(BaseModel):
ModelType.DIFFUSERS_SDXL, ModelType.DIFFUSERS_SDXL,
ModelType.DIFFUSERS_SD_INPAINT, ModelType.DIFFUSERS_SD_INPAINT,
ModelType.DIFFUSERS_SDXL_INPAINT, ModelType.DIFFUSERS_SDXL_INPAINT,
] or self.name in [Kandinsky22.name, PowerPaint.name] ] or self.name in [KANDINSKY22_NAME, POWERPAINT_NAME]
@computed_field @computed_field
@property @property
@ -86,7 +88,7 @@ class ModelInfo(BaseModel):
ModelType.DIFFUSERS_SDXL, ModelType.DIFFUSERS_SDXL,
ModelType.DIFFUSERS_SD_INPAINT, ModelType.DIFFUSERS_SD_INPAINT,
ModelType.DIFFUSERS_SDXL_INPAINT, ModelType.DIFFUSERS_SDXL_INPAINT,
] or self.name in [PowerPaint.name] ] or self.name in [POWERPAINT_NAME]
@computed_field @computed_field
@property @property
@ -96,4 +98,4 @@ class ModelInfo(BaseModel):
ModelType.DIFFUSERS_SDXL, ModelType.DIFFUSERS_SDXL,
ModelType.DIFFUSERS_SD_INPAINT, ModelType.DIFFUSERS_SD_INPAINT,
ModelType.DIFFUSERS_SDXL_INPAINT, ModelType.DIFFUSERS_SDXL_INPAINT,
] or self.name in [InstructPix2Pix.name] ] or self.name in [INSTRUCT_PIX2PIX_NAME]

View File

@ -7,7 +7,7 @@ controlnet-aux==0.0.3
fastapi==0.108.0 fastapi==0.108.0
python-socketio==5.7.2 python-socketio==5.7.2
typer typer
pydantic pydantic>=2.5.2
rich rich
loguru loguru
yacs yacs