diff --git a/iopaint/download.py b/iopaint/download.py index 18e3054..b2ed07d 100644 --- a/iopaint/download.py +++ b/iopaint/download.py @@ -2,7 +2,7 @@ import glob import json import os from functools import lru_cache -from typing import List +from typing import List, Optional from iopaint.schema import ModelType, ModelInfo from loguru import logger @@ -49,7 +49,7 @@ def folder_name_to_show_name(name: str) -> str: @lru_cache(maxsize=512) -def get_sd_model_type(model_abs_path: str) -> ModelType: +def get_sd_model_type(model_abs_path: str) -> Optional[ModelType]: if "inpaint" in Path(model_abs_path).name.lower(): model_type = ModelType.DIFFUSERS_SD_INPAINT else: @@ -65,15 +65,19 @@ def get_sd_model_type(model_abs_path: str) -> ModelType: ) model_type = ModelType.DIFFUSERS_SD_INPAINT except ValueError as e: - if "Trying to set a tensor of shape torch.Size([320, 4, 3, 3])" in str(e): + if "[320, 4, 3, 3]" in str(e): model_type = ModelType.DIFFUSERS_SD else: - logger.info(f"Ignore non sd or sdxl file: {model_abs_path}") + logger.info(f"Ignore non sdxl file: {model_abs_path}") + return + except Exception as e: + logger.error(f"Failed to load {model_abs_path}: {e}") + return return model_type @lru_cache() -def get_sdxl_model_type(model_abs_path: str) -> ModelType: +def get_sdxl_model_type(model_abs_path: str) -> Optional[ModelType]: if "inpaint" in model_abs_path: model_type = ModelType.DIFFUSERS_SDXL_INPAINT else: @@ -93,10 +97,14 @@ def get_sdxl_model_type(model_abs_path: str) -> ModelType: else: model_type = ModelType.DIFFUSERS_SDXL except ValueError as e: - if "but got torch.Size([320, 4, 3, 3])" in str(e): + if "[320, 4, 3, 3]" in str(e): model_type = ModelType.DIFFUSERS_SDXL else: - logger.info(f"Ignore non sd or sdxl file: {model_abs_path}") + logger.info(f"Ignore non sdxl file: {model_abs_path}") + return + except Exception as e: + logger.error(f"Failed to load {model_abs_path}: {e}") + return return model_type @@ -121,6 +129,9 @@ def scan_single_file_diffusion_models(cache_dir) -> List[ModelInfo]: model_type = model_type_cache.get(it.name) if model_type is None: model_type = get_sd_model_type(model_abs_path) + if model_type is None: + continue + model_type_cache[it.name] = model_type res.append( ModelInfo( @@ -152,6 +163,9 @@ def scan_single_file_diffusion_models(cache_dir) -> List[ModelInfo]: model_type = sdxl_model_type_cache.get(it.name) if model_type is None: model_type = get_sdxl_model_type(model_abs_path) + if model_type is None: + continue + sdxl_model_type_cache[it.name] = model_type if stable_diffusion_xl_dir.exists(): with open(sdxl_cache_file, "w", encoding="utf-8") as fw: diff --git a/setup.py b/setup.py index d5ba1f1..000231c 100644 --- a/setup.py +++ b/setup.py @@ -28,7 +28,7 @@ def load_requirements(): # https://setuptools.readthedocs.io/en/latest/setuptools.html#including-data-files setuptools.setup( name="IOPaint", - version="1.3.2", + version="1.3.3", author="PanicByte", author_email="cwq1913@gmail.com", description="Image inpainting, outpainting tool powered by SOTA AI Model",