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, ANYTEXT_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, ANYTEXT_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, ] or self.name in [POWERPAINT_NAME, ANYTEXT_NAME] @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]