Salastil/odysseus
1
0
Fork 0
mirror of https://github.com/pewdiepie-archdaemon/odysseus.git synced 2026-08-01 19:18:35 -04:00
Code Issues Packages Projects Releases Wiki Activity
Files
d8a2059df8e53bc7275c45339849d14c8651e73c
odysseus/services/hwfit/image_models.py
T
pewdiepie-archdaemon d8a2059df8 Merge verified Odysseus fixes
2026-07-23 14:49:02 +00:00

436 lines
15 KiB
Python
Raw Blame History

"""Image generation model registry and VRAM fitting for Cookbook."""
from __future__ import annotations
import json
import re
import time
import urllib.parse
import urllib.request
from typing import Any
# Image models are discovered from HuggingFace collections/search and local cache.
# Keep this empty: source-coded repo IDs become hidden recommendations.
IMAGE_MODEL_REGISTRY: list[dict[str, Any]] = []
HF_IMAGE_COLLECTIONS = [
"stabilityai/image",
"stabilityai/stable-diffusion-35",
"black-forest-labs/flux2",
]
HF_MLX_IMAGE_COLLECTIONS = [
"mlx-community/flux2-klein-mlx",
"mlx-community/inpainting-mlx",
"mlx-community/ddcolor-mlx",
"mlx-community/boogu-image-01-mlx",
]
HF_MLX_IMAGE_REPO_SEEDS: list[str] = []
HF_IMAGE_REPO_SEEDS: list[str] = []
_HF_COLLECTION_CACHE = {"ts": 0.0, "models": []}
_HF_COLLECTION_TTL = 30 * 60
_HF_VARIANT_CACHE: dict[str, dict[str, str]] = {}
_HF_SEARCH_DISABLED_UNTIL = 0.0
def _repo_display_name(repo_id: str) -> str:
name = str(repo_id or "").split("/")[-1]
return name.replace("-", " ").replace("_", " ").strip() or repo_id
def _provider_from_repo(repo_id: str) -> str:
owner = str(repo_id or "").split("/", 1)[0].lower()
return {
"stabilityai": "Stability AI",
"black-forest-labs": "Black Forest Labs",
"tongyi-mai": "Tongyi",
"qwen": "Qwen",
"mlx-community": "mlx-community",
}.get(owner, owner.replace("-", " ").title() if owner else "HuggingFace")
def _infer_capabilities(item: dict[str, Any], repo_id: str) -> list[str]:
tasks = set()
pipeline = str(item.get("pipeline_tag") or "").strip().lower()
if pipeline:
tasks.add(pipeline)
for provider in item.get("availableInferenceProviders") or []:
if isinstance(provider, dict) and provider.get("task"):
tasks.add(str(provider["task"]).strip().lower())
text = f"{repo_id} {' '.join(tasks)}".lower()
caps = []
if "image-to-image" in tasks or "edit" in text or "inpaint" in text:
caps.append("image-editing")
if "inpaint" in text:
caps.append("inpainting")
if "text-to-image" in tasks or not caps:
caps.append("text-to-image")
return caps
def _estimate_image_model(repo_id: str) -> dict[str, Any]:
text = str(repo_id or "").lower()
params_b = 8.0
param_match = re.search(r"(?<![\d.])(\d+(?:\.\d+)?)\s*b(?:\b|[-_])", text)
if param_match:
params_b = max(0.01, float(param_match.group(1)))
if any(k in text for k in ("mi-gan", "big-lama", "lama-")):
return {"params_b": 0.01, "bf16": 1.0, "fp8": 0.7, "q4": 0.5, "quality": 65, "speed": 98, "quant": "BF16"}
quant = "BF16"
if any(k in text for k in ("4bit", "q4", "nf4")):
quant = "Q4"
elif "fp8" in text or "8bit" in text:
quant = "FP8"
bf16 = max(1.0, round(params_b * 2.6 + 3.0, 1))
fp8 = max(0.7, round(params_b * 1.35 + 2.0, 1))
q4 = max(0.5, round(params_b * 0.8 + 1.5, 1))
speed = max(35, min(95, int(98 - params_b * 3)))
quality = max(60, min(88, int(70 + min(params_b, 18) * 0.8)))
return {"params_b": params_b, "bf16": bf16, "fp8": fp8, "q4": q4, "quality": quality, "speed": speed, "quant": quant}
def _params_b_from_item(item: dict[str, Any]) -> float | None:
raw = item.get("numParameters")
if isinstance(raw, (int, float)) and raw > 0:
return max(0.01, round(float(raw) / 1_000_000_000.0, 3))
return None
def _mlx_quantize_estimate(repo_id: str, est: dict[str, Any]) -> dict[str, Any]:
text = str(repo_id or "").lower()
out = dict(est)
if "3bit" in text or "4bit" in text or "q4" in text:
out["quant"] = "Q4"
out["bf16"] = None
out["fp8"] = None
elif "8bit" in text:
out["quant"] = "FP8"
out["bf16"] = None
elif "6bit" in text or "5bit" in text:
out["quant"] = "Q4"
out["bf16"] = None
out["fp8"] = out.get("fp8") or out.get("q4")
elif "bf16" in text or "fp16" in text:
out["quant"] = "BF16"
out["fp8"] = None
out["q4"] = None
return out
def _collection_item_to_model(item: dict[str, Any], collection_title: str = "", mlx_only: bool = False) -> dict[str, Any] | None:
repo_id = str(item.get("id") or "").strip()
if "/" not in repo_id:
return None
typ = str(item.get("type") or item.get("itemType") or "model").lower()
if typ not in {"", "model"}:
return None
est = _estimate_image_model(repo_id)
item_params_b = _params_b_from_item(item)
if item_params_b is not None:
est = {
**est,
"params_b": item_params_b,
"bf16": max(0.5, round(item_params_b * 2.4 + 0.8, 1)),
"fp8": max(0.5, round(item_params_b * 1.3 + 0.5, 1)),
"q4": max(0.4, round(item_params_b * 0.8 + 0.4, 1)),
}
if mlx_only:
est = _mlx_quantize_estimate(repo_id, est)
caps = _infer_capabilities(item, repo_id)
gated = item.get("gated")
desc_bits = []
if collection_title:
desc_bits.append(f"HF collection: {collection_title}.")
if gated:
desc_bits.append("Gated on HuggingFace.")
out = {
"id": repo_id,
"name": _repo_display_name(repo_id),
"provider": _provider_from_repo(repo_id),
"params_b": est["params_b"],
"vram_bf16": est["bf16"],
"vram_fp8": est["fp8"],
"vram_q4": est["q4"],
"default_quant": est["quant"],
"quant_repos": {},
"capabilities": caps,
"description": " ".join(desc_bits).strip() or "Imported from HuggingFace collection.",
"quality": est["quality"],
"speed": est["speed"],
"released": "",
}
if mlx_only:
out["mlx_only"] = True
out["description"] = (out["description"] + " Apple Silicon / MLX only.").strip()
return out
def _fetch_hf_image_collection_models() -> list[dict[str, Any]]:
now = time.time()
if now - float(_HF_COLLECTION_CACHE.get("ts") or 0) < _HF_COLLECTION_TTL:
return list(_HF_COLLECTION_CACHE.get("models") or [])
models: list[dict[str, Any]] = []
for slug, mlx_only in [(slug, False) for slug in HF_IMAGE_COLLECTIONS] + [(slug, True) for slug in HF_MLX_IMAGE_COLLECTIONS]:
url = f"https://huggingface.co/api/collections/{slug}"
try:
req = urllib.request.Request(url, headers={"User-Agent": "Odysseus-Cookbook/1.0"})
with urllib.request.urlopen(req, timeout=2.5) as resp:
data = json.loads(resp.read().decode("utf-8", "replace"))
except Exception:
continue
title = str(data.get("title") or slug)
for item in data.get("items") or []:
if isinstance(item, dict):
model = _collection_item_to_model(item, title, mlx_only=mlx_only)
if model:
models.append(model)
_HF_COLLECTION_CACHE["ts"] = now
_HF_COLLECTION_CACHE["models"] = models
return list(models)
def _hf_model_search(query: str, limit: int = 10) -> list[dict[str, Any]]:
global _HF_SEARCH_DISABLED_UNTIL
now = time.time()
if now < _HF_SEARCH_DISABLED_UNTIL:
return []
url = "https://huggingface.co/api/models?" + urllib.parse.urlencode({
"search": query,
"limit": str(limit),
})
try:
req = urllib.request.Request(url, headers={"User-Agent": "Odysseus-Cookbook/1.0"})
with urllib.request.urlopen(req, timeout=2.5) as resp:
data = json.loads(resp.read().decode("utf-8", "replace"))
return data if isinstance(data, list) else []
except Exception:
_HF_SEARCH_DISABLED_UNTIL = now + 10 * 60
return []
def _variant_score(candidate: dict[str, Any], base_repo: str, want: str) -> float:
rid = str(candidate.get("id") or candidate.get("modelId") or "")
text = " ".join([
rid,
str(candidate.get("library_name") or ""),
str(candidate.get("pipeline_tag") or ""),
" ".join(str(t) for t in candidate.get("tags") or []),
]).lower()
base = base_repo.lower()
base_short = base_repo.rsplit("/", 1)[-1].lower()
if want == "gguf" and "gguf" not in text:
return -1
if want == "fp8" and not any(k in text for k in ("fp8", "nvfp4", "mxfp8", "mxfp4")):
return -1
score = float(candidate.get("downloads") or 0) / 1000.0 + float(candidate.get("likes") or 0)
if f"base_model:{base}" in text or f"base_model:quantized:{base}" in text:
score += 10000
elif base_short and base_short in rid.lower():
score += 1000
else:
score -= 200
if "diffusers" in text:
score += 50
if str(candidate.get("private")).lower() == "true":
score -= 10000
return score
def _best_variant_repo(base_repo: str, want: str) -> str:
base_short = str(base_repo or "").rsplit("/", 1)[-1]
candidates = _hf_model_search(f"{base_short} {want}", limit=12)
scored = []
for item in candidates:
if not isinstance(item, dict):
continue
rid = str(item.get("id") or item.get("modelId") or "").strip()
if "/" not in rid or rid.lower() == base_repo.lower():
continue
score = _variant_score(item, base_repo, want)
if score >= 0:
scored.append((score, rid))
scored.sort(reverse=True)
return scored[0][1] if scored else ""
def _should_discover_variants(repo_id: str) -> bool:
return False
def _discover_quant_repos(repo_id: str, need_fp8: bool = True, need_gguf: bool = True) -> dict[str, str]:
key = str(repo_id or "").strip()
if not key:
return {}
cache_key = f"{key.lower()}|fp8={int(need_fp8)}|gguf={int(need_gguf)}"
if cache_key in _HF_VARIANT_CACHE:
return dict(_HF_VARIANT_CACHE[cache_key])
found: dict[str, str] = {}
if need_fp8:
fp8 = _best_variant_repo(key, "fp8")
if fp8:
found["FP8"] = fp8
if need_gguf:
gguf = _best_variant_repo(key, "gguf")
if gguf:
# The image-model fitter's smallest bucket is Q4; most HF image GGUF
# repos expose Q4/Q5/Q8 files under one repo, so use it as the low-VRAM
# download source while preserving the explicit GGUF label for callers.
found["Q4"] = gguf
found["GGUF"] = gguf
_HF_VARIANT_CACHE[cache_key] = found
return dict(found)
def _merge_quant_repos(model: dict[str, Any]) -> dict[str, Any]:
out = dict(model)
existing = dict(out.get("quant_repos") or {})
repo_id = str(out.get("id") or "")
if _should_discover_variants(repo_id):
discovered = _discover_quant_repos(
repo_id,
need_fp8="FP8" not in existing,
need_gguf="Q4" not in existing and "GGUF" not in existing,
)
for k, v in discovered.items():
existing.setdefault(k, v)
out["quant_repos"] = existing
return out
def get_image_models():
"""Return the image model registry."""
merged = [_merge_quant_repos(m) for m in IMAGE_MODEL_REGISTRY]
seen = {str(m.get("id") or "").lower() for m in merged if isinstance(m, dict)}
for model in _fetch_hf_image_collection_models():
key = str(model.get("id") or "").lower()
if key and key not in seen:
merged.append(_merge_quant_repos(model))
seen.add(key)
return merged
def _is_apple_image_system(system: dict[str, Any]) -> bool:
backend = str(system.get("backend") or "").lower()
gpu_name = str(system.get("gpu_name") or "").lower()
cpu_name = str(system.get("cpu_name") or "").lower()
platform = str(system.get("platform") or "").lower()
return (
bool(system.get("unified_memory"))
or backend in {"metal", "mps", "apple"}
or "apple" in gpu_name
or "apple" in cpu_name
or platform == "darwin"
)
def rank_image_models(system, search=None, sort="fit"):
"""Score and rank image models against detected hardware.
Returns list of models with fit info (vram needed, fits, recommended quant).
"""
if not isinstance(system, dict):
system = {}
gpu_vram = system.get("gpu_vram_gb", 0) or 0
has_gpu = system.get("has_gpu", False)
ram_gb = system.get("available_ram_gb") or system.get("total_ram_gb") or 0
budget_gb = gpu_vram if has_gpu and gpu_vram > 0 else ram_gb
budget_kind = "gpu" if has_gpu and gpu_vram > 0 else "ram"
apple_system = _is_apple_image_system(system)
results = []
for model in get_image_models():
if apple_system and not (model.get("mlx_only") or model.get("apple_ok")):
continue
if model.get("mlx_only") and not apple_system:
continue
# Filter by search
if isinstance(search, str) and search:
s = search.lower()
if s not in model["name"].lower() and s not in model["id"].lower() and s not in model.get("description", "").lower():
continue
# Determine best quant that fits
quant = None
vram_needed = None
fits = False
quant_repo = None
if budget_gb > 0:
# Try BF16 first, then FP8, then Q4
for q, vram_key in [("BF16", "vram_bf16"), ("FP8", "vram_fp8"), ("Q4", "vram_q4")]:
v = model.get(vram_key)
if v is not None and v <= budget_gb * 0.90: # 10% headroom
quant = q
vram_needed = v
fits = True
quant_repo = model.get("quant_repos", {}).get(q)
break
# If nothing fits, show what it needs
if not fits:
quant = model["default_quant"]
vram_needed = model.get("vram_bf16", 0)
# Fit label
if budget_gb <= 0:
fit = "no_gpu"
fit_label = "No GPU"
elif fits:
headroom = budget_gb - vram_needed
if headroom > budget_gb * 0.3:
fit = "perfect"
fit_label = "Perfect"
elif headroom > budget_gb * 0.1:
fit = "good"
fit_label = "Good"
else:
fit = "tight"
fit_label = "Tight"
else:
fit = "no_fit"
fit_label = "Too large"
# Score: quality * speed * fit bonus
score = model["quality"] * 0.6 + model["speed"] * 0.2
if fit == "perfect":
score += 20
elif fit == "good":
score += 10
elif fit == "tight":
score += 5
elif fit == "no_fit":
score -= 30
results.append({
"id": model["id"],
"name": model["name"],
"provider": model["provider"],
"params_b": model["params_b"],
"vram_needed": vram_needed,
"quant": quant,
"quant_repo": quant_repo,
"fits": fits,
"fit": fit,
"fit_label": fit_label,
"fit_budget": budget_kind,
"quality": model["quality"],
"speed": model["speed"],
"score": round(score, 1),
"capabilities": model["capabilities"],
"description": model["description"],
"released": model.get("released", ""),
})
# Sort
if sort == "quality":
results.sort(key=lambda x: (-x["quality"], -x["score"]))
elif sort == "speed":
results.sort(key=lambda x: (-x["speed"], -x["score"]))
elif sort == "vram":
results.sort(key=lambda x: (x["vram_needed"] or 999, -x["score"]))
else: # fit (default)
results.sort(key=lambda x: (-x["score"],))
return results
Reference in New Issue View Git Blame Copy Permalink