102 lines
2.9 KiB
Python
102 lines
2.9 KiB
Python
from typing import List
|
|
|
|
from pydantic import computed_field, BaseModel
|
|
|
|
from iopaint.const import (
|
|
SDXL_CONTROLNET_CHOICES,
|
|
SD2_CONTROLNET_CHOICES,
|
|
SD_CONTROLNET_CHOICES,
|
|
INSTRUCT_PIX2PIX_NAME,
|
|
KANDINSKY22_NAME,
|
|
POWERPAINT_NAME,
|
|
)
|
|
from iopaint.schema import ModelType
|
|
|
|
|
|
class ModelInfo(BaseModel):
|
|
name: str
|
|
path: str
|
|
model_type: ModelType
|
|
is_single_file_diffusers: bool = False
|
|
|
|
@computed_field
|
|
@property
|
|
def need_prompt(self) -> bool:
|
|
return self.model_type in [
|
|
ModelType.DIFFUSERS_SD,
|
|
ModelType.DIFFUSERS_SDXL,
|
|
ModelType.DIFFUSERS_SD_INPAINT,
|
|
ModelType.DIFFUSERS_SDXL_INPAINT,
|
|
] or self.name in [
|
|
INSTRUCT_PIX2PIX_NAME,
|
|
KANDINSKY22_NAME,
|
|
POWERPAINT_NAME,
|
|
]
|
|
|
|
@computed_field
|
|
@property
|
|
def controlnets(self) -> List[str]:
|
|
if self.model_type in [
|
|
ModelType.DIFFUSERS_SDXL,
|
|
ModelType.DIFFUSERS_SDXL_INPAINT,
|
|
]:
|
|
return SDXL_CONTROLNET_CHOICES
|
|
if self.model_type in [ModelType.DIFFUSERS_SD, ModelType.DIFFUSERS_SD_INPAINT]:
|
|
if "sd2" in self.name.lower():
|
|
return SD2_CONTROLNET_CHOICES
|
|
else:
|
|
return SD_CONTROLNET_CHOICES
|
|
if self.name == POWERPAINT_NAME:
|
|
return SD_CONTROLNET_CHOICES
|
|
return []
|
|
|
|
@computed_field
|
|
@property
|
|
def support_strength(self) -> bool:
|
|
return self.model_type in [
|
|
ModelType.DIFFUSERS_SD,
|
|
ModelType.DIFFUSERS_SDXL,
|
|
ModelType.DIFFUSERS_SD_INPAINT,
|
|
ModelType.DIFFUSERS_SDXL_INPAINT,
|
|
]
|
|
|
|
@computed_field
|
|
@property
|
|
def support_outpainting(self) -> bool:
|
|
return self.model_type in [
|
|
ModelType.DIFFUSERS_SD,
|
|
ModelType.DIFFUSERS_SDXL,
|
|
ModelType.DIFFUSERS_SD_INPAINT,
|
|
ModelType.DIFFUSERS_SDXL_INPAINT,
|
|
] or self.name in [KANDINSKY22_NAME, POWERPAINT_NAME]
|
|
|
|
@computed_field
|
|
@property
|
|
def support_lcm_lora(self) -> bool:
|
|
return self.model_type in [
|
|
ModelType.DIFFUSERS_SD,
|
|
ModelType.DIFFUSERS_SDXL,
|
|
ModelType.DIFFUSERS_SD_INPAINT,
|
|
ModelType.DIFFUSERS_SDXL_INPAINT,
|
|
]
|
|
|
|
@computed_field
|
|
@property
|
|
def support_controlnet(self) -> bool:
|
|
return self.model_type in [
|
|
ModelType.DIFFUSERS_SD,
|
|
ModelType.DIFFUSERS_SDXL,
|
|
ModelType.DIFFUSERS_SD_INPAINT,
|
|
ModelType.DIFFUSERS_SDXL_INPAINT,
|
|
]
|
|
|
|
@computed_field
|
|
@property
|
|
def support_freeu(self) -> bool:
|
|
return self.model_type in [
|
|
ModelType.DIFFUSERS_SD,
|
|
ModelType.DIFFUSERS_SDXL,
|
|
ModelType.DIFFUSERS_SD_INPAINT,
|
|
ModelType.DIFFUSERS_SDXL_INPAINT,
|
|
] or self.name in [INSTRUCT_PIX2PIX_NAME]
|