mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-22 05:32:12 -03:00
25 lines
666 B
Python
25 lines
666 B
Python
from typing import Dict, 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" |