Files
ComfyUI-Lora-Manager/py/utils/model_utils.py
2025-02-24 20:41:16 +08:00

25 lines
660 B
Python

from typing import Optional
# Base model mapping based on version string
BASE_MODEL_MAPPING = {
"sd-v1-5": "SD1.5",
"sd-v2-1": "SD2.1",
"sdxl": "SDXL",
"sd-v2": "SD2.0",
"flux1": "Flux.1 D",
"flux.1 d": "Flux.1 D",
"illustrious": "IL",
"pony": "Pony"
}
def determine_base_model(version_string: Optional[str]) -> str:
"""Determine base model from version string in safetensors metadata"""
if not version_string:
return "Unknown"
version_lower = version_string.lower()
for key, value in BASE_MODEL_MAPPING.items():
if key in version_lower:
return value
return "Unknown"