fix import issue

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"

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}...")

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
 

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

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"

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]

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