Fix model lookup by HuggingFace ID

Allow using full HuggingFace model ID (e.g., Faber8/AbyssOrangeMix2_nsfw)
as --model argument by looking up both short name and full ID in MODELS.
parent 96adc8b3
...@@ -8001,17 +8001,32 @@ def main(args): ...@@ -8001,17 +8001,32 @@ def main(args):
max_mem["cpu"] = f"{args.system_ram_limit}GiB" max_mem["cpu"] = f"{args.system_ram_limit}GiB"
if is_main: if is_main:
print(f"--- Target: {args.model.upper()} ({MODELS[args.model]['vram']}) ---") # Look up model by name or HuggingFace ID
model_key = None
for name, info in MODELS.items():
if name == args.model or info.get("id") == args.model:
model_key = name
break
if model_key:
print(f"--- Target: {model_key.upper()} ({MODELS[model_key]['vram']}) ---")
else:
print(f"--- Target: {args.model.upper()} (unknown) ---")
log_memory() log_memory()
# Validate model exists # Validate model exists (check both name and HuggingFace ID)
if args.model not in MODELS: model_key = None
for name, info in MODELS.items():
if name == args.model or info.get("id") == args.model:
model_key = name
break
if not model_key:
print(f"❌ Model '{args.model}' not found in database.") print(f"❌ Model '{args.model}' not found in database.")
print(f" Use --model-list to see available models") print(f" Use --model-list to see available models")
print(f" Or use --search-models to find models on HuggingFace") print(f" Or use --search-models to find models on HuggingFace")
sys.exit(1) sys.exit(1)
m_info = MODELS[args.model] m_info = MODELS[model_key]
# Determine task type based on arguments # Determine task type based on arguments
model_id = m_info["id"] model_id = m_info["id"]
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment